// Tempo, beat grid and downbeats from the onset envelope.
//
// This runs over the WHOLE track, which is the entire reason to analyse offline:
// a causal detector has to converge, and lags for the first several bars of every
// section. Here the grid is exact from frame zero, and phase is fitted globally.
//
// The metronome (audio/metronome.js) exists to validate this by ear. If the
// clicks don't sit on the beat, nothing downstream can be trusted — every timing
// artefact in the finished video traces back to this file.

const MIN_BPM = 60;
const MAX_BPM = 200;
const PREFERRED_BPM = 124;   // log-space centre of the prior; club-ish but broad
const PRIOR_WIDTH = 0.85;

/** Remove the slow-moving floor so autocorrelation sees onsets, not loudness. */
function whiten(envelope, fps) {
    const n = envelope.length;
    const out = new Float32Array(n);
    const halfWindow = Math.max(2, Math.round(fps * 0.35));
    let sum = 0;
    const queue = [];
    for (let i = 0; i < n; i++) {
        queue.push(envelope[i]);
        sum += envelope[i];
        if (queue.length > halfWindow * 2 + 1) sum -= queue.shift();
        const mean = sum / queue.length;
        out[i] = Math.max(0, envelope[i] - mean);
    }
    let peak = 0;
    for (let i = 0; i < n; i++) if (out[i] > peak) peak = out[i];
    if (peak > 1e-9) for (let i = 0; i < n; i++) out[i] /= peak;
    return out;
}

function autocorrelate(signal, minLag, maxLag) {
    const n = signal.length;
    const scores = new Float32Array(maxLag + 1);
    for (let lag = minLag; lag <= maxLag; lag++) {
        let sum = 0;
        const limit = n - lag;
        for (let i = 0; i < limit; i++) sum += signal[i] * signal[i + lag];
        scores[lag] = limit > 0 ? sum / limit : 0;
    }
    return scores;
}

function bpmPrior(bpm) {
    const x = Math.log2(bpm / PREFERRED_BPM) / PRIOR_WIDTH;
    return Math.exp(-0.5 * x * x);
}

/**
 * Discrete onset peaks: local maxima of the whitened envelope.
 *
 * Peak positions are refined to sub-frame precision by parabolic interpolation.
 * At 60fps a whole-frame quantisation is 16.7ms, and that error accumulates
 * through the period fit into visible drift by the end of a long track.
 */
function pickOnsetPeaks(signal) {
    const peaks = [];
    let total = 0;
    for (let i = 1; i < signal.length - 1; i++) {
        const v = signal[i];
        if (v < 0.06) continue;
        if (v < signal[i - 1] || v < signal[i + 1]) continue;

        const a = signal[i - 1], b = v, c = signal[i + 1];
        const denom = a - 2 * b + c;
        const shift = Math.abs(denom) > 1e-12 ? (0.5 * (a - c)) / denom : 0;
        peaks.push({ frame: i + Math.max(-0.5, Math.min(0.5, shift)), strength: v });
        total += v;
    }
    return { peaks, total };
}

/**
 * Least-squares fit of the grid to the onsets it already matches.
 *
 * The search above only ever tests integer offsets and a discrete set of periods,
 * which leaves up to half a frame of phase error and a small period error that
 * compounds — 64 beats into a track a 0.02-frame period error is already 20ms of
 * drift. Regressing matched onset positions against their beat indices recovers
 * both to sub-frame precision in one pass.
 */
