"""Checkpoint evaluation and ablation reporting.

Produces the numbers that go in the paper: in-domain test metrics, per-class
breakdown, calibration error, and — the measurement this work exists for — the
lyric-register gap from the probe set.

    # one checkpoint
    python -m backend.src.evaluate --model models/eumora_improved_.../final

    # ablation table across checkpoints
    python -m backend.src.evaluate \\
        --model baseline=models/eumora_baseline_.../final \\
        --model improved=models/eumora_improved_.../final \\
        --model +tapt=models/eumora_tapt_.../final \\
        --markdown results/ablation.md

`--no-test-set` skips the in-domain split and runs the probe only, which needs
no dataset download and takes seconds.
"""

from __future__ import annotations

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

import numpy as np
import torch
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
    f1_score,
)
from transformers import AutoModelForSequenceClassification, AutoTokenizer

from .calibration import expected_calibration_error, negative_log_likelihood
from .data_prep import TARGET_LABELS
from .lyrics_probe import LABELS as PROBE_LABELS, load_probe


def get_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


@torch.no_grad()
def batched_logits(
    model, tokenizer, texts: List[str], device, max_length: int = 256, batch_size: int = 32
) -> np.ndarray:
    model.eval()
    out = []
    for i in range(0, len(texts), batch_size):
        enc = tokenizer(
            texts[i : i + batch_size],
            return_tensors="pt",
            truncation=True,
            max_length=max_length,
            padding=True,
        ).to(device)
        out.append(model(**enc).logits.float().cpu().numpy())
    return np.concatenate(out, axis=0)


def softmax(z: np.ndarray) -> np.ndarray:
    z = z - z.max(axis=1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=1, keepdims=True)


# ── Uncertainty ───────────────────────────────────────────────────────────────

def bootstrap_ci(
    values: np.ndarray, n_boot: int = 10000, alpha: float = 0.05, seed: int = 0
) -> tuple:
    """Percentile bootstrap CI for the mean of a 0/1 outcome vector.

    The probe has 40 items per register, so a reported accuracy carries real
    sampling uncertainty — 11.1% on 36 items has a wide interval, and a paper
    quoting the point estimate alone overstates what was measured. Bootstrapping
    avoids a normal approximation, which is poor for proportions near 0 or 1.
    """
    if len(values) == 0:
        return (None, None)
    rng = np.random.default_rng(seed)
    idx = rng.integers(0, len(values), size=(n_boot, len(values)))
    means = values[idx].mean(axis=1)
    lo, hi = np.percentile(means, [100 * alpha / 2, 100 * (1 - alpha / 2)])
    return (round(float(lo), 4), round(float(hi), 4))


def paired_gap_ci(
    literal: np.ndarray, figurative: np.ndarray,
    n_boot: int = 10000, alpha: float = 0.05, seed: int = 0
) -> tuple:
    """CI for the register gap, resampling *pairs* rather than the two arms.

    The probe is built from minimal pairs, so the two subsets are not
    independent — a pair shares its gold label and its content. Resampling pairs
    preserves that structure; resampling the arms separately would understate
    the correlation and give a misleadingly wide interval.
    """
    if len(literal) == 0 or len(literal) != len(figurative):
        return (None, None)
    rng = np.random.default_rng(seed)
    idx = rng.integers(0, len(literal), size=(n_boot, len(literal)))
    gaps = literal[idx].mean(axis=1) - figurative[idx].mean(axis=1)
    lo, hi = np.percentile(gaps, [100 * alpha / 2, 100 * (1 - alpha / 2)])
    return (round(float(lo), 4), round(float(hi), 4))


# ── Probe evaluation ──────────────────────────────────────────────────────────

