// stft.c - STFT computation, adaptive resolution, and the FFT-size LRU cache
#include "stft.h"
#include "fft.h"
#include "render.h"   // GenerateSpectrogramTexture (used by ChangeFFTSize)

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

// ===== FFT-size cache (LRU) =====
static bool IsSTFTComplete(const StftResult* r)
{
    if (r->numSegments <= 0 || r->segments == NULL) return false;
    for (int i = 0; i < r->numSegments; i++) {
        if (r->segments[i].spectrum == NULL) return false;
    }
    return true;
}

// Deep-copy src into dst. dst is assumed to be empty (freed) beforehand.
// Handles sparse results safely: a segment with no computed spectrum is copied
// as NULL rather than dereferencing a NULL source pointer (the bug that caused
// the load crash for files long enough to use a skipFactor > 1 overview).
static void CopySTFT(StftResult* dst, const StftResult* src)
{
    dst->numSegments = src->numSegments;
    dst->sampleRate = src->sampleRate;
    dst->totalSamples = src->totalSamples;
    dst->useHannWindow = src->useHannWindow;
    dst->segments = (StftSegment*)malloc(src->numSegments * sizeof(StftSegment));
    for (int i = 0; i < src->numSegments; i++) {
        const StftSegment* s = &src->segments[i];
        StftSegment* d = &dst->segments[i];
        d->numBins = s->numBins;
        d->sampleOffset = s->sampleOffset;
        d->sampleCount = s->sampleCount;
        if (s->spectrum != NULL && s->numBins > 0) {
            d->spectrum = (FrequencyData*)malloc(s->numBins * sizeof(FrequencyData));
            memcpy(d->spectrum, s->spectrum, s->numBins * sizeof(FrequencyData));
        } else {
            d->spectrum = NULL;
        }
        if (s->derivativeSpectrum != NULL && s->numBins > 0) {
            d->derivativeSpectrum = (FrequencyData*)malloc(s->numBins * sizeof(FrequencyData));
            memcpy(d->derivativeSpectrum, s->derivativeSpectrum, s->numBins * sizeof(FrequencyData));
        } else {
            d->derivativeSpectrum = NULL;
        }
    }
}

/**
 * Free all STFT results in the cache.
 */
void FreeAllCacheEntries(FFTSizeCache* cache)
{
    for (int i = 0; i < cache->count; i++) {
        FreeSTFT(&cache->entries[i].result);
        cache->entries[i].result.sampleRate = 0;
        cache->entries[i].accessOrder = 0;
    }
    cache->count = 0;
    cache->nextOrder = 0;
}

/**
 * Look up a cache entry by FFT size. Returns NULL if not present.
 * On a hit, marks the entry as most recently used.
 */
static FFTCacheEntry* FindCacheEntry(FFTSizeCache* cache, int fftSize)
{
    for (int i = 0; i < cache->count; i++) {
        if (cache->entries[i].fftSize == fftSize) {
            cache->entries[i].accessOrder = cache->nextOrder++;
            return &cache->entries[i];
        }
    }
    return NULL;
}

/**
 * Find a cache entry for the given FFT size, or create one.
 * If the cache is full, evicts the least-recently-used entry.
 * Returns a pointer to the entry (valid until next cache access).
 */
static FFTCacheEntry* FindOrCreateCacheEntry(FFTSizeCache* cache, int fftSize, int sampleRate)
{
    FFTCacheEntry* existing = FindCacheEntry(cache, fftSize);
    if (existing) return existing;

    // Entry not found — need to create it
    if (cache->count >= FFT_CACHE_SIZE) {
        // Evict least recently used (lowest accessOrder)
        int lruIdx = 0;
        for (int i = 1; i < cache->count; i++) {
            if (cache->entries[i].accessOrder < cache->entries[lruIdx].accessOrder) {
                lruIdx = i;
            }
        }
        FreeSTFT(&cache->entries[lruIdx].result);
        // Reuse slot
        cache->entries[lruIdx].fftSize = fftSize;
        cache->entries[lruIdx].result.numSegments = 0;
        cache->entries[lruIdx].result.segments = NULL;
        cache->entries[lruIdx].accessOrder = cache->nextOrder++;
        return &cache->entries[lruIdx];
    }

    // Add new entry
    int idx = cache->count++;
    cache->entries[idx].fftSize = fftSize;
    cache->entries[idx].result.numSegments = 0;
    cache->entries[idx].result.segments = NULL;
    cache->entries[idx].result.sampleRate = sampleRate;
    cache->entries[idx].accessOrder = cache->nextOrder++;
    return &cache->entries[idx];
}

