"""GPT compact : embeddings, attention causale, MLP, connexions residuelles."""
from __future__ import annotations

import math
from dataclasses import asdict, dataclass

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


@dataclass
class ModelConfig:
    vocab_size: int
    context_length: int = 512
    n_layers: int = 6
    n_heads: int = 6
    d_model: int = 384
    dropout: float = 0.1

    def __post_init__(self):
        if min(self.vocab_size, self.context_length, self.n_layers, self.n_heads, self.d_model) <= 0:
            raise ValueError("Les dimensions du modele doivent etre positives.")
        if self.d_model % self.n_heads:
            raise ValueError("d_model doit etre divisible par n_heads.")


class CausalAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_heads = config.n_heads
        self.dropout = config.dropout
        self.qkv = nn.Linear(config.d_model, 3 * config.d_model)
        self.proj = nn.Linear(config.d_model, config.d_model)
        self.resid_dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        batch, length, width = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        def split_heads(t):
            return t.view(batch, length, self.n_heads, width // self.n_heads).transpose(1, 2)
        q, k, v = map(split_heads, (q, k, v))
        attended = F.scaled_dot_product_attention(
            q, k, v, is_causal=True, dropout_p=self.dropout if self.training else 0.0,
        )
        attended = attended.transpose(1, 2).contiguous().view(batch, length, width)
        return self.resid_dropout(self.proj(attended))


class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.d_model)
        self.attn = CausalAttention(config)
        self.ln2 = nn.LayerNorm(config.d_model)
        self.mlp = nn.Sequential(
            nn.Linear(config.d_model, 4 * config.d_model), nn.GELU(),
            nn.Linear(4 * config.d_model, config.d_model), nn.Dropout(config.dropout),
        )

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        return x + self.mlp(self.ln2(x))


class GPT(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.position_embedding = nn.Embedding(config.context_length, config.d_model)
        self.dropout = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layers)])
        self.ln_final = nn.LayerNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.lm_head.weight = self.token_embedding.weight
        self.apply(self._init_weights)
        for name, param in self.named_parameters():
            if name.endswith("attn.proj.weight") or name.endswith("mlp.2.weight"):
                nn.init.normal_(param, std=0.02 / math.sqrt(2 * config.n_layers))

    @staticmethod
    def _init_weights(module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, std=0.02)
            if isinstance(module, nn.Linear) and module.bias is not None:
                nn.init.zeros_(module.bias)

    def hidden_states(self, tokens):
        length = tokens.size(1)
        if length > self.config.context_length:
            raise ValueError("Sequence plus longue que la fenetre de contexte.")
        positions = torch.arange(length, device=tokens.device)
        x = self.dropout(self.token_embedding(tokens) + self.position_embedding(positions))
        for block in self.blocks:
            x = block(x)
        return self.ln_final(x)

    def forward(self, tokens, targets=None):
        logits = self.lm_head(self.hidden_states(tokens))
        loss = None
        if targets is not None:
            # Le decalage de +1 a deja ete fait dans le chargeur de donnees.
            loss = F.cross_entropy(logits.float().reshape(-1, self.config.vocab_size), targets.reshape(-1), ignore_index=-100)
        return logits, loss

    def parameter_count(self):
        return sum(p.numel() for p in self.parameters())

    def optimizer(self, learning_rate, weight_decay, device):
        decay, no_decay = [], []
        for param in self.parameters():
            (decay if param.ndim >= 2 else no_decay).append(param)
        return torch.optim.AdamW(
            [{"params": decay, "weight_decay": weight_decay}, {"params": no_decay, "weight_decay": 0.0}],
            lr=learning_rate, betas=(0.9, 0.95), fused=(device.type == "cuda"),
        )

    @torch.inference_mode()
    def generate(self, tokens, max_new_tokens=100, temperature=0.8, top_k=50, top_p=0.95, eos_id=2):
        if max_new_tokens < 1 or not math.isfinite(temperature) or temperature < 0:
            raise ValueError("max_new_tokens positif et temperature >= 0 attendus.")
        if top_k < 0 or not 0 < top_p <= 1:
            raise ValueError("top_k >= 0 et 0 < top_p <= 1 attendus.")
        self.eval()
        for _ in range(max_new_tokens):
            logits, _ = self(tokens[:, -self.config.context_length:])
            next_logits = logits[:, -1, :].float()
            next_logits[:, :2] = -float("inf")  # PAD et BOS ne sont pas du texte.
            if temperature == 0:
                next_token = next_logits.argmax(dim=-1, keepdim=True)
            else:
                next_logits /= temperature
                if top_k:
                    cutoff = torch.topk(next_logits, min(top_k, next_logits.size(-1))).values[:, -1:]
                    next_logits.masked_fill_(next_logits < cutoff, -float("inf"))
                if top_p < 1:
                    sorted_logits, indices = next_logits.sort(descending=True)
                    remove = sorted_logits.softmax(-1).cumsum(-1) > top_p
                    remove[:, 1:] = remove[:, :-1].clone()
                    remove[:, 0] = False
                    next_logits.scatter_(1, indices, sorted_logits.masked_fill(remove, -float("inf")))
                next_token = torch.multinomial(next_logits.softmax(-1), num_samples=1)
            tokens = torch.cat((tokens, next_token), dim=1)
            if (next_token == eos_id).all():
                break
        return tokens