#include "thermal_preprocessor.h"
#include <opencv2/imgproc.hpp>
#include <cmath>
#include <algorithm>
#include <cstring>
#include <vector>

// ── 色板颜色函数 ──

cv::Vec3b ThermalPreprocessor::jet_color(int i) {
    double t = i / 255.0;
    int b = std::min(255, std::max(0, (int)(255 * std::min(1.0, std::max(0.0, 1.5 - std::abs(4.0 * t - 1.0))))));
    int g = std::min(255, std::max(0, (int)(255 * std::min(1.0, std::max(0.0, 1.5 - std::abs(4.0 * t - 2.0))))));
    int r = std::min(255, std::max(0, (int)(255 * std::min(1.0, std::max(0.0, 1.5 - std::abs(4.0 * t - 3.0))))));
    return {(uint8_t)b, (uint8_t)g, (uint8_t)r};
}

cv::Vec3b ThermalPreprocessor::ironbow_color(int i) {
    double t = i / 255.0;
    int r, g, b;
    if (t < 0.15) {
        double s = t / 0.15;
        r = (int)(s * 60); g = 0; b = (int)(s * 100);
    } else if (t < 0.35) {
        double s = (t - 0.15) / 0.20;
        r = (int)(60 + s * 195); g = (int)(s * 25); b = (int)(100 - s * 80);
    } else if (t < 0.55) {
        double s = (t - 0.35) / 0.20;
        r = 255; g = (int)(25 + s * 155); b = (int)(20 - s * 20);
    } else if (t < 0.75) {
        double s = (t - 0.55) / 0.20;
        r = 255; g = (int)(180 + s * 75); b = 0;
    } else {
        double s = (t - 0.75) / 0.25;
        r = 255; g = 255; b = (int)(s * 255);
    }
    return {(uint8_t)std::min(b,255), (uint8_t)std::min(g,255), (uint8_t)std::min(r,255)};
}

cv::Mat ThermalPreprocessor::build_colormap_lut(cv::Vec3b (*color_fn)(int)) {
    cv::Mat lut(256, 1, CV_8UC3);
    for (int i = 0; i < 256; i++)
        lut.at<cv::Vec3b>(i, 0) = color_fn(i);
    return lut;
}

void ThermalPreprocessor::build_inverse_lut(const cv::Mat& lut_bgr, cv::Mat& inv_out, int& q_out) {
    const int Q = 8;
    q_out = Q;
}

ThermalPreprocessor::ThermalPreprocessor(double tmin, double tmax)
    : tmin_(tmin), tmax_(tmax) {
    jet_lut_ = build_colormap_lut(jet_color);
    ironbow_lut_ = build_colormap_lut(ironbow_color);
    build_inverse_lut(jet_lut_, jet_inv_lut_, inv_q_);
    build_inverse_lut(ironbow_lut_, ironbow_inv_lut_, inv_q_);
}

// ── 色板自动检测 ──

ColormapType ThermalPreprocessor::detect_colormap(const cv::Mat& image) const {
    if (image.channels() == 1)
        return detect_gray_palette(image);

    std::vector<cv::Mat> channels;
    cv::split(image, channels);
    cv::Mat b = channels[0], g = channels[1], r = channels[2];

    cv::Mat diff_rg, diff_rb;
    cv::absdiff(r, g, diff_rg);
    cv::absdiff(r, b, diff_rb);
    double mean_rg = cv::mean(diff_rg)[0];
    double mean_rb = cv::mean(diff_rb)[0];
    if (mean_rg < 8 && mean_rb < 8)
        return detect_gray_palette(cv::Mat());

    double blue_ratio = cv::countNonZero(b > 150) / (double)(b.rows * b.cols);
    double red_ratio  = cv::countNonZero(r > 150) / (double)(r.rows * r.cols);

    if (blue_ratio > 0.08 && red_ratio > 0.08)
        return ColormapType::JET;
    if (red_ratio > 0.15 && blue_ratio < 0.08)
        return ColormapType::IRONBOW;
    if (red_ratio > 0.05)
        return ColormapType::THERMAL_RGB;

    return ColormapType::GRAY;
}

ColormapType ThermalPreprocessor::detect_gray_palette(const cv::Mat& gray) const {
    if (gray.empty()) return ColormapType::WHITEHOT;
    int h = gray.rows, w = gray.cols;
    int bw = std::max(1, w / 10);
    int bh = std::max(1, h / 10);

    cv::Mat top_strip = gray.rowRange(0, bh);
    cv::Mat bot_strip = gray.rowRange(h - bh, h);
    cv::Mat left_strip = gray.colRange(0, bw).rowRange(bh, h - bh);
    cv::Mat right_strip = gray.colRange(w - bw, w).rowRange(bh, h - bh);

    double border_mean = (cv::mean(top_strip)[0] * top_strip.cols +
                          cv::mean(bot_strip)[0] * bot_strip.cols +
                          cv::mean(left_strip)[0] * left_strip.cols +
                          cv::mean(right_strip)[0] * right_strip.cols) /
                         (top_strip.total() + bot_strip.total() + left_strip.total() + right_strip.total());

    double center_mean = 128.0;
    if (h > 2 * bh && w > 2 * bw)
        center_mean = cv::mean(gray.rowRange(bh, h - bh).colRange(bw, w - bw))[0];

    if (border_mean > center_mean + 15)
        return ColormapType::BLACKHOT;
    return ColormapType::WHITEHOT;
}

