"""Transformer texte decoder-only de Cortex-0 — écrit à la main.

Assemblage : embeddings de tokens + N blocs pre-norm (attention causale RoPE +
FFN GELU) + LayerNorm finale + tête de langage (`lm_head`).

Le `lm_head` peut être **lié** aux embeddings (weight tying), ce qui réduit le
nombre de paramètres et aide l'entraînement sur petits corpus.
"""

from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn as nn

from src.model.block import TransformerBlock
from src.model.rope import build_rope_cache


@dataclass
class TextConfig:
    vocab_size: int = 32000
    n_layer: int = 6
    n_head: int = 6
    n_embd: int = 384
    ctx_len: int = 512
    dropout: float = 0.0
    bias: bool = False
    ffn_mult: int = 4
    rope_theta: float = 10000.0
    use_rope: bool = True
    tie_embeddings: bool = True
    init_std: float = 0.02


class CortexTextLM(nn.Module):
    """Modèle de langage autorégressif (decoder-only) de Cortex-0."""

    def __init__(self, config: TextConfig):
        super().__init__()
        self.config = config

        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
        self.dropout = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList(
            [
                TransformerBlock(
                    config.n_embd,
                    config.n_head,
                    dropout=config.dropout,
                    bias=config.bias,
                    ffn_mult=config.ffn_mult,
                    rope_theta=config.rope_theta,
                    use_rope=config.use_rope,
                )
                for _ in range(config.n_layer)
            ]
        )
        self.ln_f = nn.LayerNorm(config.n_embd, bias=config.bias)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        if config.tie_embeddings:
            self.lm_head.weight = self.token_embedding.weight

        self.apply(self._init_weights)
        # Initialisation plus fine sur les projections résiduelles (comme GPT-2).
        for name, param in self.named_parameters():
            if name.endswith("out_proj.weight") or name.endswith("fc2.weight"):
                nn.init.normal_(param, mean=0.0, std=config.init_std / (2 * config.n_layer) ** 0.5)

    @staticmethod
    def _init_weights(module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(
        self,
        input_ids: torch.Tensor | None = None,
        targets: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Retourne ``(logits, loss)``.

        ``input_ids`` : ``(B, T)`` ou ``inputs_embeds`` : ``(B, T, n_embd)``
        (permet d'injecter des embeddings multimodaux). ``targets`` : ``(B, T)`` ;
        les positions à ignorer (``-100``) ne comptent pas dans la perte.
        """
        hidden = self.encode(input_ids=input_ids, inputs_embeds=inputs_embeds)
        logits = self.lm_head(hidden)

        loss = None
        if targets is not None:
            loss = nn.functional.cross_entropy(
                logits.reshape(-1, logits.size(-1)), targets.reshape(-1),
                ignore_index=-100,
            )
        return logits, loss

    def encode(
        self,
        input_ids: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Retourne les états cachés finaux ``(B, T, n_embd)`` (avant la tête)."""
        if inputs_embeds is None:
            if input_ids is None:
                raise ValueError("fournir input_ids ou inputs_embeds")
            x = self.dropout(self.token_embedding(input_ids))
        else:
            x = self.dropout(inputs_embeds)
        T = x.size(1)
        if T > self.config.ctx_len:
            raise ValueError(f"longueur {T} > ctx_len {self.config.ctx_len}")

        cos, sin = build_rope_cache(T, self.config.n_embd // self.config.n_head,
                                    self.config.rope_theta, x.device, x.dtype)
        for block in self.blocks:
            x = block(x, cos, sin)
        return self.ln_f(x)

    def count_parameters(self, non_embedding: bool = False) -> int:
        """Compte les paramètres (option : hors embeddings de tokens et lm_head)."""
        n = sum(p.numel() for p in self.parameters())
        if non_embedding:
            n -= self.token_embedding.weight.numel()
            if not self.config.tie_embeddings:
                n -= self.lm_head.weight.numel()
        return n