#!/usr/bin/env python3
"""emotion_timeline — slide the CLAP emotion kernel over a track → an emotion arc.

A whole track has ONE dominant vibe but an ARC: a euphoric build, a tense breakdown,
a dark drop. So we window the master (~12 s, 50 % hop — long enough that CLAP reads a
mood, short enough to catch a section change), CLAP-embed each window, and run the
emotion_ontology kernel per window → a time series of (valence, arousal, top-emotion).
We smooth single-frame flicker, then merge contiguous frames into emotion SECTIONS.

Two modes:
  • whole-mix  — read the full mix (fast).
  • stem-aware — separate via the Foundry (vocals/bass/drums/other), read each stem,
    and FUSE: VOCALS+other carry valence, DRUMS+BASS carry arousal; vocal presence is
    a section/structure cue. The mix is muddy for CLAP — the voice read is far cleaner.

    python3 emotion_timeline.py track <file> [--win 12 --hop 6] [--json out.json]
    python3 emotion_timeline.py track <file> --stems [--slug NAME] [--json out.json]

Window scale matters: emotion reads decisively at 10-30 s (CLAP's strength), weakly on
one-shots. PLN's ear confirms feel; the machine reports the arc ([[feedback_verify_own_renders]]).
See [[project_hexa_emotion_convergence]]."""
import json
import subprocess
import sys
from collections import Counter
from pathlib import Path

import numpy as np

import emotion_ontology as EMO
import sample_classify as CLF
import sample_semantics as SEM

EMBED_BATCH = 24            # windows per CLAP forward (bounds memory on long tracks)
# Fusion weights — which stem carries which affect. Valence rides the melodic/vocal
# content; arousal rides the rhythm section. Renormalized over whatever stems exist.
VAL_W = {"vocals": 0.40, "other": 0.30, "mix": 0.20, "bass": 0.05, "drums": 0.05}
ARO_W = {"drums": 0.30, "mix": 0.25, "bass": 0.20, "vocals": 0.15, "other": 0.10}


def decode_full(path, sr=CLF.SR, max_min=25):
    """Decode any file → mono float32 @ CLF.SR, capped at max_min. None on error."""
    try:
        p = subprocess.run(
            ["ffmpeg", "-v", "quiet", "-i", str(path), "-ac", "1", "-ar", str(sr),
             "-t", str(max_min * 60), "-f", "f32le", "-"],
            capture_output=True, timeout=300)
        a = np.frombuffer(p.stdout, dtype=np.float32)
        return a if a.size >= sr else None
    except Exception:
        return None


def _median_smooth(x, k=3):
    """Median filter (odd k) — kills single-window flicker without EMA lag."""
    if k <= 1 or len(x) < k:
        return np.asarray(x, dtype="float32")
    r = k // 2
    pad = np.pad(x, r, mode="edge")
    return np.array([np.median(pad[i:i + k]) for i in range(len(x))], dtype="float32")


def _window_grid(n, win, hop):
    starts = list(range(0, max(1, n - win + 1), hop)) or [0]
    if n - starts[-1] > win * 0.4 and starts[-1] + win < n:
        starts.append(n - win)                          # trailing partial window
    return starts


def _embed_at(windows):
    """Embed a list of equal-length mono arrays → (N×512), batched."""
    rows = [SEM.embed_audio_arrays(windows[b:b + EMBED_BATCH])
            for b in range(0, len(windows), EMBED_BATCH)]
    return np.vstack(rows) if rows else np.zeros((0, 512), dtype="float32")


def _windows(y, starts, win):
    out = []
    for s in starts:
        seg = y[s:s + win]
        if len(seg) < win:
            seg = np.pad(seg, (0, win - len(seg)))
        out.append(seg)
    return out


