// fft_bench.c - Isolated benchmark of FFT() over the real mlnl_samples.wav STFT
// workload (2048-pt forward transforms, hop = fftSize/4, 75% overlap).
//
// Build: cc -O2 -march=native fft_bench.c ../src/fft.c -lm -o fft_bench
// (link whichever fft.c you want to measure; -DFFT_BENCH_REPS=N to override reps)
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <complex.h>
#include <math.h>
#include <time.h>
#include "../src/fft.h"

#define FFT_SIZE   2048
#define HOP_RATIO  4
#ifndef FFT_BENCH_REPS
#define FFT_BENCH_REPS 50   // repeat the whole STFT pass this many times
#endif

// Minimal WAV loader: 16-bit PCM, returns mono float frames (first channel).
static float* load_wav(const char* path, int* outN) {
    FILE* f = fopen(path, "rb");
    if (!f) { perror("open wav"); exit(1); }
    fseek(f, 0, SEEK_END); long sz = ftell(f); fseek(f, 0, SEEK_SET);
    uint8_t* buf = malloc(sz);
    if (fread(buf, 1, sz, f) != (size_t)sz) { fprintf(stderr, "read fail\n"); exit(1); }
    fclose(f);

    // Walk RIFF chunks to find fmt + data.
    int channels = 1, bits = 16; uint32_t dataOff = 0, dataLen = 0;
    uint32_t p = 12; // skip RIFF....WAVE
    while (p + 8 <= (uint32_t)sz) {
        uint32_t cid; memcpy(&cid, buf + p, 4);
        uint32_t clen; memcpy(&clen, buf + p + 4, 4);
        if (memcmp(buf + p, "fmt ", 4) == 0) {
            memcpy(&channels, buf + p + 8 + 2, 2);
            memcpy(&bits, buf + p + 8 + 14, 2);
        } else if (memcmp(buf + p, "data", 4) == 0) {
            dataOff = p + 8; dataLen = clen; break;
        }
        p += 8 + clen + (clen & 1);
    }
    if (!dataOff || bits != 16) { fprintf(stderr, "need 16-bit PCM wav\n"); exit(1); }

    int bytesPerSample = 2 * channels;
    int n = dataLen / bytesPerSample;
    float* out = malloc(n * sizeof(float));
    const int16_t* pcm = (const int16_t*)(buf + dataOff);
    for (int i = 0; i < n; i++) out[i] = pcm[i * channels] / 32768.0f;
    free(buf);
    *outN = n;
    return out;
}

static double now_sec(void) {
    struct timespec ts; clock_gettime(CLOCK_MONOTONIC, &ts);
    return ts.tv_sec + ts.tv_nsec * 1e-9;
}

int main(int argc, char** argv) {
    const char* path = argc > 1 ? argv[1] : "../mlnl_samples.wav";
    int n; float* samples = load_wav(path, &n);

    int hop = FFT_SIZE / HOP_RATIO;
    int frames = (n - FFT_SIZE) / hop + 1;
    if (frames < 1) { fprintf(stderr, "too short\n"); return 1; }

    float complex* in  = malloc(FFT_SIZE * sizeof(float complex));
    float complex* out = malloc(FFT_SIZE * sizeof(float complex));

    // Hann-windowed frames, exactly like stft.c feeds FFT().
    double checksum = 0.0;
    double t0 = now_sec();
    for (int rep = 0; rep < FFT_BENCH_REPS; rep++) {
        for (int fr = 0; fr < frames; fr++) {
            int off = fr * hop;
            for (int i = 0; i < FFT_SIZE; i++) {
                float w = 0.5f - 0.5f * cosf(2.0f * (float)M_PI * i / (FFT_SIZE - 1));
                in[i] = samples[off + i] * w + 0.0f * I;
            }
            FFT(in, out, FFT_SIZE, false);
            checksum += crealf(out[1]) + cimagf(out[1]);
        }
    }
    double t1 = now_sec();

    long long totalFFTs = (long long)frames * FFT_BENCH_REPS;
    double secs = t1 - t0;
    printf("samples=%d  frames=%d  reps=%d  total_FFTs=%lld\n",
           n, frames, FFT_BENCH_REPS, totalFFTs);
    printf("elapsed=%.4f s   per-FFT=%.3f us   FFTs/s=%.0f\n",
           secs, secs / totalFFTs * 1e6, totalFFTs / secs);
    printf("checksum=%.6f (anti-DCE)\n", checksum);
    return 0;
}