"""A small, explicit causal self-attention layer with a correct KV cache."""

from __future__ import annotations

import math
from dataclasses import dataclass

import torch
from torch import nn


@dataclass(frozen=True)
class KVCache:
    """Projected keys, values, and valid-key flags for one attention layer."""

    key: torch.Tensor  # (B, H, L, Dh)
    value: torch.Tensor  # (B, H, L, Dh)
    valid: torch.Tensor  # (B, L), True for a real token

    def __post_init__(self) -> None:
        if self.key.ndim != 4 or self.key.shape != self.value.shape:
            raise ValueError("cached key and value must share shape (B, H, L, Dh)")
        if self.valid.dtype != torch.bool or self.valid.shape != (
            self.key.shape[0],
            self.key.shape[2],
        ):
            raise ValueError("cached validity mask must be bool with shape (B, L)")

    @property
    def length(self) -> int:
        return self.key.shape[2]


def sinusoidal_positions(
    length: int, width: int, *, offset: int, device: torch.device, dtype: torch.dtype
) -> torch.Tensor:
    """Return absolute sinusoidal positions shaped ``(length, width)``."""
    if width % 2:
        raise ValueError("the teaching position signal needs an even model width")
    positions = torch.arange(offset, offset + length, device=device, dtype=dtype)[:, None]
    frequencies = torch.exp(
        torch.arange(0, width, 2, device=device, dtype=dtype)
        * (-math.log(10_000.0) / width)
    )
    result = torch.empty((length, width), device=device, dtype=dtype)
    result[:, 0::2] = torch.sin(positions * frequencies)
    result[:, 1::2] = torch.cos(positions * frequencies)
    return result


class CausalSelfAttention(nn.Module):
    """Manual multi-head attention for correctness experiments, not production."""

    def __init__(self, model_width: int, heads: int) -> None:
        super().__init__()
        if model_width % heads:
            raise ValueError("model_width must be divisible by heads")
        if model_width % 2:
            raise ValueError("model_width must be even for the position signal")
        self.model_width = model_width
        self.heads = heads
        self.head_width = model_width // heads
        self.query = nn.Linear(model_width, model_width, bias=False)
        self.key = nn.Linear(model_width, model_width, bias=False)
        self.value = nn.Linear(model_width, model_width, bias=False)
        self.output = nn.Linear(model_width, model_width, bias=False)

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

    def _merge(self, tensor: torch.Tensor) -> torch.Tensor:
        batch, _, length, _ = tensor.shape
        return tensor.transpose(1, 2).contiguous().reshape(batch, length, self.model_width)

    def forward(
        self,
        inputs: torch.Tensor,
        *,
        valid_tokens: torch.Tensor | None = None,
        cache: KVCache | None = None,
        use_cache: bool = False,
        query_offset: int | None = None,
        return_weights: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, KVCache] | tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, KVCache, torch.Tensor]:
        """Apply causal attention to ``inputs`` shaped ``(B, T, D)``.

        A cache fixes the first query's absolute position at ``cache.length``.
        ``query_offset`` is accepted only as an assertion of that expected value;
        a mismatch is rejected rather than silently producing different answers.
        """
        if inputs.ndim != 3 or inputs.shape[-1] != self.model_width:
            raise ValueError("inputs must have shape (B, T, model_width)")
        batch, query_length, _ = inputs.shape
        if query_length == 0:
            raise ValueError("attention requires at least one query")
        if cache is not None and not use_cache:
            raise ValueError("a supplied cache requires use_cache=True")
        expected_offset = cache.length if cache is not None else 0
        if query_offset is not None and query_offset != expected_offset:
            raise ValueError(f"query_offset must equal cache length {expected_offset}")
        offset = expected_offset

        if valid_tokens is None:
            valid_tokens = torch.ones((batch, query_length), dtype=torch.bool, device=inputs.device)
        if valid_tokens.dtype != torch.bool or valid_tokens.shape != (batch, query_length):
            raise ValueError("valid_tokens must be bool with shape (B, T)")
        if cache is not None and (
            cache.key.shape[:2] != (batch, self.heads)
            or cache.key.shape[-1] != self.head_width
        ):
            raise ValueError("cache batch, head, or width does not match this input")

        positions = sinusoidal_positions(
            query_length,
            self.model_width,
            offset=offset,
            device=inputs.device,
            dtype=inputs.dtype,
        )
        positioned = inputs + positions[None, :, :]
        q = self._split(self.query(positioned))
        new_k = self._split(self.key(positioned))
        new_v = self._split(self.value(positioned))

        if cache is None:
            k, v, valid_keys = new_k, new_v, valid_tokens
        else:
            k = torch.cat((cache.key, new_k), dim=2)
            v = torch.cat((cache.value, new_v), dim=2)
            valid_keys = torch.cat((cache.valid, valid_tokens), dim=1)

        key_length = k.shape[2]
        query_positions = torch.arange(offset, offset + query_length, device=inputs.device)
        key_positions = torch.arange(key_length, device=inputs.device)
        causal = key_positions[None, :] <= query_positions[:, None]
        allowed = valid_keys[:, None, None, :] & causal[None, None, :, :]
        allowed = allowed.expand(batch, self.heads, query_length, key_length)
        if not bool(allowed.any(dim=-1).all()):
            raise ValueError("a query has no valid key at or before its position")

        scores = q @ k.transpose(-2, -1) / math.sqrt(self.head_width)
        scores = scores.masked_fill(~allowed, torch.finfo(scores.dtype).min)
        weights = torch.softmax(scores, dim=-1)
        weights = weights.masked_fill(~allowed, 0.0)
        attended = weights @ v
        result = self.output(self._merge(attended))
        result = result * valid_tokens[..., None].to(result.dtype)

        new_cache = KVCache(k, v, valid_keys)
        if use_cache and return_weights:
            return result, new_cache, weights
        if use_cache:
            return result, new_cache
        if return_weights:
            return result, weights
        return result


def run_incrementally(
    layer: CausalSelfAttention,
    inputs: torch.Tensor,
    valid_tokens: torch.Tensor,
    chunk_sizes: tuple[int, ...],
) -> tuple[torch.Tensor, KVCache]:
    """Process a sequence in cache-aware chunks whose sizes sum to its length."""
    if sum(chunk_sizes) != inputs.shape[1]:
        raise ValueError("chunk sizes must cover the sequence exactly")
    pieces: list[torch.Tensor] = []
    cache: KVCache | None = None
    start = 0
    for size in chunk_sizes:
        stop = start + size
        piece, cache = layer(
            inputs[:, start:stop],
            valid_tokens=valid_tokens[:, start:stop],
            cache=cache,
            use_cache=True,
        )
        pieces.append(piece)
        start = stop
    assert cache is not None
    return torch.cat(pieces, dim=1), cache


def main() -> None:
    torch.manual_seed(7)
    layer = CausalSelfAttention(model_width=12, heads=3).eval()
    inputs = torch.randn(2, 6, 12)
    valid = torch.tensor(
        [[True, True, True, True, True, True], [True, True, True, True, False, False]]
    )
    with torch.no_grad():
        full, weights = layer(inputs, valid_tokens=valid, return_weights=True)
        cached, cache = run_incrementally(layer, inputs, valid, (3, 2, 1))
    print("output shape:", tuple(full.shape))
    print("weight shape:", tuple(weights.shape))
    print("cache shape:", tuple(cache.key.shape))
    print("maximum full-vs-cached error:", (full - cached).abs().max().item())
    print("future weight above diagonal:", weights[0, 0].triu(1).abs().max().item())


if __name__ == "__main__":
    main()