from dataclasses import dataclass

import torch
import torch.nn as nn


@dataclass
class TransformerConfig:
    model_dim: int  # Dimension of the model
    num_heads: int
    N_ctx: int  # Number of context tokens
    use_moe: bool  # Whether to use MoE or MLP
    num_experts : int
    active_experts : int
    shared_experts : int
    expert_width : int
    vocab_size : int
    layers : int

class MHA(nn.Module):
    def __init__(self, num_heads, model_dim):
        super().__init__()
        self.num_heads = num_heads
        self.model_dim = model_dim
        self.head_dim = model_dim // num_heads
        
        self.q = nn.Linear(in_features=model_dim, out_features=model_dim)
        self.k = nn.Linear(in_features=model_dim, out_features=model_dim)
        self.v = nn.Linear(in_features=model_dim, out_features=model_dim)
        self.out_proj = nn.Linear(in_features=model_dim, out_features=model_dim)

    def forward(self, x, mask=None):
        q = self.q(x)
        k = self.k(x)
        v = self.v(x)

        q_heads = q.reshape(x.shape[0], x.shape[1], self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        k_heads = k.reshape(x.shape[0], x.shape[1], self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        v_heads = v.reshape(x.shape[0], x.shape[1], self.num_heads, self.head_dim).permute(0, 2, 1, 3)

        attention_scores = (q_heads @ torch.transpose(k_heads, 3, 2)) / self.head_dim ** (1/2)
        if mask is not None:
            attention_scores = attention_scores.masked_fill(mask == 0, float('-inf'))

        attention_scores = torch.softmax(attention_scores, dim=-1)
        self_attention = attention_scores @ v_heads

        mha_output = self_attention.permute(0, 2, 1, 3).reshape(x.shape[0], x.shape[1], self.model_dim)
        output = self.out_proj(mha_output)
        return output

class RMSNorm(nn.Module):
    def __init__(self, model_dim):
        super().__init__()
        self.model_dim = model_dim
        self.gamma = nn.Parameter(torch.ones(model_dim))
        
    def forward(self, x):
        eps = 1e-5
        rms = (((torch.sum(torch.square(x), dim=-1, keepdim=True)) / self.model_dim) + eps) ** (1/2)
        norm = self.gamma * (x / rms)
        return norm

class SinusoidalPositionalEmbeddings(nn.Module):
    def __init__(self, model_dim, N=10000):
        super().__init__()
        self.d = model_dim
        self.N = N
        self.register_buffer('k', torch.arange(0, model_dim, step=2))

    def forward(self, x):
        t = torch.arange(x.shape[1], device=x.device).unsqueeze(1)
        sin = torch.sin(t/self.N ** (self.k/self.d))
        cos = torch.cos(t/self.N ** (self.k/self.d))
        total_angles = torch.cat([sin, cos], dim=-1).unsqueeze(0)
        return x + total_angles

class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.self_attn = MHA(config.num_heads, config.model_dim)
        self.norms = nn.ModuleList([RMSNorm(config.model_dim) for _ in range(2)])
        if config.use_moe:
            self.ffn = MoE(self.config)
        else:
            self.ffn = MLP(self.config)

    def forward(self, x, mask=None, pad_mask=None):
        x = x + self.self_attn(self.norms[0](x), mask)
        h, aux, counts = self.ffn(self.norms[1](x), pad_mask)
        return x + h, aux, counts

class Expert(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.expert = nn.Sequential(
            nn.Linear(config.model_dim, config.expert_width * config.model_dim),
            nn.GELU(),
            nn.Linear(config.expert_width * config.model_dim, config.model_dim)
        )

    def forward(self, x):
        x = self.expert(x)
        return x

class MoE(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.config = config
        self.router = nn.Sequential(
            nn.Linear(config.model_dim, config.num_experts - config.shared_experts),
            nn.Softmax(-1)
        )
        self.experts = nn.ModuleList([Expert(config) for _ in range(config.num_experts)])

    def forward(self, x, pad_mask=None):
        B, N, D = x.shape
        x_flat = x.reshape(-1, D)                      # (T, D) where T = B*N tokens
        probs = self.router(x_flat)                    # (T, num_routed)
        top_w, top_idx = torch.topk(probs, self.config.active_experts)  # both (T, k)
        out = torch.zeros_like(x_flat)                 # (T, D)        

        # shared routers
        if self.config.shared_experts:
            for i in range(self.config.shared_experts):
                out += self.experts[i](x_flat)

        real = pad_mask.reshape(-1).bool()
        r_probs = probs[real]
        r_idx = top_idx[real]
        num_routed = self.config.num_experts - self.config.shared_experts
        counts = torch.bincount(r_idx.reshape(-1), minlength=num_routed)

        for e in range(self.config.num_experts - self.config.shared_experts):
            e_mask = (top_idx == e).any(dim=-1)
            if not e_mask.any():
                continue
            e_wt = (top_w * (top_idx == e)).sum(dim=-1)
            out[e_mask] += e_wt[e_mask, None] * self.experts[self.config.shared_experts + e](x_flat[e_mask])

        f = counts / r_idx.numel()
        P = r_probs.mean(dim=0)
        aux = num_routed * (f * P).sum()
        return out.reshape(B, N, D), aux, counts
            


class MLP(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(config.model_dim, 4 * config.model_dim),
            nn.GELU(),
            nn.Linear(4 * config.model_dim, config.model_dim)
        )

    def forward(self, x, pad_mask=None):
        x = self.mlp(x)
        return x, 0, None


class Transformer(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.model_dim)
        self.pos_embedding = SinusoidalPositionalEmbeddings(config.model_dim)
        self.blocks = nn.ModuleList([Block(config) for _ in range(config.layers)])
        self.norm = RMSNorm(config.model_dim)
        self.lm_head = nn.Linear(config.model_dim, config.vocab_size, bias=False)
        self.lm_head.weight = self.token_embedding.weight
        nn.init.normal_(self.token_embedding.weight, std=0.02)

    def forward(self, input_ids, attention_mask, device, last_positions=None):
        N = input_ids.shape[-1]
        causal_mask = torch.tril(torch.ones(N, N, device=device, dtype=torch.bool))
        pad_mask = attention_mask
        attn_mask = causal_mask & pad_mask[:, None, None, :].bool()

        x = self.token_embedding(input_ids)
        x = self.pos_embedding(x)

        aux_total = 0
        all_counts = []
        for block in self.blocks:
            x, aux, counts = block(x, attn_mask, pad_mask)
            aux_total = aux_total + aux
            all_counts.append(counts)

        if last_positions is not None:      # generation: only each row's next-token logits are needed -> (B, V)
            x = x[torch.arange(x.shape[0], device=x.device), last_positions]
        logits = self.lm_head(self.norm(x))

        return logits, aux_total, all_counts