// count.cpp - 计数模式(count)核心算法与文件保存实现
#include "count.h"
#include "config.h"
#include "filter.h"
#include "io_utils.h"
#include <algorithm>
#include <cmath>
#include <fstream>
#include <iostream>
#include <iomanip>
#include <sstream>
#include <thread>
#include <unordered_map>

// ==================== 点集 → 图像 ====================
std::vector<uint8_t> points_to_image(const std::vector<CountPoint>& points,
                                     double* area_out,
                                     int* out_w, int* out_h)
{
    const int MIN_DIM = 8;     // 最小尺寸保护
    const int MAX_DIM = 20000; // 单边最大尺寸保护(防止z过大导致内存爆炸)
    if (area_out) *area_out = 0.0;
    if (out_w) *out_w = 0;
    if (out_h) *out_h = 0;
    if (points.empty()) return std::vector<uint8_t>((size_t)MIN_DIM * MIN_DIM, 0);

    // 宽: Filter_Polygon 多边形 X 最大最小值直接构成 (1mm=1px)
    double polyXMin = 0.0, polyXMax = 0.0;
    if (Filter_Poly_N > 0)
    {
        polyXMin = Filter_Polygon[0].first;
        polyXMax = Filter_Polygon[0].first;
        for (int i = 1; i < Filter_Poly_N; i++)
        {
            polyXMin = std::min(polyXMin, Filter_Polygon[i].first);
            polyXMax = std::max(polyXMax, Filter_Polygon[i].first);
        }
    }
    double spanX = polyXMax - polyXMin;
    if (spanX <= 0.0) spanX = 1.0;

    // 高: z 最大值直接构成 (Z=运动方向, 从0开始递增, 1单位=1px)
    double maxZ = 0.0;
    for (const auto& p : points) maxZ = std::max(maxZ, p.z_val);
    if (maxZ <= 0.0) maxZ = 1.0;

    int imageWidth  = (int)std::lround(spanX);
    int imageHeight = (int)std::lround(maxZ);
    if (imageWidth  < MIN_DIM) imageWidth  = MIN_DIM;
    if (imageHeight < MIN_DIM) imageHeight = MIN_DIM;
    if (imageWidth  > MAX_DIM) imageWidth  = MAX_DIM;
    if (imageHeight > MAX_DIM) imageHeight = MAX_DIM;
    if (out_w) *out_w = imageWidth;
    if (out_h) *out_h = imageHeight;
    std::vector<uint8_t> image((size_t)imageWidth * imageHeight, 0);   // 0=白底

    // X坐标相对多边形左边界偏移到0；Z从0开始
    double scaleX = (double)imageWidth / spanX;   // ≈1: 宽直接由多边形X跨度构成
    double scaleY = (double)imageHeight / maxZ;   // ≈1: 高直接由z最大值构成
    double xOffset = -polyXMin;

    // 优化：Step1 先按1x1网格标记点位置（重复点只写一次），
    //      Step2 再对标记位置膨胀 Point_Dilate×Point_Dilate
    std::vector<uint32_t> marks;
    marks.reserve(points.size() / 2);
    for (const auto& p : points)
    {
        int px = (int)((p.x_mm + xOffset) * scaleX);
        int py = (int)(p.z_val * scaleY);
        if (px < 0 || px >= imageWidth || py < 0 || py >= imageHeight) continue;
        size_t idx = (size_t)py * imageWidth + px;
        if (image[idx] == 0)          // 未标记过才记录
        {
            image[idx] = 1;
            marks.push_back((uint32_t)idx);
        }
    }
    for (uint32_t idx : marks)        // 膨胀 Point_Dilate×Point_Dilate
    {
        int px = (int)(idx % imageWidth);
        int py = (int)(idx / imageWidth);
        for (int dy = 0; dy < Point_Dilate; dy++)
            for (int dx = 0; dx < Point_Dilate; dx++)
            {
                int ix = px + dx, iy = py + dy;
                if (ix < imageWidth && iy < imageHeight)
                    image[(size_t)iy * imageWidth + ix] = 255;   // 255=黑点
            }
    }
    int totalPixels = 0;
    for (size_t i = 0; i < image.size(); i++)
        if (image[i] != 0) totalPixels++;

    double pixelToAreaRatio = spanX * maxZ / (double)(imageWidth * imageHeight);
    if (area_out) *area_out = totalPixels * pixelToAreaRatio;
    return image;
}

