"""Causal load-bearing + generalization report for the SINGLE-token-per-step diffuse organism.

One soft token z_t per CA step carries the WHOLE K-cell row c_t (contrast the K-position markov
organism). Free-running, then intervene on the T-token chain and re-read the boxed digit:
  A. baseline + wrong-class answer decodability (10-way softmax + confusion).
  B. per-step necessity: ablate z_1->prompt; corrupt step t's feedback with noise.
  C. shuffle the step order of the feedback chain (z_2..z_T permuted).
  D. donor patch (replace the chain with a different problem's) -> follows-donor vs follows-prompt.
  E. subspace patch: the soft token lives in U = span(codebook reshaped [K*10, d]); patch U vs U^perp.
  F. probe: per-step WHOLE-ROW decodability (one token holds K digits) across layers + selectivity.
  G. GENERALIZATION: free-running accuracy vs chain length m (trained on train_lengths; test longer).

    python -m latent_threads.eval_single_report --ckpt .../sg1/diffuse/best --n 400
"""
from __future__ import annotations

import argparse
import json
import os
import random

import dotenv

dotenv.load_dotenv()
_USER = os.environ.get("USER", "jbauer")
os.environ.setdefault("HF_HOME", f"/workspace-vast/{_USER}/hf")
os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
from peft import PeftModel
from tqdm.auto import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer

from latent_threads import tasks as LT
from latent_threads.single import build_single_batch, build_single_mask
from latent_threads.probe_markov import probe_all

RES = f"/workspace-vast/{_USER}/exp/latent_threads/results"