def _sections(times, labels, min_sec=8.0):
    """Contiguous equal-label runs → sections; runs shorter than min_sec fuse into the
    longer neighbour so the arc reads as scenes, not strobing."""
    if not len(labels):
        return []
    segs, s = [], 0
    for i in range(1, len(labels) + 1):
        if i == len(labels) or labels[i] != labels[s]:
            segs.append([times[s], times[i - 1], labels[s], s, i - 1])
            s = i
    changed = True
    while changed and len(segs) > 1:
        changed = False
        for j, seg in enumerate(segs):
            if seg[1] - seg[0] >= min_sec:
                continue
            left = segs[j - 1] if j > 0 else None
            right = segs[j + 1] if j < len(segs) - 1 else None
            pick = left if (right is None or (left is not None and
                            (left[1] - left[0]) >= (right[1] - right[0]))) else right
            pick[0] = min(pick[0], seg[0]); pick[1] = max(pick[1], seg[1])
            pick[3] = min(pick[3], seg[3]); pick[4] = max(pick[4], seg[4])
            segs.pop(j); changed = True
            break
    return segs


def _finalize(frames, dur, win_s, hop_s, extra_summary=None):
    """Shared back half: smooth V/A, relabel from the smoothed circumplex point, fuse
    into sections, summarize. Works for both whole-mix and stem-fused frames."""
    val = _median_smooth([f["valence"] for f in frames], 3)
    aro = _median_smooth([f["arousal"] for f in frames], 3)
    va, names = EMO._matrix()["va"], EMO._matrix()["names"]
    for i, f in enumerate(frames):
        f["valence"], f["arousal"] = round(float(val[i]), 3), round(float(aro[i]), 3)
        d = (va[:, 0] - val[i]) ** 2 + (va[:, 1] - aro[i]) ** 2
        f["emotion"] = names[int(np.argmin(d))]
    times = [f["t"] for f in frames]
    labels = [f["emotion"] for f in frames]
    has_vocal = "vocal" in frames[0]
    sections = []
    for s0, s1, lab, i0, i1 in _sections(times, labels, min_sec=max(win_s, 8.0)):
        sl = frames[i0:i1 + 1]
        sec = {"start": round(s0, 2), "end": round(min(s1 + hop_s, dur), 2), "emotion": lab,
               "valence": round(float(np.mean([f["valence"] for f in sl])), 3),
               "arousal": round(float(np.mean([f["arousal"] for f in sl])), 3)}
        if has_vocal:
            sec["vocal"] = bool(np.mean([f["vocal"] for f in sl]) >= 0.5)
        sections.append(sec)
    summary = {"dominant": Counter(labels).most_common(1)[0][0],
               "valence": round(float(np.mean(val)), 3),
               "arousal": round(float(np.mean(aro)), 3),
               "arc": "→".join(dict.fromkeys(s["emotion"] for s in sections))}
    if extra_summary:
        summary.update(extra_summary)
    return {"duration": dur, "win_s": win_s, "hop_s": hop_s, "n_frames": len(frames),
            "summary": summary, "sections": sections, "frames": frames}


def analyze(path, win_s=12.0, hop_s=6.0):
    """Whole-mix emotion arc."""
    y = decode_full(path)
    if y is None:
        return None
    sr = CLF.SR
    win, hop = int(win_s * sr), int(hop_s * sr)
    starts = _window_grid(len(y), win, hop)
    ae = _embed_at(_windows(y, starts, win))
    frames = []
    for s, row in zip(starts, ae):
        r = EMO.score(row)
        frames.append({"t": round(s / sr, 2), "valence": r["valence"],
                       "arousal": r["arousal"], "emotion": None, "confidence": r["confidence"]})
    if not frames:
        return None
    res = _finalize(frames, round(len(y) / sr, 2), win_s, hop_s)
    res["file"], res["stems"] = str(path), ["mix"]
    return res