// ==================== z轴空隙插值 ====================
int z_gap_interpolate(std::vector<CountPoint>& points)
{
    if (points.empty() || Z_Interpolate_Mode == "none")
        return 0;

    // 正常相邻两条(圈)的 z 间隔: 速度模式 = 每圈行进距离(mm); 时间模式 = 每圈耗时(ms)
    double normal_gap = (Speed_Z > 0)
        ? (Speed_Z * 10.0 / Lidar_Hz)   // cm/s ×10 = mm/s, ÷Hz = mm/圈
        : (1000.0 / Lidar_Hz);          // ms/圈
    if (normal_gap <= 0) return 0;

    // 空隙判定: 本条与上一条 z 差 > 2×正常间隔 视为断流空隙
    const double gap_threshold = 2.0 * normal_gap;
    // 条件1: 前后两条 Y 平均值差超过30mm 不插(可能跨层/换物体)
    const double max_y_mean_diff = 30.0;
    // 一个空隙最多补的行数
    const int MAX_FILL_LEVELS = 10;

    // 1. 按 z 升序排序
    std::sort(points.begin(), points.end(),
              [](const CountPoint& a, const CountPoint& b) { return a.z_val < b.z_val; });

    // 2. 聚类成"圈"(同一scan包的点 z 完全相同, 用区间 [begin,end) 索引 points)
    struct Circle { double z; size_t begin; size_t end; };
    std::vector<Circle> circles;
    {
        size_t i = 0;
        while (i < points.size())
        {
            size_t j = i + 1;
            while (j < points.size() && points[j].z_val - points[i].z_val < 1e-6) ++j;
            circles.push_back(Circle());
            circles.back().z = points[i].z_val;
            circles.back().begin = i;
            circles.back().end = j;
            i = j;
        }
    }
    if (circles.size() < 2)
        return 0;

    auto circle_mean_y = [&](const Circle& c) -> double
    {
        double s = 0;
        for (size_t p = c.begin; p < c.end; ++p) s += points[p].y_mm;
        return s / (double)(c.end - c.begin);
    };

    int added = 0;
    // 3. 逐空隙补行
    for (size_t i = 1; i < circles.size(); ++i)
    {
        double gap = circles[i].z - circles[i - 1].z;
        if (gap <= gap_threshold)      // 未超过2×正常间隔, 不算空隙
            continue;

        const Circle& cA = circles[i - 1];   // 上一条
        const Circle& cB = circles[i];       // 本条

        // 条件1: 前后两条的Y平均值几乎相当(差≤30mm)才插值
        if (std::fabs(circle_mean_y(cA) - circle_mean_y(cB)) > max_y_mean_diff)
            continue;

        // 条件2: 生成行的 X 范围 = 这两条的 X 最大和最小
        double x_lo = points[cA.begin].x_mm, x_hi = points[cA.begin].x_mm;
        for (size_t p = cA.begin; p < cA.end; ++p)
        {
            if (points[p].x_mm < x_lo) x_lo = points[p].x_mm;
            if (points[p].x_mm > x_hi) x_hi = points[p].x_mm;
        }
        for (size_t p = cB.begin; p < cB.end; ++p)
        {
            if (points[p].x_mm < x_lo) x_lo = points[p].x_mm;
            if (points[p].x_mm > x_hi) x_hi = points[p].x_mm;
        }

        // 空隙内补行数: 按正常间隔容纳行数-1(空隙两端已有), 最多10行
        int n = (int)(gap / normal_gap + 0.5) - 1;
        if (n < 1) continue;
        if (n > MAX_FILL_LEVELS) n = MAX_FILL_LEVELS;

        if (Z_Interpolate_Mode == "CopyTranslation")
        {
            // 复制"上一条"的点, z 均匀平移到空隙内的各补行
            for (int k = 1; k <= n; ++k)
            {
                double zk = cA.z + gap * (double)k / (n + 1);   // 均匀分布
                for (size_t p = cA.begin; p < cA.end; ++p)
                {
                    const CountPoint& src = points[p];
                    if (src.x_mm < x_lo || src.x_mm > x_hi) continue;   // X范围约束
                    CountPoint np = src;
                    np.z_val = zk;
                    points.push_back(np);
                    ++added;
                }
            }
        }
        else if (Z_Interpolate_Mode == "AngleBucketInterpolation")
        {
            // 角度桶插值: 每桶用前A/后B平均点线性插值(平滑过渡)
            struct Buck { double sx, sy; int cnt; Buck() : sx(0), sy(0), cnt(0) {} };
            std::vector<Buck> bA(360), bB(360);
            auto fill = [&](const Circle& c, std::vector<Buck>& out)
            {
                for (size_t p = c.begin; p < c.end; ++p)
                {
                    double deg = std::atan2(points[p].x_mm, -points[p].y_mm) * 57.29577951308232;
                    if (deg < 0) deg += 360.0;
                    int b = (int)(deg + 0.5);
                    if (b >= 360) b = 0;
                    out[b].sx += points[p].x_mm;
                    out[b].sy += points[p].y_mm;
                    out[b].cnt++;
                }
            };
            fill(cA, bA);
            fill(cB, bB);
            for (int k = 1; k <= n; ++k)
            {
                double zk = cA.z + gap * (double)k / (n + 1);   // 均匀分布
                double t = (zk - cA.z) / gap;                    // 0~1
                for (int b = 0; b < 360; ++b)
                {
                    if (bA[b].cnt <= 0 || bB[b].cnt <= 0) continue;
                    double ax = bA[b].sx / bA[b].cnt, ay = bA[b].sy / bA[b].cnt;
                    double bx = bB[b].sx / bB[b].cnt, by = bB[b].sy / bB[b].cnt;
                    double nx = ax + (bx - ax) * t;
                    double ny = ay + (by - ay) * t;
                    if (nx < x_lo || nx > x_hi) continue;        // X范围约束
                    CountPoint np;
                    np.x_mm = nx;
                    np.y_mm = ny;
                    np.z_val = zk;
                    points.push_back(np);
                    ++added;
                }
            }
        }
    }

    if (added > 0)
        std::cout << "[count] z空隙插值: 模式=" << Z_Interpolate_Mode
                  << ", 正常间隔=" << normal_gap
                  << ", 判定阈值=2×间隔=" << gap_threshold
                  << ", 新增" << added << "点" << std::endl;
    return added;
}

