#include "thermal_detector.h"
#include <opencv2/imgproc.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/highgui.hpp>
#include <algorithm>
#include <cmath>
#include <numeric>

// ── 构造 ──

ThermalDetector::ThermalDetector(const DetectorConfig& cfg) : cfg_(cfg) {}

// ── 环境温度估计 ──

double ThermalDetector::estimate_ambient(const cv::Mat& temp_map) {
    int h = temp_map.rows, w = temp_map.cols;
    int bh = std::max(1, h / 10);
    int bw = std::max(1, w / 10);

    std::vector<double> border;
    auto collect = [&](const cv::Mat& region) {
        for (int r = 0; r < region.rows; r++) {
            const double* row = region.ptr<double>(r);
            for (int c = 0; c < region.cols; c++)
                border.push_back(row[c]);
        }
    };

    collect(temp_map.rowRange(0, bh));
    collect(temp_map.rowRange(h - bh, h));
    collect(temp_map.colRange(0, bw).rowRange(bh, h - bh));
    collect(temp_map.colRange(w - bw, w).rowRange(bh, h - bh));

    std::sort(border.begin(), border.end());
    return border[border.size() / 2];
}

// ── 形状特征 ──

ThermalDetector::ShapeFeatures ThermalDetector::compute_shape(
    const cv::Mat& comp_mask, int area, int x, int y, int w, int h) {

    ShapeFeatures sf;
    sf.aspect_ratio = std::max(w, h) / (double)std::max(std::min(w, h), 1);
    sf.extent = area / (double)std::max(w * h, 1);
    sf.is_vertical = (h >= w);

    std::vector<std::vector<cv::Point>> contours;
    cv::findContours(comp_mask.clone(), contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

    if (!contours.empty()) {
        auto& contour = *std::max_element(contours.begin(), contours.end(),
            [](const auto& a, const auto& b) { return cv::contourArea(a) < cv::contourArea(b); });

        double perimeter = cv::arcLength(contour, true);
        if (perimeter > 0)
            sf.circularity = 4.0 * CV_PI * area / (perimeter * perimeter);

        std::vector<cv::Point> hull;
        cv::convexHull(contour, hull);
        double hull_area = cv::contourArea(hull);
        sf.solidity = area / std::max(hull_area, 1.0);
    }
    return sf;
}

// ── 温度特征 ──

ThermalDetector::ThermalFeatures ThermalDetector::compute_thermal(
    const cv::Mat& temp_map, const cv::Mat& comp_mask, double /*ambient*/) {

    ThermalFeatures tf;
    cv::Mat masked;
    temp_map.copyTo(masked, comp_mask > 0);

    std::vector<double> temps;
    for (int r = 0; r < masked.rows; r++) {
        const double* row = masked.ptr<double>(r);
        const uchar* mask_row = comp_mask.ptr<uchar>(r);
        for (int c = 0; c < masked.cols; c++) {
            if (mask_row[c] > 0)
                temps.push_back(row[c]);
        }
    }

    if (temps.empty()) return tf;

    tf.max_t  = *std::max_element(temps.begin(), temps.end());
    tf.min_t  = *std::min_element(temps.begin(), temps.end());
    tf.mean_t = std::accumulate(temps.begin(), temps.end(), 0.0) / temps.size();

    double sq_sum = 0;
    for (double t : temps)
        sq_sum += (t - tf.mean_t) * (t - tf.mean_t);
    tf.stddev_t = std::sqrt(sq_sum / temps.size());

    return tf;
}

// ── 辅助 ──

bool ThermalDetector::near_person(double cx, double cy,
                                   const std::vector<DetectionResult>& persons,
                                   double max_ratio) {
    for (const auto& p : persons) {
        double psize = std::max(p.bbox.width, p.bbox.height);
        double dist = std::sqrt((cx - p.centroid.x) * (cx - p.centroid.x) +
                                (cy - p.centroid.y) * (cy - p.centroid.y));
        if (dist < max_ratio * psize)
            return true;
    }
    return false;
}

double ThermalDetector::compute_iou(const cv::Rect& a, const cv::Rect& b) {
    int ix = std::max(0, std::min(a.x + a.width, b.x + b.width) - std::max(a.x, b.x));
    int iy = std::max(0, std::min(a.y + a.height, b.y + b.height) - std::max(a.y, b.y));
    double inter = ix * iy;
    double union_area = a.width * a.height + b.width * b.height - inter;
    return inter / std::max(union_area, 1.0);
}

double ThermalDetector::body_temp_score(double mean_temp, double stddev) const {
    double temp_score;
    if (mean_temp >= 34.0 && mean_temp <= 40.0) temp_score = 1.0;
    else if (mean_temp >= 30.0 && mean_temp < 34.0) temp_score = 0.6;
    else if (mean_temp > 40.0 && mean_temp <= 44.0) temp_score = 0.5;
    else temp_score = 0.1;

    double uniform_score = std::max(0.0, 1.0 - stddev / 15.0);
    return temp_score * 0.7 + uniform_score * 0.3;
}

double ThermalDetector::barrel_confidence(double max_temp, double ambient,
                                           double solidity, double ar) const {
    double ideal_ar = (cfg_.barrel_ar_min + cfg_.barrel_ar_max) / 2.0;
    double ar_score = 1.0 - std::abs(ar - ideal_ar) / (cfg_.barrel_ar_max - cfg_.barrel_ar_min);
    double temp_score = std::min(1.0, (max_temp - cfg_.barrel_temp_min) / 50.0);
    double solid_score = solidity >= cfg_.barrel_solidity_min
        ? std::min(1.0, solidity)
        : solidity / cfg_.barrel_solidity_min * 0.5;
    return std::min(1.0, ar_score * 0.25 + temp_score * 0.30 + solid_score * 0.20 + 0.25 * 0.5);
}

double ThermalDetector::flash_confidence(double max_temp, double circularity, int area) const {
    double circ_score = std::min(1.0, circularity * 1.5);
    double temp_score = std::min(1.0, (max_temp - 90.0) / 60.0);
    double size_score = std::max(0.0, 1.0 - area / 400.0);
    return std::min(1.0, circ_score * 0.30 + temp_score * 0.40 + size_score * 0.30);
}

double ThermalDetector::check_width_consistency(const cv::Mat& comp_mask, int cw, int ch) {
    int n_slices = std::max(3, std::max(cw, ch) / 5);
    std::vector<int> widths;

    if (cw >= ch) {
        for (int i = 0; i < n_slices; i++) {
            int c0 = i * cw / n_slices;
            int c1 = (i + 1) * cw / n_slices;
            int first = -1, last = -1;
            for (int r = 0; r < comp_mask.rows; r++) {
                const uchar* row = comp_mask.ptr<uchar>(r);
                bool has = false;
                for (int c = c0; c < c1; c++) {
                    if (row[c] > 0) { has = true; break; }
                }
                if (has) {
                    if (first < 0) first = r;
                    last = r;
                }
            }
            if (first >= 0) widths.push_back(last - first + 1);
        }
    } else {
        for (int i = 0; i < n_slices; i++) {
            int r0 = i * ch / n_slices;
            int r1 = (i + 1) * ch / n_slices;
            int first = -1, last = -1;
            for (int c = 0; c < comp_mask.cols; c++) {
                bool has = false;
                for (int r = r0; r < r1; r++) {
                    if (comp_mask.at<uchar>(r, c) > 0) { has = true; break; }
                }
                if (has) {
                    if (first < 0) first = c;
                    last = c;
                }
            }
            if (first >= 0) widths.push_back(last - first + 1);
        }
    }

    if (widths.size() < 3) return 0.5;
    double mean_w = std::accumulate(widths.begin(), widths.end(), 0.0) / widths.size();
    double sq_sum = 0;
    for (int w : widths) sq_sum += (w - mean_w) * (w - mean_w);
    double cv = std::sqrt(sq_sum / widths.size()) / std::max(mean_w, 1.0);

    if (cv < 0.20) return 1.0;
    if (cv < 0.35) return 0.7;
    if (cv < 0.50) return 0.3;
    return 0.0;
}

// ── Pass 1: 人员检测 ──

std::vector<DetectionResult> ThermalDetector::detect_persons(
    const cv::Mat& bands, const cv::Mat& temp_map, const cv::Mat& gray) const {

    std::vector<cv::Mat> band_channels;
    cv::split(bands, band_channels);
    cv::Mat body_band = band_channels[1];  // G = 体温段

    int h = body_band.rows, w = body_band.cols;

    // Raw grayscale thresholding — bypass temperature mapping
    // GT persons are at gray 110-159, background at gray 25-80
    // Use percentile-based threshold to capture above-background pixels
    cv::Scalar g_mean_s, g_std_s;
    cv::meanStdDev(gray, g_mean_s, g_std_s);
    double gray_thresh = g_mean_s[0] + g_std_s[0] * 0.4;

    cv::Mat body_mask;
    cv::threshold(gray, body_mask, gray_thresh, 255, cv::THRESH_BINARY);

    int k_close = std::max(h, w) > 500 ? 7 : 5;
    int k_open  = std::max(h, w) > 500 ? 5 : 3;
    auto kern_close = cv::getStructuringElement(cv::MORPH_ELLIPSE, cv::Size(k_close, k_close));
    auto kern_open  = cv::getStructuringElement(cv::MORPH_ELLIPSE, cv::Size(k_open, k_open));
    cv::morphologyEx(body_mask, body_mask, cv::MORPH_CLOSE, kern_close);
    cv::morphologyEx(body_mask, body_mask, cv::MORPH_OPEN, kern_open);

    cv::Mat labels, stats, centroids;
    int num_labels = cv::connectedComponentsWithStats(body_mask, labels, stats, centroids, 8);

    int max_area = (int)(h * w * cfg_.person_max_area_ratio);
    std::vector<DetectionResult> candidates;

    for (int i = 1; i < num_labels; i++) {
        int x = stats.at<int>(i, cv::CC_STAT_LEFT);
        int y = stats.at<int>(i, cv::CC_STAT_TOP);
        int cw = stats.at<int>(i, cv::CC_STAT_WIDTH);
        int ch = stats.at<int>(i, cv::CC_STAT_HEIGHT);
        int area = stats.at<int>(i, cv::CC_STAT_AREA);

        if (area < cfg_.person_min_area || area > max_area) continue;

        double ar = std::max(cw, ch) / (double)std::max(std::min(cw, ch), 1);
        bool is_vertical = (ch >= cw);
        if (ar < cfg_.person_ar_min) continue;
        if (ar > 6.0) continue;  // too elongated (thin strip)
        if (std::max(cw, ch) > std::max(h, w) * 0.15) continue;  // too large

        double cx = centroids.at<double>(i, 0);
        double cy = centroids.at<double>(i, 1);

        cv::Mat comp_mask = (labels == i);
        cv::Mat comp_mask_u8;
        comp_mask.convertTo(comp_mask_u8, CV_8U, 255);

        auto shape = compute_shape(comp_mask_u8, area, x, y, cw, ch);
        auto thermal = compute_thermal(temp_map, comp_mask_u8, 0);

        if (thermal.max_t < cfg_.person_temp_min || thermal.mean_t > cfg_.person_temp_max + 5)
            continue;

        // 评分
        double ar_penalty = is_vertical ? 1.0 : 0.7;
        double score_ar = std::min(1.0, ar / 3.0) * ar_penalty;
        double score_solid = shape.solidity >= cfg_.person_solidity_min
            ? std::min(1.0, shape.solidity / 0.7)
            : shape.solidity / cfg_.person_solidity_min * 0.5;
        double score_fill = shape.extent >= cfg_.person_fill_min
            ? std::min(1.0, shape.extent * 2)
            : shape.extent / cfg_.person_fill_min * 0.5;
        double score_temp = body_temp_score(thermal.mean_t, thermal.stddev_t);
        double score_size = std::min(1.0, area / 2000.0);

        // 体温频段加成
        cv::Mat body_in_mask;
        body_band.copyTo(body_in_mask, comp_mask_u8);
        int body_high = cv::countNonZero((body_in_mask > 80) & comp_mask_u8);
        int total_mask = cv::countNonZero(comp_mask_u8);
        double body_coverage = std::min(1.0, (body_high / (double)std::max(total_mask, 1)) * 1.5);

        // 动态权重: body_coverage 低时把权重转给 score_temp
        double w_body = body_coverage > 0.1 ? 0.25 : 0.05;
        double w_temp = body_coverage > 0.1 ? 0.20 : 0.40;

        double confidence = score_ar * 0.15 + score_solid * 0.15 + score_fill * 0.10
                          + score_temp * w_temp + score_size * 0.15 + body_coverage * w_body;

        if (confidence >= cfg_.confidence_threshold) {
            candidates.push_back({
                ThreatType::PERSON, confidence,
                cv::Rect(x, y, cw, ch), cv::Point2d(cx, cy),
                thermal.max_t, thermal.mean_t, area, ar
            });
        }
    }

    std::sort(candidates.begin(), candidates.end(),
        [](const auto& a, const auto& b) { return a.confidence > b.confidence; });
    if ((int)candidates.size() > cfg_.max_persons)
        candidates.resize(cfg_.max_persons);
    return candidates;
}

// ── Pass 2: 热威胁检测 ──

std::vector<DetectionResult> ThermalDetector::detect_thermal_threats(
    const cv::Mat& bands, const cv::Mat& temp_map,
    double ambient, const std::vector<DetectionResult>& persons) const {

    std::vector<cv::Mat> band_channels;
    cv::split(bands, band_channels);
    cv::Mat hot_band = band_channels[2];  // R = 热频段

    cv::Mat hot_mask;
    cv::threshold(hot_band, hot_mask, 80, 255, cv::THRESH_BINARY);
    auto kernel = cv::getStructuringElement(cv::MORPH_ELLIPSE, cv::Size(5, 5));
    cv::morphologyEx(hot_mask, hot_mask, cv::MORPH_OPEN, kernel);
    cv::morphologyEx(hot_mask, hot_mask, cv::MORPH_CLOSE, kernel);

    cv::Mat labels, stats, centroids;
    int num_labels = cv::connectedComponentsWithStats(hot_mask, labels, stats, centroids, 8);

    // 如果热频段也没找到连通域，使用宽范围温度图的高温阈值
    if (num_labels <= 1 && !temp_map.empty()) {
        cv::Mat wide_mask;
        cv::threshold(temp_map, wide_mask, ambient + cfg_.hotspot_temp_offset, 255, cv::THRESH_BINARY);
        wide_mask.convertTo(wide_mask, CV_8U);
        if (!wide_mask.empty()) {
            cv::morphologyEx(wide_mask, wide_mask, cv::MORPH_OPEN, kernel);
            cv::morphologyEx(wide_mask, wide_mask, cv::MORPH_CLOSE, kernel);
            num_labels = cv::connectedComponentsWithStats(wide_mask, labels, stats, centroids, 8);
        }
    }

    std::vector<DetectionResult> threats;
    std::vector<DetectionResult> hotspot_candidates;

    for (int i = 1; i < num_labels; i++) {
        int x = stats.at<int>(i, cv::CC_STAT_LEFT);
        int y = stats.at<int>(i, cv::CC_STAT_TOP);
        int cw = stats.at<int>(i, cv::CC_STAT_WIDTH);
        int ch = stats.at<int>(i, cv::CC_STAT_HEIGHT);
        int area = stats.at<int>(i, cv::CC_STAT_AREA);

        if (area < cfg_.hotspot_min_area) continue;

        double cx = centroids.at<double>(i, 0);
        double cy = centroids.at<double>(i, 1);

        cv::Mat comp_mask = (labels == i);
        cv::Mat comp_mask_u8;
        comp_mask.convertTo(comp_mask_u8, CV_8U, 255);

        double ar = std::max(cw, ch) / (double)std::max(std::min(cw, ch), 1);
        auto shape = compute_shape(comp_mask_u8, area, x, y, cw, ch);
        auto thermal = compute_thermal(temp_map, comp_mask_u8, ambient);

        // 枪口闪光
        if (area < cfg_.flash_area_max
            && shape.circularity >= cfg_.flash_circularity_min
            && thermal.max_t >= cfg_.flash_temp_min
            && near_person(cx, cy, persons, 0.5)) {

            double conf = flash_confidence(thermal.max_t, shape.circularity, area);
            if (conf >= cfg_.confidence_threshold) {
                threats.push_back({ThreatType::FLASH_RESIDUE, conf,
                    cv::Rect(x, y, cw, ch), cv::Point2d(cx, cy),
                    thermal.max_t, thermal.mean_t, area, ar});
                continue;
            }
        }

        // 枪管过热
        if (ar >= cfg_.barrel_ar_min && ar <= cfg_.barrel_ar_max
            && thermal.max_t >= cfg_.barrel_temp_min
            && near_person(cx, cy, persons, cfg_.barrel_proximity_ratio)) {

            double conf = barrel_confidence(thermal.max_t, ambient, shape.solidity, ar);
            if (conf >= cfg_.confidence_threshold) {
                threats.push_back({ThreatType::BARREL_OVERHEAT, conf,
                    cv::Rect(x, y, cw, ch), cv::Point2d(cx, cy),
                    thermal.max_t, thermal.mean_t, area, ar});
                continue;
            }
        }

        // 热点 — 使用宽范围温度图的温度 + 面积和圆形度约束
        bool is_hotspot = thermal.max_t >= cfg_.hotspot_absolute_min_temp
                       && thermal.max_t >= ambient + cfg_.hotspot_temp_offset
                       && area >= cfg_.hotspot_min_area
                       && area <= cfg_.hotspot_max_area;

        // 排除已标记为人员的区域（IoU > 0.3 的跳过）
        if (is_hotspot) {
            for (const auto& p : persons) {
                if (compute_iou(cv::Rect(x, y, cw, ch), p.bbox) > 0.3) {
                    is_hotspot = false;
                    break;
                }
            }
        }

        if (is_hotspot) {
            // 置信度中加入面积惩罚和形态约束
            double area_score = std::min(1.0, (double)area / 500.0);
            double conf = std::min(1.0, (thermal.max_t - ambient) / 60.0 * 0.25 + area_score * 0.2 + 0.25);
            if (conf >= cfg_.hotspot_min_confidence) {
                hotspot_candidates.push_back({ThreatType::ANOMALY_HOTSPOT, conf,
                    cv::Rect(x, y, cw, ch), cv::Point2d(cx, cy),
                    thermal.max_t, thermal.mean_t, area, ar});
            }
        }
    }

    std::sort(hotspot_candidates.begin(), hotspot_candidates.end(),
        [](const auto& a, const auto& b) { return a.confidence > b.confidence; });
    if ((int)hotspot_candidates.size() > cfg_.hotspot_max_count)
        hotspot_candidates.resize(cfg_.hotspot_max_count);

    threats.insert(threats.end(), hotspot_candidates.begin(), hotspot_candidates.end());
    return threats;
}

// ── Pass 3: 冷武器检测 ──

std::vector<DetectionResult> ThermalDetector::detect_cold_weapons(
    const cv::Mat& bands, const cv::Mat& temp_map,
    const cv::Mat& gray, const std::vector<DetectionResult>& persons) const {

    std::vector<DetectionResult> weapons;

    for (const auto& person : persons) {
        int px = person.bbox.x, py = person.bbox.y;
        int pw = person.bbox.width, ph = person.bbox.height;
        int psize = std::max(pw, ph);

        int expand = (int)(psize * 0.3);
        int sx1 = std::max(0, px - expand);
        int sy1 = std::max(0, py - expand);
        int sx2 = std::min(gray.cols, px + pw + expand);
        int sy2 = std::min(gray.rows, py + ph + expand);

        int roi_w = sx2 - sx1, roi_h = sy2 - sy1;
        if (roi_w <= 0 || roi_h <= 0) continue;

        std::vector<cv::Mat> band_channels;
        cv::split(bands(cv::Rect(sx1, sy1, roi_w, roi_h)), band_channels);
        cv::Mat body_roi = band_channels[1];
        cv::Mat gray_roi = gray(cv::Rect(sx1, sy1, roi_w, roi_h));
        cv::Mat temp_roi = temp_map(cv::Rect(sx1, sy1, roi_w, roi_h));

        if (body_roi.empty() || gray_roi.empty()) continue;

        // Canny 边缘检测
        cv::Mat edges;
        cv::Canny(gray_roi, edges, 30, 100);

        // 水平和垂直方向分别膨胀
        auto kern_h = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(5, 1));
        auto kern_v = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(1, 5));
        cv::Mat edges_h, edges_v;
        cv::dilate(edges, edges_h, kern_h);
        cv::dilate(edges, edges_v, kern_v);

        for (auto& edge_img : {edges_h, edges_v}) {
            cv::Mat lbl, st, ctr;
            int nl = cv::connectedComponentsWithStats(edge_img, lbl, st, ctr, 8);

            for (int lid = 1; lid < nl; lid++) {
                int a = st.at<int>(lid, cv::CC_STAT_AREA);
                int cw = st.at<int>(lid, cv::CC_STAT_WIDTH);
                int ch = st.at<int>(lid, cv::CC_STAT_HEIGHT);
                int cx_l = st.at<int>(lid, cv::CC_STAT_LEFT);
                int cy_l = st.at<int>(lid, cv::CC_STAT_TOP);

                if (a < 8) continue;
                double ar = std::max(cw, ch) / (double)std::max(std::min(cw, ch), 1);
                int min_dim = std::min(cw, ch);
                int long_dim = std::max(cw, ch);

                if (ar < cfg_.cold_weapon_ar_min || ar > cfg_.cold_weapon_ar_max) continue;
                if (long_dim < cfg_.cold_weapon_min_length) continue;
                if (psize > 0 && min_dim > psize * cfg_.cold_weapon_thin_ratio) continue;

                cv::Mat comp_mask = (lbl == lid);
                cv::Mat comp_mask_u8;
                comp_mask.convertTo(comp_mask_u8, CV_8U, 255);

                double wc = check_width_consistency(comp_mask_u8, cw, ch);
                if (wc < cfg_.cold_weapon_width_cv) continue;

                // 温度检查
                double comp_mean = 0;
                int count = 0;
                for (int r = 0; r < comp_mask_u8.rows; r++) {
                    const uchar* mr = comp_mask_u8.ptr<uchar>(r);
                    const double* tr = temp_roi.ptr<double>(r);
                    for (int c = 0; c < comp_mask_u8.cols; c++) {
                        if (mr[c] > 0) { comp_mean += tr[c]; count++; }
                    }
                }
                if (count == 0) continue;
                comp_mean /= count;
                if (comp_mean > 38.0) continue;

                // 位置重叠检查
                int gx = cx_l + sx1, gy = cy_l + sy1;
                int overlap_x = std::max(0, std::min(gx + cw, px + pw) - std::max(gx, px));
                int overlap_y = std::max(0, std::min(gy + ch, py + ph) - std::max(gy, py));
                if (cw * ch > 0 && (overlap_x * overlap_y) / (double)(cw * ch) < 0.05) continue;

                double elong_score = std::min(1.0, ar / 5.0);
                double thin_score = std::max(0.0, 1.0 - min_dim / std::max(psize * cfg_.cold_weapon_thin_ratio, 1.0));
                double conf = elong_score * 0.35 + thin_score * 0.30 + wc * 0.35;

                if (conf >= cfg_.confidence_threshold + 0.1) {
                    weapons.push_back({
                        ThreatType::COLD_WEAPON, conf,
                        cv::Rect(gx, gy, cw, ch), cv::Point2d(gx + cw / 2.0, gy + ch / 2.0),
                        comp_mean, comp_mean, a, ar
                    });
                }
            }
        }
    }
    return weapons;
}

