"""mtp.py — V6: Multi-Token Prediction with adaptive attention + entropy regularizer.

═══════════════════════════════════════════════════════════════════════════════
V6 MELHORIAS MATEMÁTICAS (user request)
═══════════════════════════════════════════════════════════════════════════════

1. **Ativar MTP em val**: em V5, MTP só era usado em train. V6 ativa MTP
   também em val para medir a capacidade preditiva multi-token em validação
   (métrica adicional: val_mtp_loss).

2. **Regularizador de entropia nos mtp_alphas**: o MTP gera uma distribuição
   mtp_alphas (pesos para cada passo k ∈ {0..K-1}). Em V5, esses pesos podiam
   colapsar para um único k. V6 adiciona:
        L_entropy = -β * Σ_k α_k log(α_k + ε)
   que incentiva os K passos a serem igualmente utilizados (mínimo colapso).

3. **mtp_adaptive_attention**: combina MTP com attention adaptativa para
   suportar sequências longas (chunked attention, não full O(N²)).

Matemática
──────────
  L_mtp_total = Σ_k α_k * L_k - β * H(α)

  onde:
    L_k = cross_entropy(logits_k, target_k)
    α_k = softmax(raw_alpha_k / T_alpha)
    H(α) = -Σ_k α_k log(α_k + ε)  (entropia da distribuição α)
    β = 0.01 (peso do regularizador de entropia)
    T_alpha = 1.0 (temperatura do softmax)
"""
from __future__ import annotations
import math
from dataclasses import dataclass, field
from typing import Optional, Tuple, Dict, List

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


@dataclass
class MTPConfig:
    """Configuração do MTP V6."""
    hidden_size: int = 128
    vocab_size: int = 256
    K: int = 4                  # nº de tokens futuros
    share_embeddings: bool = False
    tie_weights: bool = False
    dropout: float = 0.1
    # V6 — entropy regularizer
    entropy_beta: float = 0.01  # peso do regularizador de entropia
    alpha_temp: float = 1.0     # temperatura do softmax dos alphas
    # V6 — val mode
    active_in_val: bool = True  # ativar MTP em val
    # V6 — adaptive attention
    adaptive_attention: bool = True
    adaptive_chunk_size: int = 64  # tamanho do chunk para adaptive attention