def evaluate_probe_detailed(model, tokenizer, device, max_length: int = 256) -> Dict:
    """Full probe breakdown: register gap, per-phenomenon, per-pair failures."""
    all_records = load_probe()

    # Compare by label *name*, not index: checkpoints may order their classes
    # differently, and older variants may not have all eight. Items whose gold
    # label the model cannot express are excluded and reported, rather than
    # failing the whole run.
    model_labels = [
        model.config.id2label[i] for i in sorted(int(k) for k in model.config.id2label)
    ]
    supported = [l for l in PROBE_LABELS if l in model_labels]
    skipped = [l for l in PROBE_LABELS if l not in model_labels]

    records = [r for r in all_records if r["gold"] in supported]
    if not records:
        raise ValueError(
            f"Model labels {model_labels} share no classes with the probe {PROBE_LABELS}"
        )

    logits = batched_logits(model, tokenizer, [r["text"] for r in records], device, max_length)
    probs = softmax(logits)
    pred_names = [model_labels[i] for i in probs.argmax(axis=1)]
    correct = np.array(
        [int(p == r["gold"]) for p, r in zip(pred_names, records)], dtype=int
    )
    # Gold indices in the *model's* label space, for calibration metrics.
    gold = np.array([model_labels.index(r["gold"]) for r in records])

    def subset_acc(mask) -> Optional[float]:
        idx = [i for i, r in enumerate(records) if mask(r)]
        return round(float(correct[idx].mean()), 4) if idx else None

    literal = subset_acc(lambda r: r["register"] == "literal")
    figurative = subset_acc(lambda r: r["register"] == "figurative")

    # Pairs where the literal member is right but the figurative one is wrong —
    # the cleanest evidence of register sensitivity, controlling for content.
    by_pair: Dict[str, Dict[str, int]] = {}
    for i, r in enumerate(records):
        by_pair.setdefault(r["pair_id"], {})[r["register"]] = int(correct[i])
    broken = sorted(
        p for p, v in by_pair.items()
        if v.get("literal") == 1 and v.get("figurative") == 0
    )

    phenomena = sorted({r["phenomenon"] for r in records})
    per_class = {
        name: subset_acc(lambda r, n=name: r["gold"] == n) for name in supported
    }

    # Uncertainty. The gap is bootstrapped over pairs, since the probe's two
    # halves are matched rather than independent samples.
    def subset_vec(mask):
        return correct[[i for i, r in enumerate(records) if mask(r)]]

    lit_vec = subset_vec(lambda r: r["register"] == "literal")
    fig_vec = subset_vec(lambda r: r["register"] == "figurative")

    pair_ids = sorted({r["pair_id"] for r in records})
    lit_by_pair, fig_by_pair = [], []
    for pid in pair_ids:
        members = {r["register"]: correct[i] for i, r in enumerate(records)
                   if r["pair_id"] == pid}
        if "literal" in members and "figurative" in members:
            lit_by_pair.append(members["literal"])
            fig_by_pair.append(members["figurative"])

    # How often the model retreats to `neutral` on figurative input — the
    # specific failure mode behind the low lyric scores in production.
    figurative_non_neutral = [
        i for i, r in enumerate(records)
        if r["register"] == "figurative" and r["gold"] != "neutral"
    ]
    neutral_retreat = (
        round(float(np.mean([pred_names[i] == "neutral" for i in figurative_non_neutral])), 4)
        if figurative_non_neutral and "neutral" in model_labels
        else None
    )

    # A model that predicts one class for everything has a register gap of 0,
    # which reads as "no register sensitivity" when it actually means the model
    # never discriminates at all. The baseline arm does exactly this — it
    # collapses to `neutral` and scores 7.5% on both registers. Flag it so the
    # gap is not reported as a strength.
    distinct_predictions = len(set(pred_names))
    from collections import Counter as _Counter
    top_share = _Counter(pred_names).most_common(1)[0][1] / len(pred_names)
    degenerate = distinct_predictions <= 2 or top_share >= 0.9

    return {
        "overall": round(float(correct.mean()), 4),
        "overall_ci": bootstrap_ci(correct),
        "literal": literal,
        "literal_ci": bootstrap_ci(lit_vec),
        "figurative": figurative,
        "figurative_ci": bootstrap_ci(fig_vec),
        "register_gap_ci": paired_gap_ci(np.array(lit_by_pair), np.array(fig_by_pair)),
        "n_complete_pairs": len(lit_by_pair),
        "degenerate": degenerate,
        "distinct_predictions": distinct_predictions,
        "majority_prediction_share": round(top_share, 4),
        "register_gap": round(literal - figurative, 4) if None not in (literal, figurative) else None,
        "by_phenomenon": {p: subset_acc(lambda r, p=p: r["phenomenon"] == p) for p in phenomena},
        "by_class": per_class,
        "mean_confidence": round(float(probs.max(axis=1).mean()), 4),
        "ece": round(expected_calibration_error(probs, gold), 4),
        "broken_pairs": broken,
        "n_broken_pairs": len(broken),
        "neutral_rate_figurative": neutral_retreat,
        "n_evaluated": len(records),
        "n_total": len(all_records),
        "skipped_labels": skipped,
        # Per-item outcomes, retained so two arms can be compared on the *same*
        # items rather than by eyeballing whether their intervals overlap.
        "_pair_ids": pair_ids,
        "_literal_correct": [int(x) for x in lit_by_pair],
        "_figurative_correct": [int(x) for x in fig_by_pair],
    }


# ── In-domain evaluation ──────────────────────────────────────────────────────

