# examples/transformer.py
"""Offline reduced encoder-decoder mechanics from Vaswani et al. (2017).

Paper reference: https://arxiv.org/html/1706.03762v7, sections 3.1-3.5
and 5.4. This authored implementation follows the paper's post-norm order,
ReLU FFN, interleaved sinusoids, scaled and shared embeddings. It does not
inherit later Tensor2Tensor pre-norm defaults, train a translator, or reproduce
the optimizer schedule, label smoothing, beam search, or benchmark results.
Inputs are finite CPU tensors. True in a padding mask means blocked. A fully
masked attention row raises ValueError before softmax, including combined
causal/padding masks. Empty sequences are rejected rather than invented.
"""

import argparse
import hashlib
import json
import math
from pathlib import Path

import torch
from torch import nn
from torch.nn import functional as F

PAPER_URL = "https://arxiv.org/html/1706.03762v7"
COMPARISON_REVISION = "15ea434070b341e36b356ff7224ba7d247cac8a4"


def positive_integer(value, name):
    """Reject invalid dimensions without silently truncating them."""
    if isinstance(value, bool) or not isinstance(value, int) or value < 1:
        raise ValueError(f"{name} must be a positive integer")


def validate_features(x, name):
    """Check the deliberately CPU-only attention API."""
    if not isinstance(x, torch.Tensor) or x.ndim != 3:
        raise ValueError(f"{name} must have shape [B, N, D]")
    if min(x.shape) < 1 or x.device.type != "cpu" or not x.is_floating_point():
        raise ValueError(f"{name} must be a nonempty floating CPU tensor")
    if not torch.isfinite(x).all():
        raise ValueError(f"{name} must be finite")


