"""Stasis 1 step 12,269 vs GPT-2 small matched eval. Implements METRIC_DEFINITIONS.json v1
(sha256 579ab8c12aeb2cb01ed6f8d32028f97065139217f1a7e702bf0d3ff1afa24137). Offline; no hub, no AutoModel."""
import hashlib, json, math, os, sys, time
import numpy as np, torch
import torch.nn.functional as F

METRIC_SHA = "579ab8c12aeb2cb01ed6f8d32028f97065139217f1a7e702bf0d3ff1afa24137"
CKPT_SHAS = {"2139ca7dd8b0c21e7c4b3adce14970fceadb7e542aa61c222e375b5a13b2656b": "endpoint.model_only.pt",
             "44e66daf5e0eee079711ac96d2bdd10fff7751a0573d56eb5c7b422f958e9645": "recovery-final.pt"}
GPT2_SAFETENSORS_SHA = None  # filled from SHA256SUMS
BENCHES = ["WT2_FULL_BPB", "WT2_PANEL256_CONTINUITY", "LAMBADA_OPENAI", "HELLASWAG", "PIQA", "ARC_EASY"]
LN2 = math.log(2)

def need(ok, msg):
    if not ok: raise RuntimeError("GATE FAILED: " + msg)

def file_sha(p):
    h = hashlib.sha256()
    with open(p, "rb") as f:
        for b in iter(lambda: f.read(1 << 24), b""): h.update(b)
    return h.hexdigest()

def verify_package(root):
    need(file_sha(os.path.join(root, "METRIC_DEFINITIONS.json")) == METRIC_SHA, "metric definitions sha (frozen)")
    sums = {}
    for line in open(os.path.join(root, "SHA256SUMS")):
        s, p = line.strip().split("  ", 1); sums[p] = s
        need(file_sha(os.path.join(root, p)) == s, "package file sha " + p)
    return sums

# ---------------- models ----------------
class GPT2(torch.nn.Module):
    """Minimal exact GPT-2 small forward (gelu_new, LN eps 1e-5, tied head) over bundled safetensors."""
    def __init__(self, sd):
        super().__init__()
        self.p = torch.nn.ParameterDict({k.replace(".", "__"): torch.nn.Parameter(v, requires_grad=False) for k, v in sd.items()})
    def g(self, k): return self.p[k.replace(".", "__")]
    def forward(self, ids):
        b, t = ids.shape
        x = self.g("wte.weight")[ids] + self.g("wpe.weight")[:t]
        for i in range(12):
            h = F.layer_norm(x, (768,), self.g(f"h.{i}.ln_1.weight"), self.g(f"h.{i}.ln_1.bias"), 1e-5)
            qkv = h @ self.g(f"h.{i}.attn.c_attn.weight") + self.g(f"h.{i}.attn.c_attn.bias")
            q, k, v = qkv.view(b, t, 3, 12, 64).permute(2, 0, 3, 1, 4).unbind(0)
            y = F.scaled_dot_product_attention(q, k, v, is_causal=True).transpose(1, 2).reshape(b, t, 768)
            x = x + y @ self.g(f"h.{i}.attn.c_proj.weight") + self.g(f"h.{i}.attn.c_proj.bias")
            h = F.layer_norm(x, (768,), self.g(f"h.{i}.ln_2.weight"), self.g(f"h.{i}.ln_2.bias"), 1e-5)
            h = F.gelu(h @ self.g(f"h.{i}.mlp.c_fc.weight") + self.g(f"h.{i}.mlp.c_fc.bias"), approximate="tanh")
            x = x + h @ self.g(f"h.{i}.mlp.c_proj.weight") + self.g(f"h.{i}.mlp.c_proj.bias")
        x = F.layer_norm(x, (768,), self.g("ln_f.weight"), self.g("ln_f.bias"), 1e-5)
        return x @ self.g("wte.weight").T