def evaluate_test_set(model, tokenizer, device, max_length: int = 256, **data_kwargs) -> Dict:
    from .data_prep import build_dataset

    dataset, _ = build_dataset(**data_kwargs)
    texts = dataset["test"]["text"]
    gold = np.array(dataset["test"]["label"])

    logits = batched_logits(model, tokenizer, texts, device, max_length)
    probs = softmax(logits)
    preds = probs.argmax(axis=1)

    per_class = f1_score(
        gold, preds, average=None, zero_division=0, labels=range(len(TARGET_LABELS))
    )
    return {
        "accuracy": round(float(accuracy_score(gold, preds)), 4),
        "f1_macro": round(float(f1_score(gold, preds, average="macro", zero_division=0)), 4),
        "f1_weighted": round(float(f1_score(gold, preds, average="weighted", zero_division=0)), 4),
        "per_class_f1": {n: round(float(s), 4) for n, s in zip(TARGET_LABELS, per_class)},
        "ece": round(expected_calibration_error(probs, gold), 4),
        "nll": round(negative_log_likelihood(probs, gold), 4),
        "mean_confidence": round(float(probs.max(axis=1).mean()), 4),
        "n_samples": len(gold),
        "report": classification_report(
            gold, preds, labels=range(len(TARGET_LABELS)),
            target_names=TARGET_LABELS, zero_division=0, digits=3,
        ),
        "confusion": confusion_matrix(
            gold, preds, labels=range(len(TARGET_LABELS))
        ).tolist(),
    }


def load_env() -> None:
    """Load .env so private Hub repos resolve without exporting HF_TOKEN by hand.

    The deployed classifier lives in a private repo, so benchmarking it — the
    baseline row of the ablation table — needs a token. Reading .env here means
    `--model VivDubs/eumora-emotion-classifier` just works.
    """
    for candidate in (
        Path(__file__).resolve().parents[2] / ".env",   # repo root
        Path(__file__).resolve().parents[1] / ".env",   # backend/
    ):
        if not candidate.exists():
            continue
        for line in candidate.read_text().splitlines():
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            key, value = line.split("=", 1)
            os.environ.setdefault(key.strip(), value.strip().strip('"').strip("'"))


def evaluate_checkpoint(
    model_path: str,
    run_test_set: bool = True,
    max_length: int = 256,
) -> Dict:
    device = get_device()
    print(f"\nEvaluating {model_path}")

    token = os.environ.get("HF_TOKEN") or os.environ.get("HF_API_TOKEN")
    auth = {"token": token} if token and not Path(model_path).exists() else {}
    if not Path(model_path).exists() and not token:
        print("  (no HF_TOKEN found — private Hub repos will fail to load)")

    tokenizer = AutoTokenizer.from_pretrained(model_path, **auth)
    model = AutoModelForSequenceClassification.from_pretrained(model_path, **auth).to(device)

    result: Dict = {"model_path": str(model_path)}

    calib = getattr(model.config, "calibration", None)
    if calib:
        result["calibration_folded"] = calib
        print(f"  calibration folded in: T={calib.get('temperature')}")

    print("  lyric-register probe...")
    result["probe"] = evaluate_probe_detailed(model, tokenizer, device, max_length)
    p = result["probe"]
    if p["skipped_labels"]:
        print(f"    note: model cannot predict {p['skipped_labels']} — "
              f"scoring {p['n_evaluated']}/{p['n_total']} items")
    def ci(key):
        lo, hi = p.get(key) or (None, None)
        return f" [{lo:.0%}-{hi:.0%}]" if lo is not None else ""
    print(f"    overall {p['overall']:.1%}{ci('overall_ci')} | "
          f"literal {p['literal']:.1%}{ci('literal_ci')} | "
          f"figurative {p['figurative']:.1%}{ci('figurative_ci')}")
    gci = p.get("register_gap_ci") or (None, None)
    gap_txt = (f" [{gci[0]:+.0%} to {gci[1]:+.0%}]" if gci[0] is not None else "")
    print(f"    register gap {p['register_gap']:+.1%}{gap_txt}  "
          f"({p.get('n_complete_pairs', 0)} pairs)")
    if p.get("degenerate"):
        print(f"    DEGENERATE: predicts {p['distinct_predictions']} distinct class(es), "
              f"{p['majority_prediction_share']:.0%} one class — the gap is meaningless here")
    if p["neutral_rate_figurative"] is not None:
        print(f"    neutral-retreat on figurative: {p['neutral_rate_figurative']:.1%}")

    if run_test_set:
        print("  in-domain test set...")
        result["test"] = evaluate_test_set(model, tokenizer, device, max_length)
        t = result["test"]
        print(f"    accuracy {t['accuracy']:.1%} | macro-F1 {t['f1_macro']:.1%} | ECE {t['ece']:.4f}")

    return result