function refineByRegression(peaks, period, offset, length) {
    const tolerance = Math.max(1.5, Math.min(3, period * 0.15));
    const gridCount = Math.max(1, Math.floor((length - offset) / period));

    let sw = 0, sk = 0, st = 0, skk = 0, skt = 0;
    let matched = 0;
    for (const peak of peaks) {
        const k = Math.round((peak.frame - offset) / period);
        if (k < 0 || k > gridCount) continue;
        if (Math.abs(peak.frame - (offset + k * period)) > tolerance) continue;
        const w = peak.strength;
        sw += w; sk += w * k; st += w * peak.frame;
        skk += w * k * k; skt += w * k * peak.frame;
        matched++;
    }
    if (matched < 4 || sw < 1e-9) return { period, offset };

    const denom = sw * skk - sk * sk;
    if (Math.abs(denom) < 1e-9) return { period, offset };

    const slope = (sw * skt - sk * st) / denom;
    const intercept = (st - slope * sk) / sw;

    // Reject a fit that has wandered — that means the matching was wrong, not
    // that the tempo is unusual.
    if (!isFinite(slope) || !isFinite(intercept)) return { period, offset };
    if (Math.abs(slope - period) > period * 0.05) return { period, offset };

    return { period: slope, offset: intercept };
}

/**
 * F-measure between a candidate grid and the detected onsets.
 *
 * This is what resolves tempo octaves, and a plain "mean energy at grid points"
 * cannot. A half-tempo grid hits every other kick at FULL strength, so its mean
 * energy per beat is identical to the true grid's — it only loses on RECALL,
 * because half the onsets go unexplained. A double-tempo grid has perfect recall
 * but half its beats land on silence, so it loses on PRECISION. Combining the two
 * is the only formulation that penalises both errors.
 */
function gridFScore(peaks, totalStrength, period, offset, length) {
    if (!peaks.length || period < 2) return 0;
    const tolerance = Math.max(1.5, Math.min(3, period * 0.15));
    const gridCount = Math.max(1, Math.floor((length - offset) / period));

    const matchedGrid = new Set();
    let matchedStrength = 0;

    for (const peak of peaks) {
        const k = Math.round((peak.frame - offset) / period);
        if (k < 0 || k > gridCount) continue;
        const gridFrame = offset + k * period;
        if (Math.abs(peak.frame - gridFrame) <= tolerance) {
            matchedStrength += peak.strength;
            matchedGrid.add(k);
        }
    }

    const recall = totalStrength > 1e-9 ? matchedStrength / totalStrength : 0;
    const precision = matchedGrid.size / gridCount;
    if (precision + recall < 1e-9) return 0;
    return (2 * precision * recall) / (precision + recall);
}

/**
 * Mean matched onset strength at even vs odd grid positions.
 *
 * Unmatched grid points count as zero — a grid point that lands on silence is
 * evidence against the grid, not a missing sample.
 */
function alternation(peaks, period, offset, length) {
    const tolerance = Math.max(1.5, Math.min(3, period * 0.15));
    const gridCount = Math.max(2, Math.floor((length - offset) / period));
    const strength = new Float64Array(gridCount);

    for (const peak of peaks) {
        const k = Math.round((peak.frame - offset) / period);
        if (k < 0 || k >= gridCount) continue;
        if (Math.abs(peak.frame - (offset + k * period)) <= tolerance) {
            strength[k] = Math.max(strength[k], peak.strength);
        }
    }

    let evenSum = 0, evenCount = 0, oddSum = 0, oddCount = 0;
    for (let k = 0; k < gridCount; k++) {
        if (k % 2 === 0) { evenSum += strength[k]; evenCount++; }
        else { oddSum += strength[k]; oddCount++; }
    }
    const even = evenCount ? evenSum / evenCount : 0;
    const odd = oddCount ? oddSum / oddCount : 0;
    const hi = Math.max(even, odd);
    return { even, odd, ratio: hi > 1e-9 ? Math.min(even, odd) / hi : 1 };
}

/**
 * Correct a double-tempo reading.
 *
 * Offbeat hi-hats make a double-tempo grid score perfectly on both precision and
 * recall — every grid point genuinely has an onset — so the F-measure alone
 * cannot tell 90 BPM with hats from 180 BPM. What separates them is that the
 * onsets ALTERNATE strong/weak: kick, hat, kick, hat. A systematic alternation
 * means the real beat is every other grid point. This is the same cue a listener
 * uses, and it is why the grid is fitted first and metrically interpreted second.
 */
