#include "fusion_engine.h"
#include <cmath>
#include <algorithm>
#include <set>

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

FusionEngine::FusionEngine(ThreatCache& cache, const FusionConfig& cfg)
    : cache_(cache), cfg_(cfg) {}

std::vector<FusedThreat> FusionEngine::fuse(double now, double drone_heading, int image_width) {
    std::vector<FusedThreat> threats;

    // 收集各模态数据
    auto thermal_list  = cfg_.thermal_enabled  ? cache_.get_thermal(now)  : std::vector<ThermalCacheEntry>();
    FlashDetection flash_item;
    bool has_flash = cfg_.flash_enabled && cache_.get_latest_flash(now, flash_item);
    auto acoustic_list = cfg_.acoustic_enabled ? cache_.get_acoustic(now) : std::vector<AcousticThreatData>();

    // 用于回查的平铺列表
    std::vector<ThermalCacheEntry::Detection> all_thermal_dets;
    std::vector<FlashDetection> all_flash_items;
    std::vector<AcousticThreatData> all_acoustic_items;

    // 转为统一格式
    std::vector<UnifiedEntry> unified;

    for (auto& entry : thermal_list) {
        for (auto& det : entry.detections) {
            UnifiedEntry e;
            e.source = "thermal";
            e.azimuth = bbox_to_azimuth(det.bbox_x, det.bbox_w, image_width, drone_heading);
            e.confidence = det.confidence;
            e.level = thermal_solo_level(det.type);
            e.original_index = (int)all_thermal_dets.size();
            unified.push_back(e);
            all_thermal_dets.push_back(det);
        }
    }

    if (has_flash) {
        UnifiedEntry e;
        e.source = "flash";
        e.azimuth = sight_to_azimuth(flash_item.sight_x, drone_heading);
        e.confidence = flash_item.area > 0 ? std::min(1.0, flash_item.area / 400.0) : 0.3;
        e.level = flash_item.area > 200 ? FusionThreatLevel::WARNING : FusionThreatLevel::CAUTION;
        e.original_index = 0;
        unified.push_back(e);
        all_flash_items.push_back(flash_item);
    }

    for (auto& t : acoustic_list) {
        if (t.confidence < cfg_.acoustic_confidence_threshold) continue;
        UnifiedEntry e;
        e.source = "acoustic";
        e.azimuth = t.azimuth;
        e.confidence = t.confidence;
        e.level = acoustic_solo_level(t.sound_type);
        e.original_index = (int)all_acoustic_items.size();
        unified.push_back(e);
        all_acoustic_items.push_back(t);
    }

    if (unified.empty()) return threats;

    // 空间关联分组
    auto groups = spatial_group(unified);

    for (auto& group : groups) {
        std::vector<UnifiedEntry> group_entries;
        for (int idx : group) group_entries.push_back(unified[idx]);

        auto threat = build_threat(group_entries, all_thermal_dets, all_flash_items, all_acoustic_items, now);
        if (threat.confidence >= cfg_.fusion_confidence_threshold) {
            threats.push_back(threat);
        }
    }

    // 保守策略: 单模态高置信度告警
    for (auto& e : unified) {
        if (e.confidence >= 0.8 && FusionThreatLevelInfo::is_dangerous(e.level)) {
            bool already = false;
            for (auto& t : threats) {
                if (std::fabs(angle_diff(t.bearing, e.azimuth)) < 20) {
                    already = true;
                    break;
                }
            }
            if (!already) {
                FusedThreat t;
                t.threat_id = "FT-" + std::to_string(++threat_counter_);
                std::snprintf(nullptr, 0, "%04d", threat_counter_);
                char buf[16];
                std::snprintf(buf, sizeof(buf), "%04d", threat_counter_);
                t.threat_id = std::string("FT-") + buf;
                t.level = e.level;
                t.confidence = e.confidence;
                t.bearing = e.azimuth;
                t.sources.push_back(e.source);
                t.description = e.source + " high-confidence alert";
                t.timestamp = now;
                threats.push_back(t);
            }
        }
    }

    return threats;
}

// ── 类型映射 ──

FusionThreatLevel FusionEngine::thermal_solo_level(const std::string& type) {
    if (type == "Person")          return FusionThreatLevel::WARNING;
    if (type == "Barrel Overheat") return FusionThreatLevel::DANGER;
    if (type == "Cold Weapon")     return FusionThreatLevel::CAUTION;
    if (type == "Flash Residue")   return FusionThreatLevel::CAUTION;
    if (type == "Hotspot")         return FusionThreatLevel::CAUTION;
    return FusionThreatLevel::CAUTION;
}

