"""Item 31: synthetic memory tokens — per-layer KV prefix (Nils's design).

"The one read mechanism a frozen transformer natively possesses is
attention over the KV cache." A KVMemoryAdapter maps carried L30 states
to (k, v) entries appended to every band layer's attention — frozen
heads consume the state through their existing machinery. Keys live in
post-RoPE space (prefix-tuning convention: learned keys place themselves
where frozen queries already point).

Activation is global-flag-scoped like lora_band.LoopLoRA: a custom
attention interface "kvmem" wraps the stock eager path; when _MEM is
armed and the module is a band layer, memory kv columns are appended
(mask columns visible to every query).
"""
import torch
import torch.nn as nn

from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
from transformers.models.gemma4.modeling_gemma4 import (
    eager_attention_forward)

_MEM = {"on": False, "kv": None, "band": (14, 30), "base": None}


def arm_memory(kv):
    """kv: dict layer_idx -> (k, v), each (B, kv_heads, M, head_dim)."""
    _MEM["kv"] = kv
    _MEM["on"] = kv is not None


def kvmem_attention(module, query, key, value, attention_mask, **kwargs):
    if (_MEM["on"] and _MEM["kv"] is not None
            and getattr(module, "layer_idx", -1) in _MEM["kv"]):
        mk, mv, gate = _MEM["kv"][module.layer_idx]
        m = mk.shape[2]
        key = torch.cat([key, mk.to(dtype=key.dtype, device=key.device)], 2)
        value = torch.cat([value, mv.to(dtype=value.dtype,
                                        device=value.device)], 2)
        if attention_mask is not None:
            pad = torch.zeros(*attention_mask.shape[:-1], m,
                              dtype=attention_mask.dtype,
                              device=attention_mask.device) + gate.to(
                attention_mask.dtype)
            attention_mask = torch.cat([attention_mask, pad], -1)
    base = _MEM["base"] or eager_attention_forward
    return base(module, query, key, value, attention_mask, **kwargs)


def install(model):
    """Wrap the registered implementation IN PLACE — the config name is
    untouched, so mask construction and every other branch stays
    bit-identical; only the (q,k,v,mask) seam gains the memory append."""
    if _MEM.get("installed"):
        return
    impl = model.config._attn_implementation
    _MEM["base"] = ALL_ATTENTION_FUNCTIONS.get_interface(
        impl, eager_attention_forward)
    try:
        ALL_ATTENTION_FUNCTIONS[impl] = kvmem_attention
    except TypeError:
        ALL_ATTENTION_FUNCTIONS.register(impl, kvmem_attention)
    _MEM["installed"] = True


class KVMemoryAdapter(nn.Module):
    """Carried states (B, M, d) -> per-band-layer post-RoPE (k, v).
    Per-layer geometry (head_dim, kv heads) read off the live modules —
    gemma-4 mixes sliding (256) and global (512) head dims in the band."""

    def __init__(self, model, band=(14, 30), code=512):
        super().__init__()
        self.band = band
        tm = model.model.language_model
        d = model.config.get_text_config().hidden_size
        self.trunk = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, code),
                                   nn.GELU())
        self.geom = {}
        self.k_heads = nn.ModuleDict()
        self.v_heads = nn.ModuleDict()
        tc = model.config.get_text_config()
        for l in range(band[0], band[1] + 1):
            att = tm.layers[l].self_attn
            hd = att.head_dim
            kvh = (tc.num_global_key_value_heads
                   if getattr(att, "use_alternative_attention", False)
                   else tc.num_key_value_heads)
            self.geom[l] = (kvh, hd)
            self.k_heads[str(l)] = nn.Linear(code, kvh * hd)
            self.v_heads[str(l)] = nn.Linear(code, kvh * hd)
            nn.init.zeros_(self.v_heads[str(l)].weight)
            nn.init.zeros_(self.v_heads[str(l)].bias)
        # learnable per-layer attention gate on the memory columns:
        # init -10 -> memory is invisible until training opens it
        self.gate = nn.ParameterDict(
            {str(l): nn.Parameter(torch.tensor(-10.0))
             for l in range(band[0], band[1] + 1)})

    def forward(self, states):
        B, M, _ = states.shape
        c = self.trunk(states.float())
        kv = {}
        for l in range(self.band[0], self.band[1] + 1):
            kvh, hd = self.geom[l]
            k = self.k_heads[str(l)](c).view(B, M, kvh, hd).transpose(1, 2)
            v = self.v_heads[str(l)](c).view(B, M, kvh, hd).transpose(1, 2)
            kv[l] = (k, v, self.gate[str(l)])
        return kv