#!/usr/bin/env python3
"""
extract_audio_features.py — per-second audio features for scene-boundary detection.

Audio is often a stronger scene-boundary cue than video: music swells, silence,
and ambience changes at narrative scene transitions — exactly the coarse
boundaries Amazon X-Ray marks, and exactly what the grayscale video cut detector
misses on low-contrast films. This extracts a small per-second feature series per
film, aligned to the 1 fps timeline the embedding dumps use, so it can be fused
with the RGB-histogram features in train_scene_boundary.py.

Two-tower design: this is the AUDIO tower's input, mirroring the video tower's
per-second RGB histogram. Because the scene model is an LSTM (temporal context
comes from the recurrence, not a 2D spectrogram), each second needs only a single
log-PSD vector — one FFT over a WIN_SEC window centred on that second. The LSTM
sees the sequence of per-second PSDs and learns the boundary dynamics itself.

Per second t:
  - log-PSD over [t-WIN/2, t+WIN/2], N_BINS log-spaced frequency bins, L1-norm'd
    then log1p — the spectral shape (music vs speech vs silence vs ambience),
    which changes at scene transitions.

No new dependency: ffmpeg (CLI) decodes the whole track to mono 16 kHz WAV;
numpy does the FFT.

Writes <out_dir>/<slug>.npz with `ts` (second grid) and `feat` [T, N_BINS].

Usage:
  python scripts/scene_detector/extract_audio_features.py \
      --manifest experiments/manifests/films_LVFace_opencv5.json \
      --file-lut experiments/file-lut.json \
      --out experiments/dumps/audio_features
"""
from __future__ import annotations
import argparse, json, subprocess, sys, tempfile, os
from pathlib import Path

import numpy as np
from scipy import signal as sps
from scipy.io import wavfile

SR = 16000
HOP_SEC = 1.0          # one feature vector per second (matches 1 fps presence grid)
WIN_SEC = 4.0          # FFT window per second (centred); >HOP for temporal context
N_BINS = 64            # log-spaced frequency bins per second (the audio tower dim)


def decode_mono(path: str) -> np.ndarray:
    """Whole-file mono 16 kHz float32 PCM via ffmpeg."""
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tf:
        wav = tf.name
    try:
        subprocess.run(
            ["ffmpeg", "-v", "error", "-y", "-i", path,
             "-ac", "1", "-ar", str(SR), "-f", "wav", wav],
            check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        sr, x = wavfile.read(wav)
        if x.dtype == np.int16:
            x = x.astype(np.float32) / 32768.0
        else:
            x = x.astype(np.float32)
        return x
    finally:
        try: os.unlink(wav)
        except OSError: pass


def _logbin_edges(win_samples: int) -> np.ndarray:
    """Indices into the rfft output that bound N_BINS log-spaced freq bands."""
    nfreq = win_samples // 2 + 1
    # log-space from bin 1 (skip DC) to Nyquist; unique integer edges
    edges = np.unique(np.geomspace(1, nfreq - 1, N_BINS + 1).astype(int))
    return edges


def features(mono: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Return (ts[T], feat[T, N_BINS]) — one per-second log-PSD row.

    One FFT per second over a WIN_SEC window centred on that second. Power is
    pooled into N_BINS log-spaced frequency bands (mel-like), L1-normalised across
    bands (so loudness doesn't dominate — the SHAPE is the scene cue), then
    log1p-compressed. The LSTM downstream supplies temporal context, so no
    spectrogram/2D input is needed."""
    hop = int(SR * HOP_SEC)
    win = int(SR * WIN_SEC)
    T = len(mono) // hop
    if T == 0:
        return np.zeros(0), np.zeros((0, N_BINS), np.float32)
    edges = _logbin_edges(win)
    nb = len(edges) - 1
    hann = sps.windows.hann(win)
    feat = np.zeros((T, nb), np.float32)
    half = win // 2
    for t in range(T):
        centre = t * hop + hop // 2
        s = centre - half
        seg = mono[max(0, s): s + win]
        if len(seg) < win:                       # pad edges
            seg = np.pad(seg, (0, win - len(seg)))
        psd = np.abs(np.fft.rfft(seg * hann))**2 + 1e-12
        band = np.array([psd[edges[i]:edges[i+1]].sum() for i in range(nb)])
        band /= band.sum()                       # normalise shape, drop loudness
        feat[t] = np.log1p(band * 1e3)
    ts = np.arange(T, dtype=np.float64)
    return ts, feat


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", required=True)
    ap.add_argument("--file-lut", default="experiments/file-lut.json")
    ap.add_argument("--out", default="experiments/dumps/audio_features")
    args = ap.parse_args()
    films = json.load(open(args.manifest))
    lut = json.load(open(args.file_lut))
    Path(args.out).mkdir(parents=True, exist_ok=True)
    for f in films:
        slug = f["slug"]
        outp = Path(args.out) / f"{slug}.npz"
        if outp.exists():
            print(f"[audio] {slug}: exists, skip", file=sys.stderr); continue
        path = lut.get(slug)
        if not path or not os.path.exists(path):
            print(f"[audio] {slug}: movie missing ({path})", file=sys.stderr); continue
        try:
            mono = decode_mono(path)
            ts, feat = features(mono)
            np.savez_compressed(outp, ts=ts, feat=feat)
            print(f"[audio] {slug}: {len(ts)}s  feat{feat.shape} → {outp.name}",
                  file=sys.stderr)
        except subprocess.CalledProcessError:
            print(f"[audio] {slug}: ffmpeg decode failed", file=sys.stderr)


if __name__ == "__main__":
    main()