"""Single-cell causal patch for the single-token organism.

Overwrite ONE cell value c_i(t) inside the soft token z_t (force cell i = digit d), then CONTINUE
free-running so the change propagates through the learned CA rule, and check the final answer equals
what the *true* CA produces for the altered row. This is a surgical test (one cell, one step) that the
latent at (t,i) causally carries c_i(t) AND that the downstream computation uses it.

Patch (per the card): generate free-running to step t; read z_t = softmax(head(h_t)) in R^{K x 10};
set z_t[i] <- one-hot(d) (other cells untouched); re-embed via the codebook (fed = sum_k z_t[k].C[k])
as the input to z_{t+1}; continue free-running; read the answer from the head at z_T, cell q.
Expected = CA-evolve(c_t with cell i := d) for the remaining T-1-t steps, read at cell q.

    python -m latent_threads.eval_single_cellpatch --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")

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
from tqdm.auto import tqdm

from latent_threads import tasks as LT
from latent_threads.single import build_single_batch, build_single_mask
from latent_threads.eval_single_report import load, RES


@torch.no_grad()
def patched_answer(model, head, codebook, q_emb, tok, task, batch, im, device, t_p, i_p, d_p):
    """Free-run; at step t_p[j] overwrite cell i_p[j] of the soft token with one-hot(d_p[j]); continue.
    Returns the head's predicted digit for each example's queried cell at z_T."""
    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)
    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):
        out = model(inputs_embeds=E, attention_mask=attn, position_ids=pos, output_hidden_states=True)
        if t == m - 1:
            break
        z = torch.softmax(head(out.hidden_states[-1][bidx, zs + t]).view(B, K, 10).float(), -1)  # [B,K,10]
        for j in range(B):                                                   # patch examples whose t_p == t
            if t_p[j] == t:
                z[j, i_p[j]] = F.one_hot(torch.tensor(d_p[j], device=device), 10).float()
        fed = torch.einsum("bkv,kvd->bd", z, codebook.float())
        E = E.clone(); E[bidx, zs + t + 1] = fed.to(E.dtype)
    final = head(out.hidden_states[-1][bidx, zs + m - 1]).view(B, K, 10)      # head at z_T
    return [final[j, batch[j].q].argmax().item() for j in range(B)]


def expected_answer(task, p, t, i, d):
    row = list(p.rows[t]); row[i] = d
    for _ in range(task.m - 1 - t):
        row = task._step(row)
    return row[p.q]


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|>")
    m = cfg["task_kwargs"]["m"]; task = LT.make_task(cfg["task"], k=K, m=m)
    rng = random.Random(7)
    probs = [task.sample(rng) for _ in range(n)]
    # per-example patch: a random step t in 0..m-2, cell i, target digit d != current value
    t_p, i_p, d_p = [], [], []
    for p in probs:
        t = rng.randrange(0, m - 1); i = rng.randrange(K)
        cur = p.rows[t][i]; d = rng.choice([x for x in range(10) if x != cur])
        t_p.append(t); i_p.append(i); d_p.append(d)
    exp = np.array([expected_answer(task, probs[j], t_p[j], i_p[j], d_p[j]) for j in range(n)])
    orig = np.array([task.answer(probs[j]) for j in range(n)])

    preds = []
    bs = 16
    for s in tqdm(range(0, n, bs), desc=f"{tag} cell-patch"):
        b = probs[s:s + bs]
        preds += patched_answer(model, head, codebook, q_emb, tok, task, b, im, device,
                                t_p[s:s + bs], i_p[s:s + bs], d_p[s:s + bs])
    preds = np.array(preds)

    follows_ca = float((preds == exp).mean())
    changed = exp != orig                                                     # patch actually moves the answer
    follows_on_changed = float((preds[changed] == exp[changed]).mean())
    follows_orig_on_changed = float((preds[changed] == orig[changed]).mean())  # ignores the patch?
    dist = np.array([m - 1 - t for t in t_p])                                # propagation distance
    by_dist = {int(dd): float((preds[dist == dd] == exp[dist == dd]).mean()) for dd in sorted(set(dist.tolist()))}
    print("=" * 70 + f"\nSINGLE-CELL PATCH {tag}  K={K} m={m}  n={n}\n" + "=" * 70, flush=True)
    print(f"  follows CA-propagated answer (all patches): {follows_ca:.3f}", flush=True)
    print(f"  on patches that CHANGE the answer (n={int(changed.sum())}): follows patched CA={follows_on_changed:.3f}  "
          f"follows ORIGINAL (ignores patch)={follows_orig_on_changed:.3f}", flush=True)
    print(f"  by propagation distance (steps from patch to read): { {k: round(v,3) for k,v in by_dist.items()} }", flush=True)

    res = dict(tag=tag, K=K, m=m, n=n, follows_ca=follows_ca, n_changed=int(changed.sum()),
               follows_on_changed=follows_on_changed, follows_orig_on_changed=follows_orig_on_changed,
               by_distance=by_dist)
    json.dump(res, open(f"{RES}/single_cellpatch_{tag}.json", "w"), indent=2)

    fig, ax = plt.subplots(figsize=(6, 4))
    ds = sorted(by_dist); ax.plot(ds, [by_dist[d] for d in ds], "o-", color="#2ca02c", label="follows patched CA")
    ax.axhline(1 / 10, ls="--", color="grey", label="chance 0.10")
    ax.set_xlabel("propagation distance (CA steps from patched cell to read-out)")
    ax.set_ylabel("P(answer = CA-propagated value)"); ax.set_ylim(0, 1.02); ax.legend()
    ax.set_title(f"{tag}: single-cell patch c_i(t):=d propagates correctly")
    png = f"{RES}/single_cellpatch_{tag}.png"; fig.savefig(png, bbox_inches="tight", dpi=130)
    print(png, flush=True)
    return res


def run_sweep(ckpt, model_name, n_probs, device):
    """Systematic spot-check: patch EVERY (step t, cell i) site on a small set of problems, force a
    changed value, and print the full per-patch detail (orig answer / CA-expected / model) for review."""
    model, head, codebook, q_emb, tok, did, K, cfg, tag = load(ckpt, model_name, device)
    im = tok.convert_tokens_to_ids("<|im_end|>")
    m = cfg["task_kwargs"]["m"]; task = LT.make_task(cfg["task"], k=K, m=m)
    rng = random.Random(11)
    probs = [task.sample(rng) for _ in range(n_probs)]
    sites = [(t, i) for t in range(m - 1) for i in range(K)]                  # all positions x all cells
    batch, T, I, D, meta = [], [], [], [], []
    for pi, p in enumerate(probs):
        for (t, i) in sites:
            cur = p.rows[t][i]; d = rng.choice([x for x in range(10) if x != cur])
            batch.append(p); T.append(t); I.append(i); D.append(d)
            meta.append((pi, t, i, cur, d, p.q, task.answer(p), expected_answer(task, p, t, i, d)))
    preds = []
    for s in range(0, len(batch), 16):
        preds += patched_answer(model, head, codebook, q_emb, tok, task, batch[s:s + 16], im, device,
                                T[s:s + 16], I[s:s + 16], D[s:s + 16])
    print("=" * 92 + f"\nSINGLE-CELL PATCH SWEEP {tag}  K={K} m={m}  {len(sites)} sites x {n_probs} problems = {len(meta)} patches\n" + "=" * 92, flush=True)
    print(f"{'prob':>4} {'site':>8} {'c_i(t)->d':>10} {'q':>4} {'orig':>5} {'CA_exp':>7} {'model':>6}  match", flush=True)
    persite, ok_all = {}, []
    for k, (pi, t, i, cur, d, q, orig, exp) in enumerate(meta):
        mod = preds[k]; ok = (mod == exp); ok_all.append(ok); persite.setdefault((t, i), []).append(ok)
        print(f"{pi:>4} {f's{t+1}.c{i+1}':>8} {f'{cur}->{d}':>10} {f'c{q+1}':>4} {orig:>5} {exp:>7} {mod:>6}  {'OK' if ok else 'XX'}", flush=True)
    print("-" * 92 + "\nper-site follows-CA (over the problems):", flush=True)
    for (t, i) in sites:
        v = persite[(t, i)]; print(f"  s{t+1}.c{i+1}: {sum(v)}/{len(v)} = {sum(v)/len(v):.2f}", flush=True)
    print(f"OVERALL follows CA-propagated: {sum(ok_all)}/{len(ok_all)} = {sum(ok_all)/len(ok_all):.3f}", 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)
    ap.add_argument("--sweep", action="store_true"); ap.add_argument("--probs", type=int, default=7)
    args = ap.parse_args()
    if args.sweep:
        run_sweep(args.ckpt, args.model, args.probs, "cuda")
    else:
        run(args.ckpt, args.model, args.n, "cuda")


if __name__ == "__main__":
    main()