#!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 sys

import triton
import triton.language as tl

from task import input_t, output_t # Assuming task.py defines these types

# Global dictionary to store printed FFN config summaries to avoid re-printing
_printed_ffn_config_summaries = {}

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

@triton.autotune(
    configs=[
        # num_stages = 1
        triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 64,  'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=4, num_stages=1),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N_OUT': 64,  'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=4, num_stages=1),
        triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 128, 'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=4, num_stages=1),
        triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 64,  'BLOCK_SIZE_K_IN': 128, 'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=4, num_stages=1),
        triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 64,  'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 128, 'GROUP_SIZE_M': 4}, num_warps=4, num_stages=1),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N_OUT': 128, 'BLOCK_SIZE_K_IN': 128, 'BLOCK_SIZE_K_INTER': 128, 'GROUP_SIZE_M': 4}, num_warps=8, num_stages=1),

        # num_stages = 2
        triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 64,  'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=4, num_stages=2),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N_OUT': 64,  'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=4, num_stages=2),
        triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 128, 'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=4, num_stages=2),
        triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 64,  'BLOCK_SIZE_K_IN': 128, 'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=4, num_stages=2),
        triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 64,  'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 128, 'GROUP_SIZE_M': 4}, num_warps=4, num_stages=2),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N_OUT': 128, 'BLOCK_SIZE_K_IN': 128, 'BLOCK_SIZE_K_INTER': 128, 'GROUP_SIZE_M': 4}, num_warps=8, num_stages=2),
        # Larger block sizes with more warps/stages
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N_OUT': 128, 'BLOCK_SIZE_K_IN': 64, 'BLOCK_SIZE_K_INTER': 64, 'GROUP_SIZE_M': 4}, num_warps=8, num_stages=2),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N_OUT': 64, 'BLOCK_SIZE_K_IN': 128, 'BLOCK_SIZE_K_INTER': 128, 'GROUP_SIZE_M': 4}, num_warps=8, num_stages=2),
    ],
    key=['M', 'N_out', 'K_in', 'K_inter'], # M (tokens per expert call), N_out (d_hidden), K_in (d_hidden), K_inter (d_expert)
)
@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 (for autotune key and bounds checks)
    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 (from autotuner's triton.Config)
    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):
        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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        M, K_in_val = x.shape
        K_inter_val = self.d_expert_actual
        N_out_val = self.d_hidden

        if M == 0: # If no tokens are routed to this expert
            return torch.empty((0, N_out_val), device=x.device, dtype=x.dtype)

        x = x.contiguous()
        output = torch.empty((M, N_out_val), device=x.device, dtype=x.dtype)

        grid = lambda META: (
            triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N_out_val, META['BLOCK_SIZE_N_OUT']),
        )
        fused_ffn_kernel_v2[grid](
            x, self.W_gate, self.W_up, self.W_down, output,
            M, N_out_val, K_in_val, K_inter_val,
            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)
        )
        return output

