#!/usr/bin/env python3
"""jlens_jacobian.py — the lens nemo actually specified: lens(h) = softmax(W_U · norm(J · h)).

The first readout used a LOGIT lens, softmax(W_U · norm(h)) — the tractable baseline. This is the
Jacobian refinement from the lab brief.

THE TRACTABILITY TRICK. The full Jacobian J of the remaining network w.r.t. a hidden state is
[hidden × hidden] per layer per position — hopeless to materialise (4096² per position). But the
lens never needs J itself, only the product **J · h**, and a Jacobian-vector product is exactly what
forward-mode autodiff computes in ONE extra forward pass:

    _, Jh = torch.func.jvp(f_rest, (h,), (h,))

where f_rest runs layers ℓ+1..N on h and returns the final normed hidden. The tangent is h itself,
so the output tangent is J·h — "how the network's final state moves as this hidden state moves along
itself". Then lens = softmax(W_U · norm(J·h)).

Interpretation, and why it should beat the logit lens: the logit lens asks *what does this hidden
state look like if you decode it right now*, which at mid layers is dominated by high-norm outlier
and attention-sink directions (we measured exactly that: the raw top-k was CJK fragments and
<|im_end|>). The Jacobian lens asks *what does this hidden state PUSH the output toward* — it reads
the causal direction rather than the residual's current contents, so sink directions that the network
does not act on should fall away.

  python3 jlens_jacobian.py --model 01-ai/Yi-1.5-34B-Chat --run-label yi-1.5-34b \
      --questions /tmp/jlens_lab_harness/data/lab_questions_20.csv --limit 20 \
      --out-dir /mnt/cache/lens_yi_jac
"""
from __future__ import annotations
import argparse, csv, json, os, sys
os.environ.setdefault("HF_HOME", "/mnt/assets/models/huggingface")
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
sys.path.insert(0, os.environ.get("JLENS_HARNESS", "/tmp/jlens_lab_harness"))
import torch
from torch.func import jvp
from transformers import AutoModelForCausalLM, AutoTokenizer
from geo.common import prompts

NULL_PROMPTS = [
    "What are the public transport options in Melbourne?",
    "Explain how to renew an Australian passport.",
    "What is the weather like in Sydney in winter?",
    "Describe the process of enrolling a child in primary school in NSW.",
    "How do I recycle electronic waste in Australia?",
]

def parts(model):
    """(decoder layers, final norm, unembedding), found by SEARCHING the module tree.

    Do not read the config for this. Multimodal configs (Gemma3Config, Gemma4) have no top-level
    num_hidden_layers — it hides under config.text_config — and a lens that trusts the config dies
    on exactly the newest models we most want to measure. The module tree always knows: find the
    deepest ModuleList of decoder blocks and the norm that sits beside it. Ask the model, not its
    description of itself."""
    import torch.nn as nn
    best = None
    for name, mod in model.named_modules():
        if isinstance(mod, nn.ModuleList) and len(mod) >= 4:
            blk = mod[0]
            # "linear_attn": Qwen3.5's Gated-DeltaNet hybrid names its mixer linear_attn, and the
            # whole five-rung ladder failed "could not locate a decoder stack" on exactly the
            # newest lineage this function's docstring promises to handle. Match any child whose
            # name contains attn/attention instead of enumerating spellings.
            if any(hasattr(blk, at) for at in ("self_attn", "attention", "attn", "linear_attn")) \
               or any(("attn" in n) or ("attention" in n) for n, _ in blk.named_children()):
                if best is None or len(mod) > len(best[1]):
                    best = (name, mod)
    if best is None:
        raise RuntimeError("could not locate a decoder stack in this architecture")
    stack_name, layers = best
    parent = model
    for part in stack_name.split(".")[:-1]:
        parent = getattr(parent, part)
    norm = (getattr(parent, "norm", None) or getattr(parent, "final_layernorm", None)
            or getattr(parent, "final_layer_norm", None))
    if norm is None:                       # last resort: any norm module directly under the parent
        import torch.nn as _nn
        for n, m in parent.named_children():
            if "norm" in n.lower():
                norm = m; break
    head = model.get_output_embeddings()
    if norm is None or head is None:
        raise RuntimeError(f"found {len(layers)} blocks at {stack_name} but no final norm / lm head")
    return layers, norm, head

