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

class MoELayer(nn.Module):
    def __init__(self, input_dim, output_dim, num_experts=4, k=2):
        super().__init__()
        self.k = k
        self.num_experts = num_experts

        # Router: logits = W_r · x
        self.router = nn.Linear(input_dim, num_experts)

        # Experts: each is a small MLP (or Linear here)
        self.experts = nn.ModuleList([
            nn.Linear(input_dim, output_dim) for _ in range(num_experts)
        ])

    def forward(self, x):
        """
        x: (batch_size, input_dim)
        returns: (batch_size, output_dim)
        """
        # Compute routing logits
        logits = self.router(x)  # (B, N)
        gate_probs = F.softmax(logits, dim=-1)  # (B, N)

        # Select top-k experts per example
        topk_vals, topk_idx = torch.topk(gate_probs, self.k, dim=-1)  # (B, k)

        # Initialize output
        output = torch.zeros(x.size(0), self.experts[0].out_features, device=x.device)

        for i in range(self.k):
            idx = topk_idx[:, i]  # expert index
            weight = topk_vals[:, i].unsqueeze(1)  # (B, 1)

            # For batch-wise selection
            expert_outputs = torch.stack([
                self.experts[expert_id](x[j].unsqueeze(0))
                for j, expert_id in enumerate(idx)
            ]).squeeze(1)  # (B, output_dim)

            output += weight * expert_outputs  # weighted sum

        return output

idx = torch.argmax(logits)

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

class Router(nn.Module):
    """
    The router module determines which expert each token is sent to.
    It also computes the load balancing loss.
    """
    def __init__(self, d_model, num_experts):
        super().__init__()
        self.num_experts = num_experts
        self.gate = nn.Linear(d_model, num_experts)

    def forward(self, x):
        # x: (batch_size, seq_len, d_model)
        logits = self.gate(x)
        # logits: (batch_size, seq_len, num_experts)
        
        # Get the top-1 expert for each token
        top1_logits, top1_indices = logits.max(dim=-1)
        # top1_indices: (batch_size, seq_len)
        
        # Create a one-hot encoding of the expert indices
        # This will be used to dispatch tokens to the correct expert
        expert_mask = F.one_hot(top1_indices, self.num_experts).float()
        # expert_mask: (batch_size, seq_len, num_experts)
        
        # Calculate the load balancing loss
        # This loss encourages all experts to be used equally.
        
        # Count how many tokens are sent to each expert
        tokens_per_expert = expert_mask.sum(dim=(0, 1)) # (num_experts)
        # Total number of tokens
        total_tokens = x.size(0) * x.size(1)
        
        # Calculate the fraction of tokens sent to each expert
        fraction_tokens_per_expert = tokens_per_expert / total_tokens
        
        # Calculate the expert probabilities from the logits
        expert_probs = F.softmax(logits, dim=-1).mean(dim=(0, 1)) # (num_experts)
        
        # The load balancing loss is the dot product of these two quantities
        load_balancing_loss = self.num_experts * torch.dot(fraction_tokens_per_expert, expert_probs)
        
        return expert_mask, load_balancing_loss

class SwitchTransformerLM(nn.Module):
    def __init__(self, ntoken: int, d_model: int, nhead: int, d_ff: int, num_experts: int, num_layers: int,
                 dropout: float = 0.5):
        super().__init__()
        self.model_type = 'Transformer'
        self.d_model = d_model
        self.encoder = nn.Embedding(ntoken, d_model)
        self.pos_encoder = PositionalEncoding(d_model, dropout)

        encoder_layer = SwitchTransformerEncoderLayer(d_model, nhead, d_ff, num_experts, dropout)
        self.transformer_encoder = SwitchTransformerEncoder(encoder_layer, num_layers)

        self.decoder = nn.Linear(d_model, ntoken)