/**
 * Save the current app.stft result to the cache entry matching app.fftSize.
 * Creates/overwrites the entry and marks it as most recently used.
 */
void SaveToCache(void)
{
    // Only cache fully-computed (full-resolution) results. A sparse overview
    // contains NULL segments and isn't worth caching — and restoring one would
    // leave permanent black gaps since we'd mark it finished.
    if (!IsSTFTComplete(&app.stft)) return;

    FFTCacheEntry* entry = FindOrCreateCacheEntry(&app.fftCache, app.fftSize, app.signal.sampleRate);
    FreeSTFT(&entry->result);
    CopySTFT(&entry->result, &app.stft);
    TraceLog(LOG_INFO, "Saved STFT result to cache for FFT size %d (%d segments)",
             app.fftSize, app.stft.numSegments);
}

// ===== Per-segment STFT (shared by the overview and high-res passes) =====

// Scratch buffers reused across every segment in one pass, so we don't malloc
// per segment. Allocate once, hand to ComputeSegment, free when the pass ends.
typedef struct {
    float* windowed;
    float* derivWindowed;
    float complex* fftIn;
    float complex* fftOut;
} SegScratch;

static SegScratch AllocSegScratch(int fftSize)
{
    SegScratch sc;
    sc.windowed      = (float*)malloc(fftSize * sizeof(float));
    sc.derivWindowed = (float*)malloc(fftSize * sizeof(float));
    sc.fftIn         = (float complex*)malloc(fftSize * sizeof(float complex));
    sc.fftOut        = (float complex*)malloc(fftSize * sizeof(float complex));
    return sc;
}

static void FreeSegScratch(SegScratch* sc)
{
    free(sc->windowed);
    free(sc->derivWindowed);
    free(sc->fftIn);
    free(sc->fftOut);
}

// Compute one STFT segment (normal V_f + derivative-window V_fd spectra) into
// result->segments[seg]. Caller ensures the segment isn't already computed.
static void ComputeSegment(AudioSignal* signal, StftResult* result, int fftSize, int seg, SegScratch* sc)
{
    int hopSize = fftSize / HOP_RATIO;
    int numBins = fftSize / 2 + 1;
    int offset = seg * hopSize;
    int samplesToCopy = fftSize;
    if (offset + samplesToCopy > signal->numSamples) {
        samplesToCopy = signal->numSamples - offset;
        memset(sc->windowed, 0, fftSize * sizeof(float));
        memset(sc->derivWindowed, 0, fftSize * sizeof(float));
    } else {
        memcpy(sc->windowed, signal->samples + offset, fftSize * sizeof(float));
        memcpy(sc->derivWindowed, signal->samples + offset, fftSize * sizeof(float));
    }

    // Hann window h(t) = 0.5*(1 - cos(2πt)); derivative window h'(t) = π*sin(2πt)
    for (int i = 0; i < fftSize; i++) {
        float t = (float)i / (fftSize - 1);
        sc->windowed[i]      *= 0.5f * (1.0f - cosf(2.0f * M_PI * t));
        sc->derivWindowed[i] *= M_PI * sinf(2.0f * M_PI * t);
    }

    result->segments[seg].numBins = numBins;
    result->segments[seg].sampleOffset = offset;
    result->segments[seg].sampleCount = samplesToCopy;

    // Normal STFT (V_f)
    for (int i = 0; i < fftSize; i++) sc->fftIn[i] = sc->windowed[i] + 0.0f * I;
    FFT(sc->fftIn, sc->fftOut, fftSize, false);
    result->segments[seg].spectrum = (FrequencyData*)malloc(numBins * sizeof(FrequencyData));
    for (int bin = 0; bin < numBins; bin++) {
        result->segments[seg].spectrum[bin].frequency = (float)bin * signal->sampleRate / fftSize;
        result->segments[seg].spectrum[bin].amplitude = (bin == 0) ? cabsf(sc->fftOut[bin]) / fftSize : 2.0f * cabsf(sc->fftOut[bin]) / fftSize;
        result->segments[seg].spectrum[bin].phase = cargf(sc->fftOut[bin]);
    }

    // Derivative-window STFT (V_fd) for synchrosqueezing
    for (int i = 0; i < fftSize; i++) sc->fftIn[i] = sc->derivWindowed[i] + 0.0f * I;
    FFT(sc->fftIn, sc->fftOut, fftSize, false);
    result->segments[seg].derivativeSpectrum = (FrequencyData*)malloc(numBins * sizeof(FrequencyData));
    for (int bin = 0; bin < numBins; bin++) {
        result->segments[seg].derivativeSpectrum[bin].frequency = (float)bin * signal->sampleRate / fftSize;
        result->segments[seg].derivativeSpectrum[bin].amplitude = cabsf(sc->fftOut[bin]) / fftSize;
        result->segments[seg].derivativeSpectrum[bin].phase = cargf(sc->fftOut[bin]);
    }
}

