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 ---
if hasattr(torch, '_dynamo') and hasattr(torch._dynamo, 'config'):
    torch._dynamo.config.capture_scalar_outputs = True
else:
    print("Warning: torch._dynamo.config not found. Cannot set capture_scalar_outputs.")


# --- Model Classes ---

class Expert(nn.Module):
    """
    A single expert module in the Mixture of Experts layer.
    f(x) = W_down(SiLU(W_gate(x)) * W_up(x))
    """
    def __init__(self, config: Dict, d_expert: Optional[int] = None):
        super().__init__()
        self.act_fn = nn.SiLU()
        self.d_hidden: int = config["d_hidden"]
        # d_expert is the intermediate dimension of the FFN
        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:
        # Handle cases where an expert might receive an empty tensor
        if x.shape[0] == 0: 
            output_shape = list(x.shape)
            output_shape[-1] = self.d_hidden # Output dimension remains d_hidden
            return torch.zeros(output_shape, device=x.device, dtype=x.dtype)
        
        gate_output = self.act_fn(self.W_gate(x))
        up_output = self.W_up(x)
        expert_output = self.W_down(gate_output * up_output)
        return expert_output

class MoEGate(nn.Module):
    """
    Gating network for the Mixture of Experts layer.
    Determines which experts to route tokens to and their corresponding weights.
    """
    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) # Router linear projection

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        logits = self.W_g(x) # [bs, seq_len, num_experts]
        # Compute softmax in float32 for numerical stability, then cast back
        all_scores = F.softmax(logits, dim=-1, dtype=torch.float32).to(x.dtype)
        # Get top-k scores and their indices
        topk_scores, topk_indices = torch.topk(all_scores, k=self.top_k, dim=-1, sorted=False)
        return topk_indices, topk_scores, all_scores

class MoE_Optimized(nn.Module):
    """
    Mixture of Experts (MoE) layer.
    This version uses a debug-oriented moe_infer method.
    """
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config
        self.num_routed_experts: int = config["n_routed_experts"]
        
        # Create the pool of routed experts
        self.experts = nn.ModuleList([
            Expert(config) # Each routed expert uses config["d_expert"] as intermediate dim
            for _ in range(self.num_routed_experts)
        ])
        
        self.gating_network = MoEGate(config)
        
        # Create the shared expert(s)
        # The intermediate dimension of the shared expert is scaled by the number of shared experts
        shared_expert_intermediate_dim = config["d_expert"] * config["n_shared_experts"]
        self.shared_expert = Expert(config=config, d_expert=shared_expert_intermediate_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 1. Compute shared expert output
        shared_output = self.shared_expert(x) # [bs, seq_len, d_hidden]

        # 2. Perform routing for routed experts
        # expert_indices: [bs, seq_len, top_k]
        # expert_scores: [bs, seq_len, top_k]
        expert_indices, expert_scores, _ = self.gating_network(x) # all_router_scores is not used for inference output

        # Prepare flattened input for moe_infer_v_debug
        original_shape = x.shape # [bs, seq_len, d_hidden]
        x_flat = x.reshape(-1, original_shape[-1]) # [total_tokens, d_hidden]
                                                 # total_tokens = bs * seq_len

        # 3. Compute routed expert outputs using the debug version of moe_infer
        # This version takes x_flat and the original shaped expert_indices and expert_scores
        routed_output_flat = self.moe_infer_v_debug(x_flat,
                                                    expert_indices, 
                                                    expert_scores)
        
        # Reshape routed output back to original 3D shape
        routed_output = routed_output_flat.view(*original_shape) # [bs, seq_len, d_hidden]

        # 4. Combine shared and routed outputs
        final_output = routed_output + shared_output
        return final_output

    def moe_infer_v_debug(self,
                          x_flat: torch.Tensor, # Input tokens, flattened: [total_tokens, d_hidden]
                          expert_indices_orig_shape: torch.Tensor, # Chosen expert_ids: [bs, seq_len, top_k]
                          expert_scores_orig_shape: torch.Tensor   # Scores for chosen experts: [bs, seq_len, top_k]
                         ) -> torch.Tensor:
        """
        Debug version of moe_infer. Processes each token's expert assignments individually.
        This is closer to the reference MoE loop structure for easier debugging of numerical issues.
        It will be slower than the optimized scatter-based approach.
        """
        total_tokens, hidden_dim = x_flat.shape
        top_k = self.config["n_experts_per_token"]
        
        # Reshape expert_indices and expert_scores to [total_tokens, top_k] for easier iteration
        expert_indices_flat_per_token = expert_indices_orig_shape.reshape(total_tokens, top_k)
        expert_scores_flat_per_token = expert_scores_orig_shape.reshape(total_tokens, top_k)

        # Initialize tensor to store the final routed output for each token
        routed_output_flat = torch.zeros_like(x_flat)

        # Iterate over each token (conceptually, this replaces the outer two loops of the reference)
        for i in range(total_tokens):
            current_token_input_for_expert = x_flat[i].unsqueeze(0) # Shape [1, d_hidden]
            
            # Accumulator for the current token's output, summed over its top_k chosen experts
            current_token_output_accumulator = torch.zeros(hidden_dim, device=x_flat.device, dtype=x_flat.dtype)
            
            # Iterate over the top_k experts chosen for this specific token
            for k_idx in range(top_k):
                expert_id = expert_indices_flat_per_token[i, k_idx].item() # Get the scalar expert_id
                score = expert_scores_flat_per_token[i, k_idx] # Get the scalar score (as a 0-dim tensor or float)
                
                expert_module_to_use = self.experts[expert_id]
                
                # Pass the single token (reshaped to [1, d_hidden]) through the chosen expert
                # expert_module_to_use output will be [1, d_hidden]
                expert_activation = expert_module_to_use(current_token_input_for_expert).squeeze(0) # Shape [d_hidden]
                
                # Weight the expert's activation by its score and add to the accumulator
                current_token_output_accumulator += expert_activation * score # score will broadcast
                
            # Store the accumulated sum for the current token
            routed_output_flat[i] = current_token_output_accumulator
            
        return routed_output_flat

# --- custom_kernel: Entry point for the submission ---
input_t = Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict]
output_t_custom = torch.Tensor # Expecting a single tensor output