def load(ckpt, model_name, device):
    cfg = json.load(open(os.path.join(ckpt, "lt_cfg.json")))
    K = cfg["task_kwargs"]["k"]
    tok = AutoTokenizer.from_pretrained(model_name); tok.pad_token_id = tok.pad_token_id or tok.eos_token_id
    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16,
                                                 attn_implementation="eager", device_map="cuda")
    model = PeftModel.from_pretrained(model, ckpt).eval()
    d = model.config.hidden_size
    extra = torch.load(os.path.join(ckpt, "single_extra.pt"), map_location=device)
    head = torch.nn.Sequential(torch.nn.LayerNorm(d), torch.nn.Linear(d, d // 4), torch.nn.GELU(),
                               torch.nn.Linear(d // 4, K * 10)).to(device, torch.bfloat16)
    head.load_state_dict(extra["head"])
    codebook = extra["codebook"].to(device, torch.bfloat16)                      # [K,10,d]
    q_emb = extra["q_emb"].to(device, torch.bfloat16)
    did = torch.tensor(LT.digit_ids(tok), device=device)
    tag = f"{cfg['batch_id']}_{cfg['task']}_{cfg['tag']}"
    return model, head, codebook, q_emb, tok, did, K, cfg, tag


@torch.no_grad()
def generate(model, head, codebook, q_emb, tok, task, batch, im, device, want_hidden=False,
             ablate_first=False, corrupt_step=None, noise_rng=None):
    """Free-running; returns (E, attn, pos, z_starts, lengths, out). Optionally ablate/corrupt."""
    K, m = task.K, task.m
    input_ids, roles, zs_, ls_, lens_ = build_single_batch(tok, task, batch, im, device, with_answer=False)
    B, Lmax = input_ids.shape
    attn = build_single_mask(roles, zs_, m, ablate_first=ablate_first)
    pos = torch.arange(Lmax, device=device)[None].expand(B, Lmax)
    E = model.get_input_embeddings()(input_ids)
    bidx = torch.arange(B, device=device); zs = torch.tensor(zs_, device=device)
    E = E.clone(); E[bidx, zs] = q_emb
    out = None
    for t in range(m):
        want = want_hidden or t < m - 1
        out = model(inputs_embeds=E, attention_mask=attn, position_ids=pos, output_hidden_states=want)
        if t == m - 1:
            break
        h = out.hidden_states[-1][bidx, zs + t]
        z = torch.softmax(head(h).view(B, K, 10).float(), -1)
        fed = torch.einsum("bkv,kvd->bd", z, codebook.float())
        if corrupt_step is not None and t == corrupt_step:
            fed = torch.randn(fed.shape, generator=noise_rng, dtype=torch.float32).to(device) * float(codebook.float().std())
        E = E.clone(); E[bidx, zs + t + 1] = fed.to(E.dtype)
    return E, attn, pos, zs_, lens_, out


@torch.no_grad()
def read_answer(model, head, E, attn, pos, zs_, batch, m, K, device):
    """Queried-cell answer logits, read from the per-cell HEAD at the final token z_T (the organism's
    single readout path). Returns [B, 10] for each example's queried cell."""
    H = model(inputs_embeds=E, attention_mask=attn, position_ids=pos, output_hidden_states=True).hidden_states[-1]
    out = [head(H[j, zs_[j] + m - 1]).view(K, 10)[batch[j].q] for j in range(len(batch))]
    return torch.stack(out).float()


def get_slots(E, zs_, m):
    return torch.stack([E[j, zs_[j]:zs_[j] + m].clone() for j in range(E.shape[0])])   # [B, m, d]


def put_slots(E, slots, zs_):
    for j in range(E.shape[0]):
        E[j, zs_[j]:zs_[j] + slots.shape[1]] = slots[j]
    return E


def project(x, Q):
    return (x @ Q) @ Q.T


@torch.no_grad()
def free_acc(model, head, codebook, q_emb, tok, task, probs, im, did, device):
    correct, bs = 0, 16
    for s in range(0, len(probs), bs):
        b = probs[s:s + bs]
        E, attn, pos, zs_, lens_, _ = generate(model, head, codebook, q_emb, tok, task, b, im, device)
        pred = read_answer(model, head, E, attn, pos, zs_, b, task.m, task.K, device).argmax(-1).tolist()
        correct += sum(int(pred[j] == task.answer(b[j])) for j in range(len(b)))
    return correct / len(probs)


def run(ckpt, model_name, n, device):
    model, head, codebook, q_emb, tok, did, K, cfg, tag = load(ckpt, model_name, device)
    im = tok.convert_tokens_to_ids("<|im_end|>")
    base_m = cfg["task_kwargs"]["m"]; chance = 1.0 / 10
    task = LT.make_task(cfg["task"], k=K, m=base_m)
    print("=" * 78 + f"\nSINGLE-TOKEN REPORT  {tag}  K={K} m={base_m} chance={chance:.2f}\n" + "=" * 78, flush=True)

    # U = span of the codebook (the soft-token vocabulary), reshaped [K*10, d]
    C30 = codebook.reshape(K * 10, -1).float()
    dc = C30 - C30.mean(0, keepdim=True)
    Uq, Us, _ = torch.linalg.svd(dc.T, full_matrices=False)
    rank = int((Us > 1e-4 * Us[0]).sum())
    Q = Uq[:, :rank].contiguous()
    print(f"  codebook subspace rank={rank} (of {K*10} symbols, D={C30.shape[1]})", flush=True)

    rng = random.Random(2025)
    probs = [task.sample(rng) for _ in range(n)]
    rows = np.array([np.array(task.step_states(p)) for p in probs])               # [N, m, K]

    MODES = ["intact", "shuffle_steps", "donor_full", "donor_subspace", "donor_complement"]
    preds = {mo: [] for mo in MODES}; pmass = {mo: [] for mo in MODES}
    probe_layers = sorted(set(list(range(6, model.config.num_hidden_layers + 1, 6)) + [model.config.num_hidden_layers]))
    feats = {L: [] for L in probe_layers}
    rec_q, don_q = [], []
    gen = torch.Generator().manual_seed(20260617); bs = 16
    for s in tqdm(range(0, n, bs), desc=f"{tag} interventions"):
        b = probs[s:s + bs]; B = len(b)
        E, attn, pos, zs_, lens_, out = generate(model, head, codebook, q_emb, tok, task, b, im, device, want_hidden=True)
        for L in probe_layers:
            hl = out.hidden_states[L]
            feats[L].append(torch.stack([hl[j, zs_[j]:zs_[j] + base_m] for j in range(B)]).to(torch.float16).cpu())
        slots = get_slots(E, zs_, base_m); donor = slots.roll(1, dims=0)
        for j in range(B):
            qj = b[j].q; rec_q.append(b[j].rows[-1][qj]); don_q.append(b[(j - 1) % B].rows[-1][qj])
        fb = slice(1, base_m)                                                     # feedback steps z_2..z_m (z_1=query)
        for mo in MODES:
            sl = slots.clone()
            if mo == "shuffle_steps":
                perm = torch.randperm(base_m - 1, generator=gen) + 1
                sl[:, fb] = slots[:, perm]
            elif mo == "donor_full":
                sl[:, fb] = donor[:, fb]
            elif mo == "donor_subspace":
                sl[:, fb] = slots[:, fb] + project((donor[:, fb] - slots[:, fb]).float(), Q).to(slots.dtype)
            elif mo == "donor_complement":
                delta = (donor[:, fb] - slots[:, fb]).float()
                sl[:, fb] = slots[:, fb] + (delta - project(delta, Q)).to(slots.dtype)
            lg = read_answer(model, head, put_slots(E.clone(), sl, zs_), attn, pos, zs_, b, base_m, K, device)
            p = torch.softmax(lg, -1); preds[mo] += p.argmax(-1).tolist(); pmass[mo].append(p.cpu().numpy())
    preds = {mo: np.array(v) for mo, v in preds.items()}; pmass = {mo: np.concatenate(v) for mo, v in pmass.items()}
    rec_q, don_q = np.array(rec_q), np.array(don_q)

    acc = {mo: float((preds[mo] == rec_q).mean()) for mo in MODES}
    pm = pmass["intact"]; p_correct = float(pm[np.arange(n), rec_q].mean())
    wrong = pm.copy(); wrong[np.arange(n), rec_q] = -1
    margin = float((np.log(pm[np.arange(n), rec_q] + 1e-9) - np.log(wrong.max(1) + 1e-9)).mean())
    conf = np.zeros((10, 10))
    for t, pr in zip(rec_q, preds["intact"]):
        conf[t, pr] += 1
    print(f"  [A] intact acc={acc['intact']:.3f} P(correct)={p_correct:.3f} margin={margin:.2f}nats", flush=True)
    print(f"  [C] shuffle_steps acc={acc['shuffle_steps']:.3f}", flush=True)

    # per-step necessity (ablate + corrupt)
    eprobs = [task.sample(random.Random(13371337)) for _ in range(200)]
    abl = free_acc_intervene(model, head, codebook, q_emb, tok, task, eprobs, im, did, device, ablate_first=True)
    org = free_acc(model, head, codebook, q_emb, tok, task, eprobs, im, did, device)
    nrng = torch.Generator().manual_seed(7)
    corrupt = [free_acc_intervene(model, head, codebook, q_emb, tok, task, eprobs, im, did, device,
                                  corrupt_step=t, noise_rng=nrng) for t in range(base_m - 1)]
    print(f"  [B] organism={org:.3f} ablate={abl:.3f} corrupt(worst)={min(corrupt):.3f}", flush=True)

    dmask = don_q != rec_q
    def follows(mo):
        pr = preds[mo]
        return dict(follow_donor=float((pr[dmask] == don_q[dmask]).mean()),
                    follow_prompt=float((pr[dmask] == rec_q[dmask]).mean()), n=int(dmask.sum()))
    donor_stats = {mo: follows(mo) for mo in ["donor_full", "donor_subspace", "donor_complement"]}
    print(f"  [D/E] donor_full {donor_stats['donor_full']}", flush=True)
    print(f"        donor_subspace {donor_stats['donor_subspace']} donor_complement {donor_stats['donor_complement']}", flush=True)

    # probe: per-step WHOLE-ROW decodability (decode each of K cells from one token), across layers
    feats = {L: torch.cat(v).numpy().astype(np.float32) for L, v in feats.items()}
    cellprobe = np.zeros((len(probe_layers), base_m, K))
    for li, L in enumerate(probe_layers):
        for t in range(base_m):
            cellprobe[li, t] = probe_all(feats[L][:, t, :], rows[:, t, :])
    best_li = int(np.nanmean(cellprobe, (1, 2)).argmax())
    print(f"  [F] probe best layer L{probe_layers[best_li]} whole-row decodability={np.nanmean(cellprobe[best_li]):.3f}", flush=True)

    # generalization: free-running accuracy vs chain length
    gprobs_seed = random.Random(99)
    gen_curve = {}
    for gm in cfg.get("gen_eval_lengths", [2, 3, 4, 5, 6, 7, 8, 10, 12, 14]):
        gtask = LT.make_task(cfg["task"], k=K, m=gm)
        gen_curve[gm] = free_acc(model, head, codebook, q_emb, tok, gtask, [gtask.sample(gprobs_seed) for _ in range(200)], im, did, device)
    print(f"  [G] generalization acc vs m: { {k: round(v,3) for k,v in gen_curve.items()} }", flush=True)

    train_lengths = (list(range(2, base_m + 1)) if cfg.get("curriculum") else cfg.get("train_lengths"))
    result = dict(tag=tag, K=K, base_m=base_m, chance=chance, n=n, subspace_rank=rank,
                  train_lengths=train_lengths, acc=acc, p_correct=p_correct, margin_nats=margin,
                  confusion=conf.tolist(), organism=org, ablate=abl, corrupt_step_acc=corrupt,
                  donor=donor_stats, donor_disagree_n=int(dmask.sum()),
                  probe_layers=probe_layers, cellprobe=cellprobe.tolist(), best_layer=probe_layers[best_li],
                  gen_curve={str(k): v for k, v in gen_curve.items()})
    os.makedirs(RES, exist_ok=True)
    json.dump(result, open(f"{RES}/single_report_{tag}.json", "w"), indent=2)
    make_figures(result, tag)
    return result


@torch.no_grad()
def free_acc_intervene(model, head, codebook, q_emb, tok, task, probs, im, did, device, ablate_first=False, corrupt_step=None, noise_rng=None):
    correct, bs = 0, 16
    for s in range(0, len(probs), bs):
        b = probs[s:s + bs]
        E, attn, pos, zs_, lens_, _ = generate(model, head, codebook, q_emb, tok, task, b, im, device,
                                               ablate_first=ablate_first, corrupt_step=corrupt_step, noise_rng=noise_rng)
        pred = read_answer(model, head, E, attn, pos, zs_, b, task.m, task.K, device).argmax(-1).tolist()
        correct += sum(int(pred[j] == task.answer(b[j])) for j in range(len(b)))
    return correct / len(probs)


def make_figures(r, tag):
    m, K, chance = r["base_m"], r["K"], r["chance"]
    fig, ax = plt.subplots(1, 3, figsize=(15, 4.2))
    # (a) interventions
    bm = ["intact", "shuffle_steps", "donor_full"]; bn = ["intact", "shuffle\nsteps", "donor\nfull"]
    ax[0].bar(range(len(bm)), [r["acc"][x] for x in bm], color=["#1f77b4", "#d62728", "#ff7f0e"])
    ax[0].axhline(chance, ls="--", color="grey", label=f"chance {chance:.2f}")
    ax[0].set_xticks(range(len(bm))); ax[0].set_xticklabels(bn); ax[0].set_ylim(0, 1)
    ax[0].set_ylabel("acc vs original"); ax[0].legend(); ax[0].set_title("(a) intervene on the chain")
    # (b) donor follows
    dm = ["donor_full", "donor_subspace", "donor_complement"]; dn = ["all", f"U ({r['subspace_rank']}d)", "U-perp"]
    x = np.arange(len(dm)); w = 0.38
    ax[1].bar(x - w / 2, [r["donor"][d]["follow_donor"] for d in dm], w, label="follows DONOR", color="#2ca02c")
    ax[1].bar(x + w / 2, [r["donor"][d]["follow_prompt"] for d in dm], w, label="follows PROMPT", color="#8c564b")
    ax[1].axhline(chance, ls="--", color="grey"); ax[1].set_xticks(x); ax[1].set_xticklabels(dn); ax[1].set_ylim(0, 1)
    ax[1].legend(); ax[1].set_title("(b) donor patch")
    # (c) generalization curve
    gm = sorted(int(k) for k in r["gen_curve"]); gv = [r["gen_curve"][str(k)] for k in gm]
    ax[2].plot(gm, gv, "o-", color="#9467bd")
    tl = r.get("train_lengths") or []
    if tl:
        ax[2].axvspan(min(tl), max(tl), color="green", alpha=0.12, label=f"trained m={min(tl)}-{max(tl)}")
    ax[2].axhline(chance, ls="--", color="grey"); ax[2].set_ylim(0, 1)
    ax[2].set_xlabel("chain length m"); ax[2].set_ylabel("free-running acc"); ax[2].legend()
    ax[2].set_title("(c) generalization to longer chains")
    fig.suptitle(f"{tag}: single token per step — load-bearing + length generalization")
    fig.tight_layout(); f1 = f"{RES}/single_interventions_{tag}.png"; fig.savefig(f1, bbox_inches="tight", dpi=130); plt.close(fig)

    fig, ax = plt.subplots(1, 2, figsize=(11, 4.4))
    conf = np.array(r["confusion"]); confn = conf / conf.sum(1, keepdims=True).clip(min=1)
    im0 = ax[0].imshow(confn, vmin=0, vmax=1, cmap="magma"); ax[0].set_xlabel("predicted digit"); ax[0].set_ylabel("true digit")
    ax[0].set_title("(a) answer read-out"); fig.colorbar(im0, ax=ax[0], fraction=0.046)
    cp = np.array(r["cellprobe"]).mean(2)                                         # [layers, m] mean over cells
    im1 = ax[1].imshow(cp, vmin=0, vmax=1, cmap="viridis", aspect="auto")
    ax[1].set_xticks(range(m)); ax[1].set_xticklabels([f"z{t+1}" for t in range(m)])
    ax[1].set_yticks(range(len(r["probe_layers"]))); ax[1].set_yticklabels(r["probe_layers"]); ax[1].set_ylabel("layer")
    ax[1].set_title("(b) whole-row decodability per token"); fig.colorbar(im1, ax=ax[1], fraction=0.046)
    fig.suptitle(f"{tag}: each single token encodes the whole {K}-cell row")
    fig.tight_layout(); f2 = f"{RES}/single_decode_{tag}.png"; fig.savefig(f2, bbox_inches="tight", dpi=130); plt.close(fig)
    print(f1, flush=True); print(f2, flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True); ap.add_argument("--model", default="Qwen/Qwen3-8B")
    ap.add_argument("--n", type=int, default=400)
    args = ap.parse_args()
    run(args.ckpt, args.model, args.n, "cuda")


if __name__ == "__main__":
    main()