// ===== Background high-res computation =====
// Fill segments [startSeg, endSeg) at full resolution, skipping any already
// computed. Returns the next segment index to resume from.
int ComputeNextHighResChunk(AudioSignal* signal, StftResult* result,
                                     int fftSize, int startSeg, int endSeg)
{
    SegScratch sc = AllocSegScratch(fftSize);
    for (int seg = startSeg; seg < endSeg && seg < result->numSegments; seg++) {
        if (result->segments[seg].spectrum != NULL) continue;  // already computed
        ComputeSegment(signal, result, fftSize, seg, &sc);
    }
    FreeSegScratch(&sc);

    if (endSeg >= result->numSegments) return result->numSegments;
    return endSeg;
}

// ===== STFT computation =====
void ComputeSTFTInit(AudioSignal* signal, StftResult* result, int fftSize)
{
    FreeSTFT(result);  // release any previous result before reallocating
    int hopSize = fftSize / HOP_RATIO;  // 75% overlap
    int numSegments = (signal->numSamples - fftSize) / hopSize + 1;
    if (numSegments <= 0) numSegments = 1;

    result->numSegments = numSegments;
    result->segments = (StftSegment*)calloc(numSegments, sizeof(StftSegment));
    result->sampleRate = signal->sampleRate;
    result->totalSamples = signal->numSamples;
    result->useHannWindow = true;
}

bool ComputeSTFTIncremental(AudioSignal* signal, StftResult* result, int fftSize, int startSegment)
{
    SegScratch sc = AllocSegScratch(fftSize);
    for (int seg = startSegment; seg < result->numSegments; seg++) {
        if (seg % app.skipFactor != 0) continue;               // overview stride
        if (result->segments[seg].spectrum != NULL) continue;  // already computed
        ComputeSegment(signal, result, fftSize, seg, &sc);
    }
    FreeSegScratch(&sc);
    return true;
}

void FreeSTFT(StftResult* result)
{
    if (!result) return;
    if (result->segments) {
        for (int i = 0; i < result->numSegments; i++) {
            free(result->segments[i].spectrum);
            result->segments[i].spectrum = NULL;
            if (result->segments[i].derivativeSpectrum) {
                free(result->segments[i].derivativeSpectrum);
                result->segments[i].derivativeSpectrum = NULL;
            }
        }
        free(result->segments);
        result->segments = NULL;
    }
    result->numSegments = 0;
}

/**
 * Change the FFT size. If a fully-computed result for the new size is cached,
 * restore it directly (no recomputation). Otherwise free the current STFT and
 * let the main loop recompute it from scratch.
 */
