#!/usr/bin/env python3
"""emotion_calibrate — fit a V/A calibration from the GT corpus to de-compress CLAP.

The grader's baseline finding: CLAP's emotion read is center-compressed — predicted
valence/arousal span ~±0.4 while the training-knowledge priors span ±0.9. A
distribution-weighted mean of emotion anchors regresses to the centre, so extremes
(euphoric, aggressive, sad) get pulled toward neutral. We fix it from DATA, not by
guessing: fit a per-axis affine map  gt ≈ a·pred + b  (least squares) over every
analyzed corpus track, which learns the scale-up (a>1) that de-compresses + any offset.

Honesty: we report LEAVE-ONE-OUT cross-validated error (fit on N-1, test the held-out
one), so "calibration helps" is a generalization claim, not an in-sample illusion. We
also re-derive the emotion LABEL from the calibrated V/A (nearest circumplex anchor) and
re-grade — if de-compression is right, labels improve too.

    python3 emotion_calibrate.py fit [--stem]     # fit + LOO report → calibration.json
    python3 emotion_calibrate.py paired            # whole-mix vs stem head-to-head, SAME slugs
    python3 emotion_calibrate.py apply V A         # map one (valence, arousal) point

Feeds emotion_ontology/emotion_timeline once trusted. See [[feedback_build_katana_first]]."""
import json
import sys
from pathlib import Path

import numpy as np

import emotion_ontology as EMO

HERE = Path(__file__).resolve().parent
GT = HERE / "emotion_corpus_gt.json"
ANALYSIS = HERE / "_corpus_cache" / "analysis"
# Per-mode files so whole-mix and stem calibrations coexist for the #92 decision;
# `calibration.json` is the symlink/copy of whichever mode we ultimately trust.
OUT = HERE / "calibration.json"


def _out(stem):
    return HERE / f"calibration.{'stem' if stem else 'whole-mix'}.json"


def _pairs(stem=False):
    """(pred_v, pred_a, gt_v, gt_a, slug) for every track that's both GT and analyzed."""
    gt = {t["slug"]: t for t in json.loads(GT.read_text())["tracks"]}
    rows = []
    for slug, t in gt.items():
        f = ANALYSIS / f"{slug}{'' if stem else '_mix'}.json"
        if not f.exists():
            continue
        s = json.loads(f.read_text())["summary"]
        rows.append((s["valence"], s["arousal"], t["valence"], t["arousal"], slug))
    return rows


def _fit_affine(x, y):
    """Least-squares a,b for y ≈ a·x + b."""
    A = np.vstack([x, np.ones_like(x)]).T
    a, b = np.linalg.lstsq(A, y, rcond=None)[0]
    return float(a), float(b)


def _nearest_emotion(v, a):
    va = EMO._matrix()["va"]
    return EMO._matrix()["names"][int(np.argmin((va[:, 0] - v) ** 2 + (va[:, 1] - a) ** 2))]


def _quad(v, a):
    return ("+" if v >= 0 else "-") + ("+" if a >= 0 else "-")


def _emos():
    return {t["slug"]: set(t["emotions"]) for t in json.loads(GT.read_text())["tracks"]}


def _evaluate(pv, pa, gv, ga, slugs, emos):
    """raw + in-sample-calibrated + LOO-calibrated metrics for one mode on one track set.
    Returns {coef:(va_a,va_b,ar_a,ar_b), raw, cal, loo} where each metric is (err,quad,hit).
    LOO = leave-one-out cross-validation, so 'helps' is a generalization claim not an
    in-sample illusion. Shared by fit() and paired() so both score identically."""
    pv = pv.astype(float); pa = pa.astype(float)
    va_a, va_b = _fit_affine(pv, gv)
    ar_a, ar_b = _fit_affine(pa, ga)

    def metrics(vv, aa):
        err = float(np.mean(np.hypot(vv - gv, aa - ga)))
        quad = float(np.mean([_quad(vv[i], aa[i]) == _quad(gv[i], ga[i]) for i in range(len(vv))]))
        hit = float(np.mean([_nearest_emotion(vv[i], aa[i]) in emos[slugs[i]] for i in range(len(vv))]))
        return err, quad, hit

    raw = metrics(pv, pa)
    cal_in = metrics(va_a * pv + va_b, ar_a * pa + ar_b)

    loo_v, loo_a = np.zeros_like(pv), np.zeros_like(pa)
    for i in range(len(slugs)):
        m = np.ones(len(slugs), bool); m[i] = False
        a1, b1 = _fit_affine(pv[m], gv[m]); a2, b2 = _fit_affine(pa[m], ga[m])
        loo_v[i] = a1 * pv[i] + b1; loo_a[i] = a2 * pa[i] + b2
    cal_loo = metrics(loo_v, loo_a)
    return {"coef": (va_a, va_b, ar_a, ar_b), "raw": raw, "cal": cal_in, "loo": cal_loo}


