#!/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 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"
OUT = HERE / "calibration.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 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).")
    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])
    gt = json.loads(GT.read_text())["tracks"]
    emos = {t["slug"]: set(t["emotions"]) for t in gt}

    # in-sample fit (the shipped calibration)
    va_a, va_b = _fit_affine(pv, gv)
    ar_a, ar_b = _fit_affine(pa, ga)

    def metrics(vv, aa):
        err = np.mean(np.hypot(vv - gv, aa - ga))
        quad = np.mean([_quad(vv[i], aa[i]) == _quad(gv[i], ga[i]) for i in range(len(vv))])
        hit = np.mean([_nearest_emotion(vv[i], aa[i]) in emos[rows[i][4]] 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)

    # leave-one-out CV (honest generalization)
    loo_v, loo_a = np.zeros_like(pv), np.zeros_like(pa)
    for i in range(len(rows)):
        m = np.ones(len(rows), 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)

    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})")

    OUT.write_text(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))
    print(f"  ✓ {OUT.name}")


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] == "apply" and len(args) >= 3:
        print(apply(float(args[1]), float(args[2])))
    else:
        sys.exit("usage: emotion_calibrate.py [fit [--stem] | apply <V> <A>]")


if __name__ == "__main__":
    main()
