// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// example_fft.cpp
// FFT API demo: runs all code snippets from fft.html
//
// Build:
//   cd build && cmake --build . --config Release --target example-fft

#define _USE_MATH_DEFINES
#include <cmath>
#include <math/fft/fft.hpp>
#include <math/fft/fft_utils.hpp>
#include <algorithm>
#include <iostream>
#include <vector>
#include <complex>
#include <span>

using namespace sangi;

// ============================================================
// 1. FFT forward/inverse transform (complex)
// ============================================================
void demo_fft_roundtrip() {
    std::cout << "=== demo_fft_roundtrip ===" << std::endl;

    using C = std::complex<double>;
    std::vector<C> data = {1, 2, 3, 4};
    auto original = data;

    std::cout << "  Input: ";
    for (auto& v : data) std::cout << v.real() << " ";
    std::cout << std::endl;

    FFT<C>::fft(data);   // forward transform
    std::cout << "  After FFT:" << std::endl;
    for (size_t i = 0; i < data.size(); ++i)
        std::cout << "    X[" << i << "] = " << data[i] << std::endl;

    FFT<C>::ifft(data);  // inverse transform -- recovers original
    std::cout << "  After IFFT (should match original):" << std::endl;
    for (size_t i = 0; i < data.size(); ++i)
        std::cout << "    x[" << i << "] = " << data[i].real() << std::endl;

    std::cout << std::endl;
}

// ============================================================
// 2. RealFFT forward/inverse transform
// ============================================================
void demo_realfft_roundtrip() {
    std::cout << "=== demo_realfft_roundtrip ===" << std::endl;

    std::vector<double> signal = {1, 0, -1, 0, 1, 0, -1, 0};
    std::cout << "  Input signal: ";
    for (auto v : signal) std::cout << v << " ";
    std::cout << "(N=" << signal.size() << ")" << std::endl;

    auto spectrum = RealFFT<double>::fft(signal);       // 5 complex values
    auto recovered = RealFFT<double>::ifft(spectrum, 8); // 8 real values

    std::cout << "  Spectrum (" << spectrum.size() << " elements):" << std::endl;
    for (size_t i = 0; i < spectrum.size(); ++i)
        std::cout << "    S[" << i << "] = " << spectrum[i] << std::endl;

    std::cout << "  Recovered signal:" << std::endl;
    double max_err = 0;
    for (size_t i = 0; i < recovered.size(); ++i) {
        std::cout << "    x[" << i << "] = " << recovered[i] << std::endl;
        max_err = (std::max)(max_err, std::abs(signal[i] - recovered[i]));
    }
    std::cout << "  Max reconstruction error: " << max_err << std::endl;

    std::cout << std::endl;
}

// ============================================================
// 3. FFTEngine (instance-based, with normalization mode)
// ============================================================
void demo_fft_engine() {
    std::cout << "=== demo_fft_engine ===" << std::endl;

    const int N = 1024;
    std::cout << "  N=" << N << ", signal = sin(2*pi*10*n/N)" << std::endl;
    FFTEngine<double> engine(N, FFTDivMode::Inverse);

    using C = std::complex<double>;
    std::vector<C> input(N);
    for (int i = 0; i < N; ++i)
        input[i] = C(std::sin(2.0 * M_PI * 10.0 * i / N), 0);

    auto spectrum = engine.transform(input);
    auto recovered = engine.inverse(spectrum);

    double max_err = 0;
    for (int i = 0; i < N; ++i)
        max_err = (std::max)(max_err, std::abs(input[i].real() - recovered[i].real()));
    std::cout << "  Complex transform roundtrip error: " << max_err << std::endl;

    std::vector<double> real_input(N);
    for (int i = 0; i < N; ++i)
        real_input[i] = std::sin(2.0 * M_PI * 10.0 * i / N);

    auto real_spec = engine.real_transform(real_input);
    auto real_out = engine.real_inverse(real_spec);

    max_err = 0;
    for (int i = 0; i < N; ++i)
        max_err = (std::max)(max_err, std::abs(real_input[i] - real_out[i]));
    std::cout << "  Real transform roundtrip error: " << max_err << std::endl;

    std::cout << std::endl;
}

