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

# --- Configuration for torch.compile ---
# Attempt to prevent graph break from .item() as suggested by the error log.
if hasattr(torch, '_dynamo') and hasattr(torch._dynamo, 'config'):
    torch._dynamo.config.capture_scalar_outputs = True
    # Optional: torch._dynamo.config.suppress_warnings = True # To hide the .item() warning if capture works
else:
    print("Warning: torch._dynamo.config not found. Cannot set capture_scalar_outputs.")


# --- Reference Classes (Expert, MoEGate are used as is for structure) ---
class Expert(nn.Module):
    def __init__(self, config: Dict, d_expert: Optional[int] = None):
        super().__init__()
        self.act_fn = nn.SiLU()
        self.d_hidden: int = config["d_hidden"]
        self.d_expert: int = d_expert if d_expert is not None else config["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]:
        logits = self.W_g(x)
        scores = logits.softmax(dim=-1)
        topk_scores, topk_indices = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
        return topk_indices, topk_scores
# --- End of Reference Classes for Expert and MoEGate ---

class MoE_Optimized(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config 
        self.num_routed_experts = config["n_routed_experts"] 
        self.experts = nn.ModuleList([
            Expert(config) 
            for _ in range(self.num_routed_experts) 
        ])
        self.gating_network = MoEGate(config)
        shared_expert_dim = config["d_expert"] * config["n_shared_experts"]
        self.shared_expert = Expert(config=config, d_expert=shared_expert_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        shared_output = self.shared_expert(x)
        expert_indices, expert_scores = self.gating_network(x)
        
        orig_shape = x.shape 
        
        x_flat = x.view(-1, x.shape[-1]) 
        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)
        return routed_output + shared_output

    def moe_infer(self,
                  x_flat: torch.Tensor,
                  flat_expert_indices: torch.Tensor,
                  flat_expert_weights: torch.Tensor
                 ) -> torch.Tensor:
        expert_cache = torch.zeros_like(x_flat)
        idxs = flat_expert_indices.argsort() 
        
        counts_gpu = flat_expert_indices.bincount(minlength=self.num_routed_experts)
        tokens_per_expert_gpu_cumsum = counts_gpu.cumsum(dim=0) 
        
        num_per_tok = self.config["n_experts_per_token"] 
        original_token_indices_for_sorted_assignments = idxs // num_per_tok

        for expert_id in range(self.num_routed_experts): 
            if expert_id == 0:
                start_idx_for_assignments = 0
            else:
                start_idx_for_assignments = tokens_per_expert_gpu_cumsum[expert_id - 1].item()
            
            end_idx_for_assignments = tokens_per_expert_gpu_cumsum[expert_id].item()

            if start_idx_for_assignments == end_idx_for_assignments: 
                continue 

            expert = self.experts[expert_id]
            
            current_expert_assignment_original_indices = idxs[start_idx_for_assignments:end_idx_for_assignments]
            current_expert_token_indices_in_xflat = original_token_indices_for_sorted_assignments[start_idx_for_assignments:end_idx_for_assignments]
            
            expert_tokens_input = x_flat[current_expert_token_indices_in_xflat]
            expert_output_weighted = expert(expert_tokens_input) 
            
            weights_for_current_assignments = flat_expert_weights[current_expert_assignment_original_indices]
            expert_output_weighted.mul_(weights_for_current_assignments) 
            
            scatter_indices = current_expert_token_indices_in_xflat.view(-1, 1).repeat(1, x_flat.shape[-1])
            
            expert_cache.scatter_reduce_(
                0, 
                scatter_indices,
                expert_output_weighted,
                reduce='sum'
            )
            
        return expert_cache

# --- custom_kernel: Entry point for the submission ---
input_t = Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict]
output_t_for_custom_kernel = torch.Tensor

def custom_kernel(data: input_t) -> output_t_for_custom_kernel:
    input_tensor, weights, config = data
    
    input_tensor = input_tensor.to(device='cuda', dtype=torch.float16)

    moe_module = MoE_Optimized(config) 
    
    # Weight loading
    moe_module.gating_network.W_g.weight = nn.Parameter(
        weights['router.weight'].to(device='cuda', dtype=torch.float16)
    )
    num_routed_experts = config["n_routed_experts"]
    for i in range(num_routed_experts):
        gate_proj_weight = weights[f'experts.{i}.0.weight']
        up_proj_weight = weights[f'experts.{i}.1.weight']
        down_proj_weight = weights[f'experts.{i}.2.weight']
        moe_module.experts[i].W_gate.weight = nn.Parameter(gate_proj_weight.t().to(device='cuda', dtype=torch.float16))
        moe_module.experts[i].W_up.weight = nn.Parameter(up_proj_weight.t().to(device='cuda', dtype=torch.float16))
        moe_module.experts[i].W_down.weight = nn.Parameter(down_proj_weight.t().to(device='cuda', dtype=torch.float16))
        
    moe_module.shared_expert.W_gate.weight = nn.Parameter(
        weights['shared_experts.0.weight'].t().to(device='cuda', dtype=torch.float16)
    )
    moe_module.shared_expert.W_up.weight = nn.Parameter(
        weights['shared_experts.1.weight'].t().to(device='cuda', dtype=torch.float16)
    )
    moe_module.shared_expert.W_down.weight = nn.Parameter(
        weights['shared_experts.2.weight'].t().to(device='cuda', dtype=torch.float16)
    )
    
    moe_module = moe_module.to(device='cuda', dtype=torch.float16)
    moe_module.eval()

    if hasattr(torch, 'compile'):
        try:
            # Using "default" mode as the primary attempt
            # print("Attempting to compile MoE module with mode='default'...") # For local debugging
            moe_module = torch.compile(moe_module, mode="default") 
            # print("Compilation with 'default' mode successful.") # For local debugging
        except Exception as e_default:
            print(f"Warning: torch.compile with mode='default' failed: {e_default}. Trying 'reduce-overhead'.")
            try:
                moe_module = torch.compile(moe_module, mode="reduce-overhead")
            except Exception as e_reduce_overhead:
                 print(f"Warning: torch.compile with mode='reduce-overhead' also failed: {e_reduce_overhead}. Running uncompiled.")
    
    with torch.inference_mode():
        output_tensor = moe_module(input_tensor) 
    
    return output_tensor

python eval.py benchmark /tmp/tmppu1asygo
W0518 22:14:43.734000 747 torch/fx/experimental/symbolic_shapes.py:6903] [2/0] failed during evaluate_expr(Eq(u0, 0), hint=None, size_oblivious=False, forcing_spec=False