class SwitchTransformerEncoderLayer(nn.Module):
    def __init__(self, d_model: int, nhead: int, d_ff: int, num_experts: int, dropout: float = 0.1) -> None:
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.moe_layer = MoELayer(d_model, d_ff, num_experts)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, src: Tensor,
                src_mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]:
        # Multi-head self-attention
        src2 = self.self_attn(src, src, src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout(src2)
        src = self.norm1(src)

        # MoE layer
        src2, load_balancing_loss = self.moe_layer(src)
        src = src + self.dropout(src2)
        src = self.norm2(src)

        return src, load_balancing_loss

class MoELayer(nn.Module):
    """
    The Mixture-of-Experts (MoE) layer, which replaces the FFN layer in a standard Transformer.
    """
    def __init__(self, d_model: int, d_ff: int, num_experts: int, capacity_factor: float = 1.25) -> None:
        super().__init__()
        self.d_model = d_model
        self.num_experts = num_experts
        self.capacity_factor = capacity_factor

        self.router = Router(d_model, num_experts)
        self.experts = nn.ModuleList([Expert(d_model, d_ff) for _ in range(num_experts)])
        
    def forward(self, x: Tensor) -> Tuple[Tensor, Tensor]:
        """
        Example: Let's say we have:
        - batch_size=2, seq_len=4, d_model=8, num_experts=2
        - Input x: shape (2, 4, 8) - 2 sequences, each with 4 tokens of 8 dimensions
        """
        # x: (batch_size, seq_len, d_model)
        # Example: x.shape = (2, 4, 8)
        batch_size, seq_len, d_model = x.shape
        
        # Get the expert mask and load balancing loss from the router
        # The router decides which expert should process each token
        expert_mask, load_balancing_loss = self.router(x)
        # expert_mask: (batch_size, seq_len, num_experts)
        # Example: expert_mask.shape = (2, 4, 2)
        #   expert_mask[0, 0, :] = [1, 0] means token 0 goes to expert 0
        #   expert_mask[0, 1, :] = [0, 1] means token 1 goes to expert 1
        
        # Determine the capacity of each expert
        capacity = int((seq_len / self.num_experts) * self.capacity_factor)
        
        # Reshape tensors for easier processing - flatten batch and sequence dimensions
        x_flat = x.view(-1, d_model) # (batch_size * seq_len, d_model)
        expert_mask_flat = expert_mask.view(-1, self.num_experts) # (batch_size * seq_len, num_experts)

        # Initialize output tensor with zeros - same shape as flattened input
        final_output = torch.zeros_like(x_flat)
        # Example: final_output.shape = (8, 8)
        
        # Process each expert separately
        for i, expert in enumerate(self.experts):
            # Find which tokens should go to this expert (non-zero entries in the mask)
            expert_indices = torch.where(expert_mask_flat[:, i] > 0)[0]
            # Example for expert 0: expert_indices might be [0, 2, 4] (tokens 0, 2, 4)
            # Example for expert 1: expert_indices might be [1, 3, 5, 6, 7] (tokens 1, 3, 5, 6, 7)
            
            # Apply capacity constraint - if too many tokens assigned, keep only the first 'capacity' tokens
            if expert_indices.shape[0] > capacity:
                expert_indices = expert_indices[:capacity]
                # Example: if expert 1 has 5 tokens but capacity=2, keep only [1, 3]

            # Process tokens assigned to this expert
            if expert_indices.shape[0] > 0:
                # Extract the tokens that should go to this expert
                expert_input = x_flat[expert_indices]
                # Example: if expert_indices=[1, 3], expert_input.shape = (2, 8)
                
                # Process the tokens through the expert network
                expert_output = expert(expert_input)
                # Example: expert_output.shape = (2, 8) - same as expert_input
                
                # Ensure dtype consistency
                expert_output = expert_output.to(final_output.dtype)

                # Place the expert's output back to the final output at the correct positions
                # Switch Transformer uses exclusive routing: each token goes to exactly ONE expert
                # So we use assignment (=) not addition (+=)
                # Example: if expert_indices=[1, 3] and expert_output has 2 rows
                #   final_output[1] = expert_output[0]  # Place expert's output for token 1
                #   final_output[3] = expert_output[1]  # Place expert's output for token 3
                for j, token_idx in enumerate(expert_indices):
                    final_output[token_idx] = expert_output[j]

        # Reshape back to original dimensions
        final_output = final_output.view(batch_size, seq_len, d_model)
        # Example: final_output.shape = (2, 4, 8) - back to original shape
        
        return final_output, load_balancing_loss