FusionThreatLevel FusionEngine::acoustic_solo_level(const std::string& type) {
    if (type == "gunshot")    return FusionThreatLevel::WARNING;
    if (type == "artillery")  return FusionThreatLevel::DANGER;
    if (type == "explosion")  return FusionThreatLevel::DANGER;
    return FusionThreatLevel::SAFE;
}

// ── 方位角计算 ──

double FusionEngine::bbox_to_azimuth(int bx, int bw, int image_width, double heading) const {
    double cx = bx + bw / 2.0;
    double norm = (cx / std::max(image_width, 1) - 0.5) * cfg_.thermal_hfov;
    return std::fmod(heading + norm + 360.0, 360.0);
}

double FusionEngine::sight_to_azimuth(double sight_x, double heading) {
    return std::fmod(heading + sight_x * (180.0 / M_PI) + 360.0, 360.0);
}

// ── 空间分组 ──

std::vector<std::vector<size_t>> FusionEngine::spatial_group(
    const std::vector<UnifiedEntry>& entries, double threshold) const
{
    size_t n = entries.size();
    std::vector<bool> assigned(n, false);
    std::vector<std::vector<size_t>> groups;

    for (size_t i = 0; i < n; i++) {
        if (assigned[i]) continue;
        std::vector<size_t> group;
        group.push_back(i);
        assigned[i] = true;

        for (size_t j = i + 1; j < n; j++) {
            if (assigned[j]) continue;
            if (entries[i].source == entries[j].source) continue;
            double diff = std::fabs(angle_diff(entries[i].azimuth, entries[j].azimuth));
            if (diff < threshold) {
                group.push_back(j);
                assigned[j] = true;
            }
        }
        groups.push_back(group);
    }

    for (size_t i = 0; i < n; i++) {
        if (!assigned[i]) {
            groups.push_back({i});
        }
    }
    return groups;
}

// ── 构建融合威胁 ──

FusedThreat FusionEngine::build_threat(
    const std::vector<UnifiedEntry>& group_entries,
    const std::vector<ThermalCacheEntry::Detection>& thermal_dets,
    const std::vector<FlashDetection>& flash_items,
    const std::vector<AcousticThreatData>& acoustic_items,
    double now)
{
    FusedThreat t;
    char buf[16];
    std::snprintf(buf, sizeof(buf), "%04d", ++threat_counter_);
    t.threat_id = std::string("FT-") + buf;

    // 收集来源
    std::set<std::string> src_set;
    for (auto& e : group_entries) src_set.insert(e.source);
    t.sources.assign(src_set.begin(), src_set.end());

    // 加权置信度
    auto get_weight = [&](const std::string& s) -> double {
        if (s == "thermal")  return cfg_.weight_thermal;
        if (s == "flash")    return cfg_.weight_flash;
        if (s == "acoustic") return cfg_.weight_acoustic;
        return 0.3;
    };

    double total_w = 0, total_wc = 0;
    for (auto& e : group_entries) {
        double w = get_weight(e.source);
        total_w += w;
        total_wc += w * e.confidence;
    }
    double fused_conf = total_wc / std::max(total_w, 0.01);

    // 跨模态奖励
    if (src_set.size() >= 2) fused_conf = std::min(1.0, fused_conf + cfg_.cross_boost_2);
    if (src_set.size() >= 3) fused_conf = std::min(1.0, fused_conf + cfg_.cross_boost_3);

    t.confidence = fused_conf;

    // 方位角: 加权圆均值
    std::vector<std::pair<double, double>> bearing_weights;
    for (auto& e : group_entries) {
        double w = get_weight(e.source) * std::max(e.confidence, 0.1);
        bearing_weights.push_back({e.azimuth, w});
    }
    t.bearing = circular_weighted_mean(bearing_weights);

    // 距离: 取已知最小距离
    t.distance = -1.0;
    for (auto& e : group_entries) {
        double d = -1.0;
        if (e.source == "acoustic" && e.original_index < (int)acoustic_items.size()) {
            d = acoustic_items[e.original_index].distance;
        }
        if (d > 0 && (t.distance < 0 || d < t.distance)) t.distance = d;
    }

    // 仰角: 取声源 elevation
    t.elevation = 0.0;
    for (auto& e : group_entries) {
        if (e.source == "acoustic" && e.original_index < (int)acoustic_items.size()) {
            t.elevation = acoustic_items[e.original_index].elevation;
            break;
        }
    }

    // 威胁等级
    t.level = escalate(group_entries, thermal_dets, acoustic_items);

    // 描述
    std::string desc;
    for (auto& e : group_entries) {
        if (!desc.empty()) desc += " + ";
        if (e.source == "thermal" && e.original_index < (int)thermal_dets.size()) {
            auto& det = thermal_dets[e.original_index];
            char tmp[64];
            std::snprintf(tmp, sizeof(tmp), "%s(%.0f%%)", det.type.c_str(), e.confidence * 100);
            desc += tmp;
        } else if (e.source == "flash") {
            char tmp[64];
            std::snprintf(tmp, sizeof(tmp), "Flash(%.0fpx)", flash_items.empty() ? 0 : flash_items[0].area);
            desc += tmp;
        } else if (e.source == "acoustic" && e.original_index < (int)acoustic_items.size()) {
            auto& ac = acoustic_items[e.original_index];
            char tmp[64];
            std::snprintf(tmp, sizeof(tmp), "%s(%.0f%%)", ac.sound_type.c_str(), e.confidence * 100);
            desc += tmp;
        } else {
            desc += e.source;
        }
    }
    t.description = desc;
    t.timestamp = now;

    return t;
}

