"""A decoder-only transformer whose feed-forward block is swappable.

Attention, RoPE and RMSNorm are carried over from Assignment 1 and are held
fixed across all five ablations -- the feed-forward network is the only thing
that changes, which is what makes the comparison mean anything.
"""

import math

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

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


class RMSNorm(nn.Module):
    def __init__(self, d_model, eps=1e-6):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(d_model))
        self.eps = eps

    def forward(self, x):
        dtype = x.dtype
        x = x.float()
        rms = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return (self.gamma.float() * x * rms).to(dtype)


class RotaryEmbedding(nn.Module):
    """Rotary position embedding, applied to q and k inside attention."""

    def __init__(self, head_dim, max_len, base=10000.0):
        super().__init__()
        inv = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
        t = torch.arange(max_len).float()
        freqs = torch.outer(t, inv)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos", emb.cos()[None, None], persistent=False)
        self.register_buffer("sin", emb.sin()[None, None], persistent=False)

    @staticmethod
    def _rotate_half(x):
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat((-x2, x1), dim=-1)

    def forward(self, q, k, offset=0):
        n = q.size(-2)
        cos = self.cos[..., offset:offset + n, :].to(q.dtype)
        sin = self.sin[..., offset:offset + n, :].to(q.dtype)
        q = q * cos + self._rotate_half(q) * sin
        k = k * cos + self._rotate_half(k) * sin
        return q, k


class CausalSelfAttention(nn.Module):
    def __init__(self, config: TransformerConfig, rope: RotaryEmbedding):
        super().__init__()
        self.n_heads = config.n_heads
        self.n_kv = config.n_kv_heads or config.n_heads
        assert self.n_heads % self.n_kv == 0
        self.n_groups = self.n_heads // self.n_kv
        self.d_head = config.d_head
        self.rope = rope

        d = config.d_model
        self.wq = nn.Linear(d, self.n_heads * self.d_head, bias=False)
        self.wk = nn.Linear(d, self.n_kv * self.d_head, bias=False)
        self.wv = nn.Linear(d, self.n_kv * self.d_head, bias=False)
        self.wo = nn.Linear(self.n_heads * self.d_head, d, bias=False)
        self.dropout = config.dropout

    def forward(self, x):
        B, N, _ = x.shape
        q = self.wq(x).view(B, N, self.n_heads, self.d_head).transpose(1, 2)
        k = self.wk(x).view(B, N, self.n_kv, self.d_head).transpose(1, 2)
        v = self.wv(x).view(B, N, self.n_kv, self.d_head).transpose(1, 2)

        q, k = self.rope(q, k)
        if self.n_groups > 1:
            k = k.repeat_interleave(self.n_groups, dim=1)
            v = v.repeat_interleave(self.n_groups, dim=1)

        out = F.scaled_dot_product_attention(
            q, k, v, is_causal=True,
            dropout_p=self.dropout if self.training else 0.0)
        out = out.transpose(1, 2).reshape(B, N, -1)
        return self.wo(out)


class Block(nn.Module):
    """Pre-norm decoder block; ``self.ffn`` is the MoE/MLP drop-in point."""

    def __init__(self, config: TransformerConfig, rope: RotaryEmbedding):
        super().__init__()
        self.norm1 = RMSNorm(config.d_model)
        self.attn = CausalSelfAttention(config, rope)
        self.norm2 = RMSNorm(config.d_model)
        self.ffn = build_ffn(config)

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x


class Transformer(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.config = config
        self.embed = nn.Embedding(config.vocab_size, config.d_model)
        self.rope = RotaryEmbedding(config.d_head, config.n_ctx, config.rope_base)
        self.blocks = nn.ModuleList(
            [Block(config, self.rope) for _ in range(config.n_layers)])
        self.norm_f = RMSNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        if config.tie_embeddings:
            self.lm_head.weight = self.embed.weight
        self.apply(self._init)
        # Scale the residual-path output projections by 1/sqrt(2*L) so the
        # residual stream does not blow up with depth (GPT-2 initialisation).
        for name, p in self.named_parameters():
            if name.endswith("wo.weight") or name.endswith("fc2.weight"):
                nn.init.normal_(p, std=0.02 / math.sqrt(2 * config.n_layers))

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

    def forward(self, idx, targets=None, loss_mask=None):
        # Input shape: (B, N_ctx); logits shape: (B, N_ctx, N_vocab)
        x = self.embed(idx)
        for block in self.blocks:
            x = block(x)
        logits = self.lm_head(self.norm_f(x))

        if targets is None:
            return logits, None, self.aux_losses()

        loss = F.cross_entropy(
            logits.view(-1, logits.size(-1)).float(), targets.reshape(-1),
            ignore_index=-100, reduction="none")
        if loss_mask is not None:
            m = loss_mask.reshape(-1).float()
            loss = (loss * m).sum() / m.sum().clamp_min(1.0)
        else:
            loss = loss[targets.reshape(-1) != -100].mean()
        return logits, loss, self.aux_losses()

    def aux_losses(self):
        """Summed router losses over every MoE layer (zero for the dense model)."""
        aux = torch.zeros((), device=self.embed.weight.device)
        z = torch.zeros((), device=self.embed.weight.device)
        for block in self.blocks:
            if isinstance(block.ffn, MoE):
                aux = aux + block.ffn.aux_loss
                z = z + block.ffn.z_loss
        return aux, z

    def expert_counts(self):
        """Per-layer token counts from the last forward, or None if dense."""
        out = []
        for block in self.blocks:
            if isinstance(block.ffn, MoE):
                out.append(block.ffn.last_counts)
        return out or None

    def expert_assignments(self):
        """Per-layer (B, N, k) expert ids from the last forward."""
        out = []
        for block in self.blocks:
            if isinstance(block.ffn, MoE):
                out.append(block.ffn.last_assign)
        return out or None

    def num_params(self, non_embedding=False):
        n = sum(p.numel() for p in self.parameters())
        if non_embedding:
            n -= self.embed.weight.numel()
        return n