# moe_transformer.py
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class SparseMoE(nn.Module):
    """Sparse Mixture-of-Experts layer."""
    def __init__(self, d_model, num_experts, top_k, routing_algorithm, d_ff):
        super().__init__()
        self.d_model = d_model
        self.num_experts = num_experts
        self.top_k = top_k
        self.routing_algorithm = routing_algorithm
        
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_ff),
                nn.ReLU(),
                nn.Linear(d_ff, d_model)
            ) for _ in range(num_experts)
        ])

        if self.routing_algorithm == 'top_k':
            self.gate = nn.Linear(d_model, num_experts)
        
        self.load_balancing_loss = 0.0

    def hash_routing(self, x):
        token_hashes = x.sum(dim=-1).long().abs()
        expert_indices = token_hashes % self.num_experts
        return F.one_hot(expert_indices, num_classes=self.num_experts).float()

    def top_k_routing(self, x):
        gate_logits = self.gate(x)
        top_k_logits, top_k_indices = torch.topk(gate_logits, self.top_k, dim=-1)
        gate_scores = F.softmax(top_k_logits, dim=-1)
        
        router_mask = torch.zeros_like(gate_logits).scatter_(-1, top_k_indices, gate_scores)

        if self.training:
            probs_per_expert = gate_logits.softmax(dim=-1)
            tokens_per_batch_seq = router_mask.shape[0]
            fraction_tokens_per_expert = router_mask.sum(dim=0) / tokens_per_batch_seq
            mean_prob_per_expert = probs_per_expert.mean(dim=0)
            self.load_balancing_loss = self.num_experts * torch.sum(fraction_tokens_per_expert * mean_prob_per_expert)

        return router_mask

    def forward(self, x):
        batch_size, seq_len, _ = x.shape
        x_flat = x.view(-1, self.d_model)

        if self.routing_algorithm == 'top_k':
            router_output = self.top_k_routing(x_flat)
        elif self.routing_algorithm == 'hash':
            router_output = self.hash_routing(x_flat)
        else:
            raise ValueError(f"Unknown routing algorithm: {self.routing_algorithm}")
        
        final_output = torch.zeros_like(x_flat)
        for i, expert in enumerate(self.experts):
            expert_mask = router_output[:, i].unsqueeze(1)
            active_tokens_indices = torch.where(expert_mask.squeeze() > 0)[0]
            if active_tokens_indices.numel() > 0:
                active_tokens = x_flat[active_tokens_indices]
                expert_out = expert(active_tokens)
                weighted_out = expert_out * expert_mask[active_tokens_indices]
                final_output.index_add_(0, active_tokens_indices, weighted_out)

        return final_output.view(batch_size, seq_len, self.d_model)

class GroupedQueryAttention(nn.Module):
    """
    Implements Grouped-Query Attention (GQA).
    - MHA is a special case of GQA where num_kv_heads == num_heads.
    - MQA is a special case of GQA where num_kv_heads == 1.
    """
    def __init__(self, d_model, num_heads, num_kv_heads):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        assert num_heads % num_kv_heads == 0, "num_heads must be divisible by num_kv_heads"

        self.d_model = d_model
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.num_key_value_groups = num_heads // num_kv_heads
        self.d_k = d_model // num_heads
        
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, self.num_kv_heads * self.d_k)
        self.W_v = nn.Linear(d_model, self.num_kv_heads * self.d_k)
        self.W_o = nn.Linear(d_model, d_model)
    
    def scaled_dot_product_attention(self, Q, K, V, mask=None):
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            attn_scores = attn_scores.masked_fill(mask == 0, -1e9)
        attn_probs = F.softmax(attn_scores, dim=-1)
        output = torch.matmul(attn_probs, V)
        return output

    def forward(self, q, k, v, mask=None):
        batch_size = q.size(0)
        
        Q = self.W_q(q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(k).view(batch_size, -1, self.num_kv_heads, self.d_k).transpose(1, 2)
        V = self.W_v(v).view(batch_size, -1, self.num_kv_heads, self.d_k).transpose(1, 2)
        
        if self.num_key_value_groups > 1:
            K = K.repeat_interleave(self.num_key_value_groups, dim=1)
            V = V.repeat_interleave(self.num_key_value_groups, dim=1)
        
        context = self.scaled_dot_product_attention(Q, K, V, mask)
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        
        output = self.W_o(context)
        return output

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x):
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)

class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, num_kv_heads, d_ff, num_experts, top_k, routing_algorithm, dropout):
        super().__init__()
        self.self_attn = GroupedQueryAttention(d_model, num_heads, num_kv_heads)
        self.moe_ffn = SparseMoE(d_model, num_experts, top_k, routing_algorithm, d_ff)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask):
        attn_output = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_output))
        
        moe_output = self.moe_ffn(x)
        x = self.norm2(x + self.dropout(moe_output))
        return x

class DecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, num_kv_heads, d_ff, num_experts, top_k, routing_algorithm, dropout):
        super().__init__()
        self.self_attn = GroupedQueryAttention(d_model, num_heads, num_kv_heads)
        self.cross_attn = GroupedQueryAttention(d_model, num_heads, num_kv_heads)
        self.moe_ffn = SparseMoE(d_model, num_experts, top_k, routing_algorithm, d_ff)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, enc_output, src_mask, tgt_mask):
        attn_output = self.self_attn(x, x, x, tgt_mask)
        x = self.norm1(x + self.dropout(attn_output))
        
        cross_attn_output = self.cross_attn(x, enc_output, enc_output, src_mask)
        x = self.norm2(x + self.dropout(cross_attn_output))
        
        moe_output = self.moe_ffn(x)
        x = self.norm3(x + self.dropout(moe_output))
        return x

class MoETransformer(nn.Module):
    def __init__(self, config, vocab_size):
        super().__init__()
        self.config = config
        self.encoder_embedding = nn.Embedding(vocab_size, config['d_model'])
        self.decoder_embedding = nn.Embedding(vocab_size, config['d_model'])
        self.positional_encoding = PositionalEncoding(config['d_model'], config['dropout'])

        self.encoder_layers = nn.ModuleList([
            EncoderLayer(config['d_model'], config['num_heads'], config['num_kv_heads'], config['d_ff'], config['num_experts'], config['top_k'], config['routing_algorithm'], config['dropout']) 
            for _ in range(config['num_encoder_layers'])
        ])
        self.decoder_layers = nn.ModuleList([
            DecoderLayer(config['d_model'], config['num_heads'], config['num_kv_heads'], config['d_ff'], config['num_experts'], config['top_k'], config['routing_algorithm'], config['dropout']) 
            for _ in range(config['num_decoder_layers'])
        ])

        self.fc_out = nn.Linear(config['d_model'], vocab_size)

    def generate_mask(self, src, tgt, pad_idx):
        src_mask = (src != pad_idx).unsqueeze(1).unsqueeze(2)
        tgt_pad_mask = (tgt != pad_idx).unsqueeze(1).unsqueeze(2)
        seq_len = tgt.size(1)
        tgt_sub_mask = torch.tril(torch.ones((seq_len, seq_len), device=tgt.device)).bool()
        tgt_mask = tgt_pad_mask & tgt_sub_mask
        return src_mask, tgt_mask

    def forward(self, src, tgt, pad_idx=0):
        src_mask, tgt_mask = self.generate_mask(src, tgt, pad_idx)
        
        src_emb = self.positional_encoding(self.encoder_embedding(src) * math.sqrt(self.config['d_model']))
        tgt_emb = self.positional_encoding(self.decoder_embedding(tgt) * math.sqrt(self.config['d_model']))

        enc_output = src_emb
        for layer in self.encoder_layers:
            enc_output = layer(enc_output, src_mask)
            
        dec_output = tgt_emb
        for layer in self.decoder_layers:
            dec_output = layer(dec_output, enc_output, src_mask, tgt_mask)
            
        return self.fc_out(dec_output)

    def get_total_load_balancing_loss(self):
        total_loss = 0
        for layer in self.encoder_layers + self.decoder_layers:
            total_loss += layer.moe_ffn.load_balancing_loss
        return total_loss

    @torch.no_grad()
    def generate(self, src, max_length, start_symbol, pad_idx=0):
        self.eval()
        device = next(self.parameters()).device
        src = src.to(device)
        batch_size = src.shape[0]

        src_mask = (src != pad_idx).unsqueeze(1).unsqueeze(2)
        src_emb = self.positional_encoding(self.encoder_embedding(src) * math.sqrt(self.config['d_model']))
        enc_output = src_emb
        for layer in self.encoder_layers:
            enc_output = layer(enc_output, src_mask)

        tgt = torch.full((batch_size, 1), start_symbol, dtype=torch.long, device=device)

        for _ in range(max_length - 1):
            _, tgt_mask = self.generate_mask(src, tgt, pad_idx)
            tgt_emb = self.positional_encoding(self.decoder_embedding(tgt) * math.sqrt(self.config['d_model']))
            dec_output = tgt_emb
            for layer in self.decoder_layers:
                dec_output = layer(dec_output, enc_output, src_mask, tgt_mask)
            
            logits = self.fc_out(dec_output[:, -1])
            next_token = torch.argmax(logits, dim=-1).unsqueeze(1)
            tgt = torch.cat([tgt, next_token], dim=1)

            if (tgt == 1).any(dim=-1).all():
                break

        return tgt