#include "acoustic_analyzer/core/pipeline.h"
#include "acoustic_analyzer/core/audio_buffer.h"
#include "acoustic_analyzer/core/feature_extractor.h"
#include "acoustic_analyzer/core/gunshot_classifier.h"
#include "acoustic_analyzer/core/gcc_phat_localizer.h"
#include "acoustic_analyzer/core/distance_estimator.h"
#include "acoustic_analyzer/core/threat_tracker.h"

// yaml-cpp is optional: if absent, Pipeline still works but FromYaml returns defaults
#if defined(__has_include)
#  if __has_include(<yaml-cpp/yaml.h>)
#    define ACOUSTIC_HAS_YAML_CPP 1
#  endif
#endif

#ifdef ACOUSTIC_HAS_YAML_CPP
#include <yaml-cpp/yaml.h>
#endif

#include <cmath>
#include <algorithm>
#include <numeric>
#include <memory>

namespace acoustic {

struct Pipeline::Impl {
    PipelineConfig config;

    std::unique_ptr<AudioBuffer> audio_buffer;
    std::unique_ptr<FeatureExtractor> feature_extractor;
    std::unique_ptr<GunshotClassifier> classifier;
    std::unique_ptr<GccPhatLocalizer> localizer;
    std::unique_ptr<DistanceEstimator> distance_estimator;
    std::unique_ptr<ThreatTracker> threat_tracker;

    // Temporal smoothing state: per-class probability history
    std::vector<std::vector<float>> class_history_;
    int smoothing_window_ = 3;

    explicit Impl(const PipelineConfig& cfg) : config(cfg) {
        smoothing_window_ = static_cast<int>(cfg.classifier.smoothing_window);
        if (smoothing_window_ < 1) smoothing_window_ = 1;
        ResetModules();
    }

    void ResetModules() {
        std::size_t chunk_samples = static_cast<std::size_t>(config.sample_rate * config.chunk_duration);
        audio_buffer = std::make_unique<AudioBuffer>(chunk_samples,
                                                     static_cast<int>(config.mic_array.num_mics));

        feature_extractor = std::make_unique<FeatureExtractor>(
            static_cast<int>(config.sample_rate),
            static_cast<int>(config.n_fft),
            static_cast<int>(config.hop_length),
            static_cast<int>(config.n_mels),
            config.f_min,
            config.f_max,
            config.preemphasis);

        classifier = std::make_unique<GunshotClassifier>(config.classifier);

        if (config.mic_array.num_mics >= 2) {
            localizer = std::make_unique<GccPhatLocalizer>(
                config.mic_array,
                static_cast<int>(config.sample_rate),
                config.max_tdoa,
                config.interpolation_factor);
        }

        distance_estimator = std::make_unique<DistanceEstimator>(config.distance);
        threat_tracker = std::make_unique<ThreatTracker>(
            config.min_event_interval, 15.0f, 5);

        // Dynamic class history dimension from classifier labels
        int num_classes = static_cast<int>(classifier->Labels().size());
        if (num_classes == 0) num_classes = 2;  // fallback
        class_history_.assign(num_classes, std::vector<float>());
    }

    bool VadGate(const std::vector<float>& audio) {
        if (audio.empty()) return false;
        // Energy in dB
        float energy = 0.0f;
        for (float s : audio) energy += s * s;
        energy /= static_cast<float>(audio.size());
        float energy_db = 10.0f * std::log10(energy + 1e-10f);

        // Zero-crossing rate
        std::size_t zcr = 0;
        for (std::size_t i = 1; i < audio.size(); ++i) {
            if ((audio[i] >= 0.0f) != (audio[i - 1] >= 0.0f)) ++zcr;
        }
        float zcr_rate = static_cast<float>(zcr) / (audio.size() - 1);

        bool energy_ok = energy_db > -60.0f;
        bool zcr_ok = zcr_rate > 0.01f && zcr_rate < 0.5f;
        return energy_ok && zcr_ok;
    }

    std::vector<float> TemporalSmoothing(const std::vector<float>& probs) {
        int num_classes = static_cast<int>(probs.size());
        if (num_classes == 0) return probs;
        if (static_cast<int>(class_history_.size()) != num_classes) {
            class_history_.assign(num_classes, std::vector<float>());
        }

        std::vector<float> smoothed(num_classes);
        for (int i = 0; i < num_classes; ++i) {
            auto& hist = class_history_[i];
            hist.push_back(probs[i]);
            if (static_cast<int>(hist.size()) > smoothing_window_) {
                hist.erase(hist.begin());
            }
            float sum = 0.0f;
            for (float v : hist) sum += v;
            smoothed[i] = sum / static_cast<float>(hist.size());
        }
        return smoothed;
    }

