#!POPCORN leaderboard amd-mixture-of-experts
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple

import triton
import triton.language as tl

from task import input_t, output_t 

HARDCODED_FFN_CONFIGS = {
    (7168, 2048, 7168, 2048): (128, 128, 128, 128, 4, 8, 1),
    (7168, 2048, 7168, 8192): (256, 128, 64, 64, 4, 8, 2),
}
DEFAULT_FFN_CONFIG = (128, 64, 64, 64, 4, 4, 2) # Example default, adjust as needed


@triton.jit
def silu(x): return (x * tl.sigmoid(x.to(tl.float32))).to(tl.float16)

@triton.jit
def fused_ffn_kernel_v2(
    # Pointers to Matrices
    x_ptr, w_gate_ptr, w_up_ptr, w_down_ptr, output_ptr,
    # Matrix dimensions
    M, N_out, K_in, K_inter,
    # Strides
    stride_x_m, stride_x_k,
    stride_w_gate_k, stride_w_gate_inter,
    stride_w_up_k, stride_w_up_inter,
    stride_w_down_inter, stride_w_down_n,
    stride_out_m, stride_out_n,
    # Meta-parameters
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N_OUT: tl.constexpr,
    BLOCK_SIZE_K_IN: tl.constexpr, BLOCK_SIZE_K_INTER: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N_out, BLOCK_SIZE_N_OUT)
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + (pid % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n_out = pid_n * BLOCK_SIZE_N_OUT + tl.arange(0, BLOCK_SIZE_N_OUT)
    x_start_ptr = x_ptr + offs_m[:, None] * stride_x_m
    weight_dtype = w_gate_ptr.dtype.element_ty
    accum_dtype = tl.float32
    acc_down = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N_OUT), dtype=accum_dtype)

    for k_inter_start in range(0, tl.cdiv(K_inter, BLOCK_SIZE_K_INTER)):
        offs_k_inter = k_inter_start * BLOCK_SIZE_K_INTER + tl.arange(0, BLOCK_SIZE_K_INTER)
        acc_gate_fp32 = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_K_INTER), dtype=accum_dtype)
        acc_up_fp32 = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_K_INTER), dtype=accum_dtype)
        w_gate_k_inter_ptr = w_gate_ptr + offs_k_inter[None, :] * stride_w_gate_inter
        w_up_k_inter_ptr = w_up_ptr + offs_k_inter[None, :] * stride_w_up_inter

        for k_in_start in range(0, tl.cdiv(K_in, BLOCK_SIZE_K_IN)):
            offs_k_in = k_in_start * BLOCK_SIZE_K_IN + tl.arange(0, BLOCK_SIZE_K_IN)
            x_ptrs = x_start_ptr + offs_k_in[None, :] * stride_x_k
            x_mask = (offs_m[:, None] < M) & (offs_k_in[None, :] < K_in)
            x_block = tl.load(x_ptrs, mask=x_mask, other=0.0)
            w_gate_ptrs = w_gate_k_inter_ptr + offs_k_in[:, None] * stride_w_gate_k
            w_gate_mask = (offs_k_in[:, None] < K_in) & (offs_k_inter[None, :] < K_inter)
            w_gate_block = tl.load(w_gate_ptrs, mask=w_gate_mask, other=0.0)
            acc_gate_fp32 += tl.dot(x_block, w_gate_block, out_dtype=accum_dtype)
            w_up_ptrs = w_up_k_inter_ptr + offs_k_in[:, None] * stride_w_up_k
            w_up_mask = (offs_k_in[:, None] < K_in) & (offs_k_inter[None, :] < K_inter)
            w_up_block = tl.load(w_up_ptrs, mask=w_up_mask, other=0.0)
            acc_up_fp32 += tl.dot(x_block, w_up_block, out_dtype=accum_dtype)

        acc_gate_fp16 = acc_gate_fp32.to(weight_dtype)
        acc_up_fp16 = acc_up_fp32.to(weight_dtype)
        gate_activated_fp16 = silu(acc_gate_fp16)
        intermediate_block_fp16 = gate_activated_fp16 * acc_up_fp16
        w_down_start_ptr = w_down_ptr + offs_n_out[None, :] * stride_w_down_n
        w_down_ptrs = w_down_start_ptr + offs_k_inter[:, None] * stride_w_down_inter
        w_down_mask = (offs_k_inter[:, None] < K_inter) & (offs_n_out[None, :] < N_out)
        w_down_block = tl.load(w_down_ptrs, mask=w_down_mask, other=0.0)
        acc_down += tl.dot(intermediate_block_fp16, w_down_block, out_dtype=accum_dtype)

    output_ptrs = output_ptr + offs_m[:, None] * stride_out_m + offs_n_out[None, :] * stride_out_n
    output_mask = (offs_m[:, None] < M) & (offs_n_out[None, :] < N_out)
    tl.store(output_ptrs, acc_down.to(output_ptr.dtype.element_ty), mask=output_mask)

