"""Jacobian lens (J-lens) for gemma-4-E2B-it.

Reproduces the method of "Verbalizable Representations Form a Global
Workspace in Language Models" (transformer-circuits.pub/2026/workspace):

  J_l = E_{prompt, t, t' >= t} [ d h_{final, t'} / d h_{l, t} ]      (d x d per layer)
  lens(h_l) = softmax(W_U . finalnorm(J_l h_l))
  J-lens vector of token s at layer l:  v_s = (W_U J_l)[s, :] = W_U[s] J_l
  swap: h <- h + V (sigma(c) - c),  c = pinv(V) h,  V = [v_s; v_t]

h_l is the residual stream at the *output* of decoder layer l.
h_final is the output of the last decoder layer (pre final-RMSNorm).
"""

import os

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = os.environ.get("JLENS_MODEL", "google/gemma-4-E2B-it")


def load_model(dtype=torch.float32, device="cuda", model_id=None):
    model_id = model_id or MODEL_ID
    tok = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id, dtype=dtype, attn_implementation="eager"
    ).to(device).eval()
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    return model, tok


def _text_model(model):
    return model.model.language_model


class ResidualCapture:
    """Forward hooks that record each decoder layer's output residual stream."""

    def __init__(self, model, layers=None, detach=True):
        self.tm = _text_model(model)
        self.layers = list(range(len(self.tm.layers))) if layers is None else layers
        self.detach = detach
        self.acts = {}
        self.handles = []

    def __enter__(self):
        for i in self.layers:
            def hook(mod, inp, out, i=i):
                self.acts[i] = out.detach() if self.detach else out
            self.handles.append(self.tm.layers[i].register_forward_hook(hook))
        return self

    def __exit__(self, *a):
        for h in self.handles:
            h.remove()


@torch.no_grad()
def collect_residuals(model, input_ids):
    """Return (n_layers, T, d) residual stream stack for a single prompt."""
    with ResidualCapture(model) as cap:
        model(input_ids=input_ids)
    n = len(_text_model(model).layers)
    return torch.stack([cap.acts[i][0] for i in range(n)])


@torch.no_grad()
def generate_with_residuals(model, tok, input_ids, max_new_tokens=40):
    """Greedy-generate while capturing the residual stream at every position
    (prompt + generated). Returns (text, ids (T,), residuals (L, T, d))."""
    tm = _text_model(model)
    L = len(tm.layers)
    steps = {i: [] for i in range(L)}
    handles = []
    for i in range(L):
        def hook(mod, inp, out, i=i):
            steps[i].append(out[0].detach())
        handles.append(tm.layers[i].register_forward_hook(hook))
    try:
        out = model.generate(input_ids, max_new_tokens=max_new_tokens,
                             do_sample=False,
                             pad_token_id=tok.pad_token_id or tok.eos_token_id)
    finally:
        for h in handles:
            h.remove()
    hs = torch.stack([torch.cat(steps[i], dim=0) for i in range(L)])
    text = tok.decode(out[0, input_ids.shape[1]:], skip_special_tokens=True)
    return text, out[0], hs


def prompt_jacobian_pairsum(model, input_ids, chunk=128):
    """Sum over all pairs (t, t' >= t) of d h_{final,t'} / d h_{l,t} for one prompt.

    One forward pass; backward from S = sum_{t'} h_{final,t'} with batched
    identity cotangents. The gradient of <e_i, S> w.r.t. h_{l,t} is row i of
    sum_{t' >= t} J^{(t',t)}_l (causality zeroes t' < t), so summing the
    gradient over t gives row i of the pair-sum. Returns (fp32 (L, d, d), n_pairs).
    """
    tm = _text_model(model)
    L = len(tm.layers)
    d = model.config.get_text_config().hidden_size
    T = input_ids.shape[1]
    dev = input_ids.device

    with ResidualCapture(model, detach=False) as cap:
        with torch.enable_grad():
            model(input_ids=input_ids)
    hs = [cap.acts[i] for i in range(L)]  # each (1, T, d), in graph
    S = hs[-1].sum(dim=1)  # (1, d)

    J = torch.zeros(L, d, d, device=dev, dtype=torch.float32)
    eye = torch.eye(d, device=dev, dtype=S.dtype)
    starts = list(range(0, d, chunk))
    for ci, s in enumerate(starts):
        e = s + min(chunk, d - s)
        V = eye[s:e].unsqueeze(1)  # (B, 1, d) cotangents
        grads = torch.autograd.grad(
            outputs=S, inputs=hs, grad_outputs=V,
            retain_graph=(ci < len(starts) - 1), is_grads_batched=True,
        )
        for l in range(L):
            # grads[l]: (B, 1, T, d); row i of pair-sum = sum over t
            J[l, s:e] += grads[l].squeeze(1).sum(dim=1).float()
    n_pairs = T * (T + 1) // 2
    return J, n_pairs


