"""Original tiny LLaMA-1-style decoder: pre-RMSNorm, RoPE, MHA and SwiGLU.

One layer, width 16, four heads, vocabulary 32. No trained weights, tokenizer,
generation service, or original LLaMA parameter count/training recipe.
"""
import torch
import torch.nn as nn


class RMSNorm(nn.Module):
    def __init__(self):
        super().__init__()
        self.gain = nn.Parameter(torch.ones(16))

    def forward(self, x):
        return x * torch.rsqrt(x.square().mean(-1, keepdim=True) + 1e-6) * self.gain


class Rotary(nn.Module):
    def forward(self, x):
        positions = torch.arange(x.shape[2], device=x.device, dtype=x.dtype)
        frequencies = torch.tensor([1.0, 0.01], device=x.device, dtype=x.dtype)
        angles = positions[:, None] * frequencies[None, :]
        pairs = x.reshape(x.shape[0], 4, x.shape[2], 2, 2)
        a, b = pairs[..., 0], pairs[..., 1]
        return torch.stack(
            [a * angles.cos() - b * angles.sin(), a * angles.sin() + b * angles.cos()],
            -1,
        ).flatten(-2)


class Attention(nn.Module):
    def __init__(self):
        super().__init__()
        self.query = nn.Linear(16, 16, bias=False)
        self.key = nn.Linear(16, 16, bias=False)
        self.value = nn.Linear(16, 16, bias=False)
        self.query_rope = Rotary()
        self.key_rope = Rotary()
        self.softmax = nn.Softmax(dim=-1)
        self.project = nn.Linear(16, 16, bias=False)

    def forward(self, x):
        b, t, _ = x.shape
        q = self.query_rope(self.query(x).reshape(b, t, 4, 4).transpose(1, 2))
        k = self.key_rope(self.key(x).reshape(b, t, 4, 4).transpose(1, 2))
        v = self.value(x).reshape(b, t, 4, 4).transpose(1, 2)
        scores = q @ k.transpose(-2, -1) / 2
        mask = torch.ones(t, t, dtype=torch.bool, device=x.device).triu(1)
        weights = self.softmax(scores.masked_fill(mask, float("-inf")))
        return self.project((weights @ v).transpose(1, 2).reshape(b, t, 16))


class SwiGLU(nn.Module):
    def __init__(self):
        super().__init__()
        self.gate = nn.Linear(16, 40, bias=False)
        self.up = nn.Linear(16, 40, bias=False)
        self.silu = nn.SiLU()
        self.down = nn.Linear(40, 16, bias=False)

    def forward(self, x):
        return self.down(self.silu(self.gate(x)) * self.up(x))


class DecoderBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.attention_norm = RMSNorm()
        self.attention = Attention()
        self.ffn_norm = RMSNorm()
        self.ffn = SwiGLU()

    def forward(self, x):
        hidden = x + self.attention(self.attention_norm(x))
        return hidden + self.ffn(self.ffn_norm(hidden))


class TinyLlama(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(32, 16)
        self.block = DecoderBlock()
        self.final_norm = RMSNorm()
        self.lm_head = nn.Linear(16, 32, bias=False)

    def forward(self, tokens):
        x = self.embedding(tokens)
        hidden = self.block(x)
        return self.lm_head(self.final_norm(hidden))


def verify():
    torch.manual_seed(8)
    model = TinyLlama().double().eval()
    tokens = torch.tensor([[1, 2, 3, 4, 5, 6]])
    with torch.no_grad():
        changed = tokens.clone()
        changed[:, 4:] = 7
        torch.testing.assert_close(model(tokens)[:, :4], model(changed)[:, :4])
        for length in range(1, 7):
            torch.testing.assert_close(
                model(tokens[:, :length]), model(tokens)[:, :length]
            )
        x = model.embedding(tokens)
        model.block.attention.project.weight.zero_()
        model.block.ffn.down.weight.zero_()
        torch.testing.assert_close(model.block(x), x, rtol=0, atol=0)
        pair = torch.randn(1, 4, 6, 4, dtype=torch.float64)
        torch.testing.assert_close(Rotary()(pair).norm(dim=-1), pair.norm(dim=-1))
    return [
        "Later token changes leave earlier causal logits unchanged",
        "Zero attention/FFN output projections make the entire pre-norm block an identity",
        "RoPE preserves each head's query/key norm",
        "Prefix evaluation matches the corresponding full-sequence outputs",
    ]


def experiment():
    torch.manual_seed(8)
    model = TinyLlama().double().eval()
    tokens = torch.tensor([[1, 2, 3, 4, 5, 6]])
    cases = []
    with torch.no_grad():
        x = model.embedding(tokens)
        a = model.block.attention_norm(x)
        after_attention = x + model.block.attention(a)
        f = model.block.ffn_norm(after_attention)
        out = model.block(x)
        for identity, label, target, values in [
            ("embedding", "Token embedding", "embedding", x),
            ("norm", "Before attention: RMSNorm", "block.attention_norm", a),
            (
                "attention",
                "After attention residual",
                "block.attention",
                after_attention,
            ),
            ("ffn", "Before SwiGLU: RMSNorm", "block.ffn_norm", f),
            ("residual", "After SwiGLU residual", "block.ffn", out),
        ]:
            cases.append(
                {
                    "id": identity,
                    "label": label,
                    "target": target,
                    "note": "These are recorded activations for the final token of the same six-token untrained decoder. RMSNorm controls feature scale inside each branch; residual additions preserve a direct input path. The vocabulary IDs have no language semantics.",
                    "vectors": [
                        {
                            "label": "Final token · 16 features",
                            "values": values[0, -1].tolist(),
                        }
                    ],
                    "metrics": [
                        {
                            "label": "Feature RMS",
                            "value": values[0, -1].square().mean().sqrt().item(),
                        },
                        {"label": "Model width", "value": 16},
                        {"label": "Attention heads", "value": 4},
                    ],
                }
            )
    return {
        "kind": "vectors",
        "title": "See the components in one decoder.",
        "description": "A small synthesis of LLaMA-1 architectural ingredients. It uses ordinary multi-head attention; later LLaMA-family models vary their head-sharing and training choices.",
        "controlLabel": "Activation checkpoint",
        "cases": cases,
    }