class MTPHead(nn.Module):
    """Head de Multi-Token Prediction (V6).

    Prediz K distribuições sobre o vocabulário, uma para cada k-ésimo token
    futuro, a partir de uma única hidden state.

    V6: adiciona parâmetro mtp_alphas (K,) — pesos aprendíveis para cada
    passo k. Esses pesos são regularizados por entropia para evitar colapso.

    Args:
        hidden_size: dimensão da hidden state
        vocab_size: tamanho do vocabulário
        K: nº de tokens futuros a predizer
        share_embeddings: se True, compartilha a matriz de embedding
        tie_weights: se True, usa a mesma W para todos os K passos
        dropout: probabilidade de dropout

    Forward:
        h: (B, T, hidden_size) — sequência de hidden states
        → logits (B, T, K, vocab_size), alphas (K,)
    """
    def __init__(
        self,
        hidden_size: int,
        vocab_size: int,
        K: int = 4,
        share_embeddings: bool = False,
        tie_weights: bool = False,
        dropout: float = 0.0,
    ):
        super().__init__()
        if K < 1:
            raise ValueError(f"K deve ser >= 1, recebeu {K}")
        self.hidden_size = hidden_size
        self.vocab_size = vocab_size
        self.K = K
        self.share_embeddings = share_embeddings
        self.tie_weights = tie_weights

        if tie_weights:
            self.shared_proj = nn.Linear(hidden_size, vocab_size, bias=False)
            self.biases = nn.ParameterList([
                nn.Parameter(torch.zeros(vocab_size)) for _ in range(K)
            ])
        else:
            self.heads = nn.ModuleList([
                nn.Linear(hidden_size, vocab_size) for _ in range(K)
            ])

        # V6: mtp_alphas — pesos aprendíveis para cada passo k
        # Inicializado uniforme para que Σα_k = 1 no começo
        self.raw_alphas = nn.Parameter(torch.zeros(K))  # softmax(0) = uniforme
        self.dropout = nn.Dropout(dropout)

    @property
    def alphas(self) -> torch.Tensor:
        """Distribuição α normalizada via softmax com temperatura."""
        return F.softmax(self.raw_alphas, dim=-1)

    def forward(self, h: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """h: (B, T, hidden_size) → logits (B, T, K, V), alphas (K,)"""
        h = self.dropout(h)
        if self.tie_weights:
            # Cada passo k: shared_proj(h) + biases[k]
            logits_list = [self.shared_proj(h) + b for b in self.biases]
        else:
            logits_list = [head(h) for head in self.heads]
        # Stack: (B, T, K, V)
        logits = torch.stack(logits_list, dim=2)
        return logits, self.alphas


def mtp_entropy_regularizer(alphas: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
    """V6: regularizador de entropia nos mtp_alphas.

    H(α) = -Σ_k α_k log(α_k + ε)

    Retornamos -H (para minimizar = maximizar entropia).
    L_reg = -H(α) = Σ_k α_k log(α_k + ε)
    """
    return -(alphas * torch.log(alphas + eps)).sum()


def mtp_loss(
    logits: torch.Tensor,
    targets: torch.Tensor,
    alphas: torch.Tensor,
    entropy_beta: float = 0.01,
) -> Tuple[torch.Tensor, Dict[str, float]]:
    """Computa a loss MTP V6 com regularizador de entropia.

    L = Σ_k α_k * L_k - β * H(α)

    Args:
        logits: (B, T, K, V)
        targets: (B, T) — índices do token alinhado com o último T (next-token)
                 Será shiftado para alinhar com cada k
        alphas: (K,)
        entropy_beta: peso do regularizador de entropia

    Returns:
        loss: escalar
        metrics: dict com breakdown da loss
    """
    B, T, K, V = logits.shape
    # Computa loss por passo k
    losses_per_k = []
    for k in range(K):
        # Target para passo k: targets shifted por k
        # logits[:, t, k, :] prediz token em t+k+1 (ou apenas t+k dependendo da convenção)
        # Aqui usamos shift simples: prediz target[t+k] a partir de h[t]
        if T - k - 1 <= 0:
            losses_per_k.append(torch.tensor(0.0, device=logits.device))
            continue
        logits_k = logits[:, :T-k-1, k, :].reshape(-1, V)
        targets_k = targets[:, k+1:T].reshape(-1)
        loss_k = F.cross_entropy(logits_k, targets_k, ignore_index=-100)
        losses_per_k.append(loss_k)
    losses_tensor = torch.stack(losses_per_k)  # (K,)

    # Loss ponderada: Σ α_k L_k
    weighted_loss = (alphas * losses_tensor).sum()
    # Regularizador de entropia: -β H(α) = β * Σ α_k log(α_k + ε)
    entropy_reg = mtp_entropy_regularizer(alphas)
    total_loss = weighted_loss - entropy_beta * entropy_reg  # negativo porque reg = -H

    metrics = {
        "weighted_loss": float(weighted_loss.item()),
        "entropy_reg": float(entropy_reg.item()),
        "total_loss": float(total_loss.item()),
        "alphas": alphas.detach().cpu().tolist(),
        "per_k_loss": [float(l.item()) for l in losses_per_k],
    }
    return total_loss, metrics


class MTPAdaptiveAttention(nn.Module):
    """V6: MTP + adaptive attention para sequências longas (2M tokens).

    Estratégia:
    - Sequência longa é dividida em chunks de tamanho chunk_size
    - Cada chunk: attention full (custo O(chunk²))
    - Entre chunks: attention sparse (top-k dos vizinhos)
    - MTP prediz K tokens futuros por chunk anchor

    Args:
        vocab_size, embed_dim, hidden_size, K, num_layers, num_heads
        chunk_size: tamanho do chunk para adaptive attention
        knn_k: top-k para atenção entre chunks
    """
    def __init__(
        self,
        vocab_size: int,
        embed_dim: int = 128,
        hidden_size: int = 128,
        K: int = 4,
        num_layers: int = 2,
        num_heads: int = 4,
        chunk_size: int = 64,
        knn_k: int = 8,
        max_N: int = 2_000_000,  # 2M tokens (limite teórico)
        dropout: float = 0.1,
    ):
        super().__init__()
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim
        self.hidden_size = hidden_size
        self.K = K
        self.chunk_size = chunk_size
        self.knn_k = knn_k
        self.max_N = max_N

        self.embed = nn.Embedding(vocab_size, embed_dim)
        # Camadas de attention chunked
        self.layers = nn.ModuleList([
            nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout, batch_first=True)
            for _ in range(num_layers)
        ])
        self.norms = nn.ModuleList([nn.LayerNorm(embed_dim) for _ in range(num_layers)])
        # MTP head
        self.mtp_head = MTPHead(embed_dim, vocab_size, K=K, dropout=dropout)

    def forward(self, input_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """input_ids: (B, T) → logits (B, T, K, V), alphas (K,)"""
        B, T = input_ids.shape
        if T > self.max_N:
            raise ValueError(f"Sequence length {T} exceeds max_N={self.max_N}")
        h = self.embed(input_ids)  # (B, T, embed_dim)
        # Adaptive attention: processa em chunks
        chunk_size = min(self.chunk_size, T)
        n_chunks = (T + chunk_size - 1) // chunk_size
        outputs = []
        for c in range(n_chunks):
            start = c * chunk_size
            end = min((c + 1) * chunk_size, T)
            chunk = h[:, start:end, :]  # (B, chunk, embed_dim)
            for i, (attn, norm) in enumerate(zip(self.layers, self.norms)):
                normed = norm(chunk)
                attn_out, _ = attn(normed, normed, normed, need_weights=False)
                chunk = chunk + attn_out
            outputs.append(chunk)
        h_out = torch.cat(outputs, dim=1)  # (B, T, embed_dim)
        logits, alphas = self.mtp_head(h_out)
        return logits, alphas


__all__ = [
    "MTPConfig", "MTPHead", "MTPAdaptiveAttention",
    "mtp_loss", "mtp_entropy_regularizer",
]