#include "acoustic_analyzer/core/threat_tracker.h"
#include <sstream>
#include <iomanip>
#include <cmath>
#include <algorithm>
#include <chrono>

namespace acoustic {

ThreatTracker::ThreatTracker(float min_event_interval,
                             float azimuth_gate_deg,
                             int max_missing_frames)
    : min_event_interval_(min_event_interval),
      azimuth_gate_deg_(azimuth_gate_deg),
      max_missing_frames_(max_missing_frames),
      id_counter_(0) {}

std::string ThreatTracker::GenerateId() {
    std::ostringstream oss;
    oss << "THREAT-" << std::setw(4) << std::setfill('0') << ++id_counter_;
    return oss.str();
}

void ThreatTracker::Reset() {
    tracks_.clear();
    id_counter_ = 0;
}

std::size_t ThreatTracker::NumTracked() const noexcept {
    return tracks_.size();
}

std::vector<AcousticThreat> ThreatTracker::Update(
    const std::vector<AcousticThreat>& detected_threats) {
    auto now = std::chrono::steady_clock::now();

    // Increment missing frames for all existing tracks
    for (auto& track : tracks_) {
        track.missing_frames++;
    }

    // Nearest-neighbor data association (azimuth difference < azimuth_gate_deg_)
    std::vector<bool> det_matched(detected_threats.size(), false);
    for (auto& track : tracks_) {
        if (track.missing_frames > max_missing_frames_) continue;

        float best_diff = azimuth_gate_deg_;
        std::size_t best_idx = detected_threats.size();

        for (std::size_t d = 0; d < detected_threats.size(); ++d) {
            if (det_matched[d]) continue;
            float diff = std::abs(detected_threats[d].azimuth - track.last_azimuth);
            if (diff > 180.0f) diff = 360.0f - diff;
            if (diff < best_diff) {
                best_diff = diff;
                best_idx = d;
            }
        }

        if (best_idx < detected_threats.size()) {
            const auto& det = detected_threats[best_idx];
            track.last_azimuth = det.azimuth;
            track.last_elevation = det.elevation;
            track.sound_type = det.sound_type;
            track.confidence = det.confidence;
            track.distance = det.distance;
            track.distance_confidence = det.distance_confidence;
            track.missing_frames = 0;
            det_matched[best_idx] = true;
        }
    }

    // Initialize new tracks for unmatched detections
    for (std::size_t d = 0; d < detected_threats.size(); ++d) {
        if (det_matched[d]) continue;
        TrackedThreat track;
        track.threat_id = GenerateId();
        track.last_azimuth = detected_threats[d].azimuth;
        track.last_elevation = detected_threats[d].elevation;
        track.sound_type = detected_threats[d].sound_type;
        track.confidence = detected_threats[d].confidence;
        track.distance = detected_threats[d].distance;
        track.distance_confidence = detected_threats[d].distance_confidence;
        track.missing_frames = 0;
        track.ever_published = false;
        tracks_.push_back(std::move(track));
    }

    // Remove stale tracks (连续 5 帧未检测淘汰)
    tracks_.erase(
        std::remove_if(tracks_.begin(), tracks_.end(),
                       [this](const TrackedThreat& t) {
                           return t.missing_frames > max_missing_frames_;
                       }),
        tracks_.end());

    // Filter by min_event_interval and return active threats
    std::vector<AcousticThreat> output;
    for (auto& track : tracks_) {
        if (track.missing_frames != 0) continue;

        bool should_publish = !track.ever_published;
        if (!should_publish) {
            float elapsed = std::chrono::duration_cast<std::chrono::duration<float>>(
                now - track.last_publish_time).count();
            should_publish = (elapsed >= min_event_interval_);
        }

        if (should_publish) {
            track.ever_published = true;
            track.last_publish_time = now;
            AcousticThreat threat;
            threat.timestamp = now;
            threat.threat_id = track.threat_id;
            threat.sound_type = track.sound_type;
            threat.confidence = track.confidence;
            threat.azimuth = track.last_azimuth;
            threat.elevation = track.last_elevation;
            threat.distance = track.distance;
            threat.distance_confidence = track.distance_confidence;
            output.push_back(threat);
        }
    }

    return output;
}

}  // namespace acoustic
