"""Stage 4: evaluate Jedi as a verifier and as the gate of an escalation ladder.

Jedi's confidence is P(accept) read from its first output token (accept vs escalate), Platt-calibrated on the val
split. Reports, on held-out data:
  * verifier quality per split / agent: AUROC, false-accept rate (wrong answers it lets through), ECE, Brier;
  * the ladder glm-5.3-flash -> qwen3.8-27b -> nemotron-ultra, with Jedi checking rungs 1 and 2 and the top rung's
    answer always accepted, swept over the confidence threshold (accuracy, relative cost, silent-failure rate);
  * baselines on exactly the same tasks: every single agent, Shaktimaan v1.0 routing, a format check (escalate only
    truncated / unparseable answers), self-consistency (escalate when two cheap samples disagree) and the oracle
    ladder (a perfect verifier);
  * overhead: Jedi's own latency per check and the extra agent latency on escalated tasks.
Answers for glm / qwen are fresh samples (sample 0). Nemotron-ultra's outcome and token count come from the graded
record in the Shaktimaan-routing dataset (same task ids), as do agents Shaktimaan routes to that were not re-run.
Relative cost = sum over called agents of (agent price x completion tokens), divided by always-nemotron-ultra.
"""
import json, sys, time
from collections import defaultdict
import numpy as np
import torch
from datasets import load_dataset, load_from_disk
from transformers import AutoModelForCausalLM, AutoTokenizer
from common import LADDER, POOL, JEDI_SYSTEM, STRONGEST

path = sys.argv[1] if len(sys.argv) > 1 else "jedi-0.5b"
tok = AutoTokenizer.from_pretrained(path)
tok.padding_side = "left"
model = AutoModelForCausalLM.from_pretrained(path, dtype=torch.bfloat16).cuda().eval()
ds = load_from_disk("data/jedi_judge")
ACC, ESC = (tok.encode(w, add_special_tokens=False)[0] for w in ("accept", "escalate"))
assert ACC != ESC
CHEAP, MID = LADDER[0], LADDER[1]


def logit_gap(views, bs=16):
    """accept-vs-escalate logit gap at Jedi's first output token, for each judge input."""
    out = []
    order = np.argsort([len(v) for v in views])  # length-sorted batches: far less padding
    for i in range(0, len(views), bs):
        idx = order[i:i + bs]
        msgs = [[{"role": "system", "content": JEDI_SYSTEM}, {"role": "user", "content": views[j]}] for j in idx]
        enc = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt", padding=True,
                                      return_dict=True).to("cuda")
        with torch.no_grad():
            last = model(**enc).logits[:, -1, :].float()
        out += list(zip(idx, (last[:, ACC] - last[:, ESC]).tolist()))
    gap = np.zeros(len(views))
    for j, g in out:
        gap[j] = g
    return gap


def fit_platt(z, y):
    a, b = torch.ones(1, requires_grad=True), torch.zeros(1, requires_grad=True)
    zt, yt = torch.tensor(z, dtype=torch.float32), torch.tensor(y, dtype=torch.float32)
    opt = torch.optim.LBFGS([a, b], max_iter=200)

    def closure():
        opt.zero_grad()
        loss = torch.nn.functional.binary_cross_entropy_with_logits(a * zt + b, yt)
        loss.backward()
        return loss
    opt.step(closure)
    return a.item(), b.item()


def sigmoid(x):
    return 1 / (1 + np.exp(-x))


def auroc(p, y):
    p, y = np.asarray(p), np.asarray(y, bool)
    if y.all() or not y.any():
        return None
    ranks = np.argsort(np.argsort(p)) + 1
    return float((ranks[y].sum() - y.sum() * (y.sum() + 1) / 2) / (y.sum() * (~y).sum()))


def ece(p, y, bins=10):
    p, y = np.asarray(p), np.asarray(y, float)
    edges = np.linspace(0, 1, bins + 1)
    total, table = 0.0, []
    for lo, hi in zip(edges[:-1], edges[1:]):
        m = (p >= lo) & ((p < hi) | (hi == 1))
        if m.any():
            total += m.mean() * abs(p[m].mean() - y[m].mean())
            table.append(dict(bin=f"{lo:.1f}-{hi:.1f}", n=int(m.sum()), confidence=round(float(p[m].mean()), 3),
                              accuracy=round(float(y[m].mean()), 3)))
    return float(total), table


def verifier_metrics(p, y, tau=0.5):
    p, y = np.asarray(p), np.asarray(y, bool)
    e, _ = ece(p, y)
    return dict(n=int(len(y)), answers_wrong=round(float(1 - y.mean()), 4), auroc=round(auroc(p, y) or 0, 4),
                false_accept_rate=round(float(((p >= tau) & ~y).sum() / max(1, (~y).sum())), 4),
                false_escalate_rate=round(float(((p < tau) & y).sum() / max(1, y.sum())), 4),
                ece=round(e, 4), brier=round(float(((p - y) ** 2).mean()), 4))