    AcousticFrame Process(const std::vector<float>& audio_samples) {
        AcousticFrame frame;
        frame.timestamp = std::chrono::steady_clock::now();
        frame.is_clear = true;

        if (audio_samples.empty()) return frame;

        // Push interleaved multi-channel audio into ring buffer
        audio_buffer->Push(audio_samples);

        std::size_t chunk_samples = static_cast<std::size_t>(config.sample_rate * config.chunk_duration);

        // VAD on first channel
        std::size_t avail = audio_buffer->Size();
        std::size_t offset = (avail > chunk_samples) ? (avail - chunk_samples) : 0;
        auto latest_flat = audio_buffer->Get(offset, chunk_samples);
        std::vector<float> latest;
        latest.reserve(latest_flat.size() / config.mic_array.num_mics);
        for (std::size_t i = 0; i < latest_flat.size(); i += config.mic_array.num_mics) {
            latest.push_back(latest_flat[i]);
        }
        if (!VadGate(latest)) {
            return frame;
        }

        // Feature extraction: multi-channel Mel spectrogram
        auto mel_specs = feature_extractor->MelSpectrogramMultiChannel(
            audio_samples, config.mic_array.num_mics);

        // Average across channels in frequency domain
        Eigen::MatrixXf avg_mel;
        if (!mel_specs.empty()) {
            avg_mel = mel_specs[0];
            for (std::size_t i = 1; i < mel_specs.size(); ++i) {
                avg_mel += mel_specs[i];
            }
            if (mel_specs.size() > 1) {
                avg_mel /= static_cast<float>(mel_specs.size());
            }
        }

        // Classification
        auto [label, confidence] = classifier->Predict(avg_mel);

        // Build probability vector for temporal smoothing
        const auto& labels = classifier->Labels();
        int num_classes = static_cast<int>(labels.size());
        if (num_classes == 0) num_classes = 2;
        std::vector<float> probs(num_classes, 0.0f);
        int label_idx = -1;
        for (int i = 0; i < num_classes; ++i) {
            if (i < static_cast<int>(labels.size()) && labels[i] == label) {
                label_idx = i;
                break;
            }
        }
        if (label_idx >= 0) probs[label_idx] = confidence;
        auto smoothed = TemporalSmoothing(probs);

        // Find best after smoothing
        float best_conf = 0.0f;
        int best_idx = 0;
        for (int i = 0; i < num_classes; ++i) {
            if (smoothed[i] > best_conf) {
                best_conf = smoothed[i];
                best_idx = i;
            }
        }

        if (best_idx < static_cast<int>(labels.size()) &&
            best_conf >= config.confidence_threshold) {
            std::string best_label = labels[best_idx];

            // Localization
            float azimuth = 0.0f;
            float elevation = 0.0f;
            if (localizer && config.mic_array.num_mics >= 2) {
                Eigen::MatrixXf audio_mat(static_cast<int>(chunk_samples),
                                          static_cast<int>(config.mic_array.num_mics));
                for (std::size_t ch = 0; ch < config.mic_array.num_mics; ++ch) {
                    std::vector<float> ch_data;
                    ch_data.reserve(latest_flat.size() / config.mic_array.num_mics);
                    for (std::size_t i = ch; i < latest_flat.size(); i += config.mic_array.num_mics) {
                        ch_data.push_back(latest_flat[i]);
                    }
                    for (std::size_t i = 0; i < ch_data.size(); ++i) {
                        audio_mat(static_cast<int>(i), static_cast<int>(ch)) = ch_data[i];
                    }
                }
                auto [az, el] = localizer->Localize(audio_mat);
                azimuth = az;
                elevation = el;
            }

            // Distance estimation
            float distance = -1.0f;
            float distance_conf = 0.0f;
            if (distance_estimator) {
                float spl = distance_estimator->ComputeSpl(latest);
                float raw_distance = distance_estimator->Estimate(spl, best_label);
                distance = distance_estimator->UpdateKalman(raw_distance);
                distance_conf = 0.8f;
            }

            AcousticThreat threat;
            threat.timestamp = frame.timestamp;
            threat.threat_id = "";
            threat.sound_type = best_label;
            threat.confidence = best_conf;
            threat.azimuth = azimuth;
            threat.elevation = elevation;
            threat.distance = distance;
            threat.distance_confidence = distance_conf;

            auto filtered = threat_tracker->Update({threat});
            frame.threats = std::move(filtered);
            frame.is_clear = frame.threats.empty();
        }

        return frame;
    }
};

Pipeline::Pipeline(const PipelineConfig& config)
    : impl_(std::make_unique<Impl>(config)) {}

Pipeline::~Pipeline() = default;
Pipeline::Pipeline(Pipeline&&) noexcept = default;
Pipeline& Pipeline::operator=(Pipeline&&) noexcept = default;

AcousticFrame Pipeline::Process(const std::vector<float>& audio_samples) {
    return impl_->Process(audio_samples);
}

PipelineConfig Pipeline::FromYaml(const std::string& yaml_path) {
    PipelineConfig config;
#ifdef ACOUSTIC_HAS_YAML_CPP
    try {
        YAML::Node root = YAML::LoadFile(yaml_path);

        if (root["audio"]) {
            auto audio = root["audio"];
            if (audio["sample_rate"])
                config.sample_rate = audio["sample_rate"].as<uint32_t>();
            if (audio["chunk_duration"])
                config.chunk_duration = audio["chunk_duration"].as<float>();
            if (audio["hop_duration"])
                config.hop_duration = audio["hop_duration"].as<float>();
            if (audio["n_channels"])
                config.mic_array.num_mics = audio["n_channels"].as<uint32_t>();
        }

        if (root["source"]) {
            auto src = root["source"];
            if (src["type"]) src["type"].as<std::string>();
        }

        if (root["mic_array"]) {
            auto mic = root["mic_array"];
            if (mic["num_mics"])
                config.mic_array.num_mics = mic["num_mics"].as<uint32_t>();
            if (mic["layout"])
                config.mic_array.layout = mic["layout"].as<std::string>();
            if (mic["spacing"])
                config.mic_array.spacing = mic["spacing"].as<float>();
            if (mic["positions"]) {
                config.mic_array.positions.clear();
                for (const auto& pos : mic["positions"]) {
                    std::array<float, 3> p = {0.0f, 0.0f, 0.0f};
                    if (pos.size() > 0) p[0] = pos[0].as<float>();
                    if (pos.size() > 1) p[1] = pos[1].as<float>();
                    if (pos.size() > 2) p[2] = pos[2].as<float>();
                    config.mic_array.positions.push_back(p);
                }
            }
        }

        if (root["features"]) {
            auto feat = root["features"];
            if (feat["n_mels"]) config.n_mels = feat["n_mels"].as<uint32_t>();
            if (feat["n_fft"]) {
                config.n_fft = feat["n_fft"].as<uint32_t>();
            }
            if (feat["hop_length"]) {
                config.hop_length = feat["hop_length"].as<uint32_t>();
            }
            if (feat["f_min"]) {
                config.f_min = feat["f_min"].as<float>();
            }
            if (feat["f_max"]) {
                config.f_max = feat["f_max"].as<float>();
            }
            if (feat["preemphasis"]) {
                config.preemphasis = feat["preemphasis"].as<float>();
            }
        }

        if (root["classifier"]) {
            auto cls = root["classifier"];
            if (cls["model_path"])
                config.classifier.model_path = cls["model_path"].as<std::string>();
            if (cls["label_map_path"])
                config.classifier.label_map_path = cls["label_map_path"].as<std::string>();
            if (cls["threshold"])
                config.classifier.threshold = cls["threshold"].as<float>();
            if (cls["smoothing_window"])
                config.classifier.smoothing_window = cls["smoothing_window"].as<uint32_t>();
        }

        if (root["localization"]) {
            auto loc = root["localization"];
            if (loc["max_tdoa"]) {
                config.max_tdoa = loc["max_tdoa"].as<float>();
            }
            if (loc["interpolation_factor"]) {
                config.interpolation_factor = loc["interpolation_factor"].as<int>();
            }
        }

        if (root["distance"]) {
            auto dist = root["distance"];
            if (dist["attenuation_alpha"])
                config.distance.attenuation_alpha = dist["attenuation_alpha"].as<float>();
            if (dist["kalman_process_noise"])
                config.distance.kalman_process_noise =
                    dist["kalman_process_noise"].as<float>();
            if (dist["kalman_measurement_noise"])
                config.distance.kalman_measurement_noise =
                    dist["kalman_measurement_noise"].as<float>();
            if (dist["reference_spl"]) {
                auto ref = dist["reference_spl"];
                if (ref["threat"])
                    config.distance.ref_spl_gunshot = ref["threat"].as<float>();
                if (ref["gunshot"])
                    config.distance.ref_spl_gunshot = ref["gunshot"].as<float>();
                if (ref["artillery"])
                    config.distance.ref_spl_artillery = ref["artillery"].as<float>();
                if (ref["explosion"])
                    config.distance.ref_spl_explosion = ref["explosion"].as<float>();
            }
        }

        if (root["output"]) {
            auto out = root["output"];
            if (out["publish_rate"]) {
                // Not stored directly
            }
            if (out["min_event_interval"])
                config.min_event_interval = out["min_event_interval"].as<float>();
        }
    } catch (const YAML::Exception& e) {
        (void)e;
        // On parse error, return defaults
    }
#else
    (void)yaml_path;
#endif
    return config;
}

void Pipeline::Reset() {
    impl_->audio_buffer->Clear();
    impl_->threat_tracker->Reset();
    impl_->distance_estimator->Reset();
    for (auto& hist : impl_->class_history_) {
        hist.clear();
    }
}

const PipelineConfig& Pipeline::Config() const noexcept {
    return impl_->config;
}

}  // namespace acoustic
