#!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 # Not strictly needed now if not printing configs

import triton
import triton.language as tl

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

# --- Hardcoded Configurations for FFN Kernel ---
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)


@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)

    w_down_start_ptr = w_down_ptr + offs_n_out[None, :] * stride_w_down_n

    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_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)

    final_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) # M here is BLOCK_SIZE_M for the mask, actual M is used for pointers
    tl.store(final_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

    # Modified signature: added out_tensor_view
    def forward(self, x: torch.Tensor, out_tensor_view: torch.Tensor) -> None:
        M_tokens, K_in_dim = x.shape
        K_inter_dim = self.d_expert_actual
        N_out_dim = self.d_hidden # Should match out_tensor_view.shape[1]

        # M_tokens == 0 case is handled by the calling loop (it won't call if no tokens)
        # or by out_tensor_view having 0 rows.
        # If M_tokens > 0 but out_tensor_view.shape[0] == 0, that's an issue, but
        # this shouldn't happen with correct slicing.
        # The calling code in TritonMoE will ensure M_tokens == out_tensor_view.shape[0]

        x = x.contiguous()
        # out_tensor_view is already allocated and is where the output should be written.
        # Ensure out_tensor_view is also contiguous if the kernel expects it,
        # though FFN output strides are usually simple for contiguous.
        # Slices of contiguous tensors are often contiguous.

        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, 
            out_tensor_view, # Pass the view directly
            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),
            out_tensor_view.stride(0), out_tensor_view.stride(1), # Strides of the view
            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,
            num_warps=num_warps,
            num_stages=num_stages
        )
        # No return value needed


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)

    #@torch.compile # Optional: test with and without if issues arise
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_shape = x.shape
        x_flat = x.view(-1, self.d_hidden)
        num_total_tokens_moe_input = x_flat.shape[0] # num_tokens in the original paper's notation for MoE

        expert_indices, expert_scores = self.gating_network(x_flat)

        flat_expert_indices = expert_indices.view(-1)
        flat_expert_scores = expert_scores.view(-1, 1) # [num_total_tokens_moe_input * top_k, 1]

        token_ids_for_flat_assignments = torch.arange(num_total_tokens_moe_input, 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]

        all_expert_outputs_perm_allocated = torch.empty_like(perm_x_flat)

        expert_boundaries = torch.searchsorted(perm_expert_indices, torch.arange(self.num_experts + 1, device=x.device))
        
        current_processed_token_idx_in_perm = 0
        expert_boundaries_cpu_list = expert_boundaries.tolist()
        for expert_id in range(self.num_experts):
            start_idx_in_perm = expert_boundaries_cpu_list[expert_id]
            end_idx_in_perm = expert_boundaries_cpu_list[expert_id + 1]
            
            num_tokens_for_this_expert = end_idx_in_perm - start_idx_in_perm

            if num_tokens_for_this_expert > 0:
                current_expert_inputs = perm_x_flat[start_idx_in_perm:end_idx_in_perm]
                output_view = all_expert_outputs_perm_allocated.narrow(0, start_idx_in_perm, num_tokens_for_this_expert)
                expert_module = self.experts[expert_id]
                expert_module(current_expert_inputs, output_view)
            
            current_processed_token_idx_in_perm += num_tokens_for_this_expert
            
        weighted_expert_outputs_perm = all_expert_outputs_perm_allocated * perm_expert_scores
        
        shared_expert_output_flat = torch.empty_like(x_flat)
        self.shared_expert(x_flat, shared_expert_output_flat) # Shared expert writes to this tensor

        # Initialize final_output_flat with the shared expert's output
        final_output_flat = shared_expert_output_flat
        
        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

# Test snippet (remove or comment out for submission)
from reference import generate_input, ref_kernel 
from utils import make_match_reference 

if __name__ == '__main__':
    args_seq8192 = {
        "dhidden": 7168, "dexpert": 2048, "nroutedexperts": 8, "nsharedexperts": 1, 
        "nexpertspertoken": 2, "bs": 1, "seqlen": 8192, "seed": 123
    } # nroutedexperts was 32 in your log, but test gen might be different
    
    print("Testing with seq_len=8192 config...")
    inp_data_8192 = generate_input(**args_seq8192)
    
    print("Warmup run...")
    _ = custom_kernel(inp_data_8192)
    torch.cuda.synchronize()

    print("Profiled run...")
    with torch.profiler.profile(
        activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
        record_shapes=True, profile_memory=True, with_stack=True
    ) as prof:
        out_custom_8192 = custom_kernel(inp_data_8192)
    torch.cuda.synchronize()
    
    print(prof.key_averages(group_by_input_shape=True).table(sort_by="self_cuda_time_total", row_limit=20))
    prof.export_chrome_trace("moe_trace_no_cat.json")

    out_custom_8192 = custom_kernel(inp_data_8192)
    print(f"Output shape (seq_len=8192): {out_custom_8192.shape}")


python eval.py test /tmp/tmpa1dcqqf4
/usr/local/lib/python3.10/dist-packages/torch/_inductor/lowering.py:7078: UserWarning: 
Online softmax is disabled on the fly since Inductor decides to
split the reduction. Cut an issue to PyTorch if this is an
important use case and you want to speed it up with online
softmax.

  warnings.warn(
