"""Pre-normalized decoder-only transformer with interchangeable FFNs."""

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

from moe import TransformerConfig, MLP, MoE


class Block(nn.Module):
    """Apply pre-normalized causal attention and the selected FFN."""

    def __init__(self, config):
        """Create attention projections, normalization layers and the FFN."""
        super().__init__()
        self.norm1 = nn.LayerNorm(config.d_model)
        self.norm2 = nn.LayerNorm(config.d_model)
        self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False)
        self.proj = nn.Linear(config.d_model, config.d_model, bias=False)
        self.ffn = MLP(config) if config.num_experts == 1 else MoE(config)
        self.heads = config.n_heads
        self.dropout = config.dropout

    def forward(self, x, attention_mask, routing_mask=None):
        """Apply causal attention and the FFN to (batch, tokens, hidden) states."""
        b, t, d = x.shape
        q, k, v = self.qkv(self.norm1(x)).chunk(3, -1)
        q, k, v = [
            z.reshape(b, t, self.heads, d // self.heads).transpose(1, 2)
            for z in (q, k, v)
        ]
        mask = None
        if attention_mask is not None and not bool(attention_mask.all()):
            mask = (
                torch.ones(t, t, dtype=torch.bool, device=x.device).tril()[None, None]
                & attention_mask[:, None, None, :].bool()
            )
        attended = F.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=mask,
            is_causal=mask is None,
            dropout_p=self.dropout if self.training else 0.0,
        )
        x = x + F.dropout(
            self.proj(attended.transpose(1, 2).reshape(b, t, d)),
            self.dropout,
            self.training,
        )
        if isinstance(self.ffn, MoE):
            self.ffn.tracking_mask = (
                attention_mask if routing_mask is None else routing_mask
            )
        return x + F.dropout(self.ffn(self.norm2(x)), self.dropout, self.training)


class CausalLM(nn.Module):
    """Decode tokens using tied embeddings and dense or MoE blocks."""

    def __init__(self, config: TransformerConfig):
        """Build the decoder, initialize its weights and tie the output embedding."""
        super().__init__()
        if config.d_model % config.n_heads:
            raise ValueError("d_model must be divisible by n_heads")
        self.config = config
        self.context_length = config.context_length
        self.usage_mask = None
        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.position_embedding = nn.Embedding(config.context_length, config.d_model)
        self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layers)])
        self.norm = nn.LayerNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.apply(self._initialize)
        self.lm_head.weight = self.token_embedding.weight

    @staticmethod
    def _initialize(module):
        """Initialize projection and embedding weights normally, with zero biases."""
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, std=0.02)
            if getattr(module, "bias", None) is not None:
                nn.init.zeros_(module.bias)

    @property
    def aux_loss(self):
        """Return the mean balancing loss, or zero for a dense model."""
        losses = [
            block.ffn.aux_loss for block in self.blocks if isinstance(block.ffn, MoE)
        ]
        return (
            torch.stack(losses).mean()
            if losses
            else self.token_embedding.weight.new_zeros(())
        )

    def forward(
        self, input_ids: Tensor, attention_mask: Tensor | None = None
    ) -> Tensor:
        """Return logits of shape (batch, tokens, vocabulary) for a token sequence.

        Use learned absolute positions and causal attention. The optional mask
        marks valid input positions; sequences must fit the configured context.
        """
        if input_ids.ndim != 2 or input_ids.shape[1] > self.context_length:
            raise ValueError("Expected (batch, sequence) within context_length")
        positions = torch.arange(input_ids.shape[1], device=input_ids.device)
        hidden = (
            self.token_embedding(input_ids) + self.position_embedding(positions)[None]
        )
        for block in self.blocks:
            hidden = block(hidden, attention_mask, self.usage_mask)
        return self.lm_head(self.norm(hidden))

    def parameter_counts(self):
        """Count stored and active parameters, counting tied weights once."""
        total = sum(p.numel() for p in self.parameters())
        inactive = sum(
            sum(p.numel() for p in b.ffn.parameters()) - b.ffn.active_parameter_count()
            for b in self.blocks
            if isinstance(b.ffn, MoE)
        )
        return {"total": total, "active": total - inactive}


class DenseCausalLM(CausalLM):
    """Provide the dense decoder through the earlier constructor interface."""

    def __init__(
        self, vocab_size, d_model, n_layers, n_heads, d_ff, context_length, dropout
    ):
        """Create a decoder configuration with one dense feedforward expert."""
        super().__init__(
            TransformerConfig(
                d_model,
                d_ff,
                vocab_size,
                context_length,
                n_layers=n_layers,
                n_heads=n_heads,
                dropout=dropout,
            )
        )


def build_model(config):
    """Construct the configured dense or MoE causal language model."""
    m = config.model
    f = m.ffn
    return CausalLM(
        TransformerConfig(
            m.d_model,
            m.d_ff,
            m.vocab_size,
            m.context_length,
            f.num_experts,
            f.top_k,
            f.shared_expert,
            f.match_active_parameters,
            m.n_layers,
            m.n_heads,
            m.dropout,
        )
    )