#!/usr/bin/env python3
"""Cut the Lyria cues into their sections. Lyria writes medleys (three moods per cue, joined by
hard cuts), so the score is EDITED from sections rather than used whole. Novelty on a spectral
self-similarity curve finds the joins; each section is then written out as its own wav for audition.

    uv run --with numpy python tools/trailer/segment.py promo/trailer/audio/cue_a.mp3 ...
"""
import os, subprocess, sys
import numpy as np

SR, HOP = 22050, 11025          # half-second frames
for f in sys.argv[1:]:
    raw = subprocess.run(['ffmpeg', '-v', 'error', '-i', f, '-ac', '1', '-ar', str(SR), '-f', 's16le', '-'],
                         capture_output=True).stdout
    x = np.frombuffer(raw, np.int16).astype(float) / 32768
    n = len(x) // HOP
    fr = x[:n * HOP].reshape(n, HOP) * np.hanning(HOP)
    spec = np.abs(np.fft.rfft(fr, axis=1))
    edges = np.geomspace(40, 10000, 25)
    bins = np.fft.rfftfreq(HOP, 1 / SR)
    bands = np.stack([spec[:, (bins >= a) & (bins < b)].mean(1) for a, b in zip(edges, edges[1:])], 1)
    feat = np.log(bands + 1e-6)
    feat = (feat - feat.mean(0)) / (feat.std(0) + 1e-6)
    k = 4   # 2 s either side
    nov = np.zeros(n)
    for i in range(k, n - k):
        nov[i] = np.linalg.norm(feat[i - k:i].mean(0) - feat[i:i + k].mean(0))
    cuts = [0]
    for i in np.argsort(-nov):
        if nov[i] < 2.2: break
        if all(abs(i - c) >= 12 for c in cuts) and n - i >= 8:
            cuts.append(i)
    cuts = sorted(cuts) + [n]
    rms = 20 * np.log10(np.sqrt((fr ** 2).mean(1)) + 1e-9)
    base = os.path.splitext(f)[0]
    for j, (a, b) in enumerate(zip(cuts, cuts[1:])):
        out = f'{base}.s{j}.wav'
        subprocess.run(['ffmpeg', '-v', 'error', '-y', '-ss', str(a / 2), '-t', str((b - a) / 2), '-i', f, out], check=True)
        print(f'{os.path.basename(out)}  {a/2:6.1f}-{b/2:6.1f}s  rms {rms[a:b].mean():5.1f} dB')