"""Helpers for the hands-on labs: load the model in transformers, read what a layer computes, or change it.

Three things you can do to a running model with a hook:
  record_outputs   copy what a whole layer (attention or DeltaNet mixer + MLP) adds up to, per token
  record_mlp_acts  copy the 3,584 neuron activations that feed the MLP's down_proj, per token
  steer_down       change what down_proj writes, per token, without touching any weight
"""
from __future__ import annotations

import contextlib
import json
import os

import torch

LAB = os.environ.get("WSP_LAB", "/nvme/wsp-lab")


def model_dir(name: str) -> str:
    if name == "base":
        return os.path.join(LAB, "models", "Qwen3.5-0.8B")
    return name if name.startswith("/") else os.path.join(LAB, "models", name)


def load(name="base", dtype=torch.bfloat16, device="cuda"):
    from transformers import AutoModelForCausalLM, AutoTokenizer
    tok = AutoTokenizer.from_pretrained(model_dir(name))
    m = AutoModelForCausalLM.from_pretrained(model_dir(name), dtype=dtype).to(device).eval()
    return m, tok


def blocks(m):
    return m.model.layers


def down_proj(m, layer: int):
    return blocks(m)[layer].mlp.down_proj


def encode(tok, text: str, device="cuda"):
    return tok(text, return_tensors="pt", add_special_tokens=False).input_ids.to(device)


@contextlib.contextmanager
def record_outputs(m, layers):
    """out[l] = list of [tokens, 1024] tensors: the residual stream right after block l (one per forward call)."""
    out, hs = {l: [] for l in layers}, []
    for l in layers:
        def hook(mod, inp, o, l=l):
            o = o[0] if isinstance(o, tuple) else o
            out[l].append(o.detach()[0].float().cpu())
        hs.append(blocks(m)[l].register_forward_hook(hook))
    try:
        yield out
    finally:
        for h in hs:
            h.remove()


@contextlib.contextmanager
def record_mlp_acts(m, layers):
    """out[l] = list of [tokens, 3584] tensors: the neuron activations that down_proj reads at layer l."""
    out, hs = {l: [] for l in layers}, []
    for l in layers:
        def hook(mod, inp, l=l):
            out[l].append(inp[0].detach()[0].float().cpu())
        hs.append(down_proj(m, l).register_forward_pre_hook(hook))
    try:
        yield out
    finally:
        for h in hs:
            h.remove()


@contextlib.contextmanager
def steer_down(m, layer: int, u: torch.Tensor, alpha: float, cols=None, record=None):
    """Make down_proj at `layer` write  W a - alpha * u * s,  where  s = u . (W[:, cols] a[cols]).

    This is exactly what the weight edit W' = W - alpha * u u^T W[:, cols] does, computed on the fly.
    s is one number per token: how hard the chosen neurons are pushing along u right now.
    If `record` is a list, each call appends s (a [tokens] tensor) to it."""
    lin = down_proj(m, layer)
    W = lin.weight
    u = (u / u.norm()).to(W.device, torch.float32)
    idx = None if cols is None else torch.as_tensor(list(cols), device=W.device)
    WS = (W if idx is None else W[:, idx]).float()
    r = WS.T @ u                                   # one number per chosen neuron: how much it writes along u

    def hook(mod, inp, o):
        a = inp[0].float()
        s = (a if idx is None else a[..., idx]) @ r  # [batch, tokens]
        if record is not None:
            record.append(s.detach()[0].cpu())
        return (o.float() - alpha * s[..., None] * u).to(o.dtype)

    h = lin.register_forward_hook(hook)
    try:
        yield
    finally:
        h.remove()


def released_edit(path=None):
    """u, alpha and the 64 columns per layer, as recovered from the released weights by scripts/build_models.py."""
    path = path or os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "results",
                                "edit_info.json")
    info = json.load(open(path))
    return {int(l): {"u": torch.tensor(v["u"]), "alpha": v["alpha"], "cols": v["cols"]}
            for l, v in info["layers"].items()}