import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple, Optional
import math

class Expert(nn.Module):
    def __init__(self, config: Dict, d_expert: Optional[int] = None):
        super().__init__()
        self.config = config
        self.act_fn = nn.SiLU() 
        self.d_hidden: int = config["d_hidden"]
        self.d_expert: int = config["d_expert"] if d_expert is None else d_expert

        self.W_gate = nn.Linear(self.d_hidden, self.d_expert, bias=False)
        self.W_up = nn.Linear(self.d_hidden, self.d_expert, bias=False)
        self.W_down = nn.Linear(self.d_expert, self.d_hidden, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate = self.act_fn(self.W_gate(x))
        up_hidden = self.W_up(x)
        out = self.W_down(gate * up_hidden)
        return out

class MoEGate(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.top_k: int = config["n_experts_per_token"]
        self.num_experts: int = config["n_routed_experts"]
        self.d_hidden: int = config["d_hidden"]

        self.W_g = nn.Linear(self.d_hidden, self.num_experts, bias=False)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        logits = self.W_g(x)
        scores = F.softmax(logits, dim=-1)
        topk_scores, topk_indices = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
        return topk_indices, topk_scores, scores

class MoE(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config
        self.num_routed_experts = config["n_routed_experts"] 
        self.top_k = config["n_experts_per_token"]

        self.experts = nn.ModuleList([
            Expert(config) for _ in range(self.num_routed_experts)
        ])
        self.gating_network = MoEGate(config)
        
        shared_expert_config = config.copy() 
        n_shared_experts = config.get("n_shared_experts", 1) 
        d_expert_shared = config["d_expert"] * n_shared_experts
        
        self.shared_expert = Expert(shared_expert_config, d_expert=d_expert_shared)
        
        self.aux_loss_weight = config.get("aux_loss_weight", 0.01)

    @torch.no_grad()
    def moe_infer(self, x: torch.Tensor, flat_expert_indices: torch.Tensor,
                 flat_expert_weights: torch.Tensor) -> torch.Tensor:
        expert_cache = torch.zeros_like(x)
        
        idxs = flat_expert_indices.argsort()
        
        counts = flat_expert_indices.bincount(minlength=self.num_routed_experts).cpu().numpy()
        tokens_per_expert_cumsum = counts.cumsum()

        token_idxs_for_sorted_assignments = idxs // self.top_k


        for expert_id, end_idx_in_sorted_list in enumerate(tokens_per_expert_cumsum):
            start_idx_in_sorted_list = 0 if expert_id == 0 else tokens_per_expert_cumsum[expert_id - 1]
            if start_idx_in_sorted_list == end_idx_in_sorted_list:
                continue

            expert = self.experts[expert_id]
            
            current_expert_original_token_indices = token_idxs_for_sorted_assignments[start_idx_in_sorted_list:end_idx_in_sorted_list]
            expert_tokens_input = x[current_expert_original_token_indices]
            expert_tokens_output = expert(expert_tokens_input)
            
            weights_for_current_assignments = flat_expert_weights[idxs[start_idx_in_sorted_list:end_idx_in_sorted_list]]
            expert_tokens_output.mul_(weights_for_current_assignments)
            
            expert_cache.scatter_reduce_(
                0, 
                current_expert_original_token_indices.view(-1, 1).repeat(1, x.shape[-1]),
                expert_tokens_output,
                reduce='sum',
                include_self=True 
            )
        return expert_cache

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Dict]:
        batch_size, seq_len, hidden_dim = x.shape
        orig_shape = x.shape

        shared_output = self.shared_expert(x)
        
        expert_indices, expert_scores, all_router_scores = self.gating_network(x)
        
        x_flat = x.view(-1, hidden_dim)
        flat_expert_indices = expert_indices.view(-1)
        flat_expert_weights = expert_scores.view(-1, 1)
        
        routed_output_flat = self.moe_infer(
            x_flat, flat_expert_indices, flat_expert_weights
        )
        
        routed_output = routed_output_flat.view(*orig_shape)
        final_output = routed_output + shared_output
        
        aux_data = self.compute_aux_losses(all_router_scores, expert_indices)
        
        return final_output, aux_data

    def compute_aux_losses(self, router_probs_all_experts: torch.Tensor, 
                         topk_expert_indices: torch.Tensor) -> Dict:
       
        
        topk_expert_indices_flat = topk_expert_indices.view(-1)
        counts = torch.bincount(topk_expert_indices_flat, minlength=self.num_routed_experts).float()
        total_assignments = topk_expert_indices.numel() 
        
        if total_assignments == 0:
            fraction_assignments_per_expert = torch.zeros_like(counts)
            load_balance_loss = torch.tensor(0.0, device=counts.device)
        else:
            fraction_assignments_per_expert = counts / total_assignments
            load_balance_loss = self.num_routed_experts * torch.sum(fraction_assignments_per_expert.pow(2)) - 1.0 

        z_loss = (router_probs_all_experts.pow(2)).mean()
        total_aux_loss = self.aux_loss_weight * (load_balance_loss + z_loss)
        
        return {
            "router_probs": router_probs_all_experts,
            "load_balance_loss": load_balance_loss,
            "z_loss": z_loss,
            "total_aux_loss": total_aux_loss
        }

def custom_kernel(data: Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict]) -> torch.Tensor:
    input_tensor, weights, config = data
    n_routed_experts_config = config["n_routed_experts"] 
    
    input_tensor = input_tensor.to(dtype=torch.float16)
    
    moe_module = MoE(config)
    
    moe_module.gating_network.W_g.weight = nn.Parameter(weights['router.weight'].to(dtype=torch.float16))
    
    for i in range(n_routed_experts_config): 
        gate_weight_t = weights[f'experts.{i}.0.weight'].t().to(dtype=torch.float16)
        up_weight_t = weights[f'experts.{i}.1.weight'].t().to(dtype=torch.float16)
        down_weight_t = weights[f'experts.{i}.2.weight'].t().to(dtype=torch.float16)
        
        moe_module.experts[i].W_gate.weight = nn.Parameter(gate_weight_t)
        moe_module.experts[i].W_up.weight = nn.Parameter(up_weight_t)
        moe_module.experts[i].W_down.weight = nn.Parameter(down_weight_t)
    
    moe_module.shared_expert.W_gate.weight = nn.Parameter(weights['shared_experts.0.weight'].t().to(dtype=torch.float16))
    moe_module.shared_expert.W_up.weight = nn.Parameter(weights['shared_experts.1.weight'].t().to(dtype=torch.float16))
    moe_module.shared_expert.W_down.weight = nn.Parameter(weights['shared_experts.2.weight'].t().to(dtype=torch.float16))
    
    moe_module = moe_module.to(device='cuda', dtype=torch.float16)
    
    moe_module_compiled = torch.compile(moe_module)
    
    output_tensor, _ = moe_module_compiled(input_tensor) 
    
    return output_tensor
python eval.py test /tmp/tmpppaimne1
/usr/local/lib/python3.10/dist-packages/torch/_inductor/lowering.py:7061: UserWarning: 
Online softmax is disabled on the fly since Inductor decides to
split the reduction. Cut an issue to PyTorch if this is an
important use case and you want to speed it up with online
softmax.

  warnings.warn(
