"""Tiny CPU experiments for reading 'Attention Is All You Need'."""

from __future__ import annotations

import math

import torch


def scaled_dot_product_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    allowed: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Implement equation 1 from the paper for tensors ending in ``(T, D)``."""
    scores = query @ key.transpose(-2, -1) / math.sqrt(query.shape[-1])
    if allowed is not None:
        if allowed.dtype != torch.bool or allowed.shape != scores.shape:
            raise ValueError("allowed must be bool and have the same shape as attention scores")
        if not bool(allowed.any(dim=-1).all()):
            raise ValueError("every query needs at least one allowed key")
        scores = scores.masked_fill(~allowed, torch.finfo(scores.dtype).min)
    weights = torch.softmax(scores, dim=-1)
    if allowed is not None:
        weights = weights.masked_fill(~allowed, 0.0)
    return weights @ value, weights


def sinusoidal_encoding(length: int, model_width: int) -> torch.Tensor:
    """Implement the paper's fixed sine/cosine position encoding."""
    if model_width % 2:
        raise ValueError("model_width must be even")
    position = torch.arange(length, dtype=torch.float32)[:, None]
    frequency = torch.exp(
        torch.arange(0, model_width, 2, dtype=torch.float32)
        * (-math.log(10_000.0) / model_width)
    )
    encoding = torch.empty(length, model_width)
    encoding[:, 0::2] = torch.sin(position * frequency)
    encoding[:, 1::2] = torch.cos(position * frequency)
    return encoding


def score_variances(widths: tuple[int, ...], samples: int = 20_000) -> list[tuple[int, float, float]]:
    """Measure why dividing random dot products by ``sqrt(d_k)`` is useful."""
    results = []
    for width in widths:
        query = torch.randn(samples, width)
        key = torch.randn(samples, width)
        dots = (query * key).sum(dim=-1)
        results.append((width, dots.var().item(), (dots / math.sqrt(width)).var().item()))
    return results


def main() -> None:
    torch.manual_seed(1706)
    query = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
    key = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])
    value = torch.tensor([[10.0, 0.0], [0.0, 10.0], [6.0, 6.0]])
    output, weights = scaled_dot_product_attention(query, key, value)
    print("attention weights:\n", weights.round(decimals=4))
    print("weighted values:\n", output.round(decimals=4))
    print("score variance: width, unscaled, scaled")
    for row in score_variances((4, 16, 64)):
        print(row[0], round(row[1], 3), round(row[2], 3))
    print("position encoding shape:", tuple(sinusoidal_encoding(5, 8).shape))


if __name__ == "__main__":
    main()