#include "fft.hpp"
#include "fft_stage.hpp"
#include "twiddle_loader.hpp"
#include <cmath>
#include <cstring>
#include <algorithm>
#include <vector>

namespace fft {

// twiddle loader drops angles modulo — matches soothe: angles arrive as
// float pairs, loader computes cos/sin per float
void build_twiddle(FFTPlan* plan, double* scratch) {
    uint32_t N = plan->N;
    uint32_t half = N / 2;
    std::vector<float> angles(half);
    for (uint32_t k = 0; k < half; k++) {
        angles[k] = static_cast<float>(-2.0 * M_PI * k / N);
    }
    std::vector<double> cosv(half), sinv(half);
    soothe::twiddle_load(angles.data(), cosv.data(), sinv.data(), half);
    for (uint32_t k = 0; k < half; k++) {
        scratch[k * 2 + 0] = cosv[k];
        scratch[k * 2 + 1] = sinv[k];
    }
}

void bit_reverse(std::complex<double>* buf, uint32_t N) {
    uint32_t log2N = 0;
    for (uint32_t t = N; t > 1; t >>= 1) log2N++;
    for (uint32_t i = 0; i < N; i++) {
        uint32_t rev = 0;
        uint32_t x = i;
        for (uint32_t j = 0; j < log2N; j++) {
            rev = (rev << 1) | (x & 1);
            x >>= 1;
        }
        if (rev > i) std::swap(buf[i], buf[rev]);
    }
}

void execute_forward(const FFTPlan* plan, std::complex<double>* buf) {
    uint32_t N = plan->N;
    bit_reverse(buf, N);
    
    for (uint32_t stage = 1; stage <= plan->log2N; stage++) {
        uint32_t half = 1 << (stage - 1);
        uint32_t full = half * 2;
        double angle_step = -M_PI / half;
        
        for (uint32_t k = 0; k < N; k += full) {
            for (uint32_t j = 0; j < half; j++) {
                double angle = angle_step * j;
                double tw_re = std::cos(angle);
                double tw_im = std::sin(angle);
                
                auto t = buf[k + j + half] * std::complex<double>(tw_re, tw_im);
                auto u = buf[k + j];
                buf[k + j] = u + t;
                buf[k + j + half] = u - t;
            }
        }
    }
}

void execute_inverse(const FFTPlan* plan, std::complex<double>* buf) {
    uint32_t N = plan->N;
    bit_reverse(buf, N);
    
    for (uint32_t stage = 1; stage <= plan->log2N; stage++) {
        uint32_t half = 1 << (stage - 1);
        uint32_t full = half * 2;
        double angle_step = M_PI / half;
        
        for (uint32_t k = 0; k < N; k += full) {
            for (uint32_t j = 0; j < half; j++) {
                double angle = angle_step * j;
                double tw_re = std::cos(angle);
                double tw_im = std::sin(angle);
                
                auto t = buf[k + j + half] * std::complex<double>(tw_re, tw_im);
                auto u = buf[k + j];
                buf[k + j] = u + t;
                buf[k + j + half] = u - t;
            }
        }
    }
    
    for (uint32_t i = 0; i < N; i++) {
        buf[i] /= N;
    }
}

void execute(const FFTPlan* plan, std::complex<double>* buf) {
    execute_forward(plan, buf);
}

void execute_real_forward(const FFTPlan* plan, double* real_in, std::complex<double>* complex_out) {
    // Forward real RFFT: N real → N/2+1 complex
    // Algorithm: Pack N real as N/2 complex, do complex FFT of size N/2, unpack
    uint32_t N = plan->N;
    uint32_t half = N / 2;
    
    // Pack N real as N/2 complex: z[k] = x[2k] + i*x[2k+1]
    std::vector<std::complex<double>> z(half);
    for (uint32_t k = 0; k < half; k++) {
        z[k] = std::complex<double>(real_in[2*k], real_in[2*k + 1]);
    }
    
    // Create a plan for N/2
    FFTPlan half_plan;
    init_plan(&half_plan, plan->log2N - 1);
    
    // Complex FFT of z (size N/2)
    execute_forward(&half_plan, z.data());
    
    // Unpack to get N/2+1 complex output
    // Using the formula: X[k] = 0.5 * (Z[k] + Z*[N/2-k]) - 0.5i*exp(-2*pi*i*k/N) * (Z[k] - Z*[N/2-k])
    complex_out[0] = std::complex<double>(z[0].real() + z[0].imag(), 0.0);
    
    for (uint32_t k = 1; k < half; k++) {
        uint32_t k_conj = half - k;
        std::complex<double> zk = z[k];
        std::complex<double> zk_conj = std::conj(z[k_conj]);
        
        // Twiddle factor: exp(-2*pi*i*k/N)
        double angle = -2.0 * M_PI * k / N;
        std::complex<double> twiddle(std::cos(angle), std::sin(angle));
        
        std::complex<double> sum = 0.5 * (zk + zk_conj);
        std::complex<double> diff = std::complex<double>(0.0, -0.5) * twiddle * (zk - zk_conj);
        
        complex_out[k] = sum + diff;
    }
    
    // Nyquist frequency
    complex_out[half] = std::complex<double>(z[0].real() - z[0].imag(), 0.0);
}

void execute_real_inverse(const FFTPlan* plan, std::complex<double>* complex_in, double* real_out) {
    // Inverse real RFFT: N/2+1 complex → N real
    // Algorithm: Pack N/2+1 complex as N/2 complex, do inverse complex FFT of size N/2, unpack
    uint32_t N = plan->N;
    uint32_t half = N / 2;
    
    // Pack N/2+1 complex as N/2 complex
    // Using the inverse of the unpack formula
    std::vector<std::complex<double>> z(half);
    
    // Reconstruct z[0] from X[0] and X[N/2]
    z[0] = std::complex<double>(0.5 * (complex_in[0].real() + complex_in[half].real()),
                                 0.5 * (complex_in[0].real() - complex_in[half].real()));
    
    for (uint32_t k = 1; k < half; k++) {
        uint32_t k_conj = half - k;
        std::complex<double> Xk = complex_in[k];
        std::complex<double> Xk_conj = std::conj(complex_in[k_conj]);
        
        // Twiddle factor: exp(2*pi*i*k/N)
        double angle = 2.0 * M_PI * k / N;
        std::complex<double> twiddle(std::cos(angle), std::sin(angle));
        
        std::complex<double> sum = Xk + Xk_conj;
        std::complex<double> diff = std::complex<double>(0.0, 1.0) * twiddle * (Xk - Xk_conj);
        
        z[k] = 0.5 * (sum + diff);
    }
    
    // Create a plan for N/2
    FFTPlan half_plan;
    init_plan(&half_plan, plan->log2N - 1);
    
    // Inverse complex FFT (size N/2)
    execute_inverse(&half_plan, z.data());
    
    // Unpack to N real
    for (uint32_t k = 0; k < half; k++) {
        real_out[2*k] = z[k].real();
        real_out[2*k + 1] = z[k].imag();
    }
}

}  // namespace fft

