// 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);
}

// ===== Background high-res computation =====
int ComputeNextHighResChunk(AudioSignal* signal, StftResult* result,
                                     int fftSize, int startSeg, int endSeg)
{
    int hopSize = fftSize / HOP_RATIO;
    int numBins = fftSize / 2 + 1;
    float* windowedSamples = (float*)malloc(fftSize * sizeof(float));
    float* derivWindowedSamples = (float*)malloc(fftSize * sizeof(float));
    float complex *complexInput = (float complex*)malloc(fftSize * sizeof(float complex));
    float complex* fftOutput = (float complex*)malloc(fftSize * sizeof(float complex));

    for (int seg = startSeg; seg < endSeg && seg < result->numSegments; seg++) {
        // Skip if already computed (overview or high-res)
        if (result->segments[seg].spectrum != NULL) continue;

        int offset = seg * hopSize;
        int samplesToCopy = fftSize;
        if (offset + samplesToCopy > signal->numSamples) {
            samplesToCopy = signal->numSamples - offset;
            memset(windowedSamples, 0, fftSize * sizeof(float));
            memset(derivWindowedSamples, 0, fftSize * sizeof(float));
        } else {
            memcpy(windowedSamples, signal->samples + offset, fftSize * sizeof(float));
            memcpy(derivWindowedSamples, signal->samples + offset, fftSize * sizeof(float));
        }

        // Apply Hann window and derivative window
        for (int i = 0; i < fftSize; i++) {
            float t = (float)i / (fftSize - 1);
            float hann = 0.5f * (1.0f - cosf(2.0f * M_PI * t));
            float derivHann = M_PI * sinf(2.0f * M_PI * t);
            windowedSamples[i] *= hann;
            derivWindowedSamples[i] *= derivHann;
        }

        // Normal STFT
        for (int i = 0; i < fftSize; i++) complexInput[i] = windowedSamples[i] + 0.0f * I;
        FFT(complexInput, fftOutput, fftSize, false);

        result->segments[seg].numBins = numBins;
        result->segments[seg].sampleOffset = offset;
        result->segments[seg].sampleCount = samplesToCopy;
        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(fftOutput[bin]) / fftSize : 2.0f * cabsf(fftOutput[bin]) / fftSize;
            result->segments[seg].spectrum[bin].phase = cargf(fftOutput[bin]);
        }

        // Derivative-window STFT for synchrosqueezing
        result->segments[seg].derivativeSpectrum = (FrequencyData*)malloc(numBins * sizeof(FrequencyData));
        for (int i = 0; i < fftSize; i++) complexInput[i] = derivWindowedSamples[i] + 0.0f * I;
        FFT(complexInput, fftOutput, fftSize, false);

        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(fftOutput[bin]) / fftSize;
            result->segments[seg].derivativeSpectrum[bin].phase = cargf(fftOutput[bin]);
        }
    }

    free(windowedSamples);
    free(derivWindowedSamples);
    free(complexInput);
    free(fftOutput);

    // Return next segment to process
    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)
{
    int hopSize = fftSize / HOP_RATIO;
    int numBins = fftSize / 2 + 1;
    float* windowedSamples = (float*)malloc(fftSize * sizeof(float));
    float* derivWindowedSamples = (float*)malloc(fftSize * sizeof(float));
    float complex *complexInput = (float complex*)malloc(fftSize * sizeof(float complex));
    float complex* fftOutput = (float complex*)malloc(fftSize * sizeof(float complex));

    for (int seg = startSegment; seg < result->numSegments; seg++) {
        // Skip segments not aligned with the skip factor (overview mode)
        if (seg % app.skipFactor != 0) continue;

        // Skip if already computed as high-res
        if (result->segments[seg].spectrum != NULL) continue;
        int offset = seg * hopSize;
        int samplesToCopy = fftSize;
        if (offset + samplesToCopy > signal->numSamples) {
            samplesToCopy = signal->numSamples - offset;
            memset(windowedSamples, 0, fftSize * sizeof(float));
            memset(derivWindowedSamples, 0, fftSize * sizeof(float));
        } else {
            memcpy(windowedSamples, signal->samples + offset, fftSize * sizeof(float));
            memcpy(derivWindowedSamples, signal->samples + offset, fftSize * sizeof(float));
        }

        // Apply Hann window: h(t) = 0.5 * (1 - cos(2πt))
        // And derivative window: h'(t) = π * sin(2πt)
        for (int i = 0; i < fftSize; i++) {
            float t = (float)i / (fftSize - 1);
            float hann = 0.5f * (1.0f - cosf(2.0f * M_PI * t));
            float derivHann = M_PI * sinf(2.0f * M_PI * t);
            windowedSamples[i] *= hann;
            derivWindowedSamples[i] *= derivHann;
        }

        // Compute normal STFT (V_f)
        for (int i = 0; i < fftSize; i++) complexInput[i] = windowedSamples[i] + 0.0f * I;
        FFT(complexInput, fftOutput, fftSize, false);

        result->segments[seg].numBins = numBins;
        result->segments[seg].sampleOffset = offset;
        result->segments[seg].sampleCount = samplesToCopy;
        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(fftOutput[bin]) / fftSize : 2.0f * cabsf(fftOutput[bin]) / fftSize;
            result->segments[seg].spectrum[bin].phase = cargf(fftOutput[bin]);
        }

        // Compute derivative-window STFT (V_fd) for synchrosqueezing
        result->segments[seg].derivativeSpectrum = (FrequencyData*)malloc(numBins * sizeof(FrequencyData));
        for (int i = 0; i < fftSize; i++) complexInput[i] = derivWindowedSamples[i] + 0.0f * I;
        FFT(complexInput, fftOutput, fftSize, false);

        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(fftOutput[bin]) / fftSize;
            result->segments[seg].derivativeSpectrum[bin].phase = cargf(fftOutput[bin]);
        }
    }

    free(windowedSamples);
    free(derivWindowedSamples);
    free(complexInput);
    free(fftOutput);
    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 =====
void AutoScaleAmplitude(StftResult* stft)
{
    float maxDb = -999.0f;
    float minDb = 0.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 (db < minDb) minDb = db;
        }
    }
    // Set ceiling at the max, floor 40dB below — enough range to see structure
    // but not so wide that the signal is drowned in black
    app.amplitudeCeilingDb = maxDb;
    app.amplitudeFloorDb = maxDb - 40.0f;
}