#!/usr/bin/env python3
"""Separate DELAY from INTERFERENCE in the Stage-A positive result.

The preregistered mechanistic signature is: if identity-local writes protect a
binding, accuracy should depend on rewrites of the QUERIED binding, not on how
many unrelated bindings were written in between. The pooled slope conflates
the two because n_unrelated correlates with delay. This computes the 2-D
table acc(delay bucket x unrelated-write bucket).
"""
import json
import sys
from pathlib import Path

import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
sys.path.insert(0, str(Path(__file__).resolve().parent))
from pns.common import atomic_write_json, eval_root, shards_root  # noqa
from pns.checkpoint import load_model  # noqa: E402
from pns.model.modules import enum_legal_mask  # noqa: E402
from pns.train.loader import iter_eval_batches  # noqa: E402
from pns.world.schema import Fam  # noqa: E402
from stage_a import to_dev  # noqa: E402

SEM = (int(Fam.SEM_LATEST), int(Fam.SEM_2HOP))


@torch.no_grad()
def rows_for(run, dev, n_lifetimes=400):
    m, _, _ = load_model(run, dev)
    out = []
    for b in iter_eval_batches("e3_dev", shards_root(), 24, n_lifetimes):
        g = to_dev(b, dev)
        B, L = g["etype"].shape
        state = m.initial_state(B, dev)
        wall = torch.zeros(B, L, device=dev)      # cumulative ALL binding writes
        wown = {}                                  # cumulative writes per slot
        cum = torch.zeros(B, L, m.cfg.n_slots, device=dev)
        run_c = torch.zeros(B, m.cfg.n_slots, device=dev)
        acc_all = torch.zeros(B, device=dev)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            for t in range(L):
                state, o = m.step(state, g["tok"][:, t], g["etype"][:, t],
                                  g["dt"][:, t], g["bind_write"][:, t],
                                  g["bind_read"][:, t], g["bind_slot_ent"],
                                  g["bind_slot_attr"])
                w = g["bind_write"][:, t]
                hit = (w >= 0)
                acc_all = acc_all + hit.float()
                wall[:, t] = acc_all
                if hit.any():
                    run_c[torch.arange(B, device=dev), w.clamp(min=0)] += hit.float()
                cum[:, t] = run_c
                sel = torch.isin(g["family"][:, t], torch.tensor(SEM, device=dev))
                if sel.any():
                    legal = enum_legal_mask(g["enum_legal"][sel, t])
                    pred = o["enum"][sel].masked_fill(~legal, -1e9).argmax(-1)
                    ok = (pred == g["enum_gold"][sel, t]).float().cpu().numpy()
                    idx = torch.nonzero(sel).flatten().tolist()
                    for j, i in enumerate(idx):
                        d = int(g["delay"][i, t])
                        e0 = max(0, t - d)
                        rs = int(g["bind_read"][i, t])
                        total = float(wall[i, t] - wall[i, e0])
                        own = float(cum[i, t, rs] - cum[i, e0, rs]) if rs >= 0 else 0.0
                        out.append((float(ok[j]), d, total - own, own))
    return np.array(out) if out else np.zeros((0, 4))


def table(a, name):
    ok, d, unrel, own = a[:, 0], a[:, 1], a[:, 2], a[:, 3]
    dq = [0, 16, 64, 10**9]
    uq = np.quantile(unrel, [0.33, 0.66])
    lines = [f"### {name}", "",
             "| delay \\ unrelated writes | few | mid | many | row slope |",
             "| --- | --- | --- | --- | --- |"]
    res = {}
    for i in range(3):
        dm = (d > dq[i]) & (d <= dq[i + 1])
        cells, vals = [], []
        for lo, hi in ((-1, uq[0]), (uq[0], uq[1]), (uq[1], 1e18)):
            m = dm & (unrel > lo) & (unrel <= hi)
            if m.sum() > 25:
                v = float(ok[m].mean()); vals.append(v)
                cells.append(f"{v:.3f} (n={int(m.sum())})")
            else:
                vals.append(np.nan); cells.append("-")
        sl = (vals[2] - vals[0]) if not np.isnan(vals[2] + vals[0]) else np.nan
        res[f"delay_{dq[i]}_{dq[i+1]}"] = dict(few=vals[0], mid=vals[1],
                                               many=vals[2], slope=sl)
        lines.append(f"| d {dq[i]}-{dq[i+1] if i < 2 else 'inf'} | "
                     + " | ".join(cells) + f" | {sl:+.3f} |")
    # dependence on rewrites of the QUERIED binding, holding delay wide
    lines += ["", "| own-binding rewrites since evidence | acc |", "| --- | --- |"]
    for lo, hi, nm in ((-0.5, 0.5, "0"), (0.5, 1.5, "1"), (1.5, 1e9, "2+")):
        m = (own > lo) & (own <= hi)
        if m.sum() > 25:
            lines.append(f"| {nm} | {float(ok[m].mean()):.3f} (n={int(m.sum())}) |")
            res[f"own_{nm}"] = float(ok[m].mean())
    return lines, res


def main():
    dev = "cuda"
    runs = sys.argv[1:] or ["E3A_bind_s1", "E3A_unbound_s1"]
    rep, lines = {}, ["# Experiment 2 Stage-A: delay vs interference, decomposed", "",
                      "The pooled interference slope conflates delay with "
                      "unrelated writes because the two correlate. This holds "
                      "delay fixed within rows.", ""]
    for run in runs:
        a = rows_for(run, dev)
        L, r = table(a, run)
        lines += L + [""]
        rep[run] = r
        print(run, json.dumps(r, default=str)[:400], flush=True)
    out = Path(__file__).resolve().parent.parent / "results" / "reproduced"
    out.mkdir(parents=True, exist_ok=True)
    (out / "interference.md").write_text("\n".join(lines) + "\n")
    atomic_write_json(eval_root() / "E3_INTERFERENCE.json", rep)
    print("->", out / "interference.md")


if __name__ == "__main__":
    main()