def compare_arms(
    a_name: str, a_probe: Dict, b_name: str, b_probe: Dict,
    n_boot: int = 10000, seed: int = 0,
) -> Optional[Dict]:
    """Bootstrap the *difference* in register gap between two models.

    Two arms evaluated on the same probe are matched data, so the question is
    not whether their individual intervals overlap — overlapping intervals are
    entirely compatible with a significant paired difference, and reading them
    that way is a standard error. Resample the same pair indices for both models
    and bootstrap the difference directly.

    Returns None when the two runs did not score the same pairs, since the
    comparison is only paired if the items match.
    """
    if a_probe.get("_pair_ids") != b_probe.get("_pair_ids"):
        return None

    a_lit = np.array(a_probe["_literal_correct"])
    a_fig = np.array(a_probe["_figurative_correct"])
    b_lit = np.array(b_probe["_literal_correct"])
    b_fig = np.array(b_probe["_figurative_correct"])
    if len(a_lit) == 0:
        return None

    rng = np.random.default_rng(seed)
    idx = rng.integers(0, len(a_lit), size=(n_boot, len(a_lit)))

    gap_a = a_lit[idx].mean(axis=1) - a_fig[idx].mean(axis=1)
    gap_b = b_lit[idx].mean(axis=1) - b_fig[idx].mean(axis=1)
    diff = gap_b - gap_a                      # negative => b closed the gap

    fig_diff = b_fig[idx].mean(axis=1) - a_fig[idx].mean(axis=1)
    lit_diff = b_lit[idx].mean(axis=1) - a_lit[idx].mean(axis=1)

    def ci(v):
        lo, hi = np.percentile(v, [2.5, 97.5])
        return (round(float(lo), 4), round(float(hi), 4))

    gap_lo, gap_hi = ci(diff)
    return {
        "from": a_name, "to": b_name, "n_pairs": int(len(a_lit)),
        "gap_change": round(float(diff.mean()), 4), "gap_change_ci": (gap_lo, gap_hi),
        "figurative_change": round(float(fig_diff.mean()), 4),
        "figurative_change_ci": ci(fig_diff),
        "literal_change": round(float(lit_diff.mean()), 4),
        "literal_change_ci": ci(lit_diff),
        # Only claim a change when the interval clears zero.
        "gap_significantly_closed": bool(gap_hi < 0),
        "gap_significantly_widened": bool(gap_lo > 0),
    }


# ── Reporting ─────────────────────────────────────────────────────────────────