# ---------- score every held-out answer ----------
z_val = logit_gap(ds["val"]["judge_input"])
a, b = fit_platt(z_val, ds["val"]["correct"])
cal = dict(a=a, b=b, accept_token=ACC, escalate_token=ESC)
P = {}
report = {"calibration": cal, "verifier": {}}
for split in ("val", "test", "ood"):
    d = ds[split]
    z = z_val if split == "val" else logit_gap(d["judge_input"])
    p_raw, p = sigmoid(z), sigmoid(a * z + b)
    y = np.array(d["correct"])
    rep = {"all": verifier_metrics(p, y), "uncalibrated_ece": round(ece(p_raw, y)[0], 4)}
    for ag in sorted(set(d["agent"])):
        m = np.array(d["agent"]) == ag
        rep[ag] = verifier_metrics(p[m], y[m])
    for src in sorted(set(d["source"])):
        m = np.array(d["source"]) == src
        rep[f"source:{src}"] = verifier_metrics(p[m], y[m])
    rep["reliability"] = ece(p, y)[1]
    report["verifier"][split] = rep
    for r, pi in zip(d.select_columns(["task_id", "agent", "sample", "correct", "completion_tokens", "latency",
                                       "finish", "pred", "kind"]), p):
        P[r["task_id"], r["agent"], r["sample"]] = dict(r, p=float(pi), split=split)
    print(split, json.dumps(rep["all"]), flush=True)

# ---------- ladder ----------
shak = {r["id"]: r for sp in ("train", "test") for r in load_dataset("podstack/Shaktimaan-routing", split=sp)}


def norm_pred(s):
    return (s or "").strip().strip(".").lower()


def format_ok(ans):
    return ans["finish"] != "length" and bool(ans["pred"].strip())


def run_ladder(tasks, rungs, accept, final_record=None):
    """tasks: task ids. rungs: agent names; accept(task_id, agent) -> (bool, extra_cost). The last rung is always
    accepted; if final_record is given it supplies the last rung's (correct, tokens)."""
    correct, cost, silent, calls, extra_latency = 0, 0.0, 0, 0, 0.0
    for t in tasks:
        for i, ag in enumerate(rungs):
            last = i == len(rungs) - 1
            if last and final_record:
                c, toks, lat = final_record(t)
            else:
                ans = P[t, ag, 0]
                c, toks, lat = ans["correct"], ans["completion_tokens"], ans["latency"]
            cost += POOL[ag]["cost"] * toks
            calls += 1
            if i:
                extra_latency += lat
            if last:
                correct += c
                break
            ok, extra = accept(t, ag)
            cost += extra
            if ok:
                correct += c
                silent += not c
                break
    n = len(tasks)
    return dict(accuracy=correct / n, cost=cost, silent_failure_rate=silent / n, calls_per_task=calls / n,
                extra_agent_latency_s=extra_latency / n)


def ultra_record(t):
    r = json.loads(shak[t]["agent_results"])[STRONGEST]
    return r["correct"], r["completion_tokens"], r["latency"]


def jedi(tau):
    return lambda t, ag: (P[t, ag, 0]["p"] >= tau, 0.0)


def fmt(t, ag):
    return format_ok(P[t, ag, 0]), 0.0


def self_consistency(t, ag):
    s0, s1 = P[t, ag, 0], P.get((t, ag, 1))
    if s1 is None or s0["kind"] == "code":
        return format_ok(s0), 0.0
    return format_ok(s0) and norm_pred(s0["pred"]) == norm_pred(s1["pred"]), POOL[ag]["cost"] * s1["completion_tokens"]


def oracle(t, ag):
    return P[t, ag, 0]["correct"], 0.0


def ladder_tasks(split, sources):
    return sorted({t for (t, _, _), v in P.items() if v["split"] == split and t.split("-")[0] in sources
                   and (t, CHEAP, 0) in P and (t, MID, 0) in P})


test_tasks = [t for t in ladder_tasks("test", {"gsm8k", "math", "mmlupro", "arc", "mbpp", "humaneval"}) if t in shak]
ultra_cost = sum(POOL[STRONGEST]["cost"] * ultra_record(t)[1] for t in test_tasks)


def rel(res):
    return {"accuracy": round(res["accuracy"], 4), "relative_cost": round(res["cost"] / ultra_cost, 4),
            "silent_failure_rate": round(res["silent_failure_rate"], 4),
            "calls_per_task": round(res["calls_per_task"], 3),
            "extra_agent_latency_s": round(res["extra_agent_latency_s"], 2)}