function correctOctave(peaks, totalStrength, period, offset, length, fps) {
    let currentPeriod = period;
    let currentOffset = offset;

    for (let iteration = 0; iteration < 2; iteration++) {
        const bpm = (60 * fps) / currentPeriod;
        const halvedBpm = bpm / 2;
        if (halvedBpm < MIN_BPM) break;

        const alt = alternation(peaks, currentPeriod, currentOffset, length);
        if (alt.ratio >= 0.62) break;

        const strongIsEven = alt.even >= alt.odd;
        const nextPeriod = currentPeriod * 2;
        const nextOffset = strongIsEven ? currentOffset : currentOffset + currentPeriod;

        // Only accept if the slower grid still explains the material well: this
        // guards against halving genuinely syncopated but correctly-fitted music.
        const before = gridFScore(peaks, totalStrength, currentPeriod, currentOffset, length);
        const after = gridFScore(peaks, totalStrength, nextPeriod, nextOffset, length);
        if (after < before * 0.55) break;

        currentPeriod = nextPeriod;
        currentOffset = nextOffset;
    }

    return { period: currentPeriod, offset: currentOffset };
}

function bestOffset(peaks, totalStrength, period, length) {
    let best = 0;
    let bestScore = -1;
    const steps = Math.ceil(period);
    for (let o = 0; o < steps; o++) {
        const s = gridFScore(peaks, totalStrength, period, o, length);
        if (s > bestScore) { bestScore = s; best = o; }
    }
    return { offset: best, score: bestScore };
}

/**
 * @param {Float32Array} onsetEnvelope per-frame onset strength
 * @param {number} fps
 * @returns {{bpm, period, offset, beats, downbeats, confidence, beatsPerBar}}
 */
