#!/usr/bin/env python3
"""Measure a music track against the film's beat grid before mixing or editing it.

  python3 audio-analyze.py track.mp3 [--grid 0.5] [--sections 0,6.5,9.5,29,44,55.5,61]

Prints:
  - tempo (BPM) from onset autocorrelation
  - beat phase vs the grid in 10 s windows: the shift to apply (mix.py offset_ms) so beats land on cuts.
    A steady value means the tempo is exact; a drifting one means it isn't (time-stretch or re-generate).
  - downbeat position per section (which beat of the bar is strongest), needed to cut or repeat whole bars
  - key estimate (Krumhansl profiles), for any pitched sound you add
  - loudness per section
Windows with dialogue or effects give noisy phase readings; trust the music-only ones.
"""
import argparse, numpy as np
from common import decode

ap = argparse.ArgumentParser()
ap.add_argument("track"); ap.add_argument("--grid", type=float, default=0.5); ap.add_argument("--sections", default="")
a = ap.parse_args()
SR = 22050
x = decode(a.track, SR)
hop, n = 128, 2048
frames = np.lib.stride_tricks.sliding_window_view(x, n)[::hop] * np.hanning(n)
S = np.abs(np.fft.rfft(frames, axis=1)); fps = SR / hop; t = np.arange(len(S)) / fps


def envelope(spec):
    flux = np.concatenate([[0], np.maximum(0, np.diff(np.log1p(spec), axis=0)).sum(1)])
    return np.maximum(flux - np.convolve(flux, np.ones(int(fps)) / int(fps), "same"), 0)


env = envelope(S); low = envelope(S[:, : int(200 / (SR / n))])
L = min(len(env), int(30 * fps))
ac = np.correlate(env[:L], env[:L], "full")[L - 1:]; lags = np.arange(len(ac)) / fps
ok = (lags > 60 / 180) & (lags < 60 / 70)
print(f"tempo ~{60 / lags[ok][np.argmax(ac[ok])]:.1f} BPM (grid {60 / a.grid:.0f} BPM)")


def phase(sel):
    offs = np.arange(0, a.grid, 0.0025)
    sc = [env[sel][np.abs((t[sel] - o + a.grid / 2) % a.grid - a.grid / 2) < 0.012].mean() for o in offs]
    return ((offs[int(np.argmax(sc))] + a.grid / 2) % a.grid - a.grid / 2) * 1000


dur = len(x) / SR
print("beat phase vs grid (negative = early; delay the track by that much):",
      " | ".join(f"{s}-{s + 10}s {phase((t >= s) & (t < s + 10)):+.0f}ms" for s in range(0, int(dur) - 2, 10)))
off = np.median([phase((t >= s) & (t < s + 10)) for s in range(0, int(dur) - 2, 10)]) / 1000
cuts = [float(v) for v in a.sections.split(",") if v] or [0, dur]
for s0, s1 in zip(cuts, cuts[1:]):
    sums = [0.0] * 4
    for k in range(int(np.ceil(s0 / a.grid)), int(s1 / a.grid)):
        i = int(round((k * a.grid + off) * fps)); sums[k % 4] += low[max(0, i - 2):i + 3].max() if i < len(low) else 0
    best = int(np.argmax(sums))
    print(f"section {s0}-{s1}s: downbeats at film times where (t / {a.grid}) % 4 == {best}  (strength by beat {[round(v, 1) for v in sums]})"
          f"  loudness {20 * np.log10(np.sqrt(np.mean(x[int(s0 * SR):int(s1 * SR)] ** 2)) + 1e-9):.0f} dB")
names = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]
maj = np.array([6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88]); mnr = np.array([6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17])
m = n * 4; f = np.fft.rfftfreq(m, 1 / SR); band = (f > 60) & (f < 2000); pc = np.round(12 * np.log2(f[band] / 261.63)) % 12
chroma = np.zeros(12)
for i in range(0, len(x) - m, m // 2):
    sp = np.abs(np.fft.rfft(x[i:i + m] * np.hanning(m)))[band]
    for k in range(12): chroma[k] += sp[pc == k].sum()
cands = [(np.corrcoef(np.roll(maj, k), chroma)[0, 1], names[k] + " major") for k in range(12)] + [(np.corrcoef(np.roll(mnr, k), chroma)[0, 1], names[k] + " minor") for k in range(12)]
r, k = max(cands)
print(f"key ~{k} (r={r:.2f})")