"""A small real attention computation for CPU examples, not a pretrained VLM.

Image-token vectors are supplied explicitly; this fixture has no image encoder.
"""
import math
import torch
from torch import Tensor, nn


class TinyAttention(nn.Module):
    def __init__(self, width: int):
        super().__init__()
        self.query = nn.Linear(width, width, bias=False)
        self.key = nn.Linear(width, width, bias=False)
        self.value = nn.Linear(width, width, bias=False)
        self.scores = nn.Identity()
        self.probs = nn.Softmax(dim=-1)
        self.output = nn.Linear(width, width, bias=False)

    def forward(self, x: Tensor) -> Tensor:
        q, k, v = self.query(x), self.key(x), self.value(x)
        scores = (q @ k.transpose(-2, -1))[:, None] / math.sqrt(x.shape[-1])
        future = torch.ones(x.shape[1], x.shape[1], dtype=torch.bool, device=x.device).triu(1)
        scores = self.scores(scores.masked_fill(future, -torch.inf))
        probs = self.probs(scores)
        return self.output((probs @ v[:, None]).squeeze(1))


class TinyBlock(nn.Module):
    def __init__(self, width: int):
        super().__init__()
        self.attention = TinyAttention(width)
        self.mlp = nn.Sequential(nn.Linear(width, width), nn.GELU(), nn.Linear(width, width))

    def forward(self, x: Tensor) -> Tensor:
        x = x + self.attention(x)
        return x + self.mlp(x)


class TinyModel(nn.Module):
    def __init__(self, width: int = 6, vocab_size: int = 12):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, width)
        self.layers = nn.ModuleList([TinyBlock(width), TinyBlock(width)])
        self.norm = nn.LayerNorm(width)
        self.lm_head = nn.Linear(width, vocab_size, bias=False)

    def forward(self, input_ids: Tensor, image_tokens: Tensor) -> dict[str, Tensor]:
        hidden = torch.cat([image_tokens, self.embedding(input_ids)], dim=1)
        for layer in self.layers:
            hidden = layer(hidden)
        return {"logits": self.lm_head(self.norm(hidden))}

    def probing_adapter(self):
        """An explicit model protocol: Prober(model) uses this verified contract."""
        from vlm_probing import HeadSite, HookPoint, ModelReadout, ModelSpec, TokenLayout, TorchModelAdapter

        sites, residuals, attentions, scores, heads = {}, {}, {}, {}, {}
        path_heads, path_qkv = {}, {}
        for i in range(len(self.layers)):
            residuals[i], attentions[i] = f"residual.{i}", f"attention.{i}"
            scores[i], heads[i] = f"scores.{i}", f"heads.{i}"
            sites[residuals[i]] = HookPoint(f"layers.{i}")
            sites[attentions[i]] = HookPoint(f"layers.{i}.attention.probs")
            sites[scores[i]] = HookPoint(f"layers.{i}.attention.scores")
            sites[heads[i]] = HookPoint(f"layers.{i}.attention.output")
            path_heads[i] = HeadSite(heads[i], 1, self.norm.normalized_shape[0])
            path_qkv[i] = {}
            for kind, component in (("q", "query"), ("k", "key"), ("v", "value")):
                site = f"{kind}.{i}"
                sites[site] = HookPoint(f"layers.{i}.attention.{component}")
                path_qkv[i][kind] = HeadSite(site, 1, self.norm.normalized_shape[0])
        sites["embeddings"] = HookPoint("layers.0", kind="input", selector=0)
        sites["last_mlp"] = HookPoint(f"layers.{len(self.layers)-1}.mlp")

        def layout(inputs):
            ids = inputs["input_ids"]
            prefix = torch.zeros(inputs["image_tokens"].shape[:2], dtype=ids.dtype, device=ids.device)
            expanded = torch.cat([prefix, ids], dim=1)
            valid = torch.ones_like(expanded, dtype=torch.bool)
            visual = torch.zeros_like(valid)
            visual[:, :prefix.shape[1]] = True
            return TokenLayout(valid, visual, token_ids=expanded)

        def linear_readout(full):
            scale = (full.var(-1, keepdim=True, unbiased=False) + self.norm.eps).rsqrt()
            return self.lm_head.weight * self.norm.weight, scale, True

        spec = ModelSpec(
            residuals, attentions=attentions, attention_scores=scores, heads=heads,
            embeddings="embeddings", input_embeddings=self.embedding.weight,
            layout=layout, project_heads=lambda _, value: value.unsqueeze(-2),
            linear_readout=linear_readout, editable_attention=set(attentions),
            path_heads=path_heads, path_qkv=path_qkv, path_final=residuals[len(self.layers)-1],
            edges={"last_attention_to_residual": heads[len(self.layers)-1],
                   "last_mlp_to_residual": "last_mlp"},
            notes=("Random CPU demonstration model; image embeddings are supplied directly.",),
        )
        return TorchModelAdapter(self, sites, readout=ModelReadout(self.lm_head, norm=self.norm),
                                 model_id="tiny-demo", spec=spec)