"""Multi-Head Self-Attention causale — écrite à la main.

Aucun `torch.nn.MultiheadAttention`, aucun module `transformers` : tout est
explicite (projections, découpage en têtes, masque causal, softmax, RoPE).

Causalité : la position ``t`` ne peut pas regarder ``t+1``. On le garantit par
un masque triangulaire supérieur mis à ``-inf`` avant le softmax, donc les poids
d'attention vers le futur valent **exactement 0**.
"""

from __future__ import annotations

import math

import torch
import torch.nn as nn

from src.model.rope import apply_rope, build_rope_cache


def build_causal_mask(seq_len: int, device: torch.device | None = None) -> torch.Tensor:
    """Masque booléen ``(seq_len, seq_len)`` : True = position autorisée.

    ``mask[i, j] = j <= i`` (une position voit le passé et elle-même).
    """
    idx = torch.arange(seq_len, device=device)
    return idx[None, :] <= idx[:, None]


class MultiHeadSelfAttention(nn.Module):
    """Attention multi-têtes causale avec RoPE."""

    def __init__(
        self,
        n_embd: int,
        n_head: int,
        dropout: float = 0.0,
        bias: bool = False,
        rope_theta: float = 10000.0,
        use_rope: bool = True,
        causal: bool = True,
    ):
        super().__init__()
        if n_embd % n_head != 0:
            raise ValueError(f"n_embd ({n_embd}) doit être divisible par n_head ({n_head})")
        self.n_embd = n_embd
        self.n_head = n_head
        self.head_dim = n_embd // n_head
        self.dropout = dropout
        self.rope_theta = rope_theta
        self.use_rope = use_rope
        # causal=True pour le texte (décodage autorégressif) ; causal=False pour
        # les encodeurs (image/audio), où chaque token voit toute la séquence.
        self.causal = causal

        # Projections Q, K, V et sortie. Pas de biais par défaut (style moderne).
        self.q_proj = nn.Linear(n_embd, n_embd, bias=bias)
        self.k_proj = nn.Linear(n_embd, n_embd, bias=bias)
        self.v_proj = nn.Linear(n_embd, n_embd, bias=bias)
        self.out_proj = nn.Linear(n_embd, n_embd, bias=bias)
        self.attn_dropout = nn.Dropout(dropout)
        self.resid_dropout = nn.Dropout(dropout)

    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor | None = None,
        sin: torch.Tensor | None = None,
    ) -> torch.Tensor:
        B, T, C = x.shape

        q = self.q_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        # -> (B, n_head, T, head_dim)

        if self.use_rope:
            if cos is None or sin is None:
                cos, sin = build_rope_cache(T, self.head_dim, self.rope_theta, x.device, x.dtype)
            q, k = apply_rope(q, k, cos, sin)

        # Scores d'attention : (B, n_head, T, T)
        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        if self.causal:
            mask = build_causal_mask(T, x.device)
            att = att.masked_fill(~mask, float("-inf"))
        att = torch.softmax(att, dim=-1)
        att = self.attn_dropout(att)

        y = att @ v  # (B, n_head, T, head_dim)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_dropout(self.out_proj(y))