#include <iostream>
#include <vector>
#include <cmath>
#include <cassert>
#include <cstring>
#include "acoustic_analyzer/core/audio_buffer.h"
#include "acoustic_analyzer/core/feature_extractor.h"
#include "acoustic_analyzer/core/gcc_phat_localizer.h"
#include "acoustic_analyzer/core/distance_estimator.h"
#include "acoustic_analyzer/core/threat_tracker.h"
#include "acoustic_analyzer/core/fft_utils.h"

using namespace acoustic;

bool near(float a, float b, float eps = 1e-3f) {
    return std::abs(a - b) < eps;
}

void test_audio_buffer() {
    AudioBuffer buf(100, 2);
    std::vector<float> data(200);
    for (int i = 0; i < 200; ++i) data[i] = static_cast<float>(i);
    buf.Push(data);
    assert(buf.Size() == 100);
    assert(buf.NumChannels() == 2);

    auto popped = buf.Pop(10);
    assert(popped.size() == 10 * 2);
    assert(near(popped[0], 0.0f));   // ch0, sample 0 (oldest)
    assert(near(popped[1], 1.0f));   // ch1, sample 0
    std::cout << "[PASS] audio_buffer" << std::endl;
}

void test_preemphasis() {
    float in[] = {1.0f, 0.5f, 0.25f};
    float out[3];
    apply_preemphasis(in, out, 3, 0.97f);
    assert(near(out[0], 1.0f));
    assert(near(out[1], 0.5f - 0.97f * 1.0f, 1e-4f));
    std::cout << "[PASS] preemphasis" << std::endl;
}

void test_gcc_phat_cross_array() {
    MicArrayConfig config;
    config.num_mics = 4;
    config.layout = "cross";
    config.spacing = 0.15f;
    GccPhatLocalizer loc(config, 16000, 0.00044f, 4);

    // Simulate a signal arriving from 90 degrees (right)
    int n = 512;
    Eigen::MatrixXf audio(n, 4);
    for (int i = 0; i < n; ++i) {
        float s = std::sin(2.0f * static_cast<float>(M_PI) * 1000.0f * i / 16000.0f);
        audio(i, 0) = s;
        audio(i, 1) = s;
        audio(i, 2) = s;
        audio(i, 3) = s;
    }
    float delay_samples = 0.15f / 343.0f * 16000.0f;
    for (int i = 0; i < n; ++i) {
        int idx = static_cast<int>(i + delay_samples);
        if (idx < n) audio(idx, 1) = audio(i, 0);
        idx = static_cast<int>(i - delay_samples);
        if (idx >= 0) audio(idx, 3) = audio(i, 0);
    }

    auto [azimuth, elevation] = loc.Localize(audio);
    assert(azimuth >= 0.0f && azimuth < 360.0f);
    std::cout << "[PASS] gcc_phat_cross_array (azimuth=" << azimuth << ")" << std::endl;
}

void test_distance_estimator() {
    DistanceConfig config;
    config.ref_spl_gunshot = 150.0f;
    config.attenuation_alpha = 0.6f;
    DistanceEstimator est(config);
    float dist = est.Estimate(120.0f, "gunshot");
    assert(dist > 0.0f && dist < 5000.0f);
    std::cout << "[PASS] distance_estimator (dist=" << dist << "m)" << std::endl;
}

void test_threat_tracker() {
    ThreatTracker tracker(0.3f, 15.0f, 5);
    AcousticThreat t1;
    t1.sound_type = "gunshot";
    t1.confidence = 0.9f;
    t1.azimuth = 45.0f;
    t1.elevation = 0.0f;
    t1.distance = 100.0f;
    t1.distance_confidence = 0.8f;

    auto out = tracker.Update({t1});
    assert(!out.empty());
    assert(out[0].sound_type == "gunshot");
    assert(out[0].threat_id == "THREAT-0001");
    std::cout << "[PASS] threat_tracker" << std::endl;
}

void test_feature_extractor_shape() {
    FeatureExtractor ext(16000, 2048, 512, 64);
    std::vector<float> audio(16000 * 2, 0.0f); // 2 seconds
    for (int i = 0; i < 32000; ++i) {
        audio[i] = std::sin(2.0f * static_cast<float>(M_PI) * 500.0f * i / 16000.0f) * 0.1f;
    }
    Eigen::MatrixXf mel = ext.MelSpectrogram(audio);
    assert(mel.rows() == 64);
    assert(mel.cols() == 63); // padded/truncated to 63 frames
    std::cout << "[PASS] feature_extractor_shape (" << mel.rows() << " x " << mel.cols() << ")" << std::endl;
}

void test_audio_buffer_overflow() {
    AudioBuffer buf(10, 1);
    std::vector<float> data(25, 1.0f);
    buf.Push(data);
    assert(buf.Size() == 10); // should cap at capacity
    auto popped = buf.Pop(5);
    assert(popped.size() == 5);
    assert(buf.Size() == 5);
    std::cout << "[PASS] audio_buffer_overflow" << std::endl;
}

void test_distance_estimator_multi_class() {
    DistanceConfig config;
    config.ref_spl_gunshot = 150.0f;
    config.ref_spl_artillery = 180.0f;
    config.ref_spl_explosion = 170.0f;
    config.attenuation_alpha = 0.6f;
    DistanceEstimator est(config);
    float d1 = est.Estimate(120.0f, "gunshot");
    float d2 = est.Estimate(120.0f, "artillery");
    float d3 = est.Estimate(120.0f, "explosion");
    assert(d2 > d1); // artillery has higher ref SPL -> farther estimated distance at same SPL
    assert(d3 > d1);
    std::cout << "[PASS] distance_estimator_multi_class" << std::endl;
}

void test_threat_tracker_filtering() {
    ThreatTracker tracker(0.3f, 15.0f, 5);
    AcousticThreat t1;
    t1.sound_type = "gunshot";
    t1.confidence = 0.9f;
    t1.azimuth = 45.0f;
    t1.distance = 100.0f;

    // First update should publish
    auto out1 = tracker.Update({t1});
    assert(out1.size() == 1);

    // Immediate second update with same azimuth should be filtered (interval < 0.3s)
    auto out2 = tracker.Update({t1});
    assert(out2.empty());

    std::cout << "[PASS] threat_tracker_filtering" << std::endl;
}

int main() {
    std::cout << "Running core library tests..." << std::endl;
    test_audio_buffer();
    test_audio_buffer_overflow();
    test_preemphasis();
    test_gcc_phat_cross_array();
    test_distance_estimator();
    test_distance_estimator_multi_class();
    test_threat_tracker();
    test_threat_tracker_filtering();
    test_feature_extractor_shape();
    std::cout << "All tests passed!" << std::endl;
    return 0;
}