// ============================================================
// 4. Spectrum analysis (50 Hz + 120 Hz signal)
// ============================================================
void demo_spectrum_analysis() {
    std::cout << "=== demo_spectrum_analysis ===" << std::endl;

    // Composite signal of 50 Hz + 120 Hz (assuming sample rate = 1024 Hz)
    std::cout << "  Input: sin(2*pi*50*t) + sin(2*pi*120*t), N=1024, Fs=1024 Hz" << std::endl;
    std::vector<double> signal(1024);
    for (int i = 0; i < 1024; i++)
        signal[i] = std::sin(2 * M_PI * 50 * i / 1024)   // 50 Hz component
                  + std::sin(2 * M_PI * 120 * i / 1024);  // 120 Hz component

    auto spec = RealFFT<double>::fft(signal);
    auto amp = amplitude_spectrum<double>(std::span<const std::complex<double>>(spec));
    auto phase = phase_spectrum<double>(std::span<const std::complex<double>>(spec));

    // Peaks appear at amp[50] and amp[120]
    std::cout << "  Peaks (amplitude > 100):" << std::endl;
    for (size_t i = 0; i < amp.size(); ++i) {
        if (amp[i] > 100.0)
            std::cout << "    freq=" << i << " Hz, amplitude=" << amp[i] << std::endl;
    }
    std::cout << "  Phase at 50 Hz: " << phase[50] << " rad" << std::endl;
    std::cout << "  Phase at 120 Hz: " << phase[120] << " rad" << std::endl;

    std::cout << std::endl;
}

// ============================================================
// 5. Polynomial multiplication
// ============================================================
void demo_polynomial_multiply() {
    std::cout << "=== demo_polynomial_multiply ===" << std::endl;

    using C = std::complex<double>;
    std::vector<C> a = {1, 2, 3};   // 1 + 2x + 3x^2
    std::vector<C> b = {4, 5};      // 4 + 5x
    auto result = polynomial_multiply<C>(
        std::span<const C>(a), std::span<const C>(b));
    // result = {4, 13, 22, 15} -> 4 + 13x + 22x^2 + 15x^3

    std::cout << "  (1 + 2x + 3x^2) * (4 + 5x) = [";
    for (size_t i = 0; i < result.size(); ++i) {
        if (i > 0) std::cout << ", ";
        std::cout << result[i].real();
    }
    std::cout << "]" << std::endl;
    std::cout << "  (expected: [4, 13, 22, 15])" << std::endl;

    std::cout << std::endl;
}

// ============================================================
// 6. NTT convolution (exact integer arithmetic)
// ============================================================
void demo_ntt_convolve() {
    std::cout << "=== demo_ntt_convolve ===" << std::endl;

    constexpr int P = 998244353;  // NTT-friendly prime: P-1 = 2^23 * 119
    using Mint = ModularInt<P>;
    using NTTF = FFT<Mint>;

    // ntt/intt computes twiddle factors as primitive_root.pow((P-1)/len).
    // Therefore primitive_root must be a primitive root of (Z/PZ)* (order P-1).
    // For P = 998244353, the primitive root is 3.
    Mint root(3);
    std::cout << "  Primitive root mod " << P << ": " << root.value() << std::endl;

    // (1 + 2x + 3x^2) * (4 + 5x) mod P
    std::vector<Mint> a = {Mint(1), Mint(2), Mint(3)};
    std::vector<Mint> b = {Mint(4), Mint(5)};
    auto result = NTTF::convolve_ntt(
        std::span<const Mint>(a), std::span<const Mint>(b), root);

    std::cout << "  (1 + 2x + 3x^2) * (4 + 5x) mod " << P << " = [";
    for (size_t i = 0; i < result.size(); ++i) {
        if (i > 0) std::cout << ", ";
        std::cout << result[i].value();
    }
    std::cout << "]" << std::endl;
    std::cout << "  (expected: [4, 13, 22, 15])" << std::endl;

    std::cout << std::endl;
}