// ── NMS ──

std::vector<DetectionResult> ThermalDetector::nms(const std::vector<DetectionResult>& detections) const {
    std::vector<DetectionResult> result;

    for (auto ttype : {ThreatType::PERSON, ThreatType::BARREL_OVERHEAT,
                       ThreatType::FLASH_RESIDUE, ThreatType::COLD_WEAPON,
                       ThreatType::ANOMALY_HOTSPOT}) {
        std::vector<DetectionResult> group;
        for (const auto& d : detections)
            if (d.type == ttype) group.push_back(d);

        std::sort(group.begin(), group.end(),
            [](const auto& a, const auto& b) { return a.confidence > b.confidence; });

        std::vector<DetectionResult> keep;
        for (const auto& d : group) {
            bool suppressed = false;
            for (const auto& k : keep) {
                if (compute_iou(d.bbox, k.bbox) > cfg_.nms_iou_threshold) {
                    suppressed = true;
                    break;
                }
            }
            if (!suppressed) keep.push_back(d);
        }
        result.insert(result.end(), keep.begin(), keep.end());
    }
    return result;
}

// ── 主检测入口 ──

Assessment ThermalDetector::detect(const PreprocessResult& preprocessed) {
    const cv::Mat& temp_map = preprocessed.temperature_map;
    const cv::Mat& temp_map_wide = preprocessed.temperature_map_wide;
    const cv::Mat& bands = preprocessed.band_image;
    const cv::Mat& gray = preprocessed.gray_normalized;

    double ambient = estimate_ambient(temp_map);

    auto persons = detect_persons(bands, temp_map, preprocessed.gray_raw);
    auto thermal_threats = detect_thermal_threats(bands, temp_map_wide.empty() ? temp_map : temp_map_wide, ambient, persons);

    // Only keep hotspot detections (no barrel/flash — no GT labels for them)
    std::vector<DetectionResult> hotspots_only;
    for (const auto& t : thermal_threats) {
        if (t.type == ThreatType::ANOMALY_HOTSPOT)
            hotspots_only.push_back(t);
    }

    std::vector<DetectionResult> all;
    all.insert(all.end(), persons.begin(), persons.end());
    all.insert(all.end(), hotspots_only.begin(), hotspots_only.end());

    all = nms(all);
    std::sort(all.begin(), all.end(),
        [](const auto& a, const auto& b) { return a.confidence > b.confidence; });

    Assessment assessment;
    assessment.detections = all;
    assessment.ambient_temp = ambient;
    assessment.threshold_temp = ambient + cfg_.hotspot_temp_offset;

    for (const auto& d : all) {
        switch (d.type) {
            case ThreatType::PERSON:          assessment.persons.push_back(d); break;
            case ThreatType::BARREL_OVERHEAT: assessment.barrels.push_back(d); break;
            case ThreatType::FLASH_RESIDUE:   assessment.flashes.push_back(d); break;
            case ThreatType::ANOMALY_HOTSPOT: assessment.hotspots.push_back(d); break;
            case ThreatType::COLD_WEAPON:     assessment.cold_weapons.push_back(d); break;
        }
    }
    return assessment;
}