# region: positions_and_embeddings
def sinusoidal_positions(length, d_model, *, dtype=torch.float32):
    """Return [N, D] with paper section 3.5 even/odd coordinates."""
    positive_integer(length, "length")
    positive_integer(d_model, "d_model")
    if not dtype.is_floating_point:
        raise ValueError("positional dtype must be floating")
    positions = torch.arange(length, dtype=dtype)[:, None]
    even = torch.arange(0, d_model, 2, dtype=dtype)
    angles = positions / (10000.0 ** (even / d_model))
    result = torch.empty(length, d_model, dtype=dtype)
    result[:, 0::2] = torch.sin(angles)
    # For odd D, the last sine has no corresponding cosine coordinate.
    result[:, 1::2] = torch.cos(angles[:, :d_model // 2])
    return result


def embed_positions(tokens, embedding, dropout):
    """Scale token embeddings, add positions, then apply residual dropout."""
    x = embedding(tokens) * math.sqrt(embedding.embedding_dim)
    positions = sinusoidal_positions(x.shape[1], x.shape[2], dtype=x.dtype)
    return dropout(x + positions[None])
# endregion: positions_and_embeddings


# region: project_qkv
def project_qkv(queries, key_values, q_projection, k_projection, v_projection):
    """Project queries [B,Nq,D] and a separate K/V source [B,Nk,D]."""
    validate_features(queries, "queries")
    validate_features(key_values, "key_values")
    if queries.shape[0] != key_values.shape[0]:
        raise ValueError("query and key/value batch sizes must match")
    if queries.dtype != key_values.dtype:
        raise ValueError("query and key/value dtypes must match")
    if (queries.shape[-1] != q_projection.in_features
            or key_values.shape[-1] != k_projection.in_features
            or key_values.shape[-1] != v_projection.in_features):
        raise ValueError("feature width must match the projection inputs")
    return (q_projection(queries), k_projection(key_values),
            v_projection(key_values))
# endregion: project_qkv


# region: split_heads
def split_heads(x, num_heads):
    """Map [B,N,D] to [B,H,N,Dh] without mixing token positions."""
    validate_features(x, "projected features")
    positive_integer(num_heads, "num_heads")
    batch, length, width = x.shape
    if width % num_heads:
        raise ValueError("d_model must be divisible by num_heads")
    shaped = x.reshape(batch, length, num_heads, width // num_heads)
    return shaped.transpose(1, 2)
# endregion: split_heads


# region: scaled_scores
def scaled_scores(q, k):
    """Compute [B,H,Nq,Nk] scores with per-head, not model-width, scale."""
    if q.ndim != 4 or k.ndim != 4 or min(q.shape + k.shape) < 1:
        raise ValueError("Q and K must be nonempty [B,H,N,Dh]")
    if q.shape[:2] != k.shape[:2] or q.shape[-1] != k.shape[-1]:
        raise ValueError("Q and K batch, head count and head width must match")
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
    if not torch.isfinite(scores).all():
        raise ValueError("attention scores must be finite before masking")
    return scores
# endregion: scaled_scores


# region: apply_mask
def apply_mask(scores, key_padding=None, *, causal=False):
    """Block padded keys and future positions; reject fully masked rows."""
    if scores.ndim != 4 or min(scores.shape) < 1:
        raise ValueError("scores must be nonempty [B,H,Nq,Nk]")
    if not scores.is_floating_point() or not torch.isfinite(scores).all():
        raise ValueError("scores must be finite floating values")
    batch, _, query_length, key_length = scores.shape
    allowed = torch.ones(batch, 1, query_length, key_length, dtype=torch.bool)
    if key_padding is not None:
        if (not isinstance(key_padding, torch.Tensor)
                or key_padding.dtype != torch.bool
                or key_padding.shape != (batch, key_length)):
            raise ValueError(
                "key_padding must be boolean [B,Nk], True = blocked")
        if key_padding.device.type != "cpu":
            raise ValueError("key_padding must be a CPU tensor")
        allowed = allowed & ~key_padding[:, None, None, :]
    if causal:
        if query_length != key_length:
            raise ValueError("causal mask is for equal-length self-attention")
        # Lower triangle includes the diagonal, hence the current input token.
        triangle = torch.ones(
            query_length, key_length, dtype=torch.bool).tril()
        allowed = allowed & triangle
    if not allowed.any(dim=-1).all():
        raise ValueError("fully masked attention row has no valid softmax")
    return scores.masked_fill(~allowed, -torch.inf)
# endregion: apply_mask


# region: attend_and_merge
def attend_and_merge(masked_scores, values, output_projection):
    """Softmax, mix V, concatenate heads, then apply learned W_O."""
    if (masked_scores.ndim != 4 or values.ndim != 4
            or min(masked_scores.shape + values.shape) < 1
            or masked_scores.shape[:2] != values.shape[:2]
            or masked_scores.shape[-1] != values.shape[-2]
            or values.shape[1] * values.shape[-1]
            != output_projection.in_features):
        raise ValueError("score, value and output-projection shapes must match")
    if (torch.isnan(masked_scores).any() or torch.isposinf(masked_scores).any()
            or not torch.isfinite(masked_scores).any(dim=-1).all()
            or not torch.isfinite(values).all()):
        raise ValueError(
            "attention needs finite values and an unmasked key per row")
    weights = torch.softmax(masked_scores, dim=-1)
    context = weights @ values
    batch, heads, length, head_width = context.shape
    merged = context.transpose(1, 2).contiguous().reshape(
        batch, length, heads * head_width)
    return output_projection(merged), weights
# endregion: attend_and_merge


class MultiHeadAttention(nn.Module):
    """Reduced multi-head attention with explicit paper operations."""

    def __init__(self, d_model, num_heads):
        super().__init__()
        positive_integer(d_model, "d_model")
        positive_integer(num_heads, "num_heads")
        if d_model % num_heads:
            raise ValueError("d_model must be divisible by num_heads")
        self.num_heads = num_heads
        self.q = nn.Linear(d_model, d_model, bias=False)
        self.k = nn.Linear(d_model, d_model, bias=False)
        self.v = nn.Linear(d_model, d_model, bias=False)
        self.out = nn.Linear(d_model, d_model, bias=False)

    def forward(self, queries, key_values, key_padding=None, *, causal=False):
        """Return projected attention output and per-head probabilities."""
        q, k, v = project_qkv(queries, key_values, self.q, self.k, self.v)
        q, k, v = [split_heads(x, self.num_heads) for x in (q, k, v)]
        scores = apply_mask(scaled_scores(q, k), key_padding, causal=causal)
        return attend_and_merge(scores, v, self.out)


# region: encoder_layer
class EncoderLayer(nn.Module):
    """Original post-norm self-attention and position-wise ReLU FFN."""

    def __init__(self, d_model, heads, d_ff, dropout=0.1):
        super().__init__()
        positive_integer(d_ff, "d_ff")
        self.attention = MultiHeadAttention(d_model, heads)
        self.ffn = nn.Sequential(nn.Linear(d_model, d_ff), nn.ReLU(),
                                 nn.Linear(d_ff, d_model))
        self.norm1 = nn.LayerNorm(d_model, eps=1e-6)
        self.norm2 = nn.LayerNorm(d_model, eps=1e-6)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, source_padding=None):
        """LayerNorm is after addition, never before the attention or FFN."""
        attention, weights = self.attention(x, x, source_padding)
        x = self.norm1(x + self.dropout(attention))
        x = self.norm2(x + self.dropout(self.ffn(x)))
        return x, weights
# endregion: encoder_layer


# region: decoder_layer
class DecoderLayer(nn.Module):
    """Masked self-attention, encoder K/V cross-attention, then ReLU FFN."""

    def __init__(self, d_model, heads, d_ff, dropout=0.1):
        super().__init__()
        positive_integer(d_ff, "d_ff")
        self.self_attention = MultiHeadAttention(d_model, heads)
        self.cross_attention = MultiHeadAttention(d_model, heads)
        self.ffn = nn.Sequential(nn.Linear(d_model, d_ff), nn.ReLU(),
                                 nn.Linear(d_ff, d_model))
        self.norm1 = nn.LayerNorm(d_model, eps=1e-6)
        self.norm2 = nn.LayerNorm(d_model, eps=1e-6)
        self.norm3 = nn.LayerNorm(d_model, eps=1e-6)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, memory, target_padding=None, source_padding=None):
        """Memory length Nk may differ from decoder query length Nq."""
        attended, self_weights = self.self_attention(
            x, x, target_padding, causal=True)
        x = self.norm1(x + self.dropout(attended))
        attended, cross_weights = self.cross_attention(
            x, memory, source_padding)
        x = self.norm2(x + self.dropout(attended))
        x = self.norm3(x + self.dropout(self.ffn(x)))
        return x, self_weights, cross_weights
# endregion: decoder_layer


def validate_tokens(tokens, vocab_size, name):
    """Check synthetic token IDs without clamping malformed input."""
    if (not isinstance(tokens, torch.Tensor) or tokens.ndim != 2
            or tokens.dtype != torch.long or tokens.device.type != "cpu"
            or min(tokens.shape) < 1):
        raise ValueError(f"{name} must be nonempty CPU int64 [B,N]")
    if (tokens < 0).any() or (tokens >= vocab_size).any():
        raise ValueError(f"{name} IDs must be in [0, vocab_size)")
    if not (tokens != 0).any(dim=1).all():
        raise ValueError(f"{name} must contain a nonpadding token in each row")


# region: numerical_forward
def shift_targets(targets, bos_id=1):
    """Input at position i is BOS or target i-1, never target i."""
    bos = torch.full_like(targets[:, :1], bos_id)
    return torch.cat((bos, targets[:, :-1]), dim=1)


class ReducedTransformer(nn.Module):
    """One encoder/decoder layer, D=8, H=2, FFN=16, vocabulary=13 by default."""

    def __init__(self, vocab_size=13, d_model=8, heads=2, d_ff=16, dropout=0.1):
        super().__init__()
        positive_integer(vocab_size, "vocab_size")
        if vocab_size < 3:
            raise ValueError(
                "vocabulary needs PAD=0, BOS=1 and a content token")
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.embedding_dropout = nn.Dropout(dropout)
        self.encoder = EncoderLayer(d_model, heads, d_ff, dropout)
        self.decoder = DecoderLayer(d_model, heads, d_ff, dropout)

    def forward(self, source, targets):
        """Return states/logits and masks' probabilities for teacher forcing."""
        validate_tokens(source, self.embedding.num_embeddings, "source")
        validate_tokens(targets, self.embedding.num_embeddings, "targets")
        if source.shape[0] != targets.shape[0]:
            raise ValueError("source and target batch sizes must match")
        decoder_input = shift_targets(targets)
        source_padding = source == 0
        target_padding = decoder_input == 0
        memory, encoder_weights = self.encoder(
            embed_positions(source, self.embedding, self.embedding_dropout),
            source_padding)
        decoded, self_weights, cross_weights = self.decoder(
            embed_positions(
                decoder_input, self.embedding, self.embedding_dropout),
            memory, target_padding, source_padding)
        # One shared matrix for source embedding, target embedding and readout.
        logits = F.linear(decoded, self.embedding.weight)
        return {
            "memory": memory, "decoded": decoded, "logits": logits,
            "decoder_input": decoder_input, "encoder_weights": encoder_weights,
            "self_weights": self_weights, "cross_weights": cross_weights,
        }
# endregion: numerical_forward


def smoke(seed=7):
    """Run seeded offline forward/backward and one illustrative SGD step."""
    with torch.random.fork_rng(devices=[]):
        torch.manual_seed(seed)
        model = ReducedTransformer().eval()  # Disable dropout for exact checks.
        source = torch.tensor([[2, 3, 4, 0, 0], [5, 6, 7, 8, 0]])
        targets = torch.tensor([[9, 10, 11], [4, 5, 0]])
        states = model(source, targets)
        changed_targets = targets.clone()
        changed_targets[:, 1] = 12
        changed_source = source.clone()
        changed_source[source == 0] = 12
        # Preserve the original padding mask when perturbing the padded vectors.
        altered_vectors = embed_positions(
            changed_source, model.embedding, model.embedding_dropout)
        altered_memory, _ = model.encoder(altered_vectors, source == 0)
        altered_decoded, _, _ = model.decoder(
            embed_positions(states["decoder_input"], model.embedding,
                            model.embedding_dropout),
            altered_memory, states["decoder_input"] == 0, source == 0)
        future_error = (model(source, changed_targets)["logits"][:, :2]
                        - states["logits"][:, :2]).abs().max()
        padding_error = (altered_decoded - states["decoded"]).abs().max()
        loss = F.cross_entropy(
            states["logits"].flatten(0, 1), targets.flatten(), ignore_index=0)
        optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
        optimizer.zero_grad()
        loss.backward()
        gradients = {name: float(parameter.grad.norm())
                     for name, parameter in model.named_parameters()}
        if not all(
                parameter.grad is not None
                and torch.isfinite(parameter.grad).all()
                for parameter in model.parameters()):
            raise RuntimeError("nonfinite or missing gradient")
        before = model.embedding.weight.detach().clone()
        optimizer.step()
        return {
            "source_sha256": (
                hashlib.sha256(Path(__file__).read_bytes()).hexdigest()),
            "seed": seed, "device": "cpu", "depth_per_stack": 1,
            "d_model": 8, "heads": 2, "d_head": 4, "d_ff": 16, "vocab_size": 13,
            "dropout_configured": 0.1, "dropout_active": False,
            "norm_order": "post",
            "layer_norm_epsilon": 1e-6,
            "reductions": {
                "architecture": "one encoder and decoder layer; width 8; two heads",
                "inputs": "synthetic token IDs; vocabulary 13; CPU only",
                "evaluation": "dropout disabled for deterministic observations",
                "optimization": "one SGD step; no label smoothing or beam search",
            },
            "loss_reduction": "mean over nonpadding target tokens",
            "shapes": {key: list(value.shape) for key, value in states.items()},
            "loss": float(loss.detach()), "gradient_norms": gradients,
            "future_target_max_error": float(future_error.detach()),
            "padded_source_max_error": float(padding_error.detach()),
            "masked_cross_probability": float(
                states["cross_weights"][..., -1].max().detach()),
            "optimizer_steps": 1,
            "embedding_update_norm": float(
                (before - model.embedding.weight.detach()).norm()),
            "boundary": (
                "random reduced mechanics; no translation or learned semantics"),
        }


def main():
    """Print smoke evidence only when explicitly invoked."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true")
    parser.add_argument("--seed", type=int, default=7)
    arguments = parser.parse_args()
    if not arguments.smoke:
        parser.error("use --smoke for the offline numerical reference")
    print(json.dumps(smoke(arguments.seed), indent=2, allow_nan=False))


if __name__ == "__main__":
    main()