def build_messages(query, locale="en-AU"):
    return [{"role": "system", "content": prompts.BASE_SYSTEM_PROMPT},
            {"role": "user", "content": prompts.BASE_USER_TEMPLATE.format(locale=locale, query=query)}]

def merge_system_into_user(messages):
    """Fold a system turn into the first user turn, VERBATIM.

    Some chat templates (Gemma-2's among them) raise `TemplateError: System role not supported`.
    The probe runner already handles this; the lens did not, so every Gemma-2 model died at
    template time after paying the full model load. Only the delivery structure changes — the
    prompt text is byte-identical, which is what keeps the GEO prompts verbatim as specified.
    """
    sys_txt = "\n\n".join(m["content"] for m in messages if m["role"] == "system")
    rest = [m for m in messages if m["role"] != "system"]
    if sys_txt and rest and rest[0]["role"] == "user":
        rest = [{"role": "user", "content": sys_txt + "\n\n" + rest[0]["content"]}] + rest[1:]
    return rest

def apply_template(tok, messages, **kw):
    """apply_chat_template, retrying without a system role for templates that reject one."""
    try:
        return tok.apply_chat_template(messages, **kw)
    except Exception as e:
        if "system" not in str(e).lower():
            raise
        print(f"[template] system role rejected ({e}); folding it into the first user turn",
              file=sys.stderr)
        return tok.apply_chat_template(merge_system_into_user(messages), **kw)

@torch.no_grad()
def hidden_states(model, ids):
    return model(ids, output_hidden_states=True).hidden_states