void ChangeFFTSize(int newFFT)
{
    FFTCacheEntry* entry = FindCacheEntry(&app.fftCache, newFFT);

    if (entry != NULL && IsSTFTComplete(&entry->result)) {
        // Cache hit — restore the cached full-resolution result.
        TraceLog(LOG_INFO, "FFT size %d: cache hit", newFFT);
        FreeSTFT(&app.stft);
        CopySTFT(&app.stft, &entry->result);

        app.fftSize = newFFT;
        app.skipFactor = 1;
        app.stftComputed = true;     // already complete — skip recompute
        app.loadingPhase = 0;
        app.highResFinished = true;
        app.bgHighResSeg = app.stft.numSegments;
        app.bgFinished = true;
        app.isBgProcessing = false;
        app.visibleTextureValid = false;

        // Rebuild the displayed texture from the restored data. AutoScale here
        // mirrors the recompute path so the view looks identical either way.
        AutoScaleAmplitude(&app.stft);
        GenerateSpectrogramTexture(&app.stft, &app.spectrogramImage, &app.spectrogramTexture);
    } else {
        // Cache miss — drop the current STFT and recompute. Freeing here avoids
        // leaking it, since ComputeSTFTInit re-allocates segments unconditionally.
        TraceLog(LOG_INFO, "FFT size %d: cache miss, computing", newFFT);
        FreeSTFT(&app.stft);
        app.fftSize = newFFT;
        app.stftComputed = false;
        app.loadingPhase = 0;
        app.skipFactor = 1;
        app.highResFinished = false;
        app.bgHighResSeg = 0;
        app.bgFinished = false;
        app.isBgProcessing = false;
        app.visibleTextureValid = false;
    }
}

// ===== Adaptive resolution =====
int ComputeSkipFactor(float signalDurationSec)
{
    if (signalDurationSec <= 60.0f) return 1;      // < 1 min: full-res
    if (signalDurationSec <= 300.0f) return 2;     // 1-5 min: every 2nd
    if (signalDurationSec <= 600.0f) return 4;     // 5-10 min: every 4th
    return 8;                                       // > 10 min: every 8th
}

// Compute full-resolution segments for the range [startSeg, endSeg).
// This replaces existing overview (skipFactor-strided) segments with

// ===== Amplitude auto-scaling =====
// Derive the colorizer's floor/ceiling from the current scale mode and controls.
// Called after the STFT changes (load, background completion, FFT-size change)
// and when the user toggles the mode. Deriving the floor from dynRangeDb /
// absoluteFloorDb here is what preserves the user's setting across re-scales.
void AutoScaleAmplitude(StftResult* stft)
{
    if (app.amplitudeMode == SCALE_ABSOLUTE) {
        // Fixed dBFS scale: 0 dBFS (full-scale) ceiling, user-set absolute floor.
        app.amplitudeCeilingDb = 0.0f;
        app.amplitudeFloorDb = app.absoluteFloorDb;
        return;
    }

    // Relative: ceiling tracks the signal peak; floor sits dynRangeDb below it.
    float maxDb = -999.0f;
    for (int seg = 0; seg < stft->numSegments; seg++) {
        for (int bin = 0; bin < stft->segments[seg].numBins; bin++) {
            float db = AmplitudeToDecibels(stft->segments[seg].spectrum[bin].amplitude);
            if (db > maxDb) maxDb = db;
        }
    }
    if (maxDb < -998.0f) maxDb = 0.0f;  // no data yet — sane default
    app.amplitudeCeilingDb = maxDb;
    app.amplitudeFloorDb = maxDb - app.dynRangeDb;
}

static int CompareDouble(const void* a, const void* b)
{
    double da = *(const double*)a, db = *(const double*)b;
    return (da > db) - (da < db);
}

