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

    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)

    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']
        # Shared expert also uses the TritonExpert class and writes to a view of its output tensor
        self.shared_expert = TritonExpert(config, w_gate_shared, w_up_shared, w_down_shared, shared_expert_dim, seq_len_for_config_key)

    # If you compile TritonMoE.forward, ensure torch.compile can handle
    # the modified TritonExpert calls (passing tensor views).
    #@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] # Shape: [num_assignments, d_hidden]

        # Pre-allocate tensor for all routed expert outputs
        # This tensor will have the same ordering as perm_x_flat
        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]
                
                # Get a view into the pre-allocated tensor for this expert's output
                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) # TritonExpert now writes into output_view
            
            current_processed_token_idx_in_perm += num_tokens_for_this_expert
            
        # all_expert_outputs_perm_allocated now contains all routed expert outputs, correctly ordered
        weighted_expert_outputs_perm = all_expert_outputs_perm_allocated * perm_expert_scores
        
        # --- Shared Expert ---
        # The shared expert processes all original tokens (x_flat)
        # We need a tensor to hold its output.
        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-reduce the weighted routed expert outputs onto final_output_flat
        # perm_token_ids maps each item in weighted_expert_outputs_perm back to its original token slot
        if weighted_expert_outputs_perm.numel() > 0:
            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 # Crucial: adds src to existing values in final_output_flat
            )
        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 benchmark /tmp/tmprvh0zg95