// ── 温度提取 ──

cv::Mat ThermalPreprocessor::gray_to_temp(const cv::Mat& gray) const {
    cv::Mat temp;
    gray.convertTo(temp, CV_64F, (tmax_ - tmin_) / 255.0, tmin_);
    return temp;
}

cv::Mat ThermalPreprocessor::gray_to_temp_adaptive(const cv::Mat& gray) const {
    // 计算灰度直方图的 p5 和 p95 百分位
    int histSize = 256;
    float range[] = {0, 256};
    const float* histRange = {range};
    cv::Mat hist;
    cv::calcHist(&gray, 1, 0, cv::Mat(), hist, 1, &histSize, &histRange);

    int total = gray.total();
    int cum = 0;
    double p5_val = 0, p95_val = 255;
    for (int i = 0; i < 256; i++) {
        cum += (int)hist.at<float>(i);
        if (p5_val == 0 && cum >= total * 0.05) p5_val = i;
        if (p95_val == 255 && cum >= total * 0.95) { p95_val = i; break; }
    }

    // 映射: p5 → 10°C (环境), p95 → 42°C (体温上限)
    double adaptive_tmin = 10.0;
    double adaptive_tmax = 42.0;
    double range_span = std::max(p95_val - p5_val, 1.0);
    double scale = (adaptive_tmax - adaptive_tmin) / range_span;

    cv::Mat temp;
    gray.convertTo(temp, CV_64F);
    temp = (temp - p5_val) * scale + adaptive_tmin;
    cv::min(temp, adaptive_tmax + 30, temp);
    cv::max(temp, adaptive_tmin - 10, temp);
    return temp;
}

cv::Mat ThermalPreprocessor::lut_inverse_temp(const cv::Mat& image, const cv::Mat& inv_lut, int q) const {
    return thermal_rgb_to_temp(image);
}

cv::Mat ThermalPreprocessor::thermal_rgb_to_temp(const cv::Mat& image) const {
    std::vector<cv::Mat> channels;
    cv::split(image, channels);
    cv::Mat b_f, g_f, r_f;
    channels[0].convertTo(b_f, CV_64F, 1.0 / 255.0);
    channels[1].convertTo(g_f, CV_64F, 1.0 / 255.0);
    channels[2].convertTo(r_f, CV_64F, 1.0 / 255.0);

    cv::Mat t_norm = r_f * 0.50 + g_f * 0.30 + (1.0 - b_f) * 0.20;

    double p2, p98;
    cv::Mat sorted;
    t_norm.reshape(1).copyTo(sorted);
    cv::sort(sorted, sorted, cv::SORT_EVERY_ROW + cv::SORT_ASCENDING);
    int idx2 = (int)(sorted.total() * 0.02);
    int idx98 = (int)(sorted.total() * 0.98);
    p2 = sorted.at<double>(idx2);
    p98 = sorted.at<double>(idx98);

    if (p98 - p2 > 0.01) {
        cv::Mat tmp;
        cv::subtract(t_norm, p2, tmp);
        cv::divide(tmp, (p98 - p2), t_norm);
    }

    cv::Mat clamped;
    cv::min(t_norm, 1.0, clamped);
    cv::max(clamped, 0.0, clamped);

    cv::Mat temp;
    clamped.convertTo(temp, CV_64F, (tmax_ - tmin_), tmin_);
    return temp;
}

cv::Mat ThermalPreprocessor::extract_temperature(const cv::Mat& image, ColormapType cmap) const {
    if (cmap == ColormapType::GRAY)
        cmap = detect_colormap(image);

    cv::Mat gray;
    if (image.channels() == 1) {
        gray = image;
    } else {
        cv::cvtColor(image, gray, cv::COLOR_BGR2GRAY);
    }

    switch (cmap) {
        case ColormapType::WHITEHOT:
            return gray_to_temp_adaptive(gray);
        case ColormapType::BLACKHOT:
            return gray_to_temp_adaptive(255 - gray);
        case ColormapType::JET:
        case ColormapType::IRONBOW:
            return thermal_rgb_to_temp(image);
        case ColormapType::THERMAL_RGB:
            return thermal_rgb_to_temp(image);
        default:
            return gray_to_temp_adaptive(gray);
    }
}

// ── 温度频段分解 ──

