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"]
        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]
        all_scores = F.softmax(logits, dim=-1, dtype=torch.float32).to(x.dtype)
        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 the debug-oriented moe_infer method as the primary path.
    """
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config
        self.num_routed_experts: int = config["n_routed_experts"]
        
        self.experts = nn.ModuleList([
            Expert(config)
            for _ in range(self.num_routed_experts)
        ])
        
        self.gating_network = MoEGate(config)
        
        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:
        shared_output = self.shared_expert(x) 
        expert_indices, expert_scores, _ = self.gating_network(x)

        original_shape = x.shape
        x_flat = x.reshape(-1, original_shape[-1]) 

        # Using the debug version of moe_infer, which was numerically correct in eager mode
        routed_output_flat = self.moe_infer_v_debug(x_flat,
                                                    expert_indices, 
                                                    expert_scores)
        
        routed_output = routed_output_flat.view(*original_shape)
        final_output = routed_output + shared_output
        return final_output

    def moe_infer_v_debug(self,
                          x_flat: torch.Tensor, 
                          expert_indices_orig_shape: torch.Tensor, 
                          expert_scores_orig_shape: torch.Tensor
                         ) -> torch.Tensor:
        """
        Debug version of moe_infer. Processes each token's expert assignments individually.
        This is closer to the reference MoE loop structure.
        """
        total_tokens, hidden_dim = x_flat.shape
        top_k = self.config["n_experts_per_token"]
        
        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)

        routed_output_flat = torch.zeros_like(x_flat)

        for i in range(total_tokens):
            current_token_input_for_expert = x_flat[i].unsqueeze(0).contiguous() # Ensure contiguous
            current_token_output_accumulator = torch.zeros(hidden_dim, device=x_flat.device, dtype=x_flat.dtype)
            
            for k_idx in range(top_k):
                expert_id = expert_indices_flat_per_token[i, k_idx].item()
                score = expert_scores_flat_per_token[i, k_idx]
                
                expert_module_to_use = self.experts[expert_id]
                expert_activation = expert_module_to_use(current_token_input_for_expert).squeeze(0)
                current_token_output_accumulator += expert_activation * score
                
            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

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
    
    input_tensor = input_tensor.to(device=target_device, dtype=target_dtype)

    moe_module = MoE_Optimized(config) # This now uses moe_infer_v_debug internally

    # Load weights
    moe_module.gating_network.W_g.weight = nn.Parameter(
        weights['router.weight'].to(device=target_device, dtype=target_dtype)
    )
    num_routed_experts_config_val = config["n_routed_experts"]
    for i in range(num_routed_experts_config_val):
        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)
        )
    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()

    compiled_moe_module = moe_module 
    if hasattr(torch, 'compile'):
        print("Attempting to compile MoE module (using moe_infer_v_debug) with mode='default'...")
        try:
            compiled_moe_module = torch.compile(moe_module, mode="default") 
            print("Compilation with 'default' mode successful.")
        except Exception as e_default:
            print(f"Warning: torch.compile with mode='default' failed for moe_infer_v_debug: {e_default}. Running uncompiled.")
            compiled_moe_module = moe_module # Fallback to uncompiled
    
    moe_module_to_run = compiled_moe_module

    with torch.inference_mode():
        output_tensor = moe_module_to_run(input_tensor)
        
    return output_tensor

python eval.py benchmark /tmp/tmp_8x9z3ks
W0519 12:51:27.557000 466 torch/fx/experimental/symbolic_shapes.py:6903] [0/0] failed during evaluate_expr(u0, hint=None, size_oblivious=False, forcing_spec=False
W0519 12:51:29.997000 466 torch/fx/experimental/symbolic_shapes.py:6903] [1/0] failed during evaluate_expr(u0, hint=None, size_oblivious=False, forcing_spec=False
W0519 12:52:43.012000 466 torch/fx/experimental/symbolic_shapes.py:6903] [0/1] failed during evaluate_expr(u0, hint=None, size_oblivious=False, forcing_spec=False

Attempting to compile MoE module (using moe_infer_v_debug) with mode='default'...
Compilation with 'default' mode successful.
Attempting to compile MoE module (using moe_infer_v_debug) with mode='default'...
Compilation with 'default' mode successful.
Attempting to compile MoE module (using moe_infer_v_debug) with mode='default'...
Compilation with 'default' mode successful.
Attempting to compile MoE module (using moe_infer_v_debug) with mode='default'...
Compilation with 'default' mode successful.
Attempting to compile MoE module (using moe_infer_v_debug) with mode='default'...
Compilation with 'default' mode successful.
Attempting to compile MoE module (using moe_infer_v_debug) with mode='default'...
Compilation with 'default' mode successful.
Attempting to compile MoE module (using moe_infer_v_debug) with mode='default'...
Compilation with 'default' mode successful.
Attempting to compile MoE module (using moe_infer_v_debug) with mode='default'...
Compilation with 'default' mode successful.
Attempting to compile MoE module (using moe_infer_v_debug) with mode='default'...
Compilation with 'default' mode successful.
Attempting to compile MoE module (using moe_infer_v_debug) with mode='default'...
Compilation with 'default' mode successful.
Attempting to compile MoE module (using moe_infer_v_debug) with mode='default'...
Compilation with 'default' mode successful.
Attempting to compile MoE module (using moe_infer_v_debug) with mode='default'...
Compilation with 'default' mode successful.