class TritonExpert(nn.Module):
    def __init__(self, config: Dict, W_gate: torch.Tensor, W_up: torch.Tensor, W_down: torch.Tensor, d_expert_actual: int, seq_len_for_config_key: int):
        super().__init__()
        self.d_hidden: int = config["d_hidden"]
        self.d_expert_actual = d_expert_actual
        self.W_gate = W_gate
        self.W_up = W_up
        self.W_down = W_down
        self.seq_len_for_config_key = seq_len_for_config_key

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        M_tokens, K_in_dim = x.shape 
        K_inter_dim = self.d_expert_actual
        N_out_dim = self.d_hidden

        x = x.contiguous()
        output = torch.empty((M_tokens, N_out_dim), device=x.device, dtype=x.dtype)

        ffn_config_key = (K_in_dim, K_inter_dim, N_out_dim, self.seq_len_for_config_key)
        selected_params = HARDCODED_FFN_CONFIGS.get(ffn_config_key, DEFAULT_FFN_CONFIG)
        
        BLOCK_SIZE_M, BLOCK_SIZE_N_OUT, BLOCK_SIZE_K_IN, BLOCK_SIZE_K_INTER, GROUP_SIZE_M, num_warps, num_stages = selected_params

        grid = (
            triton.cdiv(M_tokens, BLOCK_SIZE_M) * triton.cdiv(N_out_dim, BLOCK_SIZE_N_OUT),
        )

        fused_ffn_kernel_v2[grid](
            x, self.W_gate, self.W_up, self.W_down, output,
            M_tokens, N_out_dim, K_in_dim, K_inter_dim,
            x.stride(0), x.stride(1),
            self.W_gate.stride(0), self.W_gate.stride(1),
            self.W_up.stride(0), self.W_up.stride(1),
            self.W_down.stride(0), self.W_down.stride(1),
            output.stride(0), output.stride(1),
            # Pass hardcoded constexprs
            BLOCK_SIZE_M=BLOCK_SIZE_M, BLOCK_SIZE_N_OUT=BLOCK_SIZE_N_OUT,
            BLOCK_SIZE_K_IN=BLOCK_SIZE_K_IN, BLOCK_SIZE_K_INTER=BLOCK_SIZE_K_INTER,
            GROUP_SIZE_M=GROUP_SIZE_M,
            # Pass num_warps and num_stages
            num_warps=num_warps,
            num_stages=num_stages
        )
        return output

class MoEGate(nn.Module):
    def __init__(self, config: Dict, W_g_weight: torch.Tensor):
        super().__init__()
        self.top_k: int = config["n_experts_per_token"]
        self.n_experts: int = config["n_routed_experts"]
        self.register_buffer('W_g', W_g_weight)

    @torch.compile
    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        logits = F.linear(x, self.W_g)
        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

