"""Which DeBERTa reads the listener's own text better? (HANDOFF §4 item 4)

The same weights read the user's text (Vercel `/api/emotion`) and the songs'
lyrics. A candidate text model has to win on lyrics through the whole system
(`train_fusion compare`) *and* not lose on user text. This measures the second
half, on text neither model trained on:

  tweet_eval   SemEval-2018 emotion tweets, test split (1,421). anger, joy,
               optimism and sadness; optimism counts as joy. Not part of
               dair-ai/emotion, GoEmotions or tweet_eval/irony.
  literal      the register probe's 55 literal statements, the closest thing
               to what people type into the box ("I feel so sad and empty...")

Both models are read the way the route reads them: plain softmax, argmax over
all 8 classes, so a prediction outside the gold set (e.g. `neutral`) is wrong.

    python -m backend.src.user_text_eval --live models/<live>/final --candidate models/<cand>/final --tag full

Ship rule, fixed before the first run: the candidate passes the user-text side
if, on *both* sets, the paired 95% bootstrap interval of its accuracy minus the
live model's has an upper bound >= 0 (not significantly worse). Decided on
unrounded values. Macro-F1 and VAD error are reported, not ruled on.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Dict, List

import numpy as np
import torch
from sklearn.metrics import f1_score
from transformers import AutoModelForSequenceClassification, AutoTokenizer

from .acoustic import EMOTION_VAD
from .fusion import user_vad
from .lyrics_probe import load_probe

ROOT = Path(__file__).resolve().parents[2]
RESULTS = ROOT / "results"
TWEET_EVAL_MAP = {0: "anger", 1: "joy", 2: "joy", 3: "sadness"}  # optimism -> joy
N_BOOT = 10000
SEED = 0


def user_sets() -> Dict[str, List[Dict]]:
    from datasets import load_dataset
    te = load_dataset("tweet_eval", "emotion", split="test")
    return {
        "tweet_eval": [{"text": t, "gold": TWEET_EVAL_MAP[int(l)]} for t, l in zip(te["text"], te["label"])],
        "literal": [{"text": r["text"], "gold": r["gold"]} for r in load_probe() if r["register"] == "literal"],
    }


@torch.no_grad()
def softmax_probs(path: str, texts: List[str], device: torch.device, batch: int = 32) -> List[Dict[str, float]]:
    tok = AutoTokenizer.from_pretrained(path)
    model = AutoModelForSequenceClassification.from_pretrained(path).to(device).eval()
    labels = [model.config.id2label[i] for i in range(model.config.num_labels)]
    out = []
    for i in range(0, len(texts), batch):
        enc = tok(texts[i:i + batch], return_tensors="pt", truncation=True, max_length=256, padding=True).to(device)
        p = torch.softmax(model(**enc).logits.float(), dim=-1).cpu().numpy()
        out.extend({l: float(v) for l, v in zip(labels, row)} for row in p)
    return out


def vad_error(probs: Dict[str, float], gold: str) -> float:
    """Distance from the user's VAD point to the gold emotion's; an all-neutral
    reading sits at neutral, as the app has nothing else to place it by."""
    v = user_vad(probs)
    v = np.asarray(EMOTION_VAD["neutral"]) if v is None else v
    return float(np.linalg.norm(v - np.asarray(EMOTION_VAD[gold])))


def score(probs: List[Dict[str, float]], items: List[Dict]) -> Dict:
    pred = [max(p, key=p.get) for p in probs]
    gold = [it["gold"] for it in items]
    present = sorted(set(gold))
    return {
        "correct": np.array([int(a == b) for a, b in zip(pred, gold)]),
        "accuracy": float(np.mean([a == b for a, b in zip(pred, gold)])),
        "macro_f1": float(f1_score(gold, pred, labels=present, average="macro", zero_division=0)),
        "vad_error": float(np.mean([vad_error(p, g) for p, g in zip(probs, gold)])),
        "neutral_rate": float(np.mean([x == "neutral" for x in pred])),
    }


def paired_ci(a: np.ndarray, b: np.ndarray) -> tuple:
    """95% percentile interval of mean(b) - mean(a), resampling items."""
    rng = np.random.default_rng(SEED)
    idx = rng.integers(0, len(a), size=(N_BOOT, len(a)))
    d = b[idx].mean(axis=1) - a[idx].mean(axis=1)
    return float(np.percentile(d, 2.5)), float(np.percentile(d, 97.5))


def compare(live: str, candidate: str, tag: str) -> Dict:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    report: Dict = {"live": live, "candidate": candidate, "sets": {}}
    for name, items in user_sets().items():
        texts = [it["text"] for it in items]
        a, b = score(softmax_probs(live, texts, device), items), score(softmax_probs(candidate, texts, device), items)
        lo, hi = paired_ci(a["correct"], b["correct"])
        report["sets"][name] = {
            "n": len(items),
            "live": {k: round(v, 4) for k, v in a.items() if k != "correct"},
            tag: {k: round(v, 4) for k, v in b.items() if k != "correct"},
            "accuracy_diff": round(b["accuracy"] - a["accuracy"], 4),
            "ci": [round(lo, 4), round(hi, 4)],
            "not_worse": bool(hi >= 0),
        }
    report["ship_user_text"] = all(s["not_worse"] for s in report["sets"].values())
    for name, s in report["sets"].items():
        print(f"  {name:10s} n={s['n']:5d}  live {s['live']['accuracy']:.3f}  {tag} {s[tag]['accuracy']:.3f}  "
              f"diff {s['accuracy_diff']:+.3f} [{s['ci'][0]:+.3f}, {s['ci'][1]:+.3f}]  "
              f"macro-F1 {s['live']['macro_f1']:.3f} -> {s[tag]['macro_f1']:.3f}  "
              f"VAD err {s['live']['vad_error']:.2f} -> {s[tag]['vad_error']:.2f}")
    print(f"  user-text side passes: {report['ship_user_text']}")
    out = RESULTS / f"user_text_{tag}.json"
    out.write_text(json.dumps(report, indent=2))
    print(f"wrote {out}")
    return report


def main() -> int:
    parser = argparse.ArgumentParser(description="Live vs candidate DeBERTa on the listener's own text")
    parser.add_argument("--live", required=True, help="local copy of the deployed weights")
    parser.add_argument("--candidate", required=True)
    parser.add_argument("--tag", required=True)
    args = parser.parse_args()
    compare(args.live, args.candidate, args.tag)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())