def analyze_stems(sources, win_s=12.0, hop_s=6.0):
    """Stem-aware: score each source on a shared window grid, fuse V/A by stem role,
    flag vocal presence. `sources` = {mix, vocals, bass, drums, other}→Path (any subset)."""
    ys = {n: decode_full(p) for n, p in sources.items()}
    ys = {n: y for n, y in ys.items() if y is not None}
    if not ys:
        return None
    sr = CLF.SR
    win, hop = int(win_s * sr), int(hop_s * sr)
    ref = ys.get("mix") if ys.get("mix") is not None else next(iter(ys.values()))
    starts = _window_grid(len(ref), win, hop)
    per = {n: [EMO.score(r) for r in _embed_at(_windows(y, starts, win))] for n, y in ys.items()}

    voc_pres = None
    if "vocals" in ys:
        rms = np.array([float(np.sqrt(np.mean(ys["vocals"][s:s + win] ** 2) + 1e-12))
                        for s in starts])
        voc_pres = rms / (np.percentile(rms, 90) + 1e-9) > 0.15

    frames = []
    for i, s in enumerate(starts):
        vv = vw = av = aw = 0.0
        by = {}
        for n, scs in per.items():
            sv = scs[i]
            by[n] = {"valence": sv["valence"], "arousal": sv["arousal"], "emotion": sv["top"][0][0]}
            vv += VAL_W.get(n, 0) * sv["valence"]; vw += VAL_W.get(n, 0)
            av += ARO_W.get(n, 0) * sv["arousal"]; aw += ARO_W.get(n, 0)
        fr = {"t": round(s / sr, 2), "valence": round(vv / vw if vw else 0.0, 3),
              "arousal": round(av / aw if aw else 0.0, 3), "emotion": None, "by_stem": by}
        if voc_pres is not None:
            fr["vocal"] = bool(voc_pres[i])
        frames.append(fr)

    extra = {"vocal_ratio": round(float(np.mean(voc_pres)), 3)} if voc_pres is not None else None
    res = _finalize(frames, round(len(ref) / sr, 2), win_s, hop_s, extra_summary=extra)
    res["stems"] = list(per)
    return res


# ── CLI ─────────────────────────────────────────────────────────────────────────
def _arc_print(res):
    name = Path(res.get("file", "track")).name
    s = res["summary"]
    print(f"🎬 {name}  ({res['duration']:.0f}s, {res['n_frames']} frames, stems={'+'.join(res['stems'])})")
    vr = f"  vocal {s['vocal_ratio']*100:.0f}%" if "vocal_ratio" in s else ""
    print(f"   dominant: {s['dominant']}   mean V/A {s['valence']:+.2f}/{s['arousal']:+.2f}{vr}")
    print(f"   arc: {s['arc']}\n")
    span = max(1.0, res["duration"])
    for sec in res["sections"]:
        a = int(sec["start"] / span * 40); b = max(a + 1, int(sec["end"] / span * 40))
        voc = ("🎤" if sec.get("vocal") else "  ") if "vocal" in sec else ""
        print(f"   {sec['start']:>6.0f}-{sec['end']:<6.0f} {voc} {sec['emotion']:<12} "
              f"V/A {sec['valence']:+.2f}/{sec['arousal']:+.2f}")
        print(f"   {' ' * a}{'█' * (b - a)}")


def main():
    args = sys.argv[1:]
    if not (args and args[0] == "track" and len(args) > 1):
        sys.exit("usage: emotion_timeline.py track <file> [--stems] [--slug N] [--win N] [--hop N] [--json out]")
    win = float(args[args.index("--win") + 1]) if "--win" in args else 12.0
    hop = float(args[args.index("--hop") + 1]) if "--hop" in args else 6.0
    if "--stems" in args:
        import foundry_stems as FS
        slug = args[args.index("--slug") + 1] if "--slug" in args else Path(args[1]).stem
        sources = FS.ensure_stems(args[1], slug)
        res = analyze_stems(sources, win_s=win, hop_s=hop)
        if res:
            res["file"] = args[1]
    else:
        res = analyze(args[1], win_s=win, hop_s=hop)
    if not res:
        sys.exit("decode/analyze failed")
    _arc_print(res)
    if "--json" in args:
        out = Path(args[args.index("--json") + 1])
        out.write_text(json.dumps(res, ensure_ascii=False, indent=1))
        print(f"\n✓ {out}")


if __name__ == "__main__":
    main()