// ── 威胁升级 ──

FusionThreatLevel FusionEngine::escalate(
    const std::vector<UnifiedEntry>& group_entries,
    const std::vector<ThermalCacheEntry::Detection>& thermal_dets,
    const std::vector<AcousticThreatData>& acoustic_items) const
{
    std::set<std::string> thermal_types, acoustic_types;
    bool has_flash = false;
    FusionThreatLevel max_level = FusionThreatLevel::SAFE;

    for (auto& e : group_entries) {
        if (e.source == "thermal" && e.original_index < (int)thermal_dets.size()) {
            thermal_types.insert(thermal_dets[e.original_index].type);
        }
        if (e.source == "acoustic" && e.original_index < (int)acoustic_items.size()) {
            acoustic_types.insert(acoustic_items[e.original_index].sound_type);
        }
        if (e.source == "flash") has_flash = true;
        if (e.level > max_level) max_level = e.level;
    }

    bool has_person   = thermal_types.count("Person") > 0;
    bool has_barrel   = thermal_types.count("Barrel Overheat") > 0;
    bool has_weapon   = thermal_types.count("Cold Weapon") > 0;
    bool has_gunshot  = acoustic_types.count("gunshot") > 0;
    bool has_explosion = acoustic_types.count("artillery") > 0 || acoustic_types.count("explosion") > 0;

    if (has_person && has_barrel)                          return FusionThreatLevel::CRITICAL;
    if (has_person && has_weapon)                          return FusionThreatLevel::CRITICAL;
    if (has_person && has_gunshot)                         return FusionThreatLevel::CRITICAL;
    if (has_barrel && has_gunshot)                         return FusionThreatLevel::CRITICAL;
    if (has_weapon && has_gunshot)                         return FusionThreatLevel::CRITICAL;
    if (has_person && has_flash)                           return FusionThreatLevel::DANGER;
    if (has_gunshot && has_flash)                          return FusionThreatLevel::DANGER;
    if (has_explosion)                                     return FusionThreatLevel::DANGER;
    if (has_person)                                        return FusionThreatLevel::WARNING;
    if (has_barrel)                                        return FusionThreatLevel::DANGER;
    if (has_weapon)                                        return FusionThreatLevel::CAUTION;
    if (has_gunshot)                                       return FusionThreatLevel::WARNING;

    return max_level;
}

// ── 工具函数 ──

double FusionEngine::angle_diff(double a, double b) {
    double d = std::fmod(a - b + 360.0, 360.0);
    if (d > 180.0) d -= 360.0;
    return d;
}

double FusionEngine::circular_weighted_mean(
    const std::vector<std::pair<double, double>>& azimuths_weights)
{
    double sin_sum = 0, cos_sum = 0;
    for (auto& [az, w] : azimuths_weights) {
        double rad = az * M_PI / 180.0;
        sin_sum += w * std::sin(rad);
        cos_sum += w * std::cos(rad);
    }
    return std::fmod(std::atan2(sin_sum, cos_sum) * 180.0 / M_PI + 360.0, 360.0);
}