// ============================================================
// 7. DCT (Discrete Cosine Transform)
// ============================================================
void demo_dct() {
    std::cout << "=== demo_dct ===" << std::endl;

    std::vector<double> data = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0};
    std::cout << "  Input: ";
    for (auto v : data) std::cout << v << " ";
    std::cout << std::endl;

    auto coeffs = dct<double>(std::span<const double>(data));
    auto recovered = idct<double>(std::span<const double>(coeffs));
    auto fast_coeffs = fast_dct<double>(std::span<const double>(data));

    std::cout << "  DCT coefficients:" << std::endl;
    for (size_t i = 0; i < coeffs.size(); ++i)
        std::cout << "    C[" << i << "] = " << coeffs[i] << std::endl;

    std::cout << "  fast_dct coefficients:" << std::endl;
    for (size_t i = 0; i < fast_coeffs.size(); ++i)
        std::cout << "    C[" << i << "] = " << fast_coeffs[i] << std::endl;

    double max_err = 0;
    for (size_t i = 0; i < data.size(); ++i)
        max_err = (std::max)(max_err, std::abs(data[i] - recovered[i]));
    std::cout << "  IDCT reconstruction error: " << max_err << std::endl;

    double max_diff = 0;
    for (size_t i = 0; i < coeffs.size(); ++i)
        max_diff = (std::max)(max_diff, std::abs(coeffs[i] - fast_coeffs[i]));
    std::cout << "  dct vs fast_dct max difference: " << max_diff << std::endl;

    std::cout << std::endl;
}

// ============================================================
// 8. RealFFT convolution
// ============================================================
void demo_realfft_convolve() {
    std::cout << "=== demo_realfft_convolve ===" << std::endl;

    std::vector<double> a = {1, 2, 3, 4};
    std::vector<double> b = {1, 0, -1};
    auto conv = RealFFT<double>::convolve(a, b);

    std::cout << "  [1,2,3,4] * [1,0,-1] = [";
    for (size_t i = 0; i < conv.size(); ++i) {
        if (i > 0) std::cout << ", ";
        std::cout << conv[i];
    }
    std::cout << "]" << std::endl;
    std::cout << "  (expected: [1, 2, 2, 2, -3, -4])" << std::endl;

    std::cout << std::endl;
}

// ============================================================
// 9. Spectrum reconstruction (amplitude + phase -> complex spectrum)
// ============================================================
void demo_reconstruct_spectrum() {
    std::cout << "=== demo_reconstruct_spectrum ===" << std::endl;

    std::vector<double> signal = {1, 0, -1, 0, 1, 0, -1, 0};
    std::cout << "  Input: ";
    for (auto v : signal) std::cout << v << " ";
    std::cout << std::endl;

    auto spec = RealFFT<double>::fft(signal);
    auto amp = amplitude_spectrum<double>(std::span<const std::complex<double>>(spec));
    auto ph = phase_spectrum<double>(std::span<const std::complex<double>>(spec));

    // Reconstruct complex spectrum from amplitude and phase
    auto reconstructed = reconstruct_spectrum<double>(
        std::span<const double>(amp), std::span<const double>(ph));

    double max_err = 0;
    for (size_t i = 0; i < spec.size(); ++i)
        max_err = (std::max)(max_err, std::abs(spec[i] - reconstructed[i]));
    std::cout << "  Spectrum reconstruction error: " << max_err << std::endl;

    std::cout << std::endl;
}

// ============================================================
int main() {
    std::cout << "=== sangi: FFT API Demo ===" << std::endl << std::endl;

    demo_fft_roundtrip();
    demo_realfft_roundtrip();
    demo_fft_engine();
    demo_spectrum_analysis();
    demo_polynomial_multiply();
    demo_ntt_convolve();
    demo_dct();
    demo_realfft_convolve();
    demo_reconstruct_spectrum();

    std::cout << "All demos completed." << std::endl;
    return 0;
}