// ==================== 按层统计 ====================
int calc_item_count(const std::vector<CountPoint>& points, const CountParams& box,
                    std::string* detail,
                    std::vector<uint8_t>* top_layer_img,
                    int* top_layer_img_w,
                    int* top_layer_img_h,
                    int* top_layer_level,
                    double* top_layer_area)
{
    if (points.empty() || box.l <= 0 || box.w <= 0 || box.h <= 0 || box.maxPerLayer <= 0)
        return 0;

    std::unordered_map<int, std::vector<CountPoint>> layer_map;
    for (const auto& p : points)
    {
        if (FILTER_ENABLE && !point_in_polygon(p.x_mm, p.y_mm, Filter_Polygon, Filter_Poly_N))
            continue;
        double currentHeight = Floor_Height - p.y_mm;  // 高度 = 地板值 - Y轴
        if (currentHeight < 0) continue;
        int level = (int)std::round(currentHeight / box.h);
        double tolerance = Count_Base_Tolerance + level * Count_Additional_Tolerance_Per_Level;
        if (std::abs(currentHeight - level * box.h) <= tolerance)
            layer_map[level].push_back(p);
    }

    // 每层面积/图像只算一次，详情与结果共用，避免重复构建图像
    std::ostringstream oss;
    std::unordered_map<int, double> layer_area;
    std::unordered_map<int, std::vector<uint8_t>> layer_img;
    std::unordered_map<int, std::pair<int, int>> layer_img_size;
    for (int h = 5; h > 0; h--)
    {
        auto it = layer_map.find(h);
        if (it == layer_map.end()) continue;
        oss << "  层" << h << ": 点数=" << it->second.size();
        // 高度范围: 该层算法对应的理论高度区间(每层都显示)
        double tol = Count_Base_Tolerance + h * Count_Additional_Tolerance_Per_Level;
        double h_lo = h * box.h - tol;
        double h_hi = h * box.h + tol;
        oss << ", 高度范围（" << std::setprecision(0) << h_lo << "," << h_hi << "）";
        if (it->second.size() > 1000)
        {
            double area = 0.0;
            int img_w = 0, img_h = 0;
            std::vector<uint8_t> img = points_to_image(it->second, &area, &img_w, &img_h);
            layer_area[h] = area;
            layer_img[h] = std::move(img);
            layer_img_size[h] = std::make_pair(img_w, img_h);
            double item_cnt = area / (box.l * box.w);   // 未取整的个数(面积/单件面积)
            oss << ", 面积=" << std::fixed << std::setprecision(1) << area << "mm^2"
                << ", 个数" << std::setprecision(2) << item_cnt << "个";
        }
        oss << "\n";
    }

    int layersCount = box.maxPerLayer;   // 每层最大个数
    int result = 0;
    int top_lv = 0;
    double top_area = 0.0;
    for (int h = 5; h > 0; h--)          // 最多5层，从顶层往下找
    {
        auto it = layer_map.find(h);
        if (it != layer_map.end() && it->second.size() > 1000)
        {
            double area = layer_area[h];
            if (area > 0)
            {
                int i = (int)std::min(std::round(area / (box.l * box.w)), (double)layersCount);
                if (i > 0)
                {
                    result = layersCount * (h - 1) + i;
                    top_lv = h;
                    top_area = area;
                    break;
                }
            }
        }
    }
    if (detail) *detail = oss.str();
    // 输出被计数层(最高有效层)的图像/尺寸/层号/面积
    if (top_layer_img && top_lv > 0)
    {
        auto it = layer_img.find(top_lv);
        if (it != layer_img.end()) *top_layer_img = it->second;
        auto sit = layer_img_size.find(top_lv);
        if (sit != layer_img_size.end())
        {
            if (top_layer_img_w) *top_layer_img_w = sit->second.first;
            if (top_layer_img_h) *top_layer_img_h = sit->second.second;
        }
    }
    if (top_layer_level) *top_layer_level = top_lv;
    if (top_layer_area) *top_layer_area = top_area;
    return result;
}

