"""Attention reference implementation for Lab 01.

CPU-only, no downloads, no API keys.

Conventions used throughout this lab
------------------------------------
* Activations have shape ``[..., T, d_head]``. The leading dimensions are
  usually ``[B, h]`` in the multi-head case and ``[B]`` in the single-head case.
* Masks are boolean and broadcastable to the score shape ``[..., T, T]``.
  ``True`` means "this query may attend to this key". This matches the
  convention used by lessons 4-6 and by ``F.scaled_dot_product_attention``.
* A query row with no allowed keys is *fully masked*. This implementation
  returns all-zero attention weights and an all-zero output for such rows.
  It never returns NaN. The lab's `stable_softmax` maps fully‑masked rows to
  zero as a deliberate design choice (raw `torch.softmax` would return NaN for
  0/0, and SDPA backend behavior varies). The parity tests use the MATH backend
  explicitly to ensure deterministic comparison on CPU.
"""

from __future__ import annotations

import math

import torch
import torch.nn as nn


def stable_softmax(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
    """Softmax that subtracts the row maximum and tolerates fully masked rows.

    ``torch.softmax`` is already numerically stable, but it returns NaN when a
    whole row is ``-inf`` (0/0). We map those rows to zeros because a query with
    no visible keys should contribute nothing, not poison the batch with NaN.
    """
    row_max = x.max(dim=dim, keepdim=True).values
    finite = torch.isfinite(row_max)
    safe_max = torch.where(finite, row_max, torch.zeros_like(row_max))

    shifted = x - safe_max
    # Keep -inf entries as -inf; replace any NaN produced by -inf - -inf.
    shifted = torch.where(
        torch.isfinite(shifted), shifted, torch.full_like(shifted, float("-inf"))
    )

    exp = torch.exp(shifted)
    denom = exp.sum(dim=dim, keepdim=True)
    weights = exp / denom.clamp_min(torch.finfo(exp.dtype).tiny)
    return torch.where(denom > 0, weights, torch.zeros_like(weights))


def scaled_dot_product_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    mask: torch.Tensor | None = None,
    scale: float | None = None,
    return_weights: bool = False,
):
    """Scaled dot-product attention written out step by step.

    Shapes
    ------
    q, k: ``[..., T_q, d_head]`` and ``[..., T_k, d_head]``
    v:    ``[..., T_k, d_v]``
    mask: broadcastable to ``[..., T_q, T_k]``, ``True`` = keep
    out:  ``[..., T_q, d_v]``
    """
    if q.shape[-1] != k.shape[-1]:
        raise ValueError(f"q/k head width mismatch: {q.shape[-1]} vs {k.shape[-1]}")
    if k.shape[-2] != v.shape[-2]:
        raise ValueError(f"k/v sequence length mismatch: {k.shape[-2]} vs {v.shape[-2]}")

    d_head = q.shape[-1]
    if scale is None:
        scale = 1.0 / math.sqrt(d_head)

    scores = (q @ k.transpose(-2, -1)) * scale

    if mask is not None:
        mask = mask.to(torch.bool)
        scores = scores.masked_fill(~mask, float("-inf"))

    weights = stable_softmax(scores, dim=-1)
    out = weights @ v

    if return_weights:
        return out, weights
    return out


def causal_mask(seq_len: int, device: torch.device | None = None) -> torch.Tensor:
    """Lower-triangular boolean mask ``[T, T]``. ``True`` = may attend."""
    return torch.tril(
        torch.ones(seq_len, seq_len, dtype=torch.bool, device=device)
    )


def combine_causal_and_padding(
    causal: torch.Tensor, key_padding: torch.Tensor
) -> torch.Tensor:
    """AND a ``[T, T]`` causal mask with a ``[B, T]`` key-padding mask.

    ``key_padding[b, j]`` is ``True`` when key position ``j`` is a real token.
    Returns ``[B, 1, T, T]``, broadcastable over heads.
    """
    if causal.dim() != 2 or causal.shape[0] != causal.shape[1]:
        raise ValueError(f"causal must be [T, T], got {tuple(causal.shape)}")
    if key_padding.dim() != 2 or key_padding.shape[1] != causal.shape[0]:
        raise ValueError(
            "key_padding must be [B, T] with T matching causal, got "
            f"{tuple(key_padding.shape)}"
        )
    return causal[None, None, :, :] & key_padding[:, None, None, :]


class SelfAttention(nn.Module):
    """Single-head self-attention from lesson 4.

    ``X`` is ``[B, T, d_model]``; output is ``[B, T, d_v]`` and weights are
    ``[B, T, T]``.
    """

    def __init__(self, d_model: int, d_k: int, d_v: int, bias: bool = False):
        super().__init__()
        self.W_Q = nn.Parameter(torch.randn(d_model, d_k))
        self.W_K = nn.Parameter(torch.randn(d_model, d_k))
        self.W_V = nn.Parameter(torch.randn(d_model, d_v))
        if bias:
            self.b_Q = nn.Parameter(torch.zeros(d_k))
            self.b_K = nn.Parameter(torch.zeros(d_k))
            self.b_V = nn.Parameter(torch.zeros(d_v))
        else:
            self.register_parameter("b_Q", None)
            self.register_parameter("b_K", None)
            self.register_parameter("b_V", None)

    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None):
        q = x @ self.W_Q
        k = x @ self.W_K
        v = x @ self.W_V
        if self.b_Q is not None:
            q = q + self.b_Q
            k = k + self.b_K
            v = v + self.b_V
        return scaled_dot_product_attention(q, k, v, mask=mask, return_weights=True)


class MultiHeadSelfAttention(nn.Module):
    """Multi-head self-attention with a combined in-projection (lesson 6).

    ``X`` is ``[B, T, d_model]``; output is ``[B, T, d_model]`` and weights are
    ``[B, h, T, T]``.
    """

    def __init__(self, d_model: int, num_heads: int, bias: bool = False):
        super().__init__()
        if d_model % num_heads != 0:
            raise ValueError(
                f"d_model ({d_model}) must be divisible by num_heads ({num_heads})"
            )
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_head = d_model // num_heads
        self.in_proj = nn.Linear(d_model, 3 * d_model, bias=bias)
        self.out_proj = nn.Linear(d_model, d_model, bias=bias)

    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None):
        batch, seq_len, _ = x.shape

        qkv = self.in_proj(x)  # [B, T, 3d]
        q, k, v = qkv.chunk(3, dim=-1)  # each [B, T, d]

        # [B, T, h, d_head] -> [B, h, T, d_head]
        q = q.view(batch, seq_len, self.num_heads, self.d_head).transpose(1, 2)
        k = k.view(batch, seq_len, self.num_heads, self.d_head).transpose(1, 2)
        v = v.view(batch, seq_len, self.num_heads, self.d_head).transpose(1, 2)

        out, weights = scaled_dot_product_attention(
            q, k, v, mask=mask, return_weights=True
        )  # out [B, h, T, d_head], weights [B, h, T, T]

        merged = out.transpose(1, 2).contiguous().view(batch, seq_len, self.d_model)
        return self.out_proj(merged), weights