"""Decoder-only transformer with a pluggable FFN (dense MLP or MoE variants)."""

from __future__ import annotations

import math

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

from .config import TransformerConfig
from .ffn import MoE, build_ffn


def apply_rope(x: torch.Tensor, positions: torch.Tensor, base: float) -> torch.Tensor:
    # x: (B, H, T, hd), positions: (B, T)
    hd = x.shape[-1]
    inv_freq = 1.0 / (base ** (torch.arange(0, hd, 2, device=x.device).float() / hd))
    ang = positions[:, None, :, None].float() * inv_freq  # (B, 1, T, hd/2)
    cos, sin = ang.cos().to(x.dtype), ang.sin().to(x.dtype)
    x1, x2 = x[..., 0::2], x[..., 1::2]
    return torch.stack((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1).flatten(-2)


class Attention(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        assert config.D_model % config.n_heads == 0
        self.n_heads = config.n_heads
        self.head_dim = config.D_model // config.n_heads
        self.rope_base = config.rope_base
        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.dropout = config.dropout

    def forward(self, x, positions, attn_mask=None):
        B, T, D = x.shape
        q, k, v = self.qkv(x).view(B, T, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k = apply_rope(q, positions, self.rope_base), apply_rope(k, positions, self.rope_base)
        y = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=attn_mask,
            dropout_p=self.dropout if self.training else 0.0,
            is_causal=attn_mask is None,
        )
        return self.proj(y.transpose(1, 2).reshape(B, T, D))


class Block(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.ln1 = nn.RMSNorm(config.D_model)
        self.attn = Attention(config)
        self.ln2 = nn.RMSNorm(config.D_model)
        self.ffn = build_ffn(config)

    def forward(self, x, positions, attn_mask=None, pad_mask=None):
        x = x + self.attn(self.ln1(x), positions, attn_mask)
        return x + self.ffn(self.ln2(x), pad_mask)


class Transformer(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.config = config
        self.tok_emb = nn.Embedding(config.N_vocab, config.D_model)
        self.blocks = nn.ModuleList(Block(config) for _ in range(config.n_layers))
        self.ln_f = nn.RMSNorm(config.D_model)
        self.lm_head = nn.Linear(config.D_model, config.N_vocab, bias=False)
        self.lm_head.weight = self.tok_emb.weight  # weight tying
        self.apply(self._init)
        for name, p in self.named_parameters():  # GPT-2 style residual scaling
            if name.endswith(("attn.proj.weight", "down.weight")):  # incl. experts
                nn.init.normal_(p, std=0.02 / math.sqrt(2 * config.n_layers))

    @staticmethod
    def _init(m):
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, std=0.02)

    def forward(self, ids, pad_mask=None, left_padded=False):
        # Input shape: (B, N_ctx) token ids
        # Output shape: (B, N_ctx, N_vocab)
        # pad_mask: (B, N_ctx) bool, True for real tokens. It is always handed to
        # the FFN (MoE ignores padding in its aux loss / usage stats). With right
        # padding (training/eval) plain causal attention is exact, since real
        # tokens never attend to trailing pads. For left-padded batches
        # (generation) set left_padded=True to mask pad keys and shift positions.
        B, T = ids.shape
        if not left_padded:
            positions = torch.arange(T, device=ids.device).expand(B, T)
            attn_mask = None
        else:
            positions = (pad_mask.long().cumsum(-1) - 1).clamp(min=0)
            causal = torch.ones(T, T, dtype=torch.bool, device=ids.device).tril()
            eye = torch.eye(T, dtype=torch.bool, device=ids.device)
            # `eye` keeps pad queries from having an all-masked row (NaNs)
            attn_mask = (causal & pad_mask[:, None, None, :]) | eye  # (B, 1, T, T)
        x = self.tok_emb(ids)
        for blk in self.blocks:
            x = blk(x, positions, attn_mask, pad_mask)
        return self.lm_head(self.ln_f(x))

    # ---- MoE helpers ----
    def moe_layers(self) -> list[MoE]:
        return [b.ffn for b in self.blocks if isinstance(b.ffn, MoE)]

    def aux_loss(self) -> torch.Tensor:
        layers = self.moe_layers()
        if not layers:
            return torch.zeros((), device=self.tok_emb.weight.device)
        return sum(m.aux_loss for m in layers) / len(layers)

    def set_record_routing(self, flag: bool) -> None:
        for m in self.moe_layers():
            m.record_routing = flag

    def param_counts(self) -> dict[str, int]:
        """total / active parameters (active = per-token, excluding unused experts)."""
        from src.utils import count_params

        total = count_params(self)
        ffn_total = sum(count_params(b.ffn) for b in self.blocks)
        ffn_active = sum(b.ffn.active_params() for b in self.blocks)
        return {"total": total, "active": total - ffn_total + ffn_active,
                "ffn_total": ffn_total, "ffn_active": ffn_active}