// ==================== count 结果文件后台保存 ====================
void save_count_files_async(const std::vector<CountPoint>& points,
                            const std::vector<uint8_t>& top_layer_img,
                            int top_img_w, int top_img_h,
                            const CountParams& params,
                            const std::string& layer_detail,
                            int total, int top_lv, double top_area)
{
    // C++11 兼容：按值捕获（拷贝到线程），不阻塞 TCP 响应
    std::thread([points, top_layer_img, params, layer_detail, total, top_lv, top_area,
                 top_img_w, top_img_h]() {
        // 目录: data/日期/时间_taskId（由 io_utils 统一创建）
        std::string out_dir = make_task_dir(params.task_id);
        size_t pt_cnt = points.size();

        // PCD：y 换成高度(Floor_Height - y_mm)
        std::vector<double> xs, ys, zs;
        xs.reserve(pt_cnt); ys.reserve(pt_cnt); zs.reserve(pt_cnt);
        for (const auto& p : points)
        {
            xs.push_back(p.x_mm);
            ys.push_back(Floor_Height - p.y_mm);
            zs.push_back(p.z_val);
        }
        save_pcd_ascii(out_dir + "/count_data.pcd", xs, ys, zs);
        std::cout << "[count] 点云已保存: " << out_dir << "/count_data.pcd ("
                  << pt_cnt << " 点, Y=高度)" << std::endl;

        // 图片1：全点X-Z平面图
        double all_area = 0.0;
        int all_w = 0, all_h = 0;
        std::vector<uint8_t> all_img = points_to_image(points, &all_area, &all_w, &all_h);
        if (save_jpg_gray(out_dir + "/count_points.jpg", all_w, all_h, all_img))
            std::cout << "[count] 点云图已保存: " << out_dir << "/count_points.jpg ("
                      << all_w << "x" << all_h << ", X-Z平面, 面积≈"
                      << std::fixed << std::setprecision(1)
                      << all_area << "mm^2)" << std::endl;

        // 图片2：被计数层(最高有效层)图像
        if (top_lv > 0 && !top_layer_img.empty() && top_img_w > 0 && top_img_h > 0)
        {
            if (save_jpg_gray(out_dir + "/count_layer.jpg", top_img_w, top_img_h, top_layer_img))
                std::cout << "[count] 有效层图已保存: " << out_dir << "/count_layer.jpg ("
                          << top_img_w << "x" << top_img_h << ", 层" << top_lv
                          << ", 面积≈" << std::fixed
                          << std::setprecision(1) << top_area << "mm^2)" << std::endl;
        }

        // 日志
        if (Log_Flag)
        {
            std::ofstream lf(out_dir + "/log.txt");
            if (lf.is_open())
            {
                std::ostringstream sum;
                sum << "==============================\n";
                sum << "[count] task_id: " << params.task_id << "\n";
                sum << "[count] box: l=" << params.l
                    << ", w=" << params.w
                    << ", h=" << params.h
                    << ", maxPerLayer=" << params.maxPerLayer << "\n";
                sum << "[count] Floor_Height: " << Floor_Height << " mm\n";
                sum << "[count] 采集点数: " << pt_cnt << "\n";
                sum << "[count] 计数层: 层" << top_lv
                    << " (面积=" << std::fixed << std::setprecision(1)
                    << top_area << "mm^2)\n";
                sum << "[count] 分层情况:\n" << layer_detail;
                sum << "[count] 计数结果: " << total << "\n";
                sum << "[count] 输出文件:\n";
                sum << "  " << out_dir << "/count_data.pcd\n";
                sum << "  " << out_dir << "/count_points.jpg (" << all_w << "x" << all_h << ")\n";
                if (top_lv > 0 && !top_layer_img.empty() && top_img_w > 0 && top_img_h > 0)
                    sum << "  " << out_dir << "/count_layer.jpg ("
                        << top_img_w << "x" << top_img_h << ")\n";
                sum << "==============================\n";
                lf << gbk_to_utf8(sum.str());
                lf.close();
            }
        }
        std::cout << "[count] 输出目录: " << out_dir << std::endl;
    }).detach();
}