@triton.jit
def _triton_gate_kernel( # This kernel is defined but not used by default in MoEGate
    X_ptr, Wg_ptr, Indices_ptr, Scores_ptr,
    M, N: tl.constexpr, K,
    stride_xm, stride_xk, stride_wn, stride_wk,
    stride_indices_m, stride_indices_k, stride_scores_m, stride_scores_k,
    top_k: tl.constexpr, BLOCK_K: tl.constexpr,
):
    pid_m = tl.program_id(axis=0)
    x_row_ptr = X_ptr + pid_m * stride_xm
    offs_k_loop = tl.arange(0, BLOCK_K)
    offs_n = tl.arange(0, N)
    wg_ptrs_base = Wg_ptr + offs_n[:, None] * stride_wn + offs_k_loop[None, :] * stride_wk
    accumulator = tl.zeros((N,), dtype=tl.float32)
    for k_start in range(0, tl.cdiv(K, BLOCK_K)):
        current_offs_k = k_start * BLOCK_K + offs_k_loop
        x_mask = (pid_m < M) & (current_offs_k < K)
        x_chunk = tl.load(x_row_ptr + current_offs_k * stride_xk, mask=x_mask, other=0.0)[:, None]
        wg_mask = (offs_n[:, None] < N) & (current_offs_k[None, :] < K)
        wg_chunk = tl.load(wg_ptrs_base + k_start * BLOCK_K * stride_wk, mask=wg_mask, other=0.0)
        dot_result = tl.dot(wg_chunk, x_chunk)
        dot_result_reshaped = tl.reshape(dot_result, (N,))
        accumulator += dot_result_reshaped
    if pid_m < M:
        current_scores_fp16 = tl.softmax(accumulator.to(tl.float16)).to(tl.float16)
        scores_row_ptr = Scores_ptr + pid_m * stride_scores_m
        indices_row_ptr = Indices_ptr + pid_m * stride_indices_m
        loop_scores = current_scores_fp16
        neg_inf_fp16 = (-float('inf')) # Ensure FP16 for tl.where
        for k_iter in tl.static_range(top_k):
            max_val_fp16, max_idx = tl.max(loop_scores, axis=0, return_indices=True)
            tl.store(indices_row_ptr + k_iter * stride_indices_k, max_idx.to(tl.int64))
            tl.store(scores_row_ptr + k_iter * stride_scores_k, max_val_fp16)
            loop_scores = tl.where(tl.arange(0, N) == max_idx, neg_inf_fp16.to(loop_scores.dtype), loop_scores)


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)

    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]):
        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))
        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)

    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) # Handles M=0 internally
            all_expert_outputs_list.append(current_expert_outputs)
        
        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) # Shared expert always runs on all tokens
        
        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, # PyTorch handles empty src/index correctly
            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
    triton_moe_model = TritonMoE(config, weights)
    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)

    global _printed_ffn_config_summaries
    print_key_tuple = (config['seq_len'],config["d_hidden"], config["d_expert"], config['n_routed_experts'])

    if print_key_tuple not in _printed_ffn_config_summaries and hasattr(fused_ffn_kernel_v2, 'best_config'):
        best_config = fused_ffn_kernel_v2.best_config
        if best_config: # Check if not None (tuning might not be instant on first call)
            try:
                cfg = best_config.kwargs
                m_block = cfg.get('BLOCK_SIZE_M', 'NA')
                n_out_block = cfg.get('BLOCK_SIZE_N_OUT', 'NA')
                k_in_block = cfg.get('BLOCK_SIZE_K_IN', 'NA')
                k_inter_block = cfg.get('BLOCK_SIZE_K_INTER', 'NA')
                group_m = cfg.get('GROUP_SIZE_M', 'NA')
                warps = best_config.num_warps if hasattr(best_config, 'num_warps') else 'NA'
                stages = best_config.num_stages if hasattr(best_config, 'num_stages') else 'NA'

                concise_summary = f"{m_block}-{n_out_block}-{k_in_block}-{k_inter_block}-{group_m}-{warps}-{stages}"
                # The printed message refers to the general FFN kernel's autotuned config.
                # K_in/N_out are d_hidden. K_inter is d_expert for routed, or d_expert_shared.
                # The specific M for which this `best_config` was found is part of autotuner's internal key.
                print(f"{config['d_hidden']}, K_inter_routed={config['d_expert']}, seq_len={config['seq_len']}, routed_experts={config['n_routed_experts']}): {concise_summary}", file=sys.stderr)
                _printed_ffn_config_summaries[print_key_tuple] = True
            except Exception as e:
                print(f"Error summarizing FFN config for {print_key_tuple}: {e}", file=sys.stderr)
                _printed_ffn_config_summaries[print_key_tuple] = True 
    return output

'''
# Test snippet (remove or comment out for submission)
from reference import generate_input, ref_kernel # Assuming these are available
from utils import make_match_reference # Assuming this is available

if __name__ == '__main__':
    # Small config for quick local test
    args = {"dhidden": 2048, "dexpert": 512, "nroutedexperts": 2, "nsharedexperts": 1, "nexpertspertoken": 1, "bs": 1, "seqlen": 128, "seed": 1212} 
    # Larger config closer to benchmark
    # args = {"dhidden": 7168, "dexpert": 2048, "nroutedexperts": 8, "nsharedexperts": 1, "nexpertspertoken": 4, "bs": 1, "seqlen": 1024, "seed": 1212}
    
    print("Generating input...")
    # Ensure generate_input and other utils are in the path or defined
    # For local testing, you might need to mock task.input_t, output_t if not available
    # from collections import namedtuple
    # input_t = namedtuple('input_t', ['x', 'weights', 'config']) # Mock if needed

    inp_data = generate_input(**args) # Returns (tensor, weights_dict, config_dict)
    
    print("Running custom_kernel (warmup)...")
    custom_kernel(inp_data) # Warmup, autotuning runs here
    
    print("Running custom_kernel (timed)...")
    import time
    start_time = time.time()
    out_custom = custom_kernel(inp_data) # Second run uses cached best_config
    end_time = time.time()
    print(f"Custom kernel execution time: {end_time - start_time:.4f} seconds")

    # print("Running reference kernel...")
    # ref_out_val, _ = ref_kernel(inp_data)
    # print("Comparing outputs...")
    # print(f"Output shape: {out_custom.shape}, dtype: {out_custom.dtype}, device: {out_custom.device}")
    # print(f"Reference output shape: {ref_out_val.shape}, dtype: {ref_out_val.dtype}, device: {ref_out_val.device}")
    # if torch.allclose(out_custom, ref_out_val, atol=1e-1, rtol=1e-1): 
    #     print("Outputs match reference (within tolerance).")
    # else:
    #     print("Outputs DO NOT match reference.")
    #     diff = torch.abs(out_custom - ref_out_val)
    #     print(f"Max difference: {torch.max(diff)}")
    #     print(f"Mean difference: {torch.mean(diff)}")
'''

python eval.py test /tmp/tmpd35sy0jg
7168, K_inter_routed=2048, seq_len=512, routed_experts=4): 128-128-128-128-4-8-1
7168, K_inter_routed=2048, seq_len=512, routed_experts=8): 256-128-64-64-4-8-2
7168, K_inter_routed=2048, seq_len=8192, routed_experts=8): 256-128-64-64-4-8-2