def custom_kernel(data: input_t) -> output_t_custom:
    input_tensor, weights, config = data
    
    target_device = 'cuda' if torch.cuda.is_available() else 'cpu'
    target_dtype = torch.float16 # Assuming float16 is desired for MI300
    
    input_tensor = input_tensor.to(device=target_device, dtype=target_dtype)

    moe_module = MoE_Optimized(config)

    # Load weights into the model
    # Gating network weights (router)
    moe_module.gating_network.W_g.weight = nn.Parameter(
        weights['router.weight'].to(device=target_device, dtype=target_dtype)
    )
    # Routed expert weights
    num_routed_experts_config_val = config["n_routed_experts"]
    for i in range(num_routed_experts_config_val):
        # Assuming weights are stored transposed relative to nn.Linear's expectation, hence .t()
        moe_module.experts[i].W_gate.weight = nn.Parameter(
            weights[f'experts.{i}.0.weight'].t().to(device=target_device, dtype=target_dtype)
        )
        moe_module.experts[i].W_up.weight = nn.Parameter(
            weights[f'experts.{i}.1.weight'].t().to(device=target_device, dtype=target_dtype)
        )
        moe_module.experts[i].W_down.weight = nn.Parameter(
            weights[f'experts.{i}.2.weight'].t().to(device=target_device, dtype=target_dtype)
        )
    # Shared expert weights
    moe_module.shared_expert.W_gate.weight = nn.Parameter(
        weights['shared_experts.0.weight'].t().to(device=target_device, dtype=target_dtype)
    )
    moe_module.shared_expert.W_up.weight = nn.Parameter(
        weights['shared_experts.1.weight'].t().to(device=target_device, dtype=target_dtype)
    )
    moe_module.shared_expert.W_down.weight = nn.Parameter(
        weights['shared_experts.2.weight'].t().to(device=target_device, dtype=target_dtype)
    )

    moe_module = moe_module.to(device=target_device, dtype=target_dtype)
    moe_module.eval() # Set to evaluation mode

    # --- TORCH.COMPILE IS DISABLED FOR THIS TEST RUN to isolate numerical issues ---
    print("Running MoE module in eager mode with moe_infer_v_debug.")
    moe_module_to_run = moe_module 
    # --- END OF DISABLED COMPILE ---

    with torch.inference_mode(): # Disable gradient calculations
        output_tensor = moe_module_to_run(input_tensor)
        
    return output_tensor

python eval.py benchmark /tmp/tmphcvbgouv
Running MoE module in eager mode with moe_infer_v_debug.
Running MoE module in eager mode with moe_infer_v_debug.
Running MoE module in eager mode with moe_infer_v_debug.
Running MoE module in eager mode with moe_infer_v_debug.
Running MoE module in eager mode with moe_infer_v_debug.
Running MoE module in eager mode with moe_infer_v_debug.
Running MoE module in eager mode with moe_infer_v_debug.
Running MoE module in eager mode with moe_infer_v_debug.
Running MoE module in eager mode with moe_infer_v_debug.
Running MoE module in eager mode with moe_infer_v_debug.
Running MoE module in eager mode with moe_infer_v_debug.
Running MoE module in eager mode with moe_infer_v_debug.
