#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();
    }
}

}