"""Ladder (glm -> qwen -> nemotron-ultra) on the 491 held-out tasks, gated by Jedi vs by the surface-feature baseline."""
import json, math, random, re, sys
import numpy as np, torch
import common
from datasets import load_from_disk
from transformers import AutoModelForCausalLM, AutoTokenizer

def roc_auc_score(y, p):
    y, p = np.asarray(y, bool), np.asarray(p, float)
    r = np.argsort(np.argsort(p, kind="mergesort")) + 1
    return float((r[y].sum() - y.sum() * (y.sum() + 1) / 2) / (y.sum() * (~y).sum()))

class LogisticRegression:
    def fit(self, X, y):
        X = torch.tensor(X, dtype=torch.float32); y = torch.tensor(np.asarray(y, float), dtype=torch.float32)
        self.mu, self.sd = X.mean(0), X.std(0) + 1e-6
        X = (X - self.mu) / self.sd
        self.w = torch.zeros(X.shape[1], requires_grad=True); self.b = torch.zeros(1, requires_grad=True)
        opt = torch.optim.LBFGS([self.w, self.b], max_iter=500)
        def closure():
            opt.zero_grad(); l = torch.nn.functional.binary_cross_entropy_with_logits(X @ self.w + self.b, y); l.backward(); return l
        opt.step(closure); return self
    def predict_proba(self, X):
        X = (torch.tensor(X, dtype=torch.float32) - self.mu) / self.sd
        p = torch.sigmoid(X @ self.w + self.b).detach().numpy(); return np.stack([1 - p, p], 1)
random.seed(0)
P = sys.argv[1] if len(sys.argv) > 1 else "jedi-0.5b"
cal = json.load(open(P + "/calibration.json")); tok = AutoTokenizer.from_pretrained(P); tok.padding_side = "left"
model = AutoModelForCausalLM.from_pretrained(P, dtype=torch.bfloat16).cuda().eval()
ds = load_from_disk("data/jedi_judge")

def score(views, bs=16):
    out = np.zeros(len(views)); order = np.argsort([len(v) for v in views])
    for i in range(0, len(views), bs):
        idx = order[i:i + bs]
        msgs = [[{"role": "system", "content": common.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(): l = model(**enc).logits[:, -1].float()
        g = (l[:, cal["accept_token"]] - l[:, cal["escalate_token"]]).cpu().numpy()
        out[idx] = 1 / (1 + np.exp(-(cal["a"] * g + cal["b"])))
    return out

def feats(d):
    ag = sorted(set(ds["train"]["agent"])); src = sorted(set(ds["train"]["source"]) | set(ds["ood"]["source"]))
    X = []
    for r in d:
        X.append([math.log1p(r["completion_tokens"]), r["finish"] == "length", not r["pred"].strip()]
                 + [r["agent"] == a for a in ag] + [r["source"] == s for s in src])
    return np.array(X, float)

base = ds["train"].filter(lambda r: r["sample"] == 0)  # one sample per (task, agent): no oversampling dupes in fit
lr = LogisticRegression().fit(feats(base), base["correct"])

from datasets import load_dataset
from common import LADDER, POOL, STRONGEST
CHEAP, MID = LADDER[0], LADDER[1]
d = ds["test"]
pj = score(d["judge_input"]); ps = lr.predict_proba(feats(d))[:, 1]
R = {}
for r, a, b in zip(d.select_columns(["task_id", "agent", "sample", "correct", "completion_tokens"]), pj, ps):
    R[r["task_id"], r["agent"], r["sample"]] = dict(r, pj=float(a), ps=float(b))
shak = {r["id"]: r for sp in ("train", "test") for r in load_dataset("podstack/Shaktimaan-routing", split=sp)}
srcs = {"gsm8k", "math", "mmlupro", "arc", "mbpp", "humaneval"}
tasks = sorted({t for (t, _, _) in R if t.split("-")[0] in srcs and (t, CHEAP, 0) in R and (t, MID, 0) in R and t in shak})
def ultra(t):
    x = json.loads(shak[t]["agent_results"])[STRONGEST]; return x["correct"], x["completion_tokens"]
ucost = sum(POOL[STRONGEST]["cost"] * ultra(t)[1] for t in tasks)
def ladder(key, tau):
    c = cost = 0
    for t in tasks:
        for i, ag in enumerate(LADDER):
            if i == 2:
                ok, tk = ultra(t); cost += POOL[ag]["cost"] * tk; c += ok; break
            a = R[t, ag, 0]; cost += POOL[ag]["cost"] * a["completion_tokens"]
            if a[key] >= tau: c += a["correct"]; break
    return round(c / len(tasks), 4), round(cost / ucost, 4)
print("n", len(tasks))
for tau in (0.05, 0.1, 0.2, 0.3, 0.4, 0.45, 0.5):
    print(tau, "jedi", ladder("pj", tau), "shallow", ladder("ps", tau))