#include "fft.h"
#include <math.h>

static float window[FFT_SIZE];

constexpr float M_PI = 3.14159;


// Call once before use
void init_hamming_window() {
    for (int i = 0; i < FFT_SIZE; i++) {
        window[i] = 0.54f - 0.46f * cosf(2.0f * (float)M_PI * i / (FFT_SIZE - 1));
    }
}

static unsigned int bit_reverse(unsigned int x, int log2n) {
    unsigned int n = 0;
    for (int i = 0; i < log2n; i++) {
        n <<= 1;
        n |= (x & 1);
        x >>= 1;
    }
    return n;
}

void compute_fft(float* time_data, float* freq_out) {
    static float real[FFT_SIZE];
    static float imag[FFT_SIZE];

    int log2n = 0;
    for (int t = FFT_SIZE; t > 1; t >>= 1) ++log2n;

    // Apply Hamming window
    for (int i = 0; i < FFT_SIZE; i++) {
        real[i] = time_data[i] * window[i];
        imag[i] = 0.0f;
    }

    // Bit reversal
    for (int i = 0; i < FFT_SIZE; ++i) {
        int j = bit_reverse(i, log2n);
        if (j > i) {
            float tmp_re = real[i], tmp_im = imag[i];
            real[i] = real[j]; imag[i] = imag[j];
            real[j] = tmp_re; imag[j] = tmp_im;
        }
    }

    // Cooley-Tukey FFT
    for (int s = 1; s <= log2n; ++s) {
        int m = 1 << s;
        for (int k = 0; k < FFT_SIZE; k += m) {
            for (int j = 0; j < m / 2; ++j) {
                int t = k + j;
                int u = t + m / 2;

                float angle = -2.0f * (float)M_PI * j / m;
                float w_real = cosf(angle);
                float w_imag = sinf(angle);

                float re = w_real * real[u] - w_imag * imag[u];
                float im = w_real * imag[u] + w_imag * real[u];

                real[u] = real[t] - re;
                imag[u] = imag[t] - im;
                real[t] += re;
                imag[t] += im;
            }
        }
    }

    for (int i = 0; i < FFT_SIZE / 2; ++i) {
        float mag = sqrtf(real[i] * real[i] + imag[i] * imag[i]) / FFT_SIZE;
        float db = 20.0f * log10f(mag + 1e-6f); // Decibels
        float normalized = (db + 60.0f) / 60.0f; // [0,1]
        freq_out[i] = mag;
    }
}