class TritonMoE(nn.Module):
    def __init__(self, config: Dict, weights: Dict[str, torch.Tensor], seq_len_for_config_key: int):
        super().__init__()
        self.num_experts = config["n_routed_experts"]
        self.top_k = config["n_experts_per_token"]
        self.d_hidden = config["d_hidden"]
        self.d_expert = config["d_expert"]
        self.gating_network = MoEGate(config, weights['router.weight'])
        self.experts = nn.ModuleList()
        for i in range(self.num_experts):
            w_gate = weights[f'experts.{i}.0.weight']
            w_up = weights[f'experts.{i}.1.weight']
            w_down = weights[f'experts.{i}.2.weight']
            self.experts.append(TritonExpert(config, w_gate, w_up, w_down, self.d_expert, seq_len_for_config_key))

        shared_expert_dim = config["d_expert"] * config["n_shared_experts"]
        w_gate_shared = weights['shared_experts.0.weight']
        w_up_shared = weights['shared_experts.1.weight']
        w_down_shared = weights['shared_experts.2.weight']
        self.shared_expert = TritonExpert(config, w_gate_shared, w_up_shared, w_down_shared, shared_expert_dim, seq_len_for_config_key)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_shape = x.shape
        x_flat = x.view(-1, self.d_hidden)
        num_tokens = x_flat.shape[0]
        expert_indices, expert_scores = self.gating_network(x_flat)
        flat_expert_indices = expert_indices.view(-1)
        flat_expert_scores = expert_scores.view(-1, 1)
        token_ids_for_flat_assignments = torch.arange(num_tokens, device=x.device).repeat_interleave(self.top_k)
        sort_perm = flat_expert_indices.argsort()
        perm_expert_indices = flat_expert_indices[sort_perm]
        perm_token_ids = token_ids_for_flat_assignments[sort_perm]
        perm_expert_scores = flat_expert_scores[sort_perm]
        perm_x_flat = x_flat[perm_token_ids]
        expert_boundaries = torch.searchsorted(perm_expert_indices, torch.arange(self.num_experts + 1, device=x.device))

        all_expert_outputs_list = []
        for expert_id in range(self.num_experts):
            start_idx = expert_boundaries[expert_id]
            end_idx = expert_boundaries[expert_id + 1]
            current_expert_inputs = perm_x_flat[start_idx:end_idx]
            expert_module = self.experts[expert_id]
            current_expert_outputs = expert_module(current_expert_inputs)
            all_expert_outputs_list.append(current_expert_outputs)
        
        # Ensure tensor is created on the correct device and with correct dtype, even if list is empty
        if not all_expert_outputs_list:
            all_expert_outputs_perm = torch.empty((0, self.d_hidden), device=x.device, dtype=x.dtype)
        else:
            all_expert_outputs_perm = torch.cat(all_expert_outputs_list, dim=0)
            
        weighted_expert_outputs_perm = all_expert_outputs_perm * perm_expert_scores
        final_output_flat = self.shared_expert(x_flat)
        
        if weighted_expert_outputs_perm.numel() > 0: # Only scatter if there's something to scatter
            scatter_indices = perm_token_ids.unsqueeze(-1).expand_as(weighted_expert_outputs_perm)
            final_output_flat.scatter_reduce_(
                dim=0,
                index=scatter_indices,
                src=weighted_expert_outputs_perm,
                reduce="sum",
                include_self=True
            )
        output = final_output_flat.view(orig_shape)
        return output

def custom_kernel(data: input_t) -> output_t:
    input_tensor, weights, config = data

    seq_len_for_config_key = input_tensor.shape[1]

    triton_moe_model = TritonMoE(config, weights, seq_len_for_config_key)
    input_tensor = input_tensor.to('cuda')
    triton_moe_model.to('cuda')

    for k, v in weights.items():
        weights[k] = v.contiguous()
    input_dtype = next(iter(weights.values())).dtype
    input_tensor = input_tensor.to(input_dtype)

    output = triton_moe_model(input_tensor)
    return output

python eval.py leaderboard /tmp/tmp3fj1jyig