/*
 * zebra_detect.cpp — 斑马线检测实现
 * 算法对齐 gd13.py: 透视变换 → 白色掩码 → Sobel 梯度 → 滑窗判定
 */
#define ZEBRA_DEBUG 1

#include "zebra_detect.h"
#include <cmath>
#include <cstdio>

using namespace cv;
using namespace std;

#if ZEBRA_DEBUG
#define ZLOG(fmt, ...) fprintf(stderr, "[zebra] " fmt "\n", ##__VA_ARGS__)
#else
#define ZLOG(fmt, ...) ((void)0)
#endif

// ============================================================
// 参数 (对齐 gd13.py)
// ============================================================
static const int   PROC_SIZE       = 400;
static const int   WIN_H           = 120;
static const int   WIN_W           = 300;
static const int   WIN_STEP        = 20;

static const float AMP_THRESH      = 12.0f;
static const int   COUNT_THRESH    = 150;
static const float MIN_VRATIO      = 0.45f;
static const int   MIN_VPIXELS     = 300;

// ============================================================
// 透视变换矩阵
// ============================================================
static Mat get_birdview_transform()
{
    vector<Point2f> src = {
        {30.f,  380.f},
        {370.f, 380.f},
        {320.f, 180.f},
        {80.f,  180.f}
    };
    vector<Point2f> dst = {
        {50.f,  350.f},
        {350.f, 350.f},
        {350.f, 50.f},
        {50.f,  50.f}
    };
    return getPerspectiveTransform(src, dst);
}

// ============================================================
// 白色掩码 (HSV)
// ============================================================
static Mat get_white_zebra_mask(const Mat &bgr)
{
    Mat hsv;
    cvtColor(bgr, hsv, COLOR_BGR2HSV);

    Mat white_mask;
    inRange(hsv, Scalar(0, 0, 180), Scalar(180, 30, 255), white_mask);
    ZLOG("  白掩码 H:[0,180] S:[0,30] V:[180,255] 白色像素=%d", countNonZero(white_mask));

    Mat brown_mask;
    inRange(hsv, Scalar(10, 20, 80), Scalar(30, 100, 200), brown_mask);
    ZLOG("  棕掩码 H:[10,30] S:[20,100] V:[80,200] 棕色像素=%d", countNonZero(brown_mask));

    Mat not_brown;
    bitwise_not(brown_mask, not_brown);
    bitwise_and(white_mask, not_brown, white_mask);
    ZLOG("  排除棕色后白色像素=%d", countNonZero(white_mask));

    Mat kernel = getStructuringElement(MORPH_RECT, Size(2, 2));
    morphologyEx(white_mask, white_mask, MORPH_OPEN,  kernel, Point(-1,-1), 1);
    morphologyEx(white_mask, white_mask, MORPH_CLOSE, kernel, Point(-1,-1), 1);
    ZLOG("  开闭运算后白色像素=%d", countNonZero(white_mask));

    return white_mask;
}

// ============================================================
// 预处理
// ============================================================
static Mat preprocess(const Mat &bgr)
{
    Mat white_mask = get_white_zebra_mask(bgr);
    Mat gray;
    cvtColor(bgr, gray, COLOR_BGR2GRAY);
    bitwise_and(gray, gray, gray, white_mask);

    medianBlur(gray, gray, 3);
    Mat kernel1 = getStructuringElement(MORPH_RECT, Size(3, 3));
    Mat kernel2 = getStructuringElement(MORPH_RECT, Size(5, 5));
    morphologyEx(gray, gray, MORPH_OPEN,  kernel1, Point(-1,-1), 1);
    morphologyEx(gray, gray, MORPH_CLOSE, kernel2, Point(-1,-1), 1);

    return gray;
}

