"""
CausalSelfAttention
────────────────────
Modern causal self-attention with:

1. SCALED DOT-PRODUCT ATTENTION (SDPA)
   torch.nn.functional.scaled_dot_product_attention dispatches to Flash
   Attention 2 on CUDA automatically. Flash Attention rewrites the standard
   O(N²) memory algorithm as a tiled, IO-efficient kernel — peak memory is
   O(N) instead of O(N²), enabling much longer contexts on the same GPU.

2. ROTARY POSITION EMBEDDINGS (RoPE)
   Applied to Q and K before attention. Encodes relative position as a
   rotation in head-dimension space — see positional_encoding.py.

3. GROUPED-QUERY ATTENTION (GQA)
   Standard MHA uses n_heads Q, K, V projections.
   GQA shares fewer KV heads across multiple Q heads:
     - n_kv_heads == n_heads → standard MHA (original Transformer)
     - n_kv_heads <  n_heads → GQA (Llama 3, Mistral)
     - n_kv_heads == 1       → MQA (one shared KV pair)
   GQA reduces KV-cache memory at inference by a factor of (n_heads / n_kv_heads)
   with minimal quality loss.

4. KV-CACHE
   During autoregressive inference, we only ever process one new token but
   recompute all previous K, V from scratch if we run the full sequence each
   step. The KV-cache stores previous K, V tensors and appends the new ones,
   reducing each inference step from O(T²) to O(T) compute.

5. BIAS-FREE PROJECTIONS
   All linear projections use bias=False, matching Llama/Mistral practice.
   Biases provide negligible benefit in large models and waste memory.
"""

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

from model.positional_encoding import apply_rotary_emb


class CausalSelfAttention(nn.Module):
    """
    Multi-head causal self-attention with RoPE, GQA, SDPA, and KV-cache.

    Args:
        d_model:    Model hidden dimension.
        n_heads:    Number of query attention heads.
        n_kv_heads: Number of key/value heads (≤ n_heads). For standard MHA
                    set n_kv_heads == n_heads.
        dropout:    Attention dropout probability (0.0 recommended at scale).
    """

    def __init__(
        self,
        d_model:    int,
        n_heads:    int,
        n_kv_heads: int,
        dropout:    float = 0.0,
    ):
        super().__init__()
        assert d_model % n_heads == 0, (
            f"d_model ({d_model}) must be divisible by n_heads ({n_heads})"
        )
        assert n_heads % n_kv_heads == 0, (
            f"n_heads ({n_heads}) must be divisible by n_kv_heads ({n_kv_heads})"
        )

        self.n_heads    = n_heads
        self.n_kv_heads = n_kv_heads
        self.head_dim   = d_model // n_heads
        self.dropout    = dropout
        self.kv_groups  = n_heads // n_kv_heads   # how many Q heads share each KV head

        # Separate Q/K/V projections — sizes differ for GQA
        self.q_proj = nn.Linear(d_model, n_heads    * self.head_dim, bias=False)
        self.k_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(d_model, d_model,                    bias=False)

    def forward(
        self,
        x:       torch.Tensor,
        cos:     torch.Tensor,
        sin:     torch.Tensor,
        past_kv: tuple[torch.Tensor, torch.Tensor] | None = None,
    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        """
        Args:
            x:       Input, shape [B, T, d_model].
            cos:     RoPE cosines, shape [1, 1, T, head_dim].
            sin:     RoPE sines,   shape [1, 1, T, head_dim].
            past_kv: Optional cached (k, v) from previous decode steps.
                     - None  → training / first inference step (is_causal=True)
                     - tuple → incremental decode step (append to cache)

        Returns:
            out:        Attention output, shape [B, T, d_model].
            present_kv: Updated (k, v) cache tuple for use in the next step.
        """
        B, T, _ = x.shape

        # ── Project Q, K, V ──────────────────────────────────────────────────
        q = self.q_proj(x).view(B, T, self.n_heads,    self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        # Shapes after transpose: [B, heads, T, head_dim]

        # ── Apply Rotary Embeddings ───────────────────────────────────────────
        q, k = apply_rotary_emb(q, k, cos, sin)

        # ── KV-Cache (inference only) ─────────────────────────────────────────
        # Append new K/V to the cached history so we never recompute old tokens.
        if past_kv is not None:
            k = torch.cat([past_kv[0], k], dim=2)   # grow along sequence axis
            v = torch.cat([past_kv[1], v], dim=2)

        present_kv = (k, v)   # return updated cache

        # ── GQA: broadcast KV heads to match Q heads ─────────────────────────
        # Each group of (n_heads/n_kv_heads) Q heads attends to the same KV pair.
        if self.kv_groups > 1:
            k = k.repeat_interleave(self.kv_groups, dim=1)   # [B, n_heads, S, head_dim]
            v = v.repeat_interleave(self.kv_groups, dim=1)

        # ── Scaled Dot-Product Attention ──────────────────────────────────────
        # is_causal=True  → auto-builds a causal mask (training / first step)
        # is_causal=False → no mask (KV-cache handles causality by construction)
        # PyTorch dispatches this to Flash Attention 2 on CUDA when available.
        is_causal = past_kv is None
        attn_out = F.scaled_dot_product_attention(
            q, k, v,
            dropout_p = self.dropout if self.training else 0.0,
            is_causal = is_causal,
        )

        # ── Merge heads and project ───────────────────────────────────────────
        out = attn_out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.o_proj(out), present_kv