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

namespace fft {

void build_twiddle(FFTPlan* plan, double* scratch) {
    uint32_t N = plan->N;
    uint32_t half = N / 2;
    for (uint32_t k = 0; k < half; k++) {
        double angle = -2.0 * M_PI * k / N;
        scratch[k * 2 + 0] = std::cos(angle);
        scratch[k * 2 + 1] = std::sin(angle);
    }
}

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;
    for (uint32_t i = 0; i < N; i++) {
        buf[i] = std::conj(buf[i]);
    }
    
    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 init_plan(FFTPlan* plan, uint32_t log2N) {
    plan->log2N = log2N;
    plan->N = 1U << log2N;
    plan->stage_count = log2N;
    plan->bit_reverse = 1;
    plan->xor_mask = 0;
}

}