"""Differentiable counterparts and a small encoder-decoder transformer.

Install numpy and torch. Run: python torch_core.py
No nn.Transformer wrapper: the attention and block operations are explicit.
"""
from __future__ import annotations

import math
import numpy as np
import torch
from torch import Tensor, nn
from torch.nn import functional as F

# region embeddings
def embedding_demo() -> tuple[Tensor, nn.Embedding, Tensor]:
    vocab = {"<pad>": 0, "<bos>": 1, "<eos>": 2, "the": 3,
             "small": 4, "robot": 5, "moves": 6, "<unk>": 7}
    tokens = "the small robot moves".split()
    ids = torch.tensor([vocab.get(token, vocab["<unk>"]) for token in tokens])
    # Match the NumPy weights exactly, rather than assuming RNGs match.
    weights = torch.from_numpy(np.random.default_rng(0).normal(size=(len(vocab), 6)))
    table = nn.Embedding.from_pretrained(weights, freeze=False)
    x = table(ids)
    torch.testing.assert_close(x[2], table.weight[vocab["robot"]])
    return ids, table, x
# endregion embeddings

# region positions
def positions(length: int, width: int, start: int = 0, device: torch.device | str | None = None, dtype: torch.dtype = torch.float32) -> Tensor:
    if width <= 0 or width % 2:
        raise ValueError("This example uses a positive, even model width")
    pos = torch.arange(start, start + length, device=device, dtype=dtype)[:, None]
    frequency = 10000.0 ** (-torch.arange(0, width, 2, device=device, dtype=dtype) / width)
    angles = pos * frequency[None, :]
    pe = torch.empty(length, width, device=device, dtype=dtype)
    pe[:, 0::2] = torch.sin(angles)
    pe[:, 1::2] = torch.cos(angles)
    return pe
# endregion positions

# region attention
def attention(q: Tensor, k: Tensor, v: Tensor, allowed: Tensor | None = None) -> tuple[Tensor, Tensor]:
    """(..., T, dk), (..., S, dk), (..., S, dv). True means visible."""
    if q.shape[-1] != k.shape[-1] or k.shape[-2] != v.shape[-2]:
        raise ValueError("Q/K feature widths and K/V source lengths must match")
    scores = (q @ k.transpose(-2, -1)) / math.sqrt(q.shape[-1])
    if allowed is not None:
        visible = torch.broadcast_to(allowed.to(device=q.device, dtype=torch.bool), scores.shape)
        if not visible.any(dim=-1).all():
            raise ValueError("Every query needs at least one visible key")
        scores = scores.masked_fill(~visible, float("-inf"))
    weights = scores.softmax(dim=-1)
    return weights @ v, weights
# endregion attention

# region masks
def attention_masks(source_ids: Tensor, target_ids: Tensor, pad: int = 0) -> tuple[Tensor, Tensor]:
    source_visible = (source_ids != pad)[:, None, None, :]
    target_visible = (target_ids != pad)[:, None, None, :]
    length = target_ids.shape[1]
    index = torch.arange(length, device=target_ids.device)
    causal = index[None, :] <= index[:, None]
    target_visible = target_visible & causal[None, None, :, :]
    return source_visible, target_visible
# endregion masks

# region heads
class MultiHead(nn.Module):
    def __init__(self, width: int, heads: int, seed: int = 0) -> None:
        super().__init__()
        if heads <= 0 or width % heads:
            raise ValueError("Model width must be divisible by head count")
        self.heads, self.head_width = heads, width // heads
        self.wq, self.wk, self.wv, self.wo = [nn.Linear(width, width, bias=False) for _ in range(4)]
        # Shared initialization makes the two libraries directly comparable.
        rng = np.random.default_rng(seed)
        with torch.no_grad():
            for layer in [self.wq, self.wk, self.wv, self.wo]:
                w = rng.normal(size=(width, width)) / np.sqrt(width)
                layer.weight.copy_(torch.tensor(w.T, dtype=layer.weight.dtype))

    def split(self, x: Tensor) -> Tensor:
        batch, length, _ = x.shape
        return x.reshape(batch, length, self.heads, self.head_width).transpose(1, 2)

    def forward(self, query: Tensor, context: Tensor, allowed: Tensor | None = None) -> tuple[Tensor, Tensor]:
        q, k, v = self.split(self.wq(query)), self.split(self.wk(context)), self.split(self.wv(context))
        heads, weights = attention(q, k, v, allowed)
        batch, _, length, _ = heads.shape
        joined = heads.transpose(1, 2).reshape(batch, length, -1)
        return self.wo(joined), weights
# endregion heads

# region norm_ffn
def layer_norm(x: Tensor, gamma: Tensor | None = None, beta: Tensor | None = None, eps: float = 1e-5) -> Tensor:
    # normalized_shape=(D,) means statistics over features of EACH token.
    return F.layer_norm(x, (x.shape[-1],), weight=gamma, bias=beta, eps=eps)


def feed_forward(x: Tensor, w1: Tensor, b1: Tensor, w2: Tensor, b2: Tensor) -> Tensor:
    hidden = torch.relu(x @ w1 + b1)
    return hidden @ w2 + b2
# endregion norm_ffn