cv::Mat ThermalPreprocessor::to_temperature_bands(const cv::Mat& temp_map) const {
    cv::Mat temp_f;
    temp_map.convertTo(temp_f, CV_64F);

    cv::Mat hot = cv::min((temp_f - 55.0) / std::max(tmax_ - 55.0, 1.0) * 255.0, 255.0);
    cv::max(hot, 0.0);

    cv::Mat body;
    cv::exp(-((temp_f - 36.0) / 6.0).mul((temp_f - 36.0) / 6.0), body);
    body = body * 255.0;

    cv::Mat ambient = cv::min((20.0 - temp_f) / std::max(20.0 - tmin_, 1.0) * 255.0, 255.0);
    cv::max(ambient, 0.0);

    std::vector<cv::Mat> bands = {ambient, body, hot};
    for (auto& m : bands) {
        cv::max(m, 0.0);
        cv::min(m, 255.0);
    }

    cv::Mat band_img;
    cv::merge(bands, band_img);
    band_img.convertTo(band_img, CV_8U);
    return band_img;
}

// ── 灰度转彩色 ──

cv::Mat ThermalPreprocessor::colorize(const cv::Mat& image, const std::string& palette, ColormapType cmap) const {
    if (cmap == ColormapType::GRAY)
        cmap = detect_colormap(image);

    cv::Mat norm;
    if (cmap == ColormapType::WHITEHOT || cmap == ColormapType::GRAY) {
        if (image.channels() == 1)
            norm = image;
        else
            cv::cvtColor(image, norm, cv::COLOR_BGR2GRAY);
    } else if (cmap == ColormapType::BLACKHOT) {
        if (image.channels() == 1)
            norm = 255 - image;
        else {
            cv::Mat gray;
            cv::cvtColor(image, gray, cv::COLOR_BGR2GRAY);
            norm = 255 - gray;
        }
    } else {
        if (image.channels() == 3)
            return image.clone();
        norm = image;
    }

    cv::Ptr<cv::CLAHE> clahe = cv::createCLAHE(3.0, cv::Size(8, 8));
    clahe->apply(norm, norm);

    cv::Mat colorized;
    if (palette == "ironbow") {
        cv::LUT(norm, ironbow_lut_, colorized);
    } else {
        cv::applyColorMap(norm, colorized, cv::COLORMAP_JET);
    }

    return colorized;
}

// ── 完整预处理流水线 ──

PreprocessResult ThermalPreprocessor::preprocess(const cv::Mat& image, ColormapType cmap, bool enhance, bool auto_colorize) const {
    if (cmap == ColormapType::GRAY)
        cmap = detect_colormap(image);

    // 自适应温度图（窄范围，用于人员检测和频段分解）
    cv::Mat temp_map = extract_temperature(image, cmap);

    // 宽范围温度图（用于热点/枪管等高温检测）
    // 使用更宽的自适应映射：p2→-10°C, p98→120°C
    cv::Mat gray;
    if (image.channels() == 1) gray = image;
    else cv::cvtColor(image, gray, cv::COLOR_BGR2GRAY);

    cv::Mat temp_map_wide;
    {
        cv::Mat gray_adj = (cmap == ColormapType::BLACKHOT) ? (255 - gray) : gray;
        int histSize = 256;
        float range[] = {0, 256};
        const float* histRange = {range};
        cv::Mat hist;
        cv::calcHist(&gray_adj, 1, 0, cv::Mat(), hist, 1, &histSize, &histRange);
        int total = gray_adj.total();
        int cum = 0;
        double p2_val = 0, p98_val = 255;
        for (int i = 0; i < 256; i++) {
            cum += (int)hist.at<float>(i);
            if (p2_val == 0 && cum >= total * 0.02) p2_val = i;
            if (p98_val == 255 && cum >= total * 0.98) { p98_val = i; break; }
        }
        double wide_tmin = -10.0;
        double wide_tmax = 120.0;
        double scale = (wide_tmax - wide_tmin) / std::max(p98_val - p2_val, 1.0);
        gray_adj.convertTo(temp_map_wide, CV_64F);
        temp_map_wide = (temp_map_wide - p2_val) * scale + wide_tmin;
        cv::min(temp_map_wide, wide_tmax + 30, temp_map_wide);
        cv::max(temp_map_wide, wide_tmin - 10, temp_map_wide);
    }

    // 频段分解基于自适应温度图（使人员落入 G 通道）
    cv::Mat band_image = to_temperature_bands(temp_map);

    cv::Mat gray_norm;
    temp_map.convertTo(gray_norm, CV_64F);
    double mn, mx;
    cv::minMaxLoc(temp_map, &mn, &mx);
    double span = std::max(mx - mn, 1.0);
    gray_norm = (gray_norm - mn) / span * 255.0;
    cv::max(gray_norm, 0.0);
    cv::min(gray_norm, 255.0);
    gray_norm.convertTo(gray_norm, CV_8U);

    if (enhance) {
        cv::Ptr<cv::CLAHE> clahe = cv::createCLAHE(3.0, cv::Size(8, 8));
        clahe->apply(gray_norm, gray_norm);
    }

    return {temp_map, temp_map_wide, band_image, gray_norm, gray, cmap};
}