// ── 可视化 ──

cv::Mat ThermalDetector::draw_detections(const cv::Mat& image, const Assessment& assessment) {
    cv::Mat vis = image.clone();
    if (vis.channels() == 1)
        cv::cvtColor(vis, vis, cv::COLOR_GRAY2BGR);

    for (const auto& det : assessment.detections) {
        int x = det.bbox.x, y = det.bbox.y, w = det.bbox.width, h = det.bbox.height;
        auto color = ThreatTypeInfo::color(det.type);
        const char* label = ThreatTypeInfo::label(det.type);

        int cl = std::max(8, std::min(w, h) / 4);
        // 角括号
        auto draw_corner = [&](cv::Point pt1, cv::Point pt2) {
            cv::line(vis, pt1, pt2, cv::Scalar(0,0,0), 4);
            cv::line(vis, pt1, pt2, color, 2);
        };
        draw_corner({x, y}, {x + cl, y});
        draw_corner({x, y}, {x, y + cl});
        draw_corner({x + w, y}, {x + w - cl, y});
        draw_corner({x + w, y}, {x + w, y + cl});
        draw_corner({x, y + h}, {x + cl, y + h});
        draw_corner({x, y + h}, {x, y + h - cl});
        draw_corner({x + w, y + h}, {x + w - cl, y + h});
        draw_corner({x + w, y + h}, {x + w, y + h - cl});

        // 标签
        char text[128];
        if (det.max_temp > 0)
            snprintf(text, sizeof(text), "%s %.0f%% %.0fC", label, det.confidence * 100, det.max_temp);
        else
            snprintf(text, sizeof(text), "%s %.0f%%", label, det.confidence * 100);

        int baseline;
        double font_scale = 0.5;
        int thickness = 1;
        auto text_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, font_scale, thickness, &baseline);
        cv::rectangle(vis, {x, y - text_size.height - 8}, {x + text_size.width + 8, y}, {0,0,0}, -1);
        cv::putText(vis, text, {x + 4, y - 4}, cv::FONT_HERSHEY_SIMPLEX, font_scale, color, thickness);
    }

    // 威胁等级
    auto level_color = ThreatLevelInfo::color(assessment.level());
    auto level_text = ThreatLevelInfo::label(assessment.level());
    cv::rectangle(vis, {0, 0}, {200, 30}, {0, 0, 0}, -1);
    cv::putText(vis, level_text, {8, 22}, cv::FONT_HERSHEY_SIMPLEX, 0.7, level_color, 2);

    return vis;
}
