// STFT feature extraction. One pass over the whole track, producing one row per
// output video frame.
//
// Two details that matter more than they look:
//
//  * Windows are CENTRED on the frame's timestamp, not started at it. A window
//    that starts at the timestamp reports energy that arrives up to 23ms later,
//    which reads on screen as the visuals lagging the music. Centring removes
//    that systematic offset.
//
//  * Energy features are normalised against the track's own 5th/95th percentiles
//    at the end of the pass. A quiet ambient master and a brickwalled EDM master
//    then both use the full reactive range, without anyone touching a gain knob.
//    Absolute (un-normalised) statistics are kept in `summary` for the look
//    generator, which does need to know that one track is genuinely darker.

import { FFT, hannWindow } from './fft.js';

export const FFT_SIZE = 2048;

/** Band edges in Hz. Sub is deliberately narrow — it is the kick, not the bass. */
export const BANDS = {
    bandSub: [20, 60],
    bandLow: [60, 250],
    bandMid: [250, 2000],
    bandHigh: [2000, 6000],
    bandAir: [6000, 16000],
};

export const ENERGY_FEATURES = ['rms', 'loudness', 'bandSub', 'bandLow', 'bandMid', 'bandHigh', 'bandAir', 'flux'];
export const SCALE_FREE_FEATURES = ['centroid', 'flatness', 'width'];
export const RAW_FEATURES = [...ENERGY_FEATURES, ...SCALE_FREE_FEATURES];

function percentile(sorted, p) {
    if (!sorted.length) return 0;
    const i = Math.min(sorted.length - 1, Math.max(0, Math.round((sorted.length - 1) * p)));
    return sorted[i];
}

/**
 * @param {AudioBuffer|{sampleRate,length,duration,numberOfChannels,getChannelData}} audioBuffer
 * @returns {{frameCount, fps, sampleRate, duration, raw, summary, onsetEnvelope}}
 */