def fit(stem=False):
    rows = _pairs(stem)
    if len(rows) < 6:
        sys.exit(f"only {len(rows)} analyzed tracks — run `emotion_corpus.py analyze"
                 f"{'' if stem else ' --mix'}` first (need ≥6).")
    slugs = [r[4] for r in rows]
    pv = np.array([r[0] for r in rows]); pa = np.array([r[1] for r in rows])
    gv = np.array([r[2] for r in rows]); ga = np.array([r[3] for r in rows])

    ev = _evaluate(pv, pa, gv, ga, slugs, _emos())
    va_a, va_b, ar_a, ar_b = ev["coef"]
    raw, cal_in, cal_loo = ev["raw"], ev["cal"], ev["loo"]

    print(f"⛵ V/A calibration — {len(rows)} tracks ({'stem' if stem else 'whole-mix'})\n")
    print(f"  valence:  gt ≈ {va_a:.2f}·pred {va_b:+.2f}   (×{va_a:.1f} de-compression)")
    print(f"  arousal:  gt ≈ {ar_a:.2f}·pred {ar_b:+.2f}\n")
    print(f"  {'':<16}{'V/A err':>9}{'quad hit':>10}{'emo hit':>9}")
    print(f"  {'raw':<16}{raw[0]:>9.3f}{raw[1]*100:>9.0f}%{raw[2]*100:>8.0f}%")
    print(f"  {'calibrated':<16}{cal_in[0]:>9.3f}{cal_in[1]*100:>9.0f}%{cal_in[2]*100:>8.0f}%")
    print(f"  {'calibrated (LOO)':<16}{cal_loo[0]:>9.3f}{cal_loo[1]*100:>9.0f}%{cal_loo[2]*100:>8.0f}%")
    verdict = "✓ helps" if cal_loo[0] < raw[0] else "✗ no gain — need more/cleaner data"
    print(f"\n  {verdict}  (LOO V/A err {raw[0]:.3f} → {cal_loo[0]:.3f})")

    payload = json.dumps({
        "source": "stem" if stem else "whole-mix", "n": len(rows),
        "valence": {"a": va_a, "b": va_b}, "arousal": {"a": ar_a, "b": ar_b},
        "raw": {"va_err": raw[0], "quad": raw[1], "emo_hit": raw[2]},
        "loo": {"va_err": cal_loo[0], "quad": cal_loo[1], "emo_hit": cal_loo[2]},
    }, indent=1)
    out = _out(stem)
    out.write_text(payload)
    OUT.write_text(payload)  # also publish the freshest fit as the active calibration
    print(f"  ✓ {out.name} (+ {OUT.name})")


def _paired_rows():
    """(slug, mix_v, mix_a, stem_v, stem_a, gt_v, gt_a) for tracks analyzed BOTH ways.
    The only honest whole-mix-vs-stem comparison is on the SAME tracks — fit/--stem grade
    different-sized sets (whole-mix scales to 534, stems to ~150), so their LOO numbers
    aren't directly comparable. This restricts to the common subset."""
    gt = {t["slug"]: t for t in json.loads(GT.read_text())["tracks"]}
    rows = []
    for slug, t in gt.items():
        fm = ANALYSIS / f"{slug}_mix.json"
        fs = ANALYSIS / f"{slug}.json"
        if not (fm.exists() and fs.exists()):
            continue
        m = json.loads(fm.read_text())["summary"]
        s = json.loads(fs.read_text())["summary"]
        rows.append((slug, m["valence"], m["arousal"], s["valence"], s["arousal"],
                     t["valence"], t["arousal"]))
    return rows


def paired():
    """Whole-mix vs stem-aware, head-to-head on the SAME slugs — THE #92 decision."""
    rows = _paired_rows()
    if len(rows) < 6:
        sys.exit(f"only {len(rows)} tracks analyzed both whole-mix AND stem — need ≥6. "
                 f"Run `emotion_corpus.py analyze` (stem) on more tracks first.")
    slugs = [r[0] for r in rows]
    mv = np.array([r[1] for r in rows]); ma = np.array([r[2] for r in rows])
    sv = np.array([r[3] for r in rows]); sa = np.array([r[4] for r in rows])
    gv = np.array([r[5] for r in rows]); ga = np.array([r[6] for r in rows])
    emos = _emos()
    mix = _evaluate(mv, ma, gv, ga, slugs, emos)
    stm = _evaluate(sv, sa, gv, ga, slugs, emos)

    print(f"⛵ whole-mix vs stem — {len(rows)} tracks (PAIRED, same slugs)\n")
    print(f"  {'':<22}{'V/A err':>9}{'quad':>7}{'emo':>6}")
    for label, ev in (("whole-mix", mix), ("stem", stm)):
        for k, tag in (("raw", "raw"), ("loo", "cal·LOO")):
            r = ev[k]
            print(f"  {label+' '+tag:<22}{r[0]:>9.3f}{r[1]*100:>6.0f}%{r[2]*100:>5.0f}%")
    win = "whole-mix" if mix["loo"][0] <= stm["loo"][0] else "stem"
    d = abs(mix["loo"][0] - stm["loo"][0])
    print(f"\n  → {win} wins on calibrated LOO V/A err (Δ {d:.3f})")
    # conditioning: a big in-sample→LOO gap = overfit (untrustworthy on a small set)
    for label, ev in (("whole-mix", mix), ("stem", stm)):
        gap = ev["loo"][0] - ev["cal"][0]
        flag = "  ⚠ overfit" if gap > 0.04 else ""
        print(f"    {label:<10} in-sample {ev['cal'][0]:.3f} → LOO {ev['loo'][0]:.3f}  (gap {gap:+.3f}){flag}")


def apply(v, a):
    c = json.loads(OUT.read_text())
    return (round(c["valence"]["a"] * v + c["valence"]["b"], 3),
            round(c["arousal"]["a"] * a + c["arousal"]["b"], 3))


def main():
    args = sys.argv[1:]
    if args and args[0] == "fit":
        fit(stem="--stem" in args)
    elif args and args[0] == "paired":
        paired()
    elif args and args[0] == "apply" and len(args) >= 3:
        print(apply(float(args[1]), float(args[2])))
    else:
        sys.exit("usage: emotion_calibrate.py [fit [--stem] | paired | apply <V> <A>]")


if __name__ == "__main__":
    main()
