"""Bloc Transformer pre-norm avec résidus — écrit à la main.

Forme pre-norm (justifiée dans docs/architecture.md §2.2) :

    x = x + Attention(LayerNorm(x))
    x = x + FeedForward(LayerNorm(x))

La normalisation est appliquée *avant* chaque sous-couche. Le gradient traverse
la connexion résiduelle sans être multiplié par une norme, ce qui stabilise
l'entraînement des modèles profonds.
"""

from __future__ import annotations

import torch.nn as nn

from src.model.attention import MultiHeadSelfAttention
from src.model.ffn import FeedForward


class TransformerBlock(nn.Module):
    """Un bloc = attention causale + FFN, chacun précédé d'une LayerNorm."""

    def __init__(
        self,
        n_embd: int,
        n_head: int,
        dropout: float = 0.0,
        bias: bool = False,
        ffn_mult: int = 4,
        rope_theta: float = 10000.0,
        use_rope: bool = True,
        causal: bool = True,
    ):
        super().__init__()
        self.ln_1 = nn.LayerNorm(n_embd, bias=bias)
        self.attn = MultiHeadSelfAttention(
            n_embd, n_head, dropout=dropout, bias=bias, rope_theta=rope_theta,
            use_rope=use_rope, causal=causal,
        )
        self.ln_2 = nn.LayerNorm(n_embd, bias=bias)
        self.ffn = FeedForward(n_embd, mult=ffn_mult, dropout=dropout, bias=bias)

    def forward(self, x, cos=None, sin=None):
        x = x + self.attn(self.ln_1(x), cos, sin)
        x = x + self.ffn(self.ln_2(x))
        return x