// fft.c - Radix-2 Cooley-Tukey FFT
//
// Twiddle factors and the bit-reversal permutation are precomputed per FFT size
// and cached (FFTW's "plan" idea, kept lightweight): the first call for a size
// builds the tables, every later call just reads them. This removes the per-
// butterfly cexpf() and the per-element bit-twiddling that dominated the old
// straight-line version.
#include "fft.h"

#include <math.h>
#include <complex.h>
#include <stdbool.h>
#include <stdlib.h>

// Cache of plans, keyed by transform size. FFT_SIZE_MAX is 2048, so a handful
// of distinct power-of-two sizes ever appear; a tiny fixed table is plenty.
#define FFT_MAX_PLANS 16

typedef struct {
    int n;                      // transform size (0 = empty slot)
    int* rev;                   // bit-reversal permutation, length n
    float complex* twiddle;     // forward twiddles W_n^k, length n/2
} FFTPlan;

static FFTPlan g_plans[FFT_MAX_PLANS];

static FFTPlan* GetPlan(int n)
{
    for (int i = 0; i < FFT_MAX_PLANS; i++)
        if (g_plans[i].n == n) return &g_plans[i];

    // Find a free slot. (Sizes are few and long-lived, so no eviction needed.)
    FFTPlan* slot = NULL;
    for (int i = 0; i < FFT_MAX_PLANS; i++)
        if (g_plans[i].n == 0) { slot = &g_plans[i]; break; }
    if (!slot) slot = &g_plans[0]; // pathological fallback: reuse slot 0

    free(slot->rev);
    free(slot->twiddle);

    int bits = 0, temp = n;
    while (temp > 1) { bits++; temp >>= 1; }

    slot->rev = (int*)malloc((size_t)n * sizeof(int));
    for (int i = 0; i < n; i++) {
        int j = 0, k = i;
        for (int b = 0; b < bits; b++) { j = (j << 1) | (k & 1); k >>= 1; }
        slot->rev[i] = j;
    }

    // Forward twiddles W_n^k = exp(-2*pi*i*k/n) for k in [0, n/2).
    slot->twiddle = (float complex*)malloc((size_t)(n / 2) * sizeof(float complex));
    for (int k = 0; k < n / 2; k++) {
        float ang = -2.0f * (float)M_PI * k / n;
        slot->twiddle[k] = cosf(ang) + sinf(ang) * I;
    }

    slot->n = n;
    return slot;
}

void FFT(float complex* input, float complex* output, int n, bool inverse)
{
    if (n <= 1) { output[0] = input[0]; return; }

    FFTPlan* plan = GetPlan(n);

    // Bit-reversal permutation via the precomputed table.
    for (int i = 0; i < n; i++) output[plan->rev[i]] = input[i];

    // Butterflies. The cached twiddle table holds W_n^k at full resolution;
    // for a stage of length `step`, the needed root W_step^k == W_n^(k*n/step),
    // so we index the table with a per-stage stride. For the inverse transform
    // we conjugate the same factor (W_n^-k), avoiding a second table.
    for (int step = 2; step <= n; step *= 2) {
        int half = step / 2;
        int stride = n / step;          // table index multiplier for this stage
        for (int k = 0; k < half; k++) {
            float complex w = plan->twiddle[k * stride];
            if (inverse) w = conjf(w);
            for (int i = k; i < n; i += step) {
                int j = i + half;
                float complex t = output[j] * w;
                output[j] = output[i] - t;
                output[i] = output[i] + t;
            }
        }
    }

    if (inverse) for (int i = 0; i < n; i++) output[i] /= n;
}