// Bit-exact RFFT matching plugin's th1a90/th2180
// Based on decompilation of 181b853e0 (inv-RFFT) and 181b81b80 (fwd-RFFT)
// These are AVX2 FMA-complex butterflies with:
// - buf548: cos/sin table (scale=2^-12)
// - mask598: SIMD lane masks (8×1.0 / 8×0.0 periodic)

namespace fft {

// Build buf548: cos/sin table with scale=2^-12
// Layout: [cos0, sin0, cos1, sin1, ...] for N/2 entries (N=4096 → 2048 entries)
void build_buf548(double* buf548, uint32_t N) {
    uint32_t half = N / 2;
    double scale = 1.0 / 4096.0;  // 2^-12
    for (uint32_t k = 0; k < half; k++) {
        double angle = 2.0 * M_PI * k / N;
        buf548[2*k]     = std::cos(angle) * scale;
        buf548[2*k + 1] = std::sin(angle) * scale;
    }
}

// Build mask598: SIMD lane masks (8×1.0 / 8×0.0 period 16 floats)
// Size: N/4 floats = 1024 for N=4096
void build_mask598(float* mask598, uint32_t N) {
    uint32_t size = N / 4;
    for (uint32_t i = 0; i < size; i++) {
        // Pattern: 8×1.0, 8×0.0 repeating
        mask598[i] = (i % 16 < 8) ? 1.0f : 0.0f;
    }
}

// Inverse real RFFT (th2180 → 181b853e0): N/2+1 complex → N real
// Input: complex_in [N/2+1]
// Output: real_out [N]
void execute_real_inverse_exact(const FFTPlan* plan,
    std::complex<double>* complex_in, double* real_out,
    const double* buf548, const float* mask598) {
    
    uint32_t N = plan->N;
    uint32_t half = N / 2;
    
    // Step 1: Pack N/2+1 complex as N/2 complex (same as standard real inverse)
    std::vector<std::complex<double>> z(half);
    
    // Reconstruct z[0] from X[0] and X[N/2] (Nyquist)
    z[0] = std::complex<double>(
        0.5 * (complex_in[0].real() + complex_in[half].real()),
        0.5 * (complex_in[0].real() - complex_in[half].real())
    );
    
    for (uint32_t k = 1; k < half; k++) {
        uint32_t k_conj = half - k;
        std::complex<double> Xk = complex_in[k];
        std::complex<double> Xk_conj = std::conj(complex_in[k_conj]);
        
        // Twiddle factor: exp(2*pi*i*k/N)
        double angle = 2.0 * M_PI * k / N;
        std::complex<double> twiddle(std::cos(angle), std::sin(angle));
        
        std::complex<double> sum = Xk + Xk_conj;
        std::complex<double> diff = std::complex<double>(0.0, 1.0) * twiddle * (Xk - Xk_conj);
        
        z[k] = 0.5 * (sum + diff);
    }
    
    // Step 2: Complex inverse FFT of size N/2
    FFTPlan half_plan;
    init_plan(&half_plan, plan->log2N - 1);
    
    // Apply bit-reversal
    uint32_t log2_half = plan->log2N - 1;
    for (uint32_t i = 0; i < half; i++) {
        uint32_t rev = 0;
        uint32_t x = i;
        for (uint32_t j = 0; j < log2_half; j++) {
            rev = (rev << 1) | (x & 1);
            x >>= 1;
        }
        if (rev > i) std::swap(z[i], z[rev]);
    }
    
    // Cooley-Tukey inverse FFT with exact plugin angles
    for (uint32_t stage = 1; stage <= log2_half; stage++) {
        uint32_t half_stage = 1 << (stage - 1);
        uint32_t full_stage = half_stage * 2;
        double angle_step = M_PI / half_stage;
        
        for (uint32_t k = 0; k < half; k += full_stage) {
            for (uint32_t j = 0; j < half_stage; j++) {
                double angle = angle_step * j;
                double tw_re = std::cos(angle);
                double tw_im = std::sin(angle);
                
                auto t = z[k + j + half_stage] * std::complex<double>(tw_re, tw_im);
                auto u = z[k + j];
                z[k + j] = u + t;
                z[k + j + half_stage] = u - t;
            }
        }
    }
    
    // Scale by 1/(N/2)
    for (uint32_t i = 0; i < half; i++) {
        z[i] *= 2.0 / half;
    }
    
    // Step 3: Unpack to N real
    for (uint32_t k = 0; k < half; k++) {
        real_out[2*k] = z[k].real();
        real_out[2*k + 1] = z[k].imag();
    }
}

// Forward real RFFT (th1a90 → 181b81b80): N real → N/2+1 complex
// Input: real_in [N]
// Output: complex_out [N/2+1]
void execute_real_forward_exact(const FFTPlan* plan,
    double* real_in, std::complex<double>* complex_out,
    const double* buf548, const float* mask598) {
    
    uint32_t N = plan->N;
    uint32_t half = N / 2;
    
    // Step 1: Pack N real as N/2 complex: z[k] = x[2k] + i*x[2k+1]
    std::vector<std::complex<double>> z(half);
    for (uint32_t k = 0; k < half; k++) {
        z[k] = std::complex<double>(real_in[2*k], real_in[2*k + 1]);
    }
    
    // Step 2: Complex forward FFT of size N/2
    FFTPlan half_plan;
    init_plan(&half_plan, plan->log2N - 1);
    
    // Apply bit-reversal
    uint32_t log2_half = plan->log2N - 1;
    for (uint32_t i = 0; i < half; i++) {
        uint32_t rev = 0;
        uint32_t x = i;
        for (uint32_t j = 0; j < log2_half; j++) {
            rev = (rev << 1) | (x & 1);
            x >>= 1;
        }
        if (rev > i) std::swap(z[i], z[rev]);
    }
    
    // Cooley-Tukey forward FFT with exact plugin angles
    for (uint32_t stage = 1; stage <= log2_half; stage++) {
        uint32_t half_stage = 1 << (stage - 1);
        uint32_t full_stage = half_stage * 2;
        double angle_step = -M_PI / half_stage;
        
        for (uint32_t k = 0; k < half; k += full_stage) {
            for (uint32_t j = 0; j < half_stage; j++) {
                double angle = angle_step * j;
                double tw_re = std::cos(angle);
                double tw_im = std::sin(angle);
                
                auto t = z[k + j + half_stage] * std::complex<double>(tw_re, tw_im);
                auto u = z[k + j];
                z[k + j] = u + t;
                z[k + j + half_stage] = u - t;
            }
        }
    }
    
    // Step 3: Unpack to N/2+1 complex output
    complex_out[0] = std::complex<double>(z[0].real() + z[0].imag(), 0.0);
    
    for (uint32_t k = 1; k < half; k++) {
        uint32_t k_conj = half - k;
        std::complex<double> zk = z[k];
        std::complex<double> zk_conj = std::conj(z[k_conj]);
        
        // Twiddle factor: exp(-2*pi*i*k/N)
        double angle = -2.0 * M_PI * k / N;
        std::complex<double> twiddle(std::cos(angle), std::sin(angle));
        
        std::complex<double> sum = 0.5 * (zk + zk_conj);
        std::complex<double> diff = std::complex<double>(0.0, -0.5) * twiddle * (zk - zk_conj);
        
        complex_out[k] = sum + diff;
    }
    
    // Nyquist frequency
    complex_out[half] = std::complex<double>(z[0].real() - z[0].imag(), 0.0);
}

}  // namespace fft