"""Evaluate laya checkpoints on the official typed-decisions test split (400 cases / 2,000 decisions).

Metrics follow the benchmark card exactly (accuracy / soft acc / Brier / ECE / score MAE /
KL / TV / within-1-level / latency) and add routing-quality numbers (AUROC, AURC) that the
confidence-gated workflows need. Reference numbers (Jev 1.13.0, laya-typed-decisions published,
majority, teacher ceiling) are printed alongside.

Run:
  D:\\venvs\\voxcpm-cuda\\Scripts\\python.exe d:\\QoderCN\\jev_finetune\\eval_td.py --models ours,laya-td
Options: --models comma list of {ours, laya-td, laya-root, dir:<path>}; --max-cases N for smoke.
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time

os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")
os.environ.setdefault("USE_TF", "0")
sys.stdout.reconfigure(encoding="utf-8")

WORK = os.path.dirname(os.path.abspath(__file__))
TEST_PARQUET = os.path.join(WORK, "td_all_test.parquet")
OURS_DIR = os.path.join(WORK, "output", "model")

PUBLISHED = {
    "TypeSafe Jev 1.13.0 (published)": {"acc": 0.727, "soft": 0.580, "brier": 0.148, "ece": 0.144, "mae": 0.391, "ms": 710},
    "laya-typed-decisions (published)": {"acc": 0.766, "soft": 0.471, "brier": 0.062, "ece": 0.213, "mae": 0.242, "ms": None},
    "teacher self-agreement ceiling": {"acc": 0.735, "soft": None, "brier": None, "ece": None, "mae": None, "ms": None},
    "majority class": {"acc": 0.461, "soft": None, "brier": None, "ece": None, "mae": None, "ms": None},
}


def auroc(scores, labels):
    import numpy as np

    scores, labels = np.asarray(scores), np.asarray(labels)
    pos, neg = labels == 1, labels == 0
    if pos.sum() == 0 or neg.sum() == 0:
        return float("nan")
    order = np.argsort(scores)
    ranks = np.empty(len(scores))
    ranks[order] = np.arange(1, len(scores) + 1)
    s_sorted = scores[order]
    i = 0
    while i < len(s_sorted):
        j = i
        while j + 1 < len(s_sorted) and s_sorted[j + 1] == s_sorted[i]:
            j += 1
        if j > i:
            ranks[order[i:j + 1]] = (i + j + 2) / 2.0
        i = j + 1
    return float((ranks[pos].sum() - pos.sum() * (pos.sum() + 1) / 2) / (pos.sum() * neg.sum()))


def aurc(conf, correct):
    import numpy as np

    conf, correct = np.asarray(conf), np.asarray(correct)
    order = np.argsort(-conf)
    err = 1 - correct[order]
    return float((np.cumsum(err) / np.arange(1, len(err) + 1)).mean())


def auto_decidable(conf, correct, max_err=0.05):
    """Largest share of traffic routed to the top-confidences prefix with running error <= max_err."""
    import numpy as np

    conf, correct = np.asarray(conf), np.asarray(correct)
    order = np.argsort(-conf)
    err = 1 - correct[order]
    run_err = np.cumsum(err) / np.arange(1, len(err) + 1)
    ok = run_err <= max_err
    if not ok.any() or not ok[0]:
        return 0.0
    # prefix must stay under budget: last index where all up to it are True
    bad = np.flatnonzero(~ok)
    n = bad[0] if bad.size else len(err)
    return float(n / len(err))


def load_agent(spec: str):
    import laya

    if spec == "ours":
        return laya.Agent(OURS_DIR, device="cuda")
    if spec.startswith("dir:"):
        return laya.Agent(spec[4:], device="cuda")
    if spec == "laya-td":
        return laya.load("convaiinnovations/laya", subfolder="typed-decisions", device="cuda")
    if spec == "laya-root":
        return laya.load("convaiinnovations/laya", device="cuda")
    raise SystemExit(f"unknown model spec: {spec}")


def score_predictions(predictions, latencies_ms, flips=None):
    import numpy as np
    from laya.common import ece_score

    acc, soft, brier, kl, tv, maes, within1, confs, corr = [], [], [], [], [], [], [], [], []
    per_wf, per_type = {}, {}
    for item in predictions:
        for qid, qdef in item["questions"].items():
            p_ans, g_ans = item["pred"][qid], item["gold"][qid]
            t = qdef["type"]
            is_corr = None
            conf = None
            if t == "choice":
                keys = list(qdef["criteria"].keys())
                is_corr = float(p_ans["choice"] == str(g_ans["label"]))
                p = np.array([p_ans["probabilities"].get(k, 1e-6) for k in keys], dtype=float)
                g = np.array([g_ans["probabilities"].get(k, 1e-6) for k in keys], dtype=float)
                p /= p.sum(); g /= g.sum()
                conf = float(p.max())
                soft.append(float((p * g).sum()))
                brier.append(float(((p - g) ** 2).sum()))
                tv.append(float(0.5 * np.abs(p - g).sum()))
                kl.append(float((g * np.log(np.clip(g / p, 1e-12, 1e4))).sum()))
            elif t == "noul":
                pv = p_ans["noul"]
                gv = g_ans.get("noul", g_ans.get("probabilities", {}).get("true", 0.5))
                gold_label = str(g_ans["label"]).lower()
                is_corr = float(("true" if pv >= 0.5 else "false") == gold_label)
                p = np.array([1 - pv, pv]); g = np.array([1 - gv, gv])
                conf = float(max(pv, 1 - pv))
                soft.append(float((p * g).sum()))
                brier.append(float(((p - g) ** 2).sum()))
                tv.append(float(0.5 * np.abs(p - g).sum()))
                kl.append(float((g * np.log(np.clip(g / p, 1e-12, 1e4))).sum()))
            elif t == "score":
                ps, gs = p_ans["score"], g_ans.get("score", 0.0)
                maes.append(abs(ps - gs))
                within1.append(float(abs(ps - gs) <= 1.0))
                n_levels = len(qdef.get("criteria", []))
                p = np.array([p_ans["probabilities"].get(str(i), 0.0) for i in range(n_levels)], dtype=float)
                if p.sum() > 0:
                    p /= p.sum()
                    p_lvl = int(np.argmax(p))
                    conf = float(p.max())
                else:
                    p_lvl = int(round(ps))
                    conf = 0.5
                g_lvl = int(g_ans.get("label", int(round(gs))))
                is_corr = float(p_lvl == g_lvl)
                g = np.array([1.0 if i == g_lvl else 0.0 for i in range(max(1, len(p)))])
                soft.append(float((p * g).sum()) if len(p) else float("nan"))
                brier.append(float(((p - g) ** 2).sum()))
                tv.append(float(0.5 * np.abs(p - g).sum()))
                kl.append(float((g * np.log(np.clip(g / np.clip(p, 1e-6, None), 1e-12, 1e4))).sum()))
            acc.append(is_corr)
            confs.append(conf)
            corr.append(is_corr)
            per_wf.setdefault(item["workflow"], []).append(is_corr)
            per_type.setdefault(t, []).append(is_corr)

    metrics = {
        "n_decisions": len(acc),
        "accuracy": round(float(np.mean(acc)), 4),
        "soft_accuracy": round(float(np.nanmean(soft)), 4),
        "brier": round(float(np.mean(brier)), 4),
        "kl": round(float(np.mean(kl)), 4),
        "tv": round(float(np.mean(tv)), 4),
        "ece": round(float(ece_score(np.array(confs), np.array(corr))), 4),
        "score_mae": round(float(np.mean(maes)), 4) if maes else None,
        "within_1_level": round(float(np.mean(within1)), 4) if within1 else None,
        "auroc_conf": round(auroc(confs, corr), 4),
        "aurc": round(aurc(confs, corr), 4),
        "auto_decidable_5pct": round(auto_decidable(confs, corr), 4),
        "choice_flip_rate": round(float(np.mean(flips)), 4) if flips else None,
        "latency_p50_ms": round(float(np.percentile(latencies_ms, 50)), 1),
        "latency_p95_ms": round(float(np.percentile(latencies_ms, 95)), 1),
        "per_workflow": {k: round(float(np.mean(v)), 4) for k, v in per_wf.items()},
        "per_type": {k: round(float(np.mean(v)), 4) for k, v in per_type.items()},
    }
    return metrics


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--models", default="ours")
    ap.add_argument("--max-cases", type=int, default=0, help="0 = all 400")
    args = ap.parse_args()

    import pandas as pd

    assert os.path.exists(TEST_PARQUET), "run prep_data.py first"
    df = pd.read_parquet(TEST_PARQUET)
    if args.max_cases:
        df = df.head(args.max_cases)

    report = {"benchmark": "LocalLLaMA/typed-decisions test split", "published": PUBLISHED, "measured": {}}
    for spec in [m.strip() for m in args.models.split(",") if m.strip()]:
        print(f"\n===== evaluating [{spec}] =====")
        agent = load_agent(spec)
        import random as _rnd
        predictions, latencies, flips = [], [], []
        t0 = time.time()
        for i, row in enumerate(df.itertuples(index=False)):
            state = json.loads(row.state)
            questions = json.loads(row.questions)
            gold = json.loads(row.gold)
            if i == 0:
                agent.predict(state, questions)  # warmup, not timed
                continue
            t1 = time.perf_counter()
            res = agent.predict(state, questions)
            latencies.append((time.perf_counter() - t1) * 1000)
            predictions.append({"id": row.id, "workflow": row.workflow,
                                "pred": res["answers"], "gold": gold, "questions": questions})
            # option-order consistency: re-ask choice questions with criteria keys shuffled
            cq = {qid: q for qid, q in questions.items() if q["type"] == "choice"}
            if cq:
                shuffled = {}
                for qid, q in cq.items():
                    keys = list(q["criteria"].keys())
                    _rnd.Random(i * 97 + hash(qid) % 1000).shuffle(keys)
                    q2 = dict(q)
                    q2["criteria"] = {k: q["criteria"][k] for k in keys}
                    shuffled[qid] = q2
                res2 = agent.predict(state, shuffled)
                for qid in cq:
                    c1 = predictions[-1]["pred"][qid].get("choice")
                    c2 = res2["answers"][qid].get("choice")
                    if c1 is not None and c2 is not None:
                        flips.append(float(c1 != c2))
            if (i + 1) % 100 == 0:
                print(f"  {i+1}/{len(df)} cases ({time.time()-t0:.0f}s)")
        metrics = score_predictions(predictions, latencies, flips)
        report["measured"][spec] = metrics
        print(json.dumps(metrics, indent=2, ensure_ascii=False))

    out = os.path.join(WORK, "benchmark_report.json")
    prev = {}
    if os.path.exists(out):
        try:
            prev = json.load(open(out, encoding="utf-8")).get("measured", {})
        except Exception:
            prev = {}
    prev.update(report["measured"])
    report["measured"] = prev
    with open(out, "w", encoding="utf-8") as f:
        json.dump(report, f, indent=2, ensure_ascii=False)
    print(f"\n[OK] report -> {out}")

    print("\n===== SUMMARY (acc / soft / brier / ece / mae / p50ms) =====")
    for name, m in PUBLISHED.items():
        print(f"  {name:38s} {m['acc']}  {m['soft']}  {m['brier']}  {m['ece']}  {m['mae']}  {m['ms']}")
    for name, m in prev.items():
        print(f"  {name:38s} {m['accuracy']}  {m['soft_accuracy']}  {m['brier']}  {m['ece']}  {m['score_mae']}  {m['latency_p50_ms']}")


if __name__ == "__main__":
    main()