"""
Track 1 Lab C — Cached generation and GQA.
Deterministic tiny attention model on CPU with cached and uncached paths.

Model dimensions (editable):
    D  = hidden dimension
    H  = number of query heads (= KV heads for the MHA variant)
    D_HEAD = head dimension
    L  = number of transformer layers
    V  = vocabulary size

All tensors: CPU, float32, deterministic via torch.manual_seed.
No GPU, no model download, no API key required.
"""

from __future__ import annotations

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


def set_seed(seed: int = 42) -> None:
    """Seed PyTorch for deterministic runs."""
    torch.manual_seed(seed)


class AttentionBlock(nn.Module):
    """Single decoder layer with per-layer KV cache.

    The cache holds (key, value) tensors of shape [B, H_KV, S_cached, D_HEAD].
    On each forward call with `past` provided, the new token's K/V is appended
    at the next free slot (index == current cache length). With
    `broken_write=True`, the exercise instead overwrites slot 0, so the cache
    never grows.
    """

    def __init__(self, D: int, H: int, D_HEAD: int, use_gqa: bool = False, H_KV: int | None = None):
        super().__init__()
        self.D = D
        self.H = H
        self.D_HEAD = D_HEAD
        self.use_gqa = use_gqa
        self.H_KV = H_KV if H_KV is not None else H

        self.q_proj = nn.Linear(D, H * D_HEAD, bias=False)
        self.k_proj = nn.Linear(D, self.H_KV * D_HEAD, bias=False)
        self.v_proj = nn.Linear(D, self.H_KV * D_HEAD, bias=False)
        self.o_proj = nn.Linear(H * D_HEAD, D, bias=False)

        self.ln1 = nn.LayerNorm(D)
        self.ln2 = nn.LayerNorm(D)
        self.fc1 = nn.Linear(D, 4 * D)
        self.fc2 = nn.Linear(4 * D, D)

    def forward(
        self,
        x: torch.Tensor,
        past: tuple[torch.Tensor, torch.Tensor] | None = None,
        broken_write: bool = False,
    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        B, S, _ = x.shape

        # Project Q/K/V and split into heads.
        q = self.q_proj(x).view(B, S, self.H, self.D_HEAD).transpose(1, 2)       # [B, H, S, D_HEAD]
        k = self.k_proj(x).view(B, S, self.H_KV, self.D_HEAD).transpose(1, 2)     # [B, H_KV, S, D_HEAD]
        v = self.v_proj(x).view(B, S, self.H_KV, self.D_HEAD).transpose(1, 2)     # [B, H_KV, S, D_HEAD]

        # Append new K/V to the existing cache. The write position is exactly
        # the current cache length (number of tokens already stored).
        if past is not None:
            k_cache, v_cache = past
            write_position = k_cache.shape[2]          # == current cache length
            if broken_write:
                # INTENTIONAL BUG: wrong cache write position. Instead of
                # appending the new K/V at the next free slot (write_position),
                # we overwrite slot 0. The cache never grows and the new token
                # displaces the oldest entry, so subsequent reads attend to the
                # wrong keys/values.
                k_cache = k_cache.clone()
                v_cache = v_cache.clone()
                k_cache[:, :, 0, :] = k[:, :, 0, :]      # overwrite slot 0 with new token
                v_cache[:, :, 0, :] = v[:, :, 0, :]
                k = k_cache
                v = v_cache
            else:
                k = torch.cat([k_cache, k], dim=2)
                v = torch.cat([v_cache, v], dim=2)

        new_past = (k, v)

        # Scaled dot-product attention with a causal mask over the extended cache.
        S_all = k.shape[2]
        scale = self.D_HEAD ** 0.5
        scores = q @ k.transpose(-2, -1) / scale   # [B, H, S, S_all]

        # Causal mask: query at position s attends only to keys 0 .. s.
        q_pos = torch.arange(S_all - S, S_all, device=x.device)
        k_pos = torch.arange(S_all, device=x.device)
        blocked = k_pos[None, :] > q_pos[:, None]   # [S, S_all]
        scores = scores.masked_fill(blocked[None, None, :, :], float("-inf"))

        attn = scores.softmax(dim=-1)
        out = attn @ v                              # [B, H, S, D_HEAD]
        out = out.transpose(1, 2).reshape(B, S, self.H * self.D_HEAD)
        out = self.o_proj(out)

        x = x + out                                 # residual 1
        x = x + self.fc2(F.gelu(self.fc1(self.ln2(x))))  # residual 2 (MLP)
        return x, new_past


class TinyDecoder(nn.Module):
    """Tiny decoder-only transformer with optional KV cache.

    forward(idx, past_key_values=None, broken_write=False, use_cache=True)
      -> (logits, new_past_key_values)
    """

    def __init__(self, D: int = 32, H: int = 2, D_HEAD: int = 16, L: int = 2, V: int = 50,
                 use_gqa: bool = False, H_KV: int | None = None):
        super().__init__()
        self.emb = nn.Embedding(V, D)
        self.blocks = nn.ModuleList([
            AttentionBlock(D, H, D_HEAD, use_gqa=use_gqa, H_KV=H_KV) for _ in range(L)
        ])
        self.ln_f = nn.LayerNorm(D)
        self.head = nn.Linear(D, V, bias=False)
        self.D, self.H, self.D_HEAD, self.L, self.V = D, H, D_HEAD, L, V
        self.use_gqa = use_gqa
        self.H_KV = H_KV if H_KV is not None else H

    def forward(
        self,
        idx: torch.Tensor,
        past_key_values: tuple[tuple[torch.Tensor, torch.Tensor], ...] | None = None,
        broken_write: bool = False,
        use_cache: bool = True,
    ):
        x = self.emb(idx)                        # [B, S, D]
        B, S, _ = x.shape

        new_past = []
        for i, blk in enumerate(self.blocks):
            layer_past = None if past_key_values is None else past_key_values[i]
            x, layer_past = blk(x, past=layer_past, broken_write=broken_write)
            new_past.append(layer_past)

        x = self.ln_f(x)
        logits = self.head(x)                    # [B, S, V]

        if use_cache:
            return logits, tuple(new_past)
        return logits, None


def kv_cache_bytes_per_layer(H_kv: int, D_HEAD: int, S: int, B: int, b: int = 2) -> int:
    """Raw KV-cache bytes for one layer: 2 * H_kv * D_HEAD * S * B * b.

    The factor 2 accounts for one key and one value per token per KV head.
    b = bytes per element (2 for FP16, 4 for FP32).
    """
    return 2 * H_kv * D_HEAD * S * B * b


def mha_gqa_cache_bytes(L: int, H_kv: int, D_HEAD: int, S: int, B: int, b: int = 2) -> int:
    """Raw KV-cache bytes for L layers: L times the per-layer figure."""
    return L * kv_cache_bytes_per_layer(H_kv, D_HEAD, S, B, b)


def model_summary(model: TinyDecoder) -> dict:
    """Return the model's tensor-dimension summary for memory calculations."""
    return {
        "L": model.L,
        "H_q": model.H,
        "H_kv": model.H_KV,
        "D_HEAD": model.D_HEAD,
        "V": model.V,
        "D": model.D,
    }