#include "engine/framework/audio/dsp.h"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <complex>
#include <cstdint>
#include <iostream>
#include <random>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>

namespace {

constexpr double kPi = 3.14159265358979323846264338327950288;

enum class WaveformMode {
    Fixed,
    PerRun,
};

struct WaveformProfile {
    double envelope_bias = 0.0;
    double envelope_rate_scale = 1.0;
    double envelope_phase = 0.0;
    double voiced_rate_scale = 1.0;
    double fricative_rate_scale = 1.0;
    double pause_period = 2.4;
    double pause_start = 1.95;
    double pause_end = 2.20;
};

const WaveformProfile & get_per_run_waveform_profile() {
    static const WaveformProfile profile = []() {
        const auto ticks = static_cast<uint64_t>(
            std::chrono::steady_clock::now().time_since_epoch().count());
        std::seed_seq seed{
            static_cast<uint32_t>(ticks),
            static_cast<uint32_t>(ticks >> 32),
            0x6d2b79f5u,
            0x9e3779b9u,
        };
        std::mt19937 rng(seed);
        std::uniform_real_distribution<double> envelope_bias(-0.08, 0.08);
        std::uniform_real_distribution<double> envelope_rate_scale(0.85, 1.18);
        std::uniform_real_distribution<double> envelope_phase(0.0, 2.0 * kPi);
        std::uniform_real_distribution<double> voiced_rate_scale(0.92, 1.10);
        std::uniform_real_distribution<double> fricative_rate_scale(0.88, 1.15);
        std::uniform_real_distribution<double> pause_period(2.0, 2.8);
        std::uniform_real_distribution<double> pause_start(1.6, 2.1);
        std::uniform_real_distribution<double> pause_width(0.14, 0.32);
        WaveformProfile value;
        value.envelope_bias = envelope_bias(rng);
        value.envelope_rate_scale = envelope_rate_scale(rng);
        value.envelope_phase = envelope_phase(rng);
        value.voiced_rate_scale = voiced_rate_scale(rng);
        value.fricative_rate_scale = fricative_rate_scale(rng);
        value.pause_period = pause_period(rng);
        value.pause_start = pause_start(rng);
        value.pause_end = std::min(value.pause_period - 0.05, value.pause_start + pause_width(rng));
        return value;
    }();
    return profile;
}

void require(bool condition, const std::string & message) {
    if (!condition) {
        throw std::runtime_error(message);
    }
}

void require_shape_equal(
    const std::vector<int64_t> & actual,
    const std::vector<int64_t> & expected,
    const std::string & label) {
    if (actual != expected) {
        std::ostringstream oss;
        oss << label << " shape mismatch";
        throw std::runtime_error(oss.str());
    }
}

void require_close(
    const std::vector<float> & actual,
    const std::vector<float> & expected,
    float max_allowed,
    double mean_allowed,
    const std::string & label) {
    require(actual.size() == expected.size(), label + " size mismatch");
    float max_diff = 0.0f;
    size_t max_index = 0;
    double mean_diff = 0.0;
    for (size_t i = 0; i < actual.size(); ++i) {
        const float diff = std::fabs(actual[i] - expected[i]);
        mean_diff += diff;
        if (diff > max_diff) {
            max_diff = diff;
            max_index = i;
        }
    }
    mean_diff /= static_cast<double>(actual.size());
    if (max_diff > max_allowed || mean_diff > mean_allowed) {
        std::ostringstream oss;
        oss << label << " drift too large: max_diff=" << max_diff
            << " mean_diff=" << mean_diff
            << " at index " << max_index
            << " expected=" << expected[max_index]
            << " actual=" << actual[max_index];
        throw std::runtime_error(oss.str());
    }
}

std::vector<float> make_speech_like_waveform(
    int64_t sample_rate,
    double seconds,
    WaveformMode mode = WaveformMode::Fixed,
    int variant = 0) {
    const int64_t samples = static_cast<int64_t>(sample_rate * seconds);
    std::vector<float> waveform(static_cast<size_t>(samples), 0.0f);
    const WaveformProfile & profile =
        mode == WaveformMode::PerRun
            ? get_per_run_waveform_profile()
            : WaveformProfile{};
    const double variant_scale = 1.0 + 0.035 * static_cast<double>(variant);
    const double envelope_phase = profile.envelope_phase + 0.47 * static_cast<double>(variant);
    const double voiced_scale = profile.voiced_rate_scale * variant_scale;
    const double fricative_scale = profile.fricative_rate_scale * (1.0 + 0.02 * static_cast<double>(variant));
    const double pause_period = std::max(1.5, profile.pause_period + 0.06 * static_cast<double>(variant));
    const double pause_start = std::min(pause_period - 0.12, profile.pause_start + 0.03 * static_cast<double>(variant));
    const double pause_end = std::min(pause_period - 0.04, profile.pause_end + 0.03 * static_cast<double>(variant));
    for (int64_t i = 0; i < samples; ++i) {
        const double t = static_cast<double>(i) / static_cast<double>(sample_rate);
        const double envelope =
            0.55 + profile.envelope_bias +
            0.30 * std::sin(2.0 * kPi * 0.7 * profile.envelope_rate_scale * t + envelope_phase) +
            0.10 * std::sin(2.0 * kPi * 0.11 * profile.envelope_rate_scale * t + 0.3 + envelope_phase * 0.5);
        const double voiced =
            0.55 * std::sin(2.0 * kPi * 180.0 * voiced_scale * t) +
            0.25 * std::sin(2.0 * kPi * 360.0 * voiced_scale * t + 0.2 + envelope_phase * 0.25) +
            0.12 * std::sin(2.0 * kPi * 520.0 * voiced_scale * t + 0.5 + envelope_phase * 0.4);
        const double fricative =
            0.03 * std::sin(2.0 * kPi * 1900.0 * fricative_scale * t) +
            0.02 * std::sin(2.0 * kPi * 2600.0 * fricative_scale * t + 0.1 + envelope_phase * 0.15);
        const double pause_phase = std::fmod(t, pause_period);
        const double pause_gate = (pause_phase > pause_start && pause_phase < pause_end) ? 0.15 : 1.0;
        waveform[static_cast<size_t>(i)] = static_cast<float>(pause_gate * envelope * (voiced + fricative));
    }
    return waveform;
}

engine::audio::AudioTensor naive_istft_reference(
    const std::vector<float> & complex_spec,
    const std::vector<float> & window,
    int64_t batch,
    int64_t freq_bins,
    int64_t frames,
    int64_t samples,
    const engine::audio::STFTConfig & config) {
    const int64_t pad = config.n_fft / 2;
    const int64_t padded_samples = samples + 2 * pad;
    std::vector<float> accum(static_cast<size_t>(batch * padded_samples), 0.0f);
    std::vector<float> window_sums(static_cast<size_t>(batch * padded_samples), 0.0f);

    for (int64_t b = 0; b < batch; ++b) {
        for (int64_t frame_index = 0; frame_index < frames; ++frame_index) {
            std::vector<std::complex<float>> spectrum(static_cast<size_t>(config.n_fft), 0.0f);
            for (int64_t k = 0; k < freq_bins; ++k) {
                const size_t base = static_cast<size_t>((((b * freq_bins) + k) * frames + frame_index) * 2);
                spectrum[static_cast<size_t>(k)] = {complex_spec[base], complex_spec[base + 1]};
            }
            for (int64_t k = 1; k < freq_bins - 1; ++k) {
                spectrum[static_cast<size_t>(config.n_fft - k)] = std::conj(spectrum[static_cast<size_t>(k)]);
            }

            const int64_t start = frame_index * config.hop_length;
            for (int64_t n = 0; n < config.n_fft; ++n) {
                std::complex<float> sum = 0.0f;
                for (int64_t k = 0; k < config.n_fft; ++k) {
                    const float angle = 2.0f * kPi * static_cast<float>(k * n) / static_cast<float>(config.n_fft);
                    sum += spectrum[static_cast<size_t>(k)] * std::complex<float>(std::cos(angle), std::sin(angle));
                }
                float value = sum.real() / static_cast<float>(config.n_fft);
                if (n < config.win_length) {
                    value *= window[static_cast<size_t>(n)];
                } else {
                    value = 0.0f;
                }
                accum[static_cast<size_t>(b * padded_samples + start + n)] += value;
                if (n < config.win_length) {
                    const float w = window[static_cast<size_t>(n)];
                    window_sums[static_cast<size_t>(b * padded_samples + start + n)] += w * w;
                }
            }
        }
    }

    engine::audio::AudioTensor result;
    result.shape = {batch, samples};
    result.values.assign(static_cast<size_t>(batch * samples), 0.0f);
    for (int64_t b = 0; b < batch; ++b) {
        for (int64_t i = 0; i < samples; ++i) {
            const size_t padded_index = static_cast<size_t>(b * padded_samples + i + pad);
            const float denom = window_sums[padded_index] > 1e-8f ? window_sums[padded_index] : 1.0f;
            result.values[static_cast<size_t>(b * samples + i)] = accum[padded_index] / denom;
        }
    }
    return result;
}

void test_log_mel_matches_reference_pipeline() {
    constexpr int64_t kSampleRate = 16000;
    constexpr float kLogZeroGuard = 5.960464477539063e-8f;  // 2^-24
    const auto waveform = make_speech_like_waveform(kSampleRate, 10.0);
    const engine::audio::STFTConfig config{
        512,
        160,
        400,
        true,
        engine::audio::STFTPadMode::Reflect,
    };

    auto reference = engine::audio::MelSpectrogram().compute(
        waveform,
        1,
        static_cast<int64_t>(waveform.size()),
        kSampleRate,
        config,
        80);
    for (float & value : reference.values) {
        value = std::log(value + kLogZeroGuard);
    }
    const auto candidate = engine::audio::LogMelSpectrogram().compute(
        waveform,
        1,
        static_cast<int64_t>(waveform.size()),
        kSampleRate,
        config,
        80);

    require_shape_equal(candidate.shape, reference.shape, "log_mel");
    require_close(candidate.values, reference.values, 2.0e-5f, 2.0e-6, "log_mel");
}

void test_stft_istft_round_trip_matches_waveform() {
    constexpr int64_t kSampleRate = 24000;
    const auto waveform = make_speech_like_waveform(kSampleRate, 3.0);
    const engine::audio::STFTConfig config{
        20,
        5,
        20,
        true,
        engine::audio::STFTPadMode::Reflect,
        engine::audio::STFTFamily::Kokoro,
    };

    const auto & window = engine::audio::get_cached_stft_window(config);
    const auto complex = engine::audio::STFT().compute_complex(
        waveform,
        window,
        1,
        static_cast<int64_t>(waveform.size()),
        config);

    require_shape_equal(
        complex.shape,
        {1, (config.n_fft / 2) + 1, 1 + static_cast<int64_t>(waveform.size()) / config.hop_length, 2},
        "stft_complex");

    const auto reconstructed = engine::audio::ISTFT().compute(
        complex.values,
        window,
        1,
        complex.shape[1],
        complex.shape[2],
        static_cast<int64_t>(waveform.size()),
        config);

    require_shape_equal(
        reconstructed.shape,
        {1, static_cast<int64_t>(waveform.size())},
        "istft_waveform");
    require_close(
        reconstructed.values,
        waveform,
        2.0e-4f,
        2.0e-5,
        "stft_istft_round_trip");
}

void test_istft_matches_reference_across_configs_and_variants() {
    const std::vector<engine::audio::STFTConfig> configs{
        engine::audio::STFTConfig{
            20,
            5,
            20,
            true,
            engine::audio::STFTPadMode::Reflect,
            engine::audio::STFTFamily::Kokoro,
        },
        engine::audio::STFTConfig{
            512,
            160,
            400,
            true,
            engine::audio::STFTPadMode::Reflect,
            engine::audio::STFTFamily::Default,
        },
    };

    for (const auto & config : configs) {
        for (int variant = 0; variant < 4; ++variant) {
            const auto waveform = make_speech_like_waveform(
                24000,
                2.0,
            WaveformMode::PerRun,
            variant);
            const auto & window = engine::audio::get_cached_stft_window(config);
            const auto complex = engine::audio::STFT().compute_complex(
                waveform,
                window,
                1,
                static_cast<int64_t>(waveform.size()),
                config);

            const auto reference = naive_istft_reference(
                complex.values,
                window,
                1,
                complex.shape[1],
                complex.shape[2],
                static_cast<int64_t>(waveform.size()),
                config);
            const auto reconstructed = engine::audio::ISTFT().compute(
                complex.values,
                window,
                1,
                complex.shape[1],
                complex.shape[2],
                static_cast<int64_t>(waveform.size()),
                config);

            require_shape_equal(reconstructed.shape, reference.shape, "istft_variant_shape");
            require_close(
                reconstructed.values,
                reference.values,
                2.0e-5f,
                2.0e-6,
                "istft_variant_parity");
        }
    }
}

}  // namespace

int main() {
    try {
        test_log_mel_matches_reference_pipeline();
        test_stft_istft_round_trip_matches_waveform();
        test_istft_matches_reference_across_configs_and_variants();
        std::cout << "audio_dsp_test: ok\n";
        return 0;
    } catch (const std::exception & ex) {
        std::cerr << "audio_dsp_test: " << ex.what() << "\n";
        return 1;
    }
}