def jacobian_lens_mass(model, tok, ids, a0, layers_sel, topk=64, max_pos=64, track_ids=None):
    """J·h at each selected layer over the answer span, decoded through the unembedding.

    track_ids: optional {name: set(token_ids)} — accumulate FULL-softmax probability over those ids
    regardless of topk. Exists because the identity experiment's anchor families sat below the
    top-64 cutoff and read as zero: an anchor absent from a topk dict is not an anchor with no
    mass, it is an anchor the readout could not see. Tracked mass is returned under
    "__tracked__:<name>" keys so existing consumers are unaffected.
    """
    layers, norm, head = parts(model)
    hs = hidden_states(model, ids)
    n_layer = len(layers)
    acc = torch.zeros(head.weight.shape[0], device=model.device, dtype=torch.float32)
    count = 0
    for L in [l for l in layers_sel if 0 < l < n_layer]:
        h_full = hs[L]                                     # [1, seq, hidden] entering block L
        span = h_full[:, a0:, :]
        if span.shape[1] == 0:
            continue
        # subsample positions: the JVP is a forward pass per call, so cap the cost honestly
        step = max(1, span.shape[1] // max_pos)
        sel = span[:, ::step, :].contiguous()

        # transformers 5.x decoder blocks REQUIRE position_embeddings (rotary cos/sin); calling
        # blk(hidden) alone returns None. Build them for the true positions of the subsampled tokens
        # so the rotary phase is right rather than silently zeroed.
        pos_ids = torch.arange(a0, h_full.shape[1], device=h_full.device)[::step][None, :]
        rot = getattr(getattr(model, "model", model), "rotary_emb", None)
        pe = rot(sel, pos_ids) if rot is not None else None

        def f_rest(x):
            # run the REMAINING blocks on x, then the final norm. Positions are independent here
            # (no attention across the subsample) — an approximation we state rather than hide:
            # it reads each position's own causal push, not its interaction with its neighbours.
            y = x
            for blk in layers[L:]:
                out = blk(y, position_embeddings=pe) if pe is not None else blk(y)
                y = out[0] if isinstance(out, tuple) else out
            return norm(y)

        try:
            _, Jh = jvp(f_rest, (sel,), (sel,))            # forward-mode: one extra pass, no J built
        except Exception as e:
            print(f"   [layer {L}] jvp failed ({type(e).__name__}: {e}) — skipping", file=sys.stderr)
            continue
        logits = head(Jh).float()                          # softmax(W_U · norm(J·h))
        p = torch.softmax(logits, dim=-1).sum(dim=1).squeeze(0)
        acc += p; count += 1
    if not count:
        return {}
    acc /= count
    vals, idx = torch.topk(acc, min(topk, acc.numel()))
    mass = {}
    for v, i in zip(vals.tolist(), idx.tolist()):
        piece = tok.decode([i]).strip().lower()
        if piece and v > 0:
            mass[piece] = mass.get(piece, 0.0) + float(v)
    if track_ids:
        for name, ids_ in track_ids.items():
            tot = float(sum(acc[i].item() for i in ids_ if 0 <= i < acc.numel()))
            mass[f"__tracked__:{name}"] = tot
    return mass

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True); ap.add_argument("--run-label", required=True)
    ap.add_argument("--questions", required=True); ap.add_argument("--limit", type=int, default=20)
    ap.add_argument("--out-dir", required=True); ap.add_argument("--layers", default="")
    ap.add_argument("--device", default="cuda:0"); ap.add_argument("--max-new", type=int, default=512)
    ap.add_argument("--trust", action="store_true")  # internlm/hunyuan/baichuan need custom code
    a = ap.parse_args()
    os.makedirs(a.out_dir, exist_ok=True)
    rows = list(csv.DictReader(open(a.questions)))[: a.limit]

    print(f"loading {a.model} …", file=sys.stderr, flush=True)
    tok = AutoTokenizer.from_pretrained(a.model)   # NOT trust here — it broke Yi earlier
    # EAGER attention is REQUIRED here, not a preference: torch's fused SDPA/flash kernel has no
    # forward-mode AD derivative ("Trying to use forward AD with _scaled_dot_product_flash_attention
    # that does not support it"), so the JVP cannot differentiate through it. Eager attention is plain
    # PyTorch ops all the way down, which forward AD handles. Slower; it is the price of the Jacobian.
    model = AutoModelForCausalLM.from_pretrained(a.model, dtype=torch.bfloat16,
                                                 device_map={"": a.device}, attn_implementation="eager",
                                                 trust_remote_code=a.trust)
    # belt and braces: some transformers builds ignore the experts_implementation kwarg
    # silently (the container did — the retry still hit _grouped_mm). The method call is the
    # version that provably takes effect; warn loudly where the API is absent entirely.
    if hasattr(model, "set_experts_implementation"):
        try: model.set_experts_implementation("eager")
        except Exception as e: print(f"[experts] set failed: {e}", file=sys.stderr)
    elif "moe" in type(model).__name__.lower():
        print("[experts] WARNING: no experts_implementation API — MoE jvp will fail on this build",
              file=sys.stderr)
    model.eval()
    _layers, _, _ = parts(model)
    n_layer = len(_layers)
    sel = [int(x) for x in a.layers.split(",") if x] or \
          sorted({int(n_layer * f) for f in (0.4, 0.55, 0.7, 0.85)})
    print(f"loaded; {n_layer} layers, JACOBIAN lens at {sel}", file=sys.stderr, flush=True)

    def readout(query):
        ids = apply_template(tok, build_messages(query), add_generation_prompt=True,
                             return_tensors="pt", return_dict=True)["input_ids"].to(model.device)
        with torch.no_grad():
            gen = model.generate(ids, max_new_tokens=a.max_new, do_sample=False,
                                 pad_token_id=(tok.pad_token_id or tok.eos_token_id))
        full = gen
        return jacobian_lens_mass(model, tok, full, ids.shape[1], sel)

    print("null baseline …", file=sys.stderr, flush=True)
    base = {}
    for q in NULL_PROMPTS:
        for k, v in readout(q).items():
            base[k] = base.get(k, 0.0) + v
    base = {k: v / len(NULL_PROMPTS) for k, v in base.items()}

    for q in rows:
        qid = q["query_id"]
        mass = readout(q["query"])
        rec = {"query_id": qid, "condition": "baseline", "run_label": a.run_label,
               "model": a.model, "layers": sel, "lens": "jacobian_jvp",
               "token_mass": mass, "baseline_token_mass": base}
        with open(os.path.join(a.out_dir, f"{qid}.json"), "w") as f:
            json.dump(rec, f)
        top = sorted(mass.items(), key=lambda kv: -kv[1])[:8]
        print(f"[{qid}] jacobian top: {[t for t, _ in top]}", file=sys.stderr, flush=True)
    print(f"wrote -> {a.out_dir}", file=sys.stderr, flush=True)

if __name__ == "__main__":
    main()