def read_safetensors(path):
    """Dependency-free safetensors reader (F32 only), so no safetensors/transformers install is needed."""
    import struct
    raw = open(path, "rb").read(); n = struct.unpack("<Q", raw[:8])[0]; hdr = json.loads(raw[8:8 + n]); base = 8 + n; out = {}
    for k, v in hdr.items():
        if k == "__metadata__": continue
        need(v["dtype"] == "F32", "F32 tensor " + k); a, b = v["data_offsets"]
        out[k] = torch.from_numpy(np.frombuffer(raw, dtype=np.float32, count=(b - a) // 4, offset=base + a).copy()).view(v["shape"])
    return out

def load_gpt2(root, device):
    sd = read_safetensors(os.path.join(root, "gpt2", "model.safetensors"))
    keys = {"wte.weight", "wpe.weight", "ln_f.weight", "ln_f.bias"}
    for i in range(12):
        for m in ["ln_1", "attn.c_attn", "attn.c_proj", "ln_2", "mlp.c_fc", "mlp.c_proj"]:
            keys |= {f"h.{i}.{m}.weight", f"h.{i}.{m}.bias"}
    masks = {k for k in sd if k.endswith(".attn.bias") or k.endswith(".attn.masked_bias")}
    need(set(sd) - masks == keys, "exact GPT-2 tensor inventory (strict)")
    sd = {k: sd[k].float() for k in keys}
    need(sum(v.numel() for v in sd.values()) == 124439808, "GPT-2 param count")
    need(all(bool(torch.isfinite(v).all()) for v in sd.values()), "finite GPT-2")
    return GPT2(sd).to(device).eval()

class StasisWrap(torch.nn.Module):
    def __init__(self, m): super().__init__(); self.m = m
    def forward(self, ids): return self.m(ids)[0]

def load_stasis(root, ckpt_path, device, ckpt_sha=None):
    sys.path.insert(0, os.path.join(root, "code"))
    from frozen_stasis.config import ModelConfig
    from frozen_stasis.model import StasisDecoder
    sha = ckpt_sha or file_sha(ckpt_path)
    need(sha in CKPT_SHAS, "step-12269 checkpoint sha256 (got %s)" % sha)
    payload = torch.load(ckpt_path, map_location="cpu", weights_only=True, mmap=True)
    need(payload["cursor"]["optimizer_step"] == 12269, "optimizer_step 12269")
    sd = payload["model"]
    m = StasisDecoder(ModelConfig())
    need(list(sd) == list(m.state_dict()), "Stasis key set/order")
    need(all(v.dtype == torch.float32 and bool(torch.isfinite(v).all()) for v in sd.values()), "finite fp32 Stasis")
    m.load_state_dict(sd, strict=True)
    need(sum(p.numel() for p in m.parameters()) == 109529856, "Stasis param count")
    del payload, sd
    return StasisWrap(m).to(device).eval(), sha

# ---------------- scoring ----------------
def score_seqs(model, seqs, device, tok_budget):
    """seqs: list of (ctx_ids_with_bos, tgt_ids). Returns list of (target logprobs np.float64, greedy_all_match)."""
    order = sorted(range(len(seqs)), key=lambda i: -(len(seqs[i][0]) + len(seqs[i][1])))
    out = [None] * len(seqs); i = 0
    while i < len(order):
        L = len(seqs[order[i]][0]) + len(seqs[order[i]][1]) - 1
        need(L <= 1024, "sequence > 1024 tokens")
        nb = max(1, tok_budget // L); batch = order[i:i + nb]; i += nb
        x = torch.zeros((len(batch), L), dtype=torch.long)
        for r, j in enumerate(batch):
            full = seqs[j][0] + seqs[j][1]; x[r, :len(full) - 1] = torch.tensor(full[:-1])
        with torch.inference_mode():
            logits = model(x.to(device))
            for r, j in enumerate(batch):
                c, t = seqs[j]; s = len(c) - 1
                lg = logits[r, s:s + len(t)].float()
                need(bool(torch.isfinite(lg).all()), "finite logits")
                lp = torch.log_softmax(lg, -1)
                tt = torch.tensor(t, device=lg.device)
                out[j] = (lp.gather(1, tt[:, None])[:, 0].double().cpu().numpy(), bool((lg.argmax(-1) == tt).all()))
        del logits
    return out

def load_items(root, bench, limit=None):
    rows = [json.loads(l) for l in open(os.path.join(root, "data", bench + ".jsonl"))]
    if limit:
        if bench.startswith("WT2"):
            units = list(dict.fromkeys(r["unit"] for r in rows))[:limit]; rows = [r for r in rows if r["unit"] in set(units)]
        else: rows = rows[:limit]
    return rows

def run_model(name, model, root, device, tok_budget, limit, outdir, log):
    bos = 2 if name == "STASIS_12269" else 50256
    res = {}
    for bench in BENCHES:
        t0 = time.time(); rows = load_items(root, bench, limit)
        if bench in ("HELLASWAG", "PIQA", "ARC_EASY"):
            seqs = [([bos] + c["ids"][name]["ctx"], c["ids"][name]["tgt"]) for r in rows for c in r["choices"]]
            sc = score_seqs(model, seqs, device, tok_budget); k = 0; per = []
            for r in rows:
                lls, nb = [], []
                for c in r["choices"]: lls.append(float(sc[k][0].sum())); nb.append(c["tgt_bytes"]); k += 1
                a = int(np.argmax(lls)); an = int(np.argmax(np.array(lls) / np.array(nb)))
                per.append({"unit": r["unit"], "gold": r["gold"], "ll": lls, "acc": int(a == r["gold"]), "acc_norm": int(an == r["gold"])})
        else:
            live = [r for r in rows if r["ids"] is not None]
            sc = score_seqs(model, [([bos] + r["ids"][name]["ctx"], r["ids"][name]["tgt"]) for r in live], device, tok_budget)
            per = []
            for r, (lp, g) in zip(live, sc):
                per.append({"unit": r["unit"], "start": r.get("start", 0), "bytes": r["tgt_bytes"], "nll": float(-lp.sum()), "tokens": len(lp), "greedy": int(g)})
            per += [{"unit": r["unit"], "start": 0, "bytes": 0, "nll": 0.0, "tokens": 0, "greedy": 0} for r in rows if r["ids"] is None]
        with open(os.path.join(outdir, f"ROWS_{bench}_{name}.jsonl"), "w") as f:
            for p in per: f.write(json.dumps(p) + "\n")
        res[bench] = per; log(f"{name} {bench}: {len(rows)} items in {time.time()-t0:.1f}s")
    return res

# ---------------- stats ----------------
def boot_idx(n, draws=10000, seed=20261016, chunk=500):
    rng = np.random.default_rng(seed)
    for s in range(0, draws, chunk): yield rng.integers(0, n, (min(chunk, draws - s), n))

def bpb_units(per):
    u = {}
    for p in per:
        a = u.setdefault(p["unit"], [0.0, 0]); a[0] += p["nll"]; a[1] += p["bytes"]
    return u

def compare_bpb(S, G):
    us, ug = bpb_units(S), bpb_units(G); keys = sorted(us)
    need(keys == sorted(ug) and all(us[k][1] == ug[k][1] for k in keys), "paired identical bytes")
    ns = np.array([us[k][0] for k in keys]); ng = np.array([ug[k][0] for k in keys]); b = np.array([us[k][1] for k in keys], float)
    bs, bg = ns.sum() / (LN2 * b.sum()), ng.sum() / (LN2 * b.sum())
    d, rel = [], []
    for idx in boot_idx(len(keys)):
        B = b[idx].sum(1); x = ns[idx].sum(1) / (LN2 * B); y = ng[idx].sum(1) / (LN2 * B)
        d.append(x - y); rel.append(100 * (x / y - 1))
    d, rel = np.concatenate(d), np.concatenate(rel)
    ci = [float(np.percentile(d, 2.5)), float(np.percentile(d, 97.5))]
    return {"STASIS_12269_bpb": bs, "GPT2_SMALL_bpb": bg, "delta_bpb": bs - bg, "delta_ci95": ci,
            "relative_deficit_pct": 100 * (bs / bg - 1), "relative_deficit_ci95": [float(np.percentile(rel, 2.5)), float(np.percentile(rel, 97.5))],
            "bytes": int(b.sum()), "units": len(keys), "stasis_not_worse": ci[1] < 0}

def compare_acc(S, G, key):
    need([p["unit"] for p in S] == [p["unit"] for p in G], "paired items")
    a = np.array([p[key] for p in S], float); g = np.array([p[key] for p in G], float); dd = a - g
    d = np.concatenate([dd[idx].mean(1) for idx in boot_idx(len(dd))])
    ci = [float(np.percentile(d, 2.5)), float(np.percentile(d, 97.5))]
    return {"STASIS_12269": float(a.mean()), "GPT2_SMALL": float(g.mean()), "delta": float(dd.mean()), "delta_ci95": ci, "n": len(dd), "stasis_not_worse": ci[0] > 0}

def summarize(rs, rg):
    out = {}
    for bench in ["WT2_FULL_BPB", "WT2_PANEL256_CONTINUITY"]: out[bench] = {"bpb": compare_bpb(rs[bench], rg[bench])}
    out["LAMBADA_OPENAI"] = {"acc": compare_acc(rs["LAMBADA_OPENAI"], rg["LAMBADA_OPENAI"], "greedy"),
                             "bpb": compare_bpb(rs["LAMBADA_OPENAI"], rg["LAMBADA_OPENAI"])}
    for bench in ["HELLASWAG", "PIQA", "ARC_EASY"]:
        out[bench] = {"acc_norm": compare_acc(rs[bench], rg[bench], "acc_norm"), "acc": compare_acc(rs[bench], rg[bench], "acc")}
    return out