policies = {
    f"always:{CHEAP}": rel(run_ladder(test_tasks, [CHEAP], None)),
    f"always:{MID}": rel(run_ladder(test_tasks, [MID], None)),
    f"always:{STRONGEST}": rel(run_ladder(test_tasks, [STRONGEST], None, ultra_record)),
    "ladder:format-check": rel(run_ladder(test_tasks, LADDER, fmt, ultra_record)),
    "ladder:self-consistency": rel(run_ladder(test_tasks, LADDER, self_consistency, ultra_record)),
    "ladder:oracle": rel(run_ladder(test_tasks, LADDER, oracle, ultra_record)),
}

# Shaktimaan v1.0 on the same tasks: route with the published router, then score its choice.
rtok = AutoTokenizer.from_pretrained("podstack/Shaktimaan")
rtok.padding_side = "left"
router = AutoModelForCausalLM.from_pretrained("podstack/Shaktimaan", dtype=torch.bfloat16).cuda().eval()
routes = []
for i in range(0, len(test_tasks), 32):
    msgs = [shak[t]["messages"][:2] for t in test_tasks[i:i + 32]]
    enc = rtok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt", padding=True,
                                   return_dict=True).to("cuda")
    with torch.no_grad():
        out = router.generate(**enc, max_new_tokens=96, do_sample=False)
    routes += rtok.batch_decode(out[:, enc["input_ids"].shape[1]:], skip_special_tokens=True)
del router
s_correct, s_cost = 0, 0.0
for t, rt in zip(test_tasks, routes):
    try:
        agents = [x for x in json.loads(rt)["agents"] if x in POOL] or [STRONGEST]
    except Exception:  # noqa: BLE001
        agents = [STRONGEST]
    rec = json.loads(shak[t]["agent_results"])
    for ag in agents:  # fresh sample where one exists, otherwise the recorded run
        fresh = P.get((t, ag, 0))
        s_cost += POOL[ag]["cost"] * (fresh["completion_tokens"] if fresh else rec[ag]["completion_tokens"])
    fin = P.get((t, agents[-1], 0))  # ensembles are scored as the aggregating agent's own answer (conservative)
    s_correct += fin["correct"] if fin else rec[agents[-1]]["correct"]
policies["shaktimaan-v1.0"] = {"accuracy": round(s_correct / len(test_tasks), 4),
                               "relative_cost": round(s_cost / ultra_cost, 4)}

curve = []
for tau in [0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.85, 0.9, 0.95, 0.98]:
    curve.append(dict(threshold=tau, **rel(run_ladder(test_tasks, LADDER, jedi(tau), ultra_record))))

# two-rung ladder (no Ultra record) on the new domains and on unseen BIG-Bench Hard
two_rung = {}
for name, split, srcs in [("new-domains (aqua, csqa)", "test", {"aqua", "csqa"}), ("ood: bbh", "ood", {"bbh"})]:
    ids = ladder_tasks(split, srcs)
    base = sum(POOL[MID]["cost"] * P[t, MID, 0]["completion_tokens"] for t in ids)
    r = lambda res: {"accuracy": round(res["accuracy"], 4), "cost_vs_always_mid": round(res["cost"] / base, 4),
                     "silent_failure_rate": round(res["silent_failure_rate"], 4)}
    two_rung[name] = {"n_tasks": len(ids),
                      f"always:{CHEAP}": r(run_ladder(ids, [CHEAP], None)),
                      f"always:{MID}": r(run_ladder(ids, [MID], None)),
                      "oracle": r(run_ladder(ids, [CHEAP, MID], oracle)),
                      "self-consistency": r(run_ladder(ids, [CHEAP, MID], self_consistency)),
                      **{f"jedi@{tau}": r(run_ladder(ids, [CHEAP, MID], jedi(tau))) for tau in (0.5, 0.7, 0.9)}}

# overhead: one check at a time, as a live ladder would issue it
views = ds["test"]["judge_input"][:200]
torch.cuda.synchronize()
t0 = time.time()
for v in views:
    logit_gap([v], bs=1)
torch.cuda.synchronize()
n_tok = np.mean([len(tok(tok.apply_chat_template([{"role": "system", "content": JEDI_SYSTEM},
                                                  {"role": "user", "content": v}], add_generation_prompt=True,
                                                 tokenize=False)).input_ids) for v in views])

report.update({"ladder": LADDER, "n_ladder_tasks": len(test_tasks), "policies": policies, "jedi_curve": curve,
               "two_rung": two_rung,
               "overhead": {"jedi_ms_per_check_batch1": round((time.time() - t0) / len(views) * 1000, 1),
                            "jedi_input_tokens_mean": round(float(n_tok)), "jedi_output_tokens": 1}})
print(json.dumps({k: v for k, v in report.items() if k != "verifier"}, indent=2))
json.dump(report, open(f"{path}/eval_results.json", "w"), indent=2)
json.dump(cal, open(f"{path}/calibration.json", "w"), indent=2)