export function detectTempo(onsetEnvelope, fps, { beatsPerBar = 4 } = {}) {
    const signal = whiten(onsetEnvelope, fps);
    const minLag = Math.floor((60 * fps) / MAX_BPM);
    const maxLag = Math.ceil((60 * fps) / MIN_BPM);
    const acf = autocorrelate(signal, minLag, maxLag);

    // Weight autocorrelation by the tempo prior, then take the peak.
    let bestLag = minLag;
    let bestValue = -Infinity;
    for (let lag = minLag; lag <= maxLag; lag++) {
        const bpm = (60 * fps) / lag;
        const v = acf[lag] * bpmPrior(bpm);
        if (v > bestValue) { bestValue = v; bestLag = lag; }
    }

    // Parabolic refinement around the peak for sub-frame period accuracy.
    let period = bestLag;
    if (bestLag > minLag && bestLag < maxLag) {
        const a = acf[bestLag - 1], b = acf[bestLag], c = acf[bestLag + 1];
        const denom = a - 2 * b + c;
        if (Math.abs(denom) > 1e-12) period = bestLag - (0.5 * (c - a)) / denom;
    }

    const { peaks, total: totalStrength } = pickOnsetPeaks(signal);
    const length = signal.length;

    // Octave resolution. The prior only breaks near-ties; the F-measure does the
    // actual work, which is why a 174 BPM track no longer reads as 87.
    const candidates = [period / 4, period / 2, period, period * 2, period * 4].filter((p) => {
        const bpm = (60 * fps) / p;
        return bpm >= MIN_BPM && bpm <= MAX_BPM && p >= 2;
    });

    let chosen = { period, offset: 0, score: -1 };
    for (const p of candidates) {
        const { offset, score } = bestOffset(peaks, totalStrength, p, length);
        const adjusted = score * (0.6 + 0.4 * bpmPrior((60 * fps) / p));
        if (adjusted > chosen.score) chosen = { period: p, offset, score: adjusted };
    }

    period = chosen.period;
    let offset = chosen.offset;

    // Local refinement of period and offset together — catches a grid that is
    // right at the start and drifts by the end of a six-minute track.
    let bestRefined = gridFScore(peaks, totalStrength, period, offset, length);
    for (let dp = -0.5; dp <= 0.5001; dp += 0.02) {
        const p = period + dp;
        if (p < 2) continue;
        const { offset: o, score } = bestOffset(peaks, totalStrength, p, length);
        if (score > bestRefined) { bestRefined = score; period = p; offset = o; }
    }

    // Metrical interpretation, after the grid itself is fitted.
    ({ period, offset } = correctOctave(peaks, totalStrength, period, offset, length, fps));

    // Sub-frame fit last, so it refines the grid we actually committed to.
    ({ period, offset } = refineByRegression(peaks, period, offset, length));

    const bpm = (60 * fps) / period;

    const beats = [];
    for (let t = offset; t < onsetEnvelope.length; t += period) beats.push(t / fps);

    // Downbeat: of the `beatsPerBar` possible bar phases, the one whose beats
    // carry the most energy.
    //
    // Scored on the RAW envelope, not the whitened one. Whitening subtracts the
    // local mean, which is exactly the accent information that distinguishes beat
    // one from the other three — a four-to-the-floor pattern has a kick on every
    // beat and the only cue is that one of them is louder.
    let bestPhase = 0;
    let bestPhaseScore = -1;
    for (let phase = 0; phase < beatsPerBar; phase++) {
        let sum = 0;
        let count = 0;
        for (let b = phase; b < beats.length; b += beatsPerBar) {
            const i = Math.round(beats[b] * fps);
            if (i >= 1 && i < onsetEnvelope.length - 1) {
                sum += Math.max(onsetEnvelope[i - 1], onsetEnvelope[i], onsetEnvelope[i + 1]);
                count++;
            }
        }
        const score = count ? sum / count : 0;
        if (score > bestPhaseScore) { bestPhaseScore = score; bestPhase = phase; }
    }
    const downbeats = [];
    for (let b = bestPhase; b < beats.length; b += beatsPerBar) downbeats.push(beats[b]);

    // Confidence is the grid's own F-measure: how much of the onset energy the
    // chosen grid explains, and how many of its beats are actually occupied.
    const confidence = gridFScore(peaks, totalStrength, period, offset, length);

    return { bpm, period, offset, beats, downbeats, beatsPerBar, barPhase: bestPhase, confidence };
}

/**
 * Per-frame phase tracks from the grid.
 *
 * `beat` is a decaying spike, but gated by a smoothed loudness envelope so a
 * drumless breakdown doesn't strobe on a grid that is technically still running.
 */
export function buildBeatTracks(tempo, frameCount, fps, loudness) {
    const beatPhase = new Float32Array(frameCount);
    const barPhase = new Float32Array(frameCount);
    const phrasePhase = new Float32Array(frameCount);
    const beat = new Float32Array(frameCount);

    const { period, offset, beatsPerBar } = tempo;
    const barPeriod = period * beatsPerBar;
    const phrasePeriod = barPeriod * 8;
    const barOffset = offset + tempo.barPhase * period;

    // Smoothed loudness gate.
    const gate = new Float32Array(frameCount);
    const attack = 0.25, release = 0.02;
    let g = 0;
    for (let f = 0; f < frameCount; f++) {
        const target = loudness ? loudness[f] : 1;
        g += (target - g) * (target > g ? attack : release);
        gate[f] = Math.min(1, g * 1.6);
    }

    const decay = Math.max(1e-3, period * 0.28);
    for (let f = 0; f < frameCount; f++) {
        const sinceBeat = ((f - offset) % period + period) % period;
        beatPhase[f] = sinceBeat / period;
        barPhase[f] = (((f - barOffset) % barPeriod) + barPeriod) % barPeriod / barPeriod;
        phrasePhase[f] = (((f - barOffset) % phrasePeriod) + phrasePeriod) % phrasePeriod / phrasePeriod;
        beat[f] = Math.exp(-sinceBeat / decay) * gate[f];
    }

    return { beatPhase, barPhase, phrasePhase, beat };
}
