"""Experiment 4: locate sensor / workspace / motor regimes by depth.

Per-layer diagnostics averaged over a prompt set:
  sensor  = frac. of positions where top J-lens token == current input token
  motor   = frac. where top J-lens token == NEXT input token (teacher-forced)
  persist = mean Jaccard overlap of top-10 lens tokens at adjacent positions
            (workspace content should persist across positions)
  content = frac. of positions whose top lens token is a content word
            (alphabetic, len>=3, not in a junk list)

Ignition test (paper: ambiguous inputs produce sharp binary commitment at
workspace onset): replace one token's embedding (main + per-layer) with a
w-mixture of two concept embeddings and track lens commitment
C = (P1-P2)/(P1+P2) by layer at a downstream position.
"""

import os, sys
from pathlib import Path

import torch

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from jlens.core import JLens, chat_ids, collect_residuals, load_model

RES = Path(os.environ.get("JLENS_RESULTS", "results"))

PROMPTS = [
    "The animal that spins webs has how many legs? Answer with just a number.",
    "What is the capital of France? Answer with just the city name.",
    "Write one sentence about the ocean.",
    "Explain photosynthesis in one sentence.",
    "Name a famous river in Egypt.",
    "What language is spoken in Brazil? Answer with one word.",
    "Summarize the plot of Romeo and Juliet in one sentence.",
    "If I have 3 apples and eat one, how many are left?",
]

JUNK = set("·.<>|/\\()[]{}:;,!?\"'`~*#-_=+ \n\t")


def regime_profile(jl, tok, model):
    L = len(jl.tm.layers)
    sensor = torch.zeros(L)
    motor = torch.zeros(L)
    persist = torch.zeros(L)
    content = torch.zeros(L)
    n_pos = 0
    n_adj = 0
    for p in PROMPTS:
        ids = chat_ids(tok, p)
        hs = collect_residuals(model, ids)
        T = ids.shape[1]
        sl = slice(4, T - 5)  # skip bos/turn tokens and trailing template
        cur = ids[0, sl].cuda()
        nxt = ids[0, 4 + 1:T - 4].cuda()
        for l in range(L):
            idx, _ = jl.read(hs[l, sl], l, topk=10)  # (T', 10)
            top1 = idx[:, 0]
            sensor[l] += (top1 == cur).sum().item()
            motor[l] += (top1 == nxt).sum().item()
            sets = [set(r.tolist()) for r in idx]
            for a, b in zip(sets, sets[1:]):
                persist[l] += len(a & b) / len(a | b)
            for t in top1.tolist():
                s = tok.decode([t]).strip()
                content[l] += (len(s) >= 3 and s.isalpha())
        n_pos += cur.numel()
        n_adj += cur.numel() - 1
    return sensor / n_pos, motor / n_pos, persist / n_adj, content / n_pos


def ignition(jl, tok, model, pairs, ws=(0.0, 0.25, 0.5, 0.75, 1.0)):
    """Mix two concept embeddings at one position; commitment by layer."""
    tm = jl.tm
    L = len(tm.layers)
    template = ("My favorite thing in the world is the X ."
                " I think about it every single day because")
    out = {}
    for w1, w2 in pairs:
        id1 = tok.encode(" " + w1, add_special_tokens=False)[0]
        id2 = tok.encode(" " + w2, add_special_tokens=False)[0]
        ids = tok(template, return_tensors="pt")["input_ids"].cuda()
        pos = (ids[0] == tok.encode(" X", add_special_tokens=False)[0]) \
            .nonzero()[0].item()
        has_ple = getattr(tm, "embed_tokens_per_layer", None) is not None
        with torch.no_grad():
            pair_ids = torch.tensor([[id1, id2]], device="cuda")
            rows_main = tm.embed_tokens(pair_ids)[0].detach()
            rows_ple = tm.embed_tokens_per_layer(pair_ids)[0].detach() if has_ple else None
        curves = torch.zeros(len(ws), L)
        for wi, w in enumerate(ws):
            def make_hook(rows, w=w, pos=pos):
                def mix_hook(mod, inp, out_e):
                    e = out_e.clone()
                    e[0, pos] = w * rows[0] + (1 - w) * rows[1]
                    return e
                return mix_hook
            h1 = tm.embed_tokens.register_forward_hook(make_hook(rows_main))
            h2 = tm.embed_tokens_per_layer.register_forward_hook(make_hook(rows_ple)) if has_ple else None
            try:
                hs = collect_residuals(model, ids)
            finally:
                h1.remove()
                if h2: h2.remove()
            rpos = pos + 3  # downstream read position
            for l in range(L):
                logits = jl._readout(hs[l, rpos].float() @ jl.Jbar[l].T)
                probs = torch.softmax(logits.float(), -1)
                p1, p2 = probs[id1].item(), probs[id2].item()
                curves[wi, l] = (p1 - p2) / (p1 + p2 + 1e-12)
        out[(w1, w2)] = curves
    return out


def main():
    model, tok = load_model(dtype=torch.bfloat16)
    ck = torch.load(sys.argv[1] if len(sys.argv) > 1 else "results/jbar.pt",
                    map_location="cuda")
    jl = JLens(model, tok, ck["Jbar"].cuda())

    sensor, motor, persist, content = regime_profile(jl, tok, model)
    print("layer | sensor(top==cur) motor(top==next) persist(top10 Jaccard) content-word")
    for l in range(len(sensor)):
        bars = lambda x: "#" * int(20 * x)
        print(f"L{l:>2} | {sensor[l]:.2f} {bars(sensor[l]):<20} | "
              f"{motor[l]:.2f} {bars(motor[l]):<20} | "
              f"{persist[l]:.2f} | {content[l]:.2f}")

    pairs = [("dog", "piano"), ("ocean", "violin"), ("dragon", "bicycle")]
    ign = ignition(jl, tok, model, pairs)
    print("\nignition: commitment C=(P1-P2)/(P1+P2) at read pos, by layer")
    print("pair            | w:   " + "  ".join(f"{w:4.2f}" for w in (0.0, .25, .5, .75, 1.0)))
    for (w1, w2), curves in ign.items():
        for l in range(2, len(sensor), 4):
            print(f"{w1}/{w2:<9} L{l:>2} | " +
                  "  ".join(f"{curves[wi, l]:+.2f}" for wi in range(5)))
    torch.save({"sensor": sensor, "motor": motor, "persist": persist,
                "content": content, "ignition": ign}, RES / "regimes.pt")


if __name__ == "__main__":
    main()