export function analyzeBuffer(audioBuffer, { fps = 60, fftSize = FFT_SIZE, onProgress = null } = {}) {
    const sampleRate = audioBuffer.sampleRate;
    const duration = audioBuffer.duration;
    const hop = Math.max(1, Math.round(sampleRate / fps));
    const frameCount = Math.max(1, Math.ceil(duration * fps));

    const channels = audioBuffer.numberOfChannels;
    const left = audioBuffer.getChannelData(0);
    const right = channels > 1 ? audioBuffer.getChannelData(1) : left;
    const totalSamples = left.length;

    const fft = new FFT(fftSize);
    const window = hannWindow(fftSize);
    const frameBuffer = new Float32Array(fftSize);
    const halfFft = fftSize >> 1;

    // Precompute band bin ranges.
    const binHz = sampleRate / fftSize;
    const bandRanges = {};
    for (const [name, [lo, hi]] of Object.entries(BANDS)) {
        bandRanges[name] = [
            Math.max(1, Math.floor(lo / binHz)),
            Math.min(halfFft - 1, Math.ceil(hi / binHz)),
        ];
    }
    const binFrequencies = new Float32Array(halfFft);
    for (let i = 0; i < halfFft; i++) binFrequencies[i] = i * binHz;

    const raw = {};
    for (const name of RAW_FEATURES) raw[name] = new Float32Array(frameCount);

    let prevMagnitude = new Float32Array(halfFft);
    const progressEvery = Math.max(1, Math.floor(frameCount / 50));

    for (let f = 0; f < frameCount; f++) {
        // Centred window.
        const centre = f * hop;
        const start = centre - (fftSize >> 1);

        let sumSq = 0;
        let midSq = 0;
        let sideSq = 0;
        for (let i = 0; i < fftSize; i++) {
            const s = start + i;
            let l = 0, r = 0;
            if (s >= 0 && s < totalSamples) { l = left[s]; r = right[s]; }
            const mono = (l + r) * 0.5;
            const side = (l - r) * 0.5;
            frameBuffer[i] = mono * window[i];
            sumSq += mono * mono;
            midSq += mono * mono;
            sideSq += side * side;
        }

        raw.rms[f] = Math.sqrt(sumSq / fftSize);
        raw.width[f] = Math.sqrt(sideSq / fftSize) / (Math.sqrt(midSq / fftSize) + Math.sqrt(sideSq / fftSize) + 1e-9);

        const mag = fft.forward(frameBuffer);

        let total = 0;
        let weighted = 0;
        let logSum = 0;
        let flux = 0;
        for (let i = 1; i < halfFft; i++) {
            const m = mag[i];
            total += m;
            weighted += m * binFrequencies[i];
            logSum += Math.log(m + 1e-10);
            const d = m - prevMagnitude[i];
            if (d > 0) flux += d;
        }

        raw.loudness[f] = total / halfFft;
        raw.flux[f] = flux / halfFft;

        // Spectral centroid, mapped to a log-frequency 0..1 so it matches how
        // brightness is actually perceived.
        const centroidHz = total > 1e-9 ? weighted / total : 0;
        raw.centroid[f] = centroidHz > 20
            ? Math.min(1, Math.max(0, Math.log2(centroidHz / 20) / Math.log2(20000 / 20)))
            : 0;

        const arithmeticMean = total / (halfFft - 1);
        const geometricMean = Math.exp(logSum / (halfFft - 1));
        raw.flatness[f] = arithmeticMean > 1e-9 ? Math.min(1, geometricMean / arithmeticMean) : 0;

        for (const [name, [lo, hi]] of Object.entries(bandRanges)) {
            let sum = 0;
            for (let i = lo; i <= hi; i++) sum += mag[i];
            raw[name][f] = sum / (hi - lo + 1);
        }

        prevMagnitude.set(mag);

        if (onProgress && f % progressEvery === 0) onProgress(f / frameCount);
    }

    // ---------------------------------------------------------------- summary
    // Computed from RAW values, before normalisation flattens them out.
    const sortedLoudness = Float32Array.from(raw.loudness).sort();
    const loudFloor = percentile(sortedLoudness, 0.1);
    const activeFrames = [];
    for (let f = 0; f < frameCount; f++) if (raw.loudness[f] > loudFloor) activeFrames.push(f);
    const activeCount = activeFrames.length || 1;

    const meanOf = (name) => {
        let s = 0;
        for (const f of activeFrames) s += raw[name][f];
        return s / activeCount;
    };

    const sortedRms = Float32Array.from(raw.rms).sort();
    const summary = {
        duration,
        sampleRate,
        fps,
        frameCount,
        meanCentroid: meanOf('centroid'),
        meanFlatness: meanOf('flatness'),
        meanWidth: meanOf('width'),
        meanLoudness: meanOf('loudness'),
        // Crest-ish: how much room there is between typical and peak level.
        // High on dynamic ambient, low on limitered club masters.
        dynamicRange: percentile(sortedRms, 0.95) > 1e-9
            ? 1 - percentile(sortedRms, 0.4) / percentile(sortedRms, 0.95)
            : 0,
        bandBalance: {
            sub: meanOf('bandSub'), low: meanOf('bandLow'), mid: meanOf('bandMid'),
            high: meanOf('bandHigh'), air: meanOf('bandAir'),
        },
    };

    // ------------------------------------------------------------ normalise
    const normalization = {};
    for (const name of ENERGY_FEATURES) {
        const sorted = Float32Array.from(raw[name]).sort();
        const lo = percentile(sorted, 0.05);
        const hi = percentile(sorted, 0.95);
        normalization[name] = { lo, hi };
        const span = hi - lo;
        const arr = raw[name];
        if (span > 1e-12) {
            for (let f = 0; f < frameCount; f++) {
                arr[f] = Math.min(1, Math.max(0, (arr[f] - lo) / span));
            }
        } else {
            arr.fill(0);
        }
    }
    for (const name of SCALE_FREE_FEATURES) {
        const arr = raw[name];
        for (let f = 0; f < frameCount; f++) arr[f] = Math.min(1, Math.max(0, arr[f]));
    }

    // The onset envelope drives tempo detection. Kept separate from the
    // normalised flux because tempo wants raw contrast, not a clipped range.
    const onsetEnvelope = new Float32Array(frameCount);
    for (let f = 0; f < frameCount; f++) {
        // Weight the low band up: in this material the kick is the clock.
        onsetEnvelope[f] = raw.flux[f] * 0.6 + raw.bandSub[f] * 0.25 + raw.bandLow[f] * 0.15;
    }

    if (onProgress) onProgress(1);

    return { frameCount, fps, sampleRate, duration, raw, summary, normalization, onsetEnvelope };
}