# region encoder
class EncoderBlock(nn.Module):
    def __init__(self, width: int = 32, heads: int = 4, inner: int = 64, seed: int = 10) -> None:
        super().__init__()
        self.attn = MultiHead(width, heads, seed)
        self.norm1, self.norm2 = nn.LayerNorm(width), nn.LayerNorm(width)
        self.ffn = nn.Sequential(nn.Linear(width, inner), nn.ReLU(), nn.Linear(inner, width))

    def forward(self, x: Tensor, source_visible: Tensor) -> Tensor:
        normalized = self.norm1(x)
        x = x + self.attn(normalized, normalized, source_visible)[0]
        return x + self.ffn(self.norm2(x))
# endregion encoder

# region decoder
class DecoderBlock(nn.Module):
    def __init__(self, width: int = 32, heads: int = 4, inner: int = 64, seed: int = 20) -> None:
        super().__init__()
        self.self_attn = MultiHead(width, heads, seed)
        self.cross_attn = MultiHead(width, heads, seed + 1)
        self.norm1, self.norm2, self.norm3 = [nn.LayerNorm(width) for _ in range(3)]
        self.ffn = nn.Sequential(nn.Linear(width, inner), nn.ReLU(), nn.Linear(inner, width))

    def forward(self, y: Tensor, memory: Tensor, target_visible: Tensor, source_visible: Tensor) -> Tensor:
        normalized = self.norm1(y)
        y = y + self.self_attn(normalized, normalized, target_visible)[0]
        y = y + self.cross_attn(self.norm2(y), memory, source_visible)[0]
        return y + self.ffn(self.norm3(y))
# endregion decoder

# region architectures
def architecture_demo() -> tuple[Tensor, Tensor]:
    source = torch.tensor([[4, 7, 9], [6, 8, 0]])
    prefix = torch.tensor([[1, 4, 7], [1, 6, 8]])
    width, vocabulary = 32, 12
    embedding = nn.Embedding(vocabulary, width)
    source_mask, target_mask = attention_masks(source, prefix)
    x = embedding(source) + positions(source.shape[1], width)
    y = embedding(prefix) + positions(prefix.shape[1], width)

    # Encoder-only: pool only real source positions for classification.
    encoded = layer_norm(EncoderBlock()(x, source_mask))
    real = (source != 0)[..., None]
    pooled = (encoded * real).sum(1) / real.sum(1)
    class_logits = nn.Linear(width, 3)(pooled)

    # Decoder-only: same two-sublayer structure, with CAUSAL visibility.
    causal_hidden = layer_norm(EncoderBlock(seed=30)(y, target_mask))
    next_token_logits = nn.Linear(width, vocabulary)(causal_hidden)
    return class_logits, next_token_logits  # (2, 3), (2, 3, 12)
# endregion architectures

# region model
class TinyTransformer(nn.Module):
    def __init__(self, vocabulary: int = 12, width: int = 32, heads: int = 4, inner: int = 64) -> None:
        super().__init__()
        self.embedding = nn.Embedding(vocabulary, width, padding_idx=0)
        self.encoder = EncoderBlock(width, heads, inner)
        self.decoder = DecoderBlock(width, heads, inner)
        self.encoder_norm, self.decoder_norm = nn.LayerNorm(width), nn.LayerNorm(width)
        self.readout = nn.Linear(width, vocabulary)

    def embed(self, ids: Tensor) -> Tensor:
        x = self.embedding(ids)
        return x + positions(ids.shape[1], x.shape[-1], device=x.device, dtype=x.dtype)

    def encode(self, source: Tensor) -> tuple[Tensor, Tensor]:
        visible = (source != 0)[:, None, None, :]
        memory = self.encoder_norm(self.encoder(self.embed(source), visible))
        return memory, visible

    def decode(self, target_prefix: Tensor, memory: Tensor, source_visible: Tensor) -> Tensor:
        _, target_visible = attention_masks(target_prefix, target_prefix)
        y = self.decoder(self.embed(target_prefix), memory, target_visible, source_visible)
        return self.readout(self.decoder_norm(y))

    def forward(self, source: Tensor, target_prefix: Tensor) -> Tensor:
        memory, visible = self.encode(source)
        return self.decode(target_prefix, memory, visible)
# endregion model

# region generate
@torch.no_grad()
def generate(model: TinyTransformer, source: Tensor, max_new_tokens: int = 4, bos: int = 1, eos: int = 2, pad: int = 0) -> Tensor:
    was_training = model.training
    model.eval()
    try:
        memory, visible = model.encode(source)
        prefix = torch.full((source.shape[0], 1), bos, dtype=torch.long, device=source.device)
        finished = torch.zeros(source.shape[0], dtype=torch.bool, device=source.device)
        for _ in range(max_new_tokens):
            logits = model.decode(prefix, memory, visible)[:, -1, :]
            token = logits.argmax(dim=-1)  # Greedy decoding for this experiment.
            token = torch.where(finished, pad, token)
            prefix = torch.cat([prefix, token[:, None]], dim=1)
            finished |= token == eos
            if finished.all():
                break
        return prefix[:, 1:]
    finally:
        model.train(was_training)
# endregion generate

if __name__ == "__main__":
    torch.manual_seed(7)
    model = TinyTransformer()
    source = torch.tensor([[4, 7, 9], [6, 8, 0]])
    prefix = torch.tensor([[1, 4, 7], [1, 6, 8]])
    print("Vocabulary logits:", tuple(model(source, prefix).shape))