// ============================================================
// 梯度计算
// ============================================================
static void compute_gradient(const Mat &gray, Mat &amplitude, Mat &theta)
{
    Mat sobelx, sobely;
    Sobel(gray, sobelx, CV_32F, 1, 0, 3);
    Sobel(gray, sobely, CV_32F, 0, 1, 3);

    phase(sobelx, sobely, theta, true);
    for (int r = 0; r < theta.rows; ++r)
        for (int c = 0; c < theta.cols; ++c)
            theta.at<float>(r, c) = fmod(theta.at<float>(r, c), 180.f);

    magnitude(sobelx, sobely, amplitude);

    // 统计过滤前
    int before = countNonZero(amplitude);
    for (int r = 0; r < amplitude.rows; ++r)
        for (int c = 0; c < amplitude.cols; ++c)
            if (amplitude.at<float>(r, c) < 10.f)
                amplitude.at<float>(r, c) = 0.f;

    int after = countNonZero(amplitude);
    double amp_min, amp_max;
    minMaxLoc(amplitude, &amp_min, &amp_max);
    ZLOG("  梯度: 过滤前=%d, 过滤后(amp>10)=%d, 幅值范围=[%.1f, %.1f]",
         before, after, amp_min, amp_max);
}

// ============================================================
// 单窗口判定 (无 vector 分配, 避免 OOM)
// ============================================================
static bool is_zebra_window(const Mat &amp_win, const Mat &ang_win, int win_top)
{
    ZLOG("--- 窗口 top=%d [%dx%d] ---", win_top, amp_win.cols, amp_win.rows);

    int total     = 0;      // amplitude > AMP_THRESH 的像素总数
    int hist[4]   = {0};    // 水平边缘直方图 [80-85,85-90,90-95,95-100)
    int v_count   = 0;      // 垂直边缘计数 [0-10]∪[170-180]

    // 单次遍历: 同时统计 total, hist, v_count
    for (int r = 0; r < amp_win.rows; ++r)
    {
        for (int c = 0; c < amp_win.cols; ++c)
        {
            float amp = amp_win.at<float>(r, c);
            if (amp <= AMP_THRESH)
                continue;

            ++total;
            float a = ang_win.at<float>(r, c);

            if (a >= 80.f && a < 100.f)
            {
                int bin = (int)(a - 80.f) / 5;
                if (bin >= 0 && bin < 4)
                    hist[bin]++;
            }

            if (a <= 10.f || a >= 170.f)
                ++v_count;
        }
    }

    ZLOG("  梯度像素(amplitude>%.0f)=%d", AMP_THRESH, total);
    if (total == 0)
    {
        ZLOG("  -> 丢弃: 无有效梯度像素");
        return false;
    }

    // 找峰值 bin
    int peak = 0, peak_bin = 0;
    for (int i = 0; i < 4; ++i)
        if (hist[i] > peak) { peak = hist[i]; peak_bin = i; }

    // 二次遍历: 统计峰值方向的像素数 (需要知道具体区间)
    int h_count = 0;
    float low_a  = 80.f + peak_bin * 5.f;
    float high_a = low_a + 5.f;
    for (int r = 0; r < amp_win.rows; ++r)
    {
        for (int c = 0; c < amp_win.cols; ++c)
        {
            if (amp_win.at<float>(r, c) <= AMP_THRESH)
                continue;
            float a = ang_win.at<float>(r, c);
            if (a >= low_a && a <= high_a)
                ++h_count;
        }
    }

    ZLOG("  水平边缘[80-100): hist=[%d,%d,%d,%d], 峰值bin=%d(%.0f-%.0f度), "
         "峰值像素=%d, 阈值=%d",
         hist[0], hist[1], hist[2], hist[3],
         peak_bin, low_a, high_a, h_count, COUNT_THRESH);

    if (h_count < COUNT_THRESH)
    {
        ZLOG("  -> 丢弃: 水平边缘像素不足 (%d < %d)", h_count, COUNT_THRESH);
        return false;
    }

    float v_ratio = (float)v_count / (float)(total + 1e-6);
    ZLOG("  垂直边缘[0-10|170-180]: 像素=%d, 占比=%.3f, 要求: >=%d && 占比>=%.2f",
         v_count, v_ratio, MIN_VPIXELS, MIN_VRATIO);

    if (v_count < MIN_VPIXELS)
    {
        ZLOG("  -> 丢弃: 垂直像素不足 (%d < %d)", v_count, MIN_VPIXELS);
        return false;
    }
    if (v_ratio < MIN_VRATIO)
    {
        ZLOG("  -> 丢弃: 垂直占比不足 (%.3f < %.2f)", v_ratio, MIN_VRATIO);
        return false;
    }

    ZLOG("  -> 判定: 斑马线窗口!");
    return true;
}