SpectralStats ComputeSpectralStats(const StftResult* stft,
                                   float t0, float t1, float f0, float f1)
{
    SpectralStats st = { 0 };
    if (!stft || stft->numSegments <= 0 || stft->sampleRate <= 0) return st;

    // Normalize + clamp the box to [0,1].
    if (t1 < t0) { float tmp = t0; t0 = t1; t1 = tmp; }
    if (f1 < f0) { float tmp = f0; f0 = f1; f1 = tmp; }
    t0 = fmaxf(0.0f, t0); t1 = fminf(1.0f, t1);
    f0 = fmaxf(0.0f, f0); f1 = fminf(1.0f, f1);

    const float nyquist = stft->sampleRate * 0.5f;
    const float freqLow = f0 * nyquist;
    const float freqHigh = f1 * nyquist;

    int segStart = (int)(t0 * stft->numSegments);
    int segEnd   = (int)(t1 * stft->numSegments);
    if (segStart < 0) segStart = 0;
    if (segEnd > stft->numSegments) segEnd = stft->numSegments;
    if (segEnd <= segStart) segEnd = (segStart < stft->numSegments) ? segStart + 1 : segStart;

    // Learn the bin count from the first computed segment in range.
    int nbins = 0;
    for (int s = segStart; s < segEnd; s++) {
        if (stft->segments[s].spectrum && stft->segments[s].numBins > 0) {
            nbins = stft->segments[s].numBins; break;
        }
    }
    if (nbins < 2) return st;

    // Mean power per bin over the selected time span (skip uncomputed segments).
    double* power = (double*)calloc(nbins, sizeof(double));
    if (!power) return st;
    int counted = 0;
    for (int s = segStart; s < segEnd; s++) {
        const StftSegment* seg = &stft->segments[s];
        if (!seg->spectrum || seg->numBins < nbins) continue;
        for (int b = 0; b < nbins; b++) {
            float a = seg->spectrum[b].amplitude;
            power[b] += (double)a * a;
        }
        counted++;
    }
    if (counted == 0) { free(power); return st; }
    for (int b = 0; b < nbins; b++) power[b] /= counted;

    const float binHz = nyquist / (float)(nbins - 1);  // = sampleRate / fftSize
    int binLow  = (int)ceilf(freqLow / binHz);
    int binHigh = (int)floorf(freqHigh / binHz);
    if (binLow < 0) binLow = 0;
    if (binHigh > nbins - 1) binHigh = nbins - 1;
    if (binHigh < binLow) { free(power); return st; }  // band narrower than a bin

    // Peak, centroid, total in-band power.
    double sumP = 0.0, sumFP = 0.0, peakP = -1.0;
    int peakBin = binLow;
    for (int b = binLow; b <= binHigh; b++) {
        double p = power[b];
        double f = (double)b * binHz;
        sumP += p; sumFP += f * p;
        if (p > peakP) { peakP = p; peakBin = b; }
    }
    int K = binHigh - binLow + 1;
    (void)peakP;
    st.valid = true;
    st.peakFreqHz = (float)(peakBin * binHz);
    st.centroidHz = (sumP > 0.0) ? (float)(sumFP / sumP) : st.peakFreqHz;
    st.inBandLevelDb = (sumP > 0.0) ? 10.0f * log10f((float)(sumP / K) + 1e-20f) : -200.0f;

    // Robust noise floor: median power of the out-of-band bins (skip DC). Used
    // for both the occupied-bandwidth threshold and the SNR estimate.
    double noiseDensity = 0.0;
    double* out = (double*)malloc(nbins * sizeof(double));
    if (out) {
        int outCount = 0;
        for (int b = 1; b < nbins; b++) {
            if (b < binLow || b > binHigh) out[outCount++] = power[b];
        }
        if (outCount > 0) {
            qsort(out, outCount, sizeof(double), CompareDouble);
            noiseDensity = out[outCount / 2];
        }
        free(out);
    }

    // Occupied bandwidth: span of the in-band region sitting >3 dB over noise.
    // Robust for both pure tones (narrow) and noise-like bursts (wide), unlike
    // a -3 dB-around-peak walk which collapses to one bin on rough spectra.
    double thresh = noiseDensity * 2.0;  // +3 dB
    int lo = -1, hi = -1;
    for (int b = binLow; b <= binHigh; b++) {
        if (power[b] >= thresh) { if (lo < 0) lo = b; hi = b; }
    }
    st.bandwidthHz = (lo >= 0) ? (float)((hi - lo + 1) * binHz) : 0.0f;

    // SNR: in-band power above the noise floor scaled to the in-band bin count.
    double noiseInBand = noiseDensity * K;
    double sig = sumP - noiseInBand;
    if (sig < 1e-20) sig = 1e-20;
    if (noiseInBand < 1e-20) noiseInBand = 1e-20;
    st.snrDb = 10.0f * log10f((float)(sig / noiseInBand));

    free(power);
    return st;
}