# Find a "do I know this person?" knob inside NVIDIA Nemotron 3.5 Bolt (E4B) and turn it.
#   python interp.py <step> [json args]
#   steps: load (smoke test), baseline (hallucination checks), collect (record the inner state for famous vs made-up names),
#          steer (add the direction while it reads / answers). Needs a CUDA GPU with ~16 GB.
import json, os, sys, time, random, math
import torch

MODEL = "./models/NVIDIA-Nemotron-3.5-Bolt-E4B-BF16"   # download huggingface.co/nvidia/NVIDIA-Nemotron-3.5-Bolt-E4B-BF16 here
OUT = "./out"
os.makedirs(OUT, exist_ok=True)
DEV = "cuda"

KNOWN = ["Michael Jordan", "Taylor Swift", "Albert Einstein", "Barack Obama", "Lionel Messi", "Elon Musk",
         "Oprah Winfrey", "Leonardo DiCaprio", "Serena Williams", "Stephen Hawking", "William Shakespeare",
         "Marie Curie", "Tom Hanks", "LeBron James", "Bill Gates", "Steve Jobs", "Mark Zuckerberg",
         "Jackie Chan", "Bruce Lee", "Yao Ming", "Jeff Bezos", "Lady Gaga", "Cristiano Ronaldo", "Roger Federer",
         "Abraham Lincoln", "Napoleon Bonaparte", "Isaac Newton", "Charles Darwin", "Pablo Picasso",
         "Vincent van Gogh", "Elvis Presley", "Michael Jackson", "Kobe Bryant", "Tiger Woods", "Usain Bolt",
         "Muhammad Ali", "Nelson Mandela", "Keanu Reeves", "Hayao Miyazaki", "Winston Churchill",
         "Mahatma Gandhi", "Frida Kahlo", "Rihanna", "Morgan Freeman", "Warren Buffett", "Neil Armstrong",
         "Ludwig van Beethoven", "Walt Disney"]
UNKNOWN = ["Michael Batkin", "Daniel Brennick", "Sarah Wexley", "Thomas Kelbrook", "Laura Penfold",
           "Gregory Haslam", "Megan Strick", "Peter Vandrell", "Olivia Crandon", "Simon Treadwell",
           "Rachel Ormsby", "Kevin Dallimore", "Hannah Feltwell", "Marcus Pilbeam", "Emily Rowcliffe",
           "Jonathan Speake", "Natalie Brisco", "Andrew Kettering", "Chloe Mabbott", "Brian Lessard",
           "Fiona Gatley", "Patrick Sowerby", "Joanna Treloar", "Victor Hamblin", "Lucy Ambrey",
           "Derek Fennimore", "Alice Wetherall", "Nathan Corrigan", "Sophie Brayfield", "Colin Ashdown",
           "Melissa Tunstall", "Gavin Pressley", "Karen Holtby", "Ian Margetson", "Tessa Larkworthy",
           "Howard Bexley", "Jenna Calloway", "Russell Penhaligon", "Wendy Stallard", "Oscar Meldrum",
           "Diana Ruskell", "Felix Harrowby", "Grace Linwood", "Martin Clowes", "Paula Westbrook",
           "Edwin Saltmarsh", "Ruby Hensall", "Stuart Bramwell"]


def load():
    from transformers import AutoTokenizer
    try:
        from transformers import Gemma4ForConditionalGeneration as M
    except Exception:
        from transformers import AutoModelForImageTextToText as M
    tok = AutoTokenizer.from_pretrained(MODEL)
    t0 = time.time()
    model = M.from_pretrained(MODEL, dtype=torch.bfloat16, device_map=None, low_cpu_mem_usage=True)
    model.eval()
    # the per-layer embedding table is huge: keep it in system RAM and look it up there
    ple = None
    for name, mod in model.named_modules():
        if name.endswith("embed_tokens_per_layer"):
            ple = (name, mod)
    lm = None
    for name, mod in model.named_modules():
        if name.endswith("language_model") and hasattr(mod, "layers"):
            lm = mod
            break
    assert lm is not None, "language_model not found"
    # everything except that table (vision/audio towers are small) goes to the GPU
    keep = ple[1] if ple else None
    def move(mod):
        for name, child in mod.named_children():
            if child is keep:
                continue
            if keep is not None and any(m is keep for m in child.modules()):
                move(child)
                for n, p in list(child._parameters.items()):
                    if p is not None:
                        p.data = p.data.to(DEV)
                for n, b in list(child._buffers.items()):
                    if b is not None:
                        child._buffers[n] = b.to(DEV)
            else:
                child.to(DEV)
    move(model)
    for n, p in list(model._parameters.items()):
        if p is not None:
            p.data = p.data.to(DEV)
    type(model).device = property(lambda self: torch.device(DEV))
    if ple:
        emb = ple[1]
        orig = emb.forward
        emb.forward = lambda ids, *a, **k: orig(ids.to("cpu"), *a, **k).to(DEV)
    # leftover params directly on language_model (norms etc.)
    for n, p in lm.named_parameters(recurse=False):
        p.data = p.data.to(DEV)
    for n, b in list(lm._buffers.items()):
        if b is not None:
            lm._buffers[n] = b.to(DEV)
    print(f"loaded in {time.time()-t0:.0f}s; layers={len(lm.layers)}; ple={'cpu' if ple else 'none'}; "
          f"gpu={torch.cuda.memory_allocated()/1e9:.1f}GB", flush=True)
    return tok, model, lm