def markdown_table(results: Dict[str, Dict]) -> str:
    """Ablation table, ready to paste into the paper."""
    has_test = any("test" in r for r in results.values())

    lines = ["# EUMORA ablation results", ""]
    header = ["Variant", "Probe overall", "Literal", "Figurative", "Register gap", "Neutral retreat"]
    if has_test:
        header += ["Test acc", "Test macro-F1", "ECE"]
    lines.append("| " + " | ".join(header) + " |")
    lines.append("|" + "|".join(["---"] * len(header)) + "|")

    def pct(v, signed=False):
        if v is None:
            return "—"
        return f"{v:+.1%}" if signed else f"{v:.1%}"

    for name, r in results.items():
        p = r["probe"]
        label = name + ("*" if p.get("skipped_labels") else "") + ("†" if p.get("degenerate") else "")
        row = [
            label,
            pct(p["overall"]),
            pct(p["literal"]),
            pct(p["figurative"]),
            (pct(p["register_gap"], signed=True) +
             (f" [{p['register_gap_ci'][0]:+.0%}, {p['register_gap_ci'][1]:+.0%}]"
              if p.get("register_gap_ci") and p["register_gap_ci"][0] is not None else "")),
            pct(p["neutral_rate_figurative"]),
        ]
        if has_test:
            t = r.get("test")
            row += (
                [f"{t['accuracy']:.1%}", f"{t['f1_macro']:.1%}", f"{t['ece']:.4f}"]
                if t else ["—", "—", "—"]
            )
        lines.append("| " + " | ".join(row) + " |")

    lines += [
        "",
        "**Register gap** = literal accuracy − figurative accuracy on matched minimal",
        "pairs. Lower is better; it isolates sensitivity to lyric phrasing while holding",
        "content and label distribution fixed.",
        "",
        "**Neutral retreat** = share of non-neutral figurative items predicted `neutral`.",
        "This is the failure mode behind the low lyric scores in production.",
        "",
    ]

    degen = {n: r["probe"] for n, r in results.items() if r["probe"].get("degenerate")}
    for n, pr in degen.items():
        lines.append(
            f"† `{n}` predicts only {pr['distinct_predictions']} distinct class(es) "
            f"({pr['majority_prediction_share']:.0%} a single class). Its register gap is "
            f"an artefact of uniform failure, not evidence of register robustness."
        )
    if degen:
        lines.append("")

    partial = {
        n: r["probe"]["skipped_labels"]
        for n, r in results.items() if r["probe"].get("skipped_labels")
    }
    for n, sk in partial.items():
        r = results[n]["probe"]
        lines.append(
            f"\\* `{n}` cannot predict {', '.join(sk)}; scored on "
            f"{r['n_evaluated']}/{r['n_total']} probe items."
        )
    if partial:
        lines.append("")

    names = list(results)
    if len(names) > 1:
        pairwise = []
        for i in range(len(names) - 1):
            c = compare_arms(names[i], results[names[i]]["probe"],
                             names[i + 1], results[names[i + 1]]["probe"])
            if c:
                pairwise.append(c)
        if pairwise:
            lines += ["## Paired comparison between arms", "",
                      "Differences are bootstrapped over the *same* probe pairs for both",
                      "models. Overlapping individual intervals are compatible with a",
                      "significant paired difference, so this is the comparison to read.", "",
                      "| change | register gap | figurative | literal |",
                      "|---|---|---|---|"]
            for c in pairwise:
                verdict = ("**closed**" if c["gap_significantly_closed"]
                           else "**widened**" if c["gap_significantly_widened"]
                           else "not significant")
                lines.append(
                    f"| {c['from']} -> {c['to']} | "
                    f"{c['gap_change']:+.1%} [{c['gap_change_ci'][0]:+.0%}, {c['gap_change_ci'][1]:+.0%}] {verdict} | "
                    f"{c['figurative_change']:+.1%} [{c['figurative_change_ci'][0]:+.0%}, {c['figurative_change_ci'][1]:+.0%}] | "
                    f"{c['literal_change']:+.1%} [{c['literal_change_ci'][0]:+.0%}, {c['literal_change_ci'][1]:+.0%}] |"
                )
            lines.append("")

    phenomena = sorted({p for r in results.values() for p in r["probe"]["by_phenomenon"]})
    lines += ["## Accuracy by linguistic phenomenon", ""]
    lines.append("| Variant | " + " | ".join(phenomena) + " |")
    lines.append("|" + "|".join(["---"] * (len(phenomena) + 1)) + "|")
    for name, r in results.items():
        vals = r["probe"]["by_phenomenon"]
        lines.append(
            "| " + name + " | "
            + " | ".join(f"{vals[p]:.1%}" if vals.get(p) is not None else "—" for p in phenomena)
            + " |"
        )
    return "\n".join(lines) + "\n"


def main():
    load_env()
    parser = argparse.ArgumentParser(description="Evaluate EUMORA checkpoints")
    parser.add_argument(
        "--model", action="append", required=True,
        help="Path, or name=path to label it in the table. Repeatable.",
    )
    parser.add_argument("--no-test-set", action="store_true",
                        help="Probe only — no dataset download needed")
    parser.add_argument("--max-length", type=int, default=256)
    parser.add_argument("--json", type=str, default=None, help="Write full results JSON here")
    parser.add_argument("--markdown", type=str, default=None, help="Write the ablation table here")
    parser.add_argument("--per-class", action="store_true", help="Print the sklearn report")
    args = parser.parse_args()

    results: Dict[str, Dict] = {}
    for spec in args.model:
        name, path = spec.split("=", 1) if "=" in spec else (Path(spec).parent.name, spec)
        results[name] = evaluate_checkpoint(path, not args.no_test_set, args.max_length)

    if args.per_class:
        for name, r in results.items():
            if "test" in r:
                print(f"\n=== {name} — per-class (in-domain test) ===")
                print(r["test"]["report"])

    print("\n" + markdown_table(results))

    if args.json:
        out = Path(args.json)
        out.parent.mkdir(parents=True, exist_ok=True)
        serializable = {
            n: {k: v for k, v in r.items() if k != "test"} | (
                {"test": {k: v for k, v in r["test"].items() if k != "report"}}
                if "test" in r else {}
            )
            for n, r in results.items()
        }
        out.write_text(json.dumps(serializable, indent=2, default=float))
        print(f"JSON  -> {out}")

    if args.markdown:
        out = Path(args.markdown)
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(markdown_table(results))
        print(f"Table -> {out}")


if __name__ == "__main__":
    main()