#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
"""Find a song's tempo, beat grid and first downbeat with numpy.

    python3 beats.py track.mp3 [--out beats.js] [--min-bpm 80 --max-bpm 160]

Writes window.BEATS = {bpm, first, beats:[...], downbeats:[...]} for the page (or JSON
when --out ends in .json). Use beat(n) in motion.js to put each action on a beat.
Method: spectral-flux onset envelope -> autocorrelation tempo -> grid phase that
maximizes onset energy -> bar phase whose beats carry the most low-end (kick) energy.
Check the result by ear: mux a tick on every beat (see SKILL.md) before animating.
"""

import argparse
import json
import subprocess
from pathlib import Path

import numpy as np

SR = 22050
HOP = 256


def decode(path: Path) -> np.ndarray:
    raw = subprocess.run(["ffmpeg", "-v", "error", "-i", str(path), "-ac", "1", "-ar", str(SR), "-f", "f32le", "-"],
                         capture_output=True, check=True).stdout
    return np.frombuffer(raw, dtype=np.float32)


def stft_mag(x: np.ndarray, n_fft: int = 1024) -> np.ndarray:
    win = np.hanning(n_fft).astype(np.float32)
    frames = np.lib.stride_tricks.sliding_window_view(np.pad(x, (n_fft // 2, n_fft // 2)), n_fft)[::HOP]
    return np.abs(np.fft.rfft(frames * win, axis=1))


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("track", type=Path)
    ap.add_argument("--out", type=Path)
    ap.add_argument("--min-bpm", type=float, default=80)
    ap.add_argument("--max-bpm", type=float, default=160)
    ap.add_argument("--beats-per-bar", type=int, default=4)
    args = ap.parse_args()

    x = decode(args.track)
    mag = np.log1p(100 * stft_mag(x))
    flux = np.maximum(0, np.diff(mag, axis=0)).sum(axis=1)
    flux = np.concatenate([[0], flux])
    flux -= np.convolve(flux, np.ones(16) / 16, mode="same")  # remove slow loudness drift
    onset = np.maximum(flux, 0)
    fps = SR / HOP

    # Tempo: autocorrelation peak in the allowed range, weighted toward ~120 BPM.
    ac = np.correlate(onset, onset, mode="full")[len(onset) - 1:]
    lags = np.arange(len(ac))
    bpm_of = 60 * fps / np.maximum(lags, 1)
    ok = (bpm_of >= args.min_bpm) & (bpm_of <= args.max_bpm)
    weight = np.exp(-0.5 * (np.log2(bpm_of / 120) / 0.9) ** 2)
    lag = lags[ok][np.argmax((ac * weight)[ok])]
    # Refine with parabolic interpolation around the integer lag.
    if 1 <= lag < len(ac) - 1:
        a, b, c = ac[lag - 1], ac[lag], ac[lag + 1]
        lag = lag + 0.5 * (a - c) / (a - 2 * b + c + 1e-9)
    period = float(lag)
    bpm = 60 * fps / period

    # Phase: grid offset (in frames) with the most onset energy on its beats.
    n_beats = int((len(onset) - 1) / period)
    best, phase = -1.0, 0.0
    for p in np.linspace(0, period, 64, endpoint=False):
        idx = np.round(p + period * np.arange(n_beats)).astype(int)
        idx = idx[idx < len(onset)]
        score = onset[idx].sum()
        if score > best:
            best, phase = score, p
    # Refine: snap each grid beat to its nearby onset peak, then fit a straight line
    # (one tempo, one phase) so small tempo errors do not drift over a long track.
    grid = phase + period * np.arange(n_beats)
    w = max(1, int(period / 4))
    peaks = np.array([max(0, int(g) - w) + np.argmax(onset[max(0, int(g) - w):int(g) + w + 1]) for g in grid])
    keep = onset[np.minimum(peaks, len(onset) - 1)] > np.percentile(onset, 75)
    if keep.sum() >= 8:
        slope, icpt = np.polyfit(np.arange(n_beats)[keep], peaks[keep], 1)
        period, phase = float(slope), float(icpt)
        bpm = 60 * fps / period
    beats = (phase + period * np.arange(n_beats)) / fps
    beats = beats[beats >= 0]

    # Downbeat: which beat in the bar has the strongest low end (kick).
    freqs = np.fft.rfftfreq(1024, 1 / SR)
    low = stft_mag(x)[:, freqs < 150].sum(axis=1)
    bar = args.beats_per_bar
    idx = np.clip(np.round(beats * fps).astype(int), 0, len(low) - 1)
    strength = [low[idx[k::bar]].mean() if len(idx[k::bar]) else 0 for k in range(bar)]
    k0 = int(np.argmax(strength))
    downbeats = beats[k0::bar]
    first = float(downbeats[0])

    result = {"bpm": round(bpm, 3), "first": round(first, 4),
              "beats": [round(float(b), 4) for b in beats if b >= first - 1e-6],
              "downbeats": [round(float(b), 4) for b in downbeats]}
    print(f"{args.track.name}: {bpm:.2f} BPM, first downbeat {first:.3f}s, {len(result['beats'])} beats")
    if args.out:
        if args.out.suffix == ".json":
            args.out.write_text(json.dumps(result, indent=1))
        else:
            args.out.write_text(f"// Generated by video/scripts/beats.py\nwindow.BEATS = {json.dumps(result)};\n")


if __name__ == "__main__":
    main()