"""Original TensorViz reduced causal decoder, inspired by Vaswani et al. (2017).

One post-norm block, width 16, four heads, FFN width 64, vocabulary 32. Sinusoidal
positions and causal masking are retained. There is no encoder/cross-attention,
dropout, tokenizer, trained checkpoint, or claim of language quality.
"""
import math
import torch
import torch.nn as nn


class SinusoidalPositions(nn.Module):
    def forward(self, x):
        positions = torch.arange(x.shape[1], device=x.device, dtype=x.dtype).unsqueeze(1)
        frequencies = torch.exp(torch.arange(0, 16, 2, device=x.device, dtype=x.dtype) * (-math.log(10000.) / 16))
        angles = positions * frequencies
        encoding = torch.stack([angles.sin(), angles.cos()], dim=-1).flatten(-2)
        return x * 4 + encoding


class CausalScores(nn.Module):
    def forward(self, q, k):
        scores = torch.matmul(q, k.transpose(-2, -1)) / 2
        future = torch.ones(scores.shape[-2:], dtype=torch.bool, device=q.device).triu(1)
        return scores.masked_fill(future, torch.finfo(scores.dtype).min)


class CausalAttention(nn.Module):
    def __init__(self):
        super().__init__()
        self.query = nn.Linear(16, 16)
        self.key = nn.Linear(16, 16)
        self.value = nn.Linear(16, 16)
        self.scores = CausalScores()
        self.softmax = nn.Softmax(dim=-1)
        self.project = nn.Linear(16, 16)

    def forward(self, x):
        q = self.query(x).reshape(x.shape[0], x.shape[1], 4, 4).transpose(1, 2)
        k = self.key(x).reshape(x.shape[0], x.shape[1], 4, 4).transpose(1, 2)
        v = self.value(x).reshape(x.shape[0], x.shape[1], 4, 4).transpose(1, 2)
        weights = self.softmax(self.scores(q, k))
        mixed = torch.matmul(weights, v).transpose(1, 2).reshape(x.shape)
        return self.project(mixed)


class FeedForward(nn.Module):
    def __init__(self):
        super().__init__()
        self.expand = nn.Linear(16, 64)
        self.relu = nn.ReLU()
        self.project = nn.Linear(64, 16)

    def forward(self, x):
        return self.project(self.relu(self.expand(x)))


class DecoderBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.attention = CausalAttention()
        self.norm_attention = nn.LayerNorm(16)
        self.ffn = FeedForward()
        self.norm_ffn = nn.LayerNorm(16)

    def forward(self, x):
        attended = self.attention(x)
        x = self.norm_attention(x + attended)
        transformed = self.ffn(x)
        return self.norm_ffn(x + transformed)


class TinyDecoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(32, 16)
        self.positions = SinusoidalPositions()
        self.block = DecoderBlock()
        self.vocabulary = nn.Linear(16, 32)
        self.probabilities = nn.Softmax(dim=-1)

    def forward(self, tokens):
        x = self.embedding(tokens)
        x = self.positions(x)
        x = self.block(x)
        logits = self.vocabulary(x)
        return self.probabilities(logits)


def attention_weights(model, tokens):
    x = model.positions(model.embedding(tokens))
    attention = model.block.attention
    q = attention.query(x).reshape(1, tokens.shape[1], 4, 4).transpose(1, 2)
    k = attention.key(x).reshape(1, tokens.shape[1], 4, 4).transpose(1, 2)
    return attention.softmax(attention.scores(q, k))


def verify():
    torch.manual_seed(0)
    model = TinyDecoder().double().eval()
    tokens = torch.tensor([[1, 2, 3, 4, 5, 6]])
    changed = torch.tensor([[1, 2, 3, 12, 13, 14]])
    with torch.no_grad():
        output, other = model(tokens), model(changed)
        torch.testing.assert_close(output[:, :3], other[:, :3], atol=1e-10, rtol=1e-8)
        assert not torch.allclose(output[:, 3:], other[:, 3:])
        torch.testing.assert_close(output.sum(-1), torch.ones(1, 6, dtype=torch.float64))
        weights = attention_weights(model, tokens)
        assert torch.equal(weights.triu(1), torch.zeros_like(weights))
        torch.testing.assert_close(weights.sum(-1), torch.ones(1, 4, 6, dtype=torch.float64))
    return ["Changing future token IDs leaves earlier predictions unchanged", "Every attention head gives future positions exactly zero weight", "Attention rows and vocabulary distributions each sum to one"]


def experiment():
    torch.manual_seed(0)
    model = TinyDecoder().double().eval()
    tokens = torch.tensor([[1, 2, 3, 4, 5, 6]])
    with torch.no_grad():
        weights = attention_weights(model, tokens)[0, 0]
        predictions = model(tokens)[0]
    cases = []
    for position in range(6):
        cases.append({"id": f"token-{position}", "label": f"Query position {position}", "target": "block.attention.softmax",
            "note": f"Head 0 at position {position} can consult positions 0 through {position}. Later keys receive zero weight. Token IDs identify synthetic vocabulary entries; this model has no learned language ability.",
            "vectors": [{"label": "Attention over key positions 0–5", "values": weights[position].tolist()}, {"label": "Probabilities for vocabulary IDs 0–7 (of 32)", "values": predictions[position, :8].tolist()}],
            "metrics": [{"label": "Visible prefix length", "value": position + 1}, {"label": "Attention row sum", "value": weights[position].sum().item()}, {"label": "All 32 vocabulary probabilities sum to", "value": predictions[position].sum().item()}]})
    return {"kind": "vectors", "title": "A token can look back, but not ahead.", "description": "Recorded head-0 attention and vocabulary probabilities for the fixed token IDs [1, 2, 3, 4, 5, 6], using seed-0 untrained weights.", "controlLabel": "Query position", "cases": cases}