// ============================================================
// detect_zebra_crossing — 主检测函数
// ============================================================
ZebraResult detect_zebra_crossing(const Mat &bgr_frame)
{
    ZebraResult result;
    result.detected    = false;
    result.distance_px = -1;

    ZLOG("========================================");
    ZLOG("新帧开始: 输入尺寸=%dx%d", bgr_frame.cols, bgr_frame.rows);

    if (bgr_frame.empty())
    {
        ZLOG("错误: 输入帧为空");
        return result;
    }

    static Mat M = get_birdview_transform();
    static Mat proc, birdview, roi, gray, amplitude, theta;

    // 1. 缩放 → 透视变换 → 裁剪 ROI(顶部30px) → 恢复 400×400
    resize(bgr_frame, proc, Size(PROC_SIZE, PROC_SIZE));
    ZLOG("步骤1-缩放: %dx%d -> %dx%d", bgr_frame.cols, bgr_frame.rows,
         PROC_SIZE, PROC_SIZE);

    warpPerspective(proc, birdview, M, Size(PROC_SIZE, PROC_SIZE));
    roi = birdview(Rect(0, 30, PROC_SIZE, PROC_SIZE - 30));
    resize(roi, proc, Size(PROC_SIZE, PROC_SIZE));
    ZLOG("步骤1-透视+ROI: 裁剪顶部30px, 鸟瞰输出=%dx%d", proc.cols, proc.rows);

    // 2. 预处理 + 梯度
    ZLOG("步骤2-预处理开始");
    gray = preprocess(proc);
    ZLOG("步骤2-预处理完成, 灰度非零像素=%d", countNonZero(gray));

    ZLOG("步骤3-梯度计算开始");
    compute_gradient(gray, amplitude, theta);

    int amp_nz = countNonZero(amplitude);
    ZLOG("步骤3-梯度完成, 总非零梯度像素=%d", amp_nz);

    if (amp_nz == 0)
    {
        ZLOG("结果: 无梯度 → 未检测到斑马线");
        return result;
    }

    // 3. 滑动窗口
    ZLOG("步骤4-滑动窗口: 窗口=%dx%d, 步长=%d, 遍历范围=[0,%d]",
         WIN_W, WIN_H, WIN_STEP, proc.rows - WIN_H);

    int z_ymin = proc.rows;
    int z_ymax = 0;
    int valid = 0;
    int total_wins = 0;

    for (int top = 0; top <= proc.rows - WIN_H; top += WIN_STEP)
    {
        ++total_wins;
        Rect win(0, top, WIN_W, WIN_H);
        if (is_zebra_window(amplitude(win), theta(win), top))
        {
            ++valid;
            if (top < z_ymin)          z_ymin = top;
            if (top + WIN_H > z_ymax)  z_ymax = top + WIN_H;
        }
    }

    // 4. 判定
    int span = z_ymax - z_ymin;
    ZLOG("步骤5-汇总: 总窗口=%d, 有效=%d, "
         "斑马线区域: top=%d~%d, 跨度=%d, 判定阈值=%d",
         total_wins, valid, z_ymin, z_ymax, span, WIN_H);

    if (span > WIN_H)
    {
        result.detected    = true;
        result.distance_px = PROC_SIZE - z_ymax;
        ZLOG("结果: 检测到斑马线! 下沿距底部=%dpx (ymax=%d)",
             result.distance_px, z_ymax);
    }
    else
    {
        ZLOG("结果: 未检测到斑马线 (跨度%d <= 阈值%d)", span, WIN_H);
    }

    return result;
}