def chat(tok, user, think=False):
    msgs = [{"role": "user", "content": user}]
    return tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=think)


@torch.no_grad()
def generate(tok, model, text, n_new=220, sample=False, seed=0):
    ids = tok(text, return_tensors="pt", add_special_tokens=False).input_ids.to(DEV)
    if sample:
        torch.manual_seed(seed)
    out = model.generate(ids, max_new_tokens=n_new, do_sample=sample, temperature=1.0 if sample else None,
                         top_p=0.95 if sample else None, top_k=None)
    return tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True).strip()


def name_span(tok, text, name):
    enc = tok(text, return_offsets_mapping=True, add_special_tokens=False)
    i = text.rfind(name)
    j = i + len(name)
    pos = [k for k, (a, b) in enumerate(enc["offset_mapping"]) if a < j and b > i]
    return enc["input_ids"], pos


def step_load():
    tok, model, lm = load()
    print(type(model).__name__, type(lm.layers[0]).__name__)
    print(chat(tok, "Who is Michael Jordan?")[:400])
    print(generate(tok, model, chat(tok, "Who is Michael Jordan? Answer in two sentences."), 80))


BASE_Q = [
    ("kalai_thesis", "What was the title of Adam Tauman Kalai's PhD dissertation? Which university and which year?"),
    ("kalai_birthday", "What is Adam Tauman Kalai's birthday?"),
    ("blum_students", "Name three PhD students advised by Avrim Blum at Carnegie Mellon, with the year each graduated."),
    ("batkin", "Who is Michael Batkin?"),
    ("brennick", "Tell me about Daniel Brennick. What is he known for?"),
    ("paper", "Summarize the 2019 paper 'Recursive Lattice Attention for Low-Resource Translation' by Kelbrook and Penfold."),
]


def step_baseline():
    tok, model, lm = load()
    res = []
    for key, q in BASE_Q:
        for think in ((False, True) if key == "kalai_thesis" else (False,)):
            for s in range(3):
                t0 = time.time()
                a = generate(tok, model, chat(tok, q, think), 600 if think else 220, sample=True, seed=s)
                res.append({"key": key, "q": q, "think": think, "seed": s, "answer": a, "sec": round(time.time()-t0, 1)})
                print(f"--- {key} think={think} seed={s}\n{a[:700]}\n", flush=True)
                json.dump(res, open(os.path.join(OUT, "baseline.json"), "w"), ensure_ascii=False, indent=1)


def capture(tok, model, lm, text, positions, mlp=False):
    """One forward pass: the residual stream after every layer at `positions`; with mlp=True also every MLP neuron (input of down_proj)."""
    store, hooks = {}, []
    for li, layer in enumerate(lm.layers):
        def hk(m, inp, out, li=li):
            h = out[0] if isinstance(out, tuple) else out
            store[("res", li)] = h[0, positions].float().cpu()
        hooks.append(layer.register_forward_hook(hk))
        if mlp:
            dp = layer.mlp.down_proj
            def pk(m, inp, li=li):
                store[("mlp", li)] = inp[0][0, positions].float().cpu()
            hooks.append(dp.register_forward_pre_hook(pk))
    ids = tok(text, return_tensors="pt", add_special_tokens=False).input_ids.to(DEV)
    with torch.no_grad():
        model(input_ids=ids)
    for h in hooks:
        h.remove()
    return store