class JLens:
    """Averaged-Jacobian lens: reading, J-lens vectors, swap interventions."""

    def __init__(self, model, tok, Jbar):
        """Jbar: (L, d, d) fp32 averaged Jacobian per layer."""
        self.model, self.tok = model, tok
        self.Jbar = Jbar
        self.tm = _text_model(model)
        self.softcap = model.config.get_text_config().final_logit_softcapping

    def _readout(self, x):
        """finalnorm + unembed (+softcap) of residual-space vectors x (..., d)."""
        x = self.tm.norm(x.to(self.tm.norm.weight.dtype))
        logits = self.model.lm_head(x)
        if self.softcap:
            logits = self.softcap * torch.tanh(logits / self.softcap)
        return logits

    @torch.no_grad()
    def read(self, h, layer, topk=10):
        """J-lens reading of residual vector(s) h (..., d) at `layer`.

        Returns (topk token ids, topk probs)."""
        proj = h.float() @ self.Jbar[layer].T
        logits = self._readout(proj)
        probs = torch.softmax(logits.float(), dim=-1)
        p, idx = probs.topk(topk, dim=-1)
        return idx, p

    @torch.no_grad()
    def read_prompt(self, input_ids, layers, topk=8):
        """Full J-lens table for one prompt: {layer: (ids (T,k), probs (T,k))}."""
        hs = collect_residuals(self.model, input_ids)
        return {l: self.read(hs[l], l, topk=topk) for l in layers}

    @torch.no_grad()
    def concept_prob(self, input_ids, token_ids, layers):
        """P(token) under the lens for each (layer, position). (L, T, n_tokens)."""
        hs = collect_residuals(self.model, input_ids)
        out = []
        for l in layers:
            logits = self._readout(hs[l].float() @ self.Jbar[l].T)
            probs = torch.softmax(logits.float(), dim=-1)
            out.append(probs[:, token_ids])
        return torch.stack(out)

    def jlens_vector(self, token_id, layer):
        """v_s = W_U[s] J_l : residual-stream direction at `layer` for token s."""
        wu = self.model.lm_head.weight[token_id].float()
        return wu @ self.Jbar[layer]

    def swap_hooks(self, pairs, layers=None, thr=0.005, alpha=1.0,
                   write="embed"):
        """Gated concept swap. `pairs` is a list of (src_str, tgt_str) tokens.

        At each (layer, position) where the J-lens reads any source token with
        probability > thr (the workspace "holds" the concept), transfer the
        activation's content from source to target:

        - write="embed" (default): move h's projection on the unit source
          embedding onto the unit target embedding (tied-embedding write basis;
          the J-lens supplies localization). This is the variant that works on
          gemma-4-E2B.
        - write="jlens": the paper's h <- h + V(sigma(c) - c), c = pinv(V) h,
          with V rows = J-lens vectors of source+target tokens.

        Returns (handles, fired) where fired collects (layer, n_positions)."""
        layers = range(4, len(self.tm.layers) - 1) if layers is None else layers
        enc = lambda s: self.tok.encode(s, add_special_tokens=False)[0]
        s_ids = torch.tensor([enc(s) for s, _ in pairs], device="cuda")
        t_ids = torch.tensor([enc(t) for _, t in pairs], device="cuda")
        W = self.model.lm_head.weight
        E_s = torch.nn.functional.normalize(W[s_ids].float(), dim=-1)
        E_t = torch.nn.functional.normalize(W[t_ids].float(), dim=-1)
        handles, fired = [], []
        for l in layers:
            V = pinv = None
            m = len(pairs)
            if write == "jlens":
                V = torch.cat([
                    torch.stack([self.jlens_vector(int(s), l) for s in s_ids]),
                    torch.stack([self.jlens_vector(int(t), l) for t in t_ids])])
                pinv = torch.linalg.pinv(V)

            def hook(mod, inp, out, l=l, V=V, pinv=pinv, m=m):
                h = out.float()
                probs = torch.softmax(
                    self._readout(h @ self.Jbar[l].T).float(), -1)
                gate = probs[..., s_ids].sum(-1) > thr  # (B, T)
                if not gate.any():
                    return None
                if write == "embed":
                    proj = h @ E_s.T
                    delta = alpha * (proj @ E_t - proj @ E_s)
                else:
                    c = h @ pinv
                    sig = torch.cat([c[..., m:], c[..., :m]], dim=-1)
                    delta = alpha * (sig - c) @ V
                fired.append((l, int(gate.sum())))
                return (h + delta * gate.unsqueeze(-1)).to(out.dtype)

            handles.append(self.tm.layers[l].register_forward_hook(hook))
        return handles, fired

    @torch.no_grad()
    def generate(self, input_ids, max_new_tokens=20):
        out = self.model.generate(
            input_ids, max_new_tokens=max_new_tokens, do_sample=False,
            pad_token_id=self.tok.pad_token_id or self.tok.eos_token_id,
        )
        return self.tok.decode(out[0, input_ids.shape[1]:], skip_special_tokens=True)

    @torch.no_grad()
    def generate_swapped(self, input_ids, pairs, layers=None, thr=0.005,
                         alpha=1.0, write="embed", max_new_tokens=20):
        handles, fired = self.swap_hooks(pairs, layers, thr, alpha, write)
        try:
            return self.generate(input_ids, max_new_tokens), fired
        finally:
            for h in handles:
                h.remove()


def chat_ids(tok, user_msg, device="cuda", assistant_prefix=None):
    enc = tok.apply_chat_template(
        [{"role": "user", "content": user_msg}],
        add_generation_prompt=True, return_tensors="pt", return_dict=True)
    ids = enc["input_ids"]
    if assistant_prefix:
        pre = tok(assistant_prefix, add_special_tokens=False, return_tensors="pt")
        ids = torch.cat([ids, pre["input_ids"]], dim=1)
    return ids.to(device)
