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

from .rope import apply_rope


class CausalSelfAttention(nn.Module):
    """Full multi-head causal self-attention (Section 4.4).

    Deliberately NOT using Grouped Query Attention (GQA) — the doc is explicit
    that at this scale, GQA's memory savings are negligible and it can quietly
    cost quality. Every head gets its own independent K/V projections.
    """

    def __init__(self, hidden_dim: int, n_heads: int, dropout: float = 0.0):
        super().__init__()
        assert hidden_dim % n_heads == 0
        self.n_heads = n_heads
        self.head_dim = hidden_dim // n_heads
        self.dropout = dropout

        # separate q, k, v projections -- no sharing across heads (full attention)
        self.q_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.k_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.v_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.out_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)

    def forward(self, x: torch.Tensor, rope_freqs: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape

        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_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)

        q = apply_rope(q, rope_freqs[:T])
        k = apply_rope(k, rope_freqs[:T])

        # scaled dot-product attention with causal masking (built-in flash-attention
        # kernel when running on a CUDA GPU; falls back to a math kernel on CPU)
        out = F.scaled_dot_product_attention(
            q, k, v,
            is_causal=True,
            dropout_p=self.dropout if self.training else 0.0,
        )

        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.out_proj(out)