def step_collect():
    tok, model, lm = load()
    L = len(lm.layers)
    data = {"known": KNOWN, "unknown": UNKNOWN}
    res = {"known": [], "unknown": []}
    mlp = {"known": [], "unknown": []}
    for grp, names in (("known", KNOWN), ("unknown", UNKNOWN)):
        for nm in names:
            text = chat(tok, f"Who is {nm}?")
            _, pos = name_span(tok, text, nm)
            st = capture(tok, model, lm, text, [pos[-1]], mlp=True)
            res[grp].append(torch.stack([st[("res", l)][0] for l in range(L)]))   # [L, d]
            mlp[grp].append(torch.stack([st[("mlp", l)][0] for l in range(L)]))   # [L, 10240]
        print(grp, "done", flush=True)
    R = {g: torch.stack(v) for g, v in res.items()}   # [N, L, d]
    M = {g: torch.stack(v) for g, v in mlp.items()}   # [N, L, f]
    torch.save({"R": R, "M": M, **data}, os.path.join(OUT, "collect.pt"))
    # per layer: hold out half the names and check the mean-difference direction separates them
    rnd = random.Random(0)
    idx_k = list(range(len(KNOWN))); idx_u = list(range(len(UNKNOWN)))
    rnd.shuffle(idx_k); rnd.shuffle(idx_u)
    hk, hu = len(idx_k)//2, len(idx_u)//2
    rows = []
    for l in range(L):
        k_tr, k_te = R["known"][idx_k[:hk], l], R["known"][idx_k[hk:], l]
        u_tr, u_te = R["unknown"][idx_u[:hu], l], R["unknown"][idx_u[hu:], l]
        d = k_tr.mean(0) - u_tr.mean(0)
        dn = d / d.norm()
        mid = ((k_tr.mean(0) + u_tr.mean(0)) / 2 @ dn).item()
        pk = (k_te @ dn); pu = (u_te @ dn)
        acc = (((pk > mid).float().sum() + (pu <= mid).float().sum()) / (len(pk) + len(pu))).item()
        cohen = ((pk.mean() - pu.mean()) / torch.sqrt((pk.var() + pu.var()) / 2)).item()
        rows.append({"layer": l, "acc": round(acc, 3), "cohen": round(cohen, 2), "dnorm": round(d.norm().item(), 1),
                     "resnorm": round(R["known"][:, l].norm(dim=-1).mean().item(), 1)})
        print(rows[-1], flush=True)
    # single neurons that differ most between famous and made-up names
    neu = []
    for l in range(L):
        a, b = M["known"][:, l], M["unknown"][:, l]
        diff = a.mean(0) - b.mean(0)
        sd = torch.sqrt((a.var(0) + b.var(0)) / 2) + 1e-4
        z = diff / sd
        for sign in (1, -1):
            top = torch.topk(sign * z, 3)
            for v, i in zip(top.values, top.indices):
                i = int(i)
                acc = (((a[:, i] - b[:, i].mean()) * sign > 0).float().mean().item() + ((b[:, i] - a[:, i].mean()) * sign < 0).float().mean().item()) / 2
                neu.append({"layer": l, "idx": i, "sign": sign, "z": round(float(v), 2), "known_mean": round(a[:, i].mean().item(), 3),
                            "unknown_mean": round(b[:, i].mean().item(), 3)})
    neu.sort(key=lambda r: -r["z"])
    json.dump({"layers": rows, "neurons": neu[:40]}, open(os.path.join(OUT, "collect.json"), "w"), indent=1)
    print(json.dumps(neu[:12], indent=0))


class Steer:
    def __init__(self, layer_mod, vec, mode, positions=None):
        self.vec, self.mode, self.positions = vec, mode, positions
        self.h = layer_mod.register_forward_hook(self.hook)

    def hook(self, m, inp, out):
        h = out[0] if isinstance(out, tuple) else out
        if self.mode == "name":
            if h.shape[1] > 1:
                h[:, self.positions] += self.vec.to(h.dtype)
        else:  # all: every position from the name on, plus every generated token
            if h.shape[1] > 1:
                h[:, self.positions[0]:] += self.vec.to(h.dtype)
            else:
                h += self.vec.to(h.dtype)
        return out

    def remove(self):
        self.h.remove()


def step_steer(layers=None, coefs=None, names=None, mode="name", signed=False, tag=""):
    tok, model, lm = load()
    C = torch.load(os.path.join(OUT, "collect.pt"))
    R = C["R"]
    layers = layers or [8, 12, 16, 20]
    coefs = coefs or [0, 1, 2, 4, 6]
    names = names or ["Michael Jordan", "Albert Einstein", "Michael Batkin", "Sarah Wexley"]
    out_path = os.path.join(OUT, f"steer_{mode}{tag}.json")
    res = json.load(open(out_path)) if os.path.exists(out_path) else []
    done = {(r["name"], r["layer"], r["coef"]) for r in res}
    for l in layers:
        d = (R["known"][:, l].mean(0) - R["unknown"][:, l].mean(0)).to(DEV)
        for nm in names:
            known = nm in KNOWN
            for c in coefs:
                if (nm, l, c) in done:
                    continue
                vec = (c if signed else (-c if known else c)) * d   # signed: positive = toward "known"; otherwise push famous names toward "stranger" and made-up ones toward "known"
                text = chat(tok, f"Who is {nm}?")
                _, pos = name_span(tok, text, nm)
                s = Steer(lm.layers[l], vec, mode, pos) if c else None
                a = generate(tok, model, text, 160)
                if s:
                    s.remove()
                res.append({"name": nm, "known": known, "layer": l, "coef": c, "mode": mode, "answer": a})
                print(f"--- {nm} L{l} c={c} {mode}\n{a[:400]}\n", flush=True)
                json.dump(res, open(out_path, "w"), ensure_ascii=False, indent=1)


if __name__ == "__main__":
    step = sys.argv[1]
    if step == "load":
        step_load()
    elif step == "baseline":
        step_baseline()
    elif step == "collect":
        step_collect()
    elif step == "steer":
        kw = json.loads(sys.argv[2]) if len(sys.argv) > 2 else {}
        step_steer(**kw)