"""Does Jedi read the answer content, or surface cues? (1) shallow-feature baselines (2) final-answer swap test."""
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"])
res = {}
for sp in ("test", "ood"):
    d = ds[sp]; y = np.array(d["correct"])
    p = score(d["judge_input"])
    trunc = np.array([f == "length" for f in d["finish"]])
    res[sp] = dict(
        jedi=roc_auc_score(y, p),
        shallow_logreg=roc_auc_score(y, lr.predict_proba(feats(d))[:, 1]),
        length_only=roc_auc_score(y, -np.array(d["completion_tokens"])),
        jedi_not_truncated=roc_auc_score(y[~trunc], p[~trunc]),
        truncated_share=float(trunc.mean()))
    # within-task: tasks with both a right and a wrong answer -> does Jedi rank the right one higher?
    by = {}
    for i, r in enumerate(d):
        by.setdefault(r["task_id"], []).append(i)
    wins = tot = 0
    for ids in by.values():
        for i in ids:
            for j in ids:
                if y[i] and not y[j]:
                    tot += 1; wins += (p[i] > p[j]) + 0.5 * (p[i] == p[j])
    res[sp]["within_task_pairwise_acc"] = wins / tot; res[sp]["within_task_pairs"] = tot
    if sp == "test":
        # swap test: replace the 'Final answer:' value. wrong -> gold, right -> a wrong value.
        orig, swapped, ys = [], [], []
        for r in d:
            if r["kind"] not in ("mc", "num") or r["finish"] == "length": continue
            m = list(re.finditer(r"(Final answer\s*[:：]\s*)(.+)", r["judge_input"], flags=re.I))
            if not m: continue
            m = m[-1]; gold = r["gold"].strip()
            if r["correct"]:
                new = random.choice([c for c in "ABCDEFGHIJ"[:5] if c != gold.upper()]) if r["kind"] == "mc" else str(int(float(gold.replace(',', ''))) + 7) if re.fullmatch(r"-?[\d,]+(\.0+)?", gold) else None
            else:
                new = gold
            if new is None: continue
            v = r["judge_input"][:m.start(2)] + new + r["judge_input"][m.end(2):]
            orig.append(r["judge_input"]); swapped.append(v); ys.append(r["correct"])
        po, ps = score(orig), score(swapped); ys = np.array(ys)
        res["swap"] = dict(n=len(ys), n_right_made_wrong=int(ys.sum()), n_wrong_made_right=int((~ys).sum()),
            right_orig=float(po[ys].mean()), right_after_swap_to_wrong=float(ps[ys].mean()),
            wrong_orig=float(po[~ys].mean()), wrong_after_swap_to_gold=float(ps[~ys].mean()),
            auroc_on_swapped=roc_auc_score(~ys, ps))  # label of swapped = flipped correctness
print(json.dumps(res, indent=1))