import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple, Optional, List
from task import input_t, output_t
import triton
import triton.language as tl
from copy import deepcopy
import time
@triton.autotune(
    configs=[
        triton.Config(
            {
                "BLOCK_SIZE": 512,
            },
            num_warps=4,
            num_stages=1,
        ),    
    ],
    key=[],
)
@triton.jit
def calculate_expert2token_count_and_within_expert_index(
    topk_indices,
    expert2token_count,
    within_expert_index,
    token_count,
    n_experts_per_token,
    n_routed_experts,
    BLOCK_SIZE: tl.constexpr,
):
    expert_idx = tl.program_id(0)
    
    total_count = 0
    for idx in range(0, token_count * n_experts_per_token, BLOCK_SIZE):
        offsets = tl.arange(0, BLOCK_SIZE) + idx
        mask = offsets < token_count * n_experts_per_token
        topk_indices_block = tl.load(topk_indices + offsets, mask=mask, other=-1)
        is_current_expert = tl.where(topk_indices_block == expert_idx, 1, 0) 
        cur_cumsum = tl.cumsum(is_current_expert, axis=0) + total_count - 1 
        total_count += tl.sum(is_current_expert)
        
        within_expert_index_block = tl.where(is_current_expert, cur_cumsum, -1)
        tl.store(
            within_expert_index + idx + tl.arange(0, BLOCK_SIZE), within_expert_index_block, 
            mask= (idx + tl.arange(0, BLOCK_SIZE) < token_count * n_experts_per_token) and (is_current_expert.to(tl.int1))
        )
        
    tl.store(expert2token_count + expert_idx, total_count)

@triton.autotune(
    configs=[
        triton.Config(
            {
                "BLOCK_SIZE": 512,
            },
            num_warps=4,
            num_stages=1,
        ),    
    ],
    key=[],
)
@triton.jit
def calculate_token2index(
    topk_indices,
    expert2token_count,
    within_expert_index,
    token2index,
    expert_token_count_cumsum,
    token_count,
    n_experts_per_token,
    n_routed_experts,
    BLOCK_SIZE: tl.constexpr,
):
    expert2token_count_block = tl.load(
        expert2token_count + tl.arange(0, BLOCK_SIZE),
        mask=tl.arange(0, BLOCK_SIZE) < n_routed_experts,
        other=0,
    ) 
    expert_token_count_cumsum_block = tl.cumsum(
        expert2token_count_block,
        axis=0,
    ) 
    if tl.program_id(0) == 0:
        tl.store(expert_token_count_cumsum + tl.arange(0, BLOCK_SIZE) + 1, expert_token_count_cumsum_block, mask=tl.arange(0, BLOCK_SIZE) < n_routed_experts)
    
    topk_indices_block = tl.load(
        topk_indices + tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE),
        mask=tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) < token_count * n_experts_per_token,
        other=0,
    ) 
    within_expert_index_block = tl.load(
        within_expert_index + tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE),
        mask=tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) < token_count * n_experts_per_token,
        other=0,
    ) 
    is_zero = topk_indices_block == 0
    expert_offset = tl.gather(
        expert_token_count_cumsum_block ,
        tl.where(is_zero, 0, topk_indices_block - 1),
        axis=0,
    ) 
    expert_offset = tl.where(is_zero, 0, expert_offset) 
    token2index_block = expert_offset + within_expert_index_block
    tl.store(
        token2index + tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE),
        token2index_block,
        mask=tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) < token_count * n_experts_per_token,
    )
@triton.autotune(
    configs=[
        triton.Config(
            {
                "BLOCK_SIZE": 1024,
            },
            num_warps=4,
            num_stages=1,
        ),
    ],
    key=[],
)
@triton.jit
def scatter(
    input_tensor,
    token2index,
    repeated_tokens,
    batch_size,
    seq_len,
    n_experts_per_token,
    d_hidden, 
    BLOCK_SIZE: tl.constexpr,
):
    expert_idx = tl.program_id(0)
    seq_idx = tl.program_id(1)
    batch_idx = tl.program_id(2)
    
    token_idx = batch_idx * seq_len + seq_idx
    dst_idx = tl.load(token2index + token_idx * n_experts_per_token + expert_idx)
    for i in range(0, d_hidden, BLOCK_SIZE):
        offset = tl.arange(0, BLOCK_SIZE) + i
        mask = offset < d_hidden
        input_tensor_block = tl.load(
            input_tensor + token_idx * d_hidden + offset,
            mask=mask,
        )
        tl.store(
            repeated_tokens + dst_idx * d_hidden + offset,
            input_tensor_block,
            mask=mask,
        )
@triton.autotune(
    configs=[
        triton.Config(
            {
                "BLOCK_M": 128,
                "BLOCK_N": 128,
                "BLOCK_K": 128,
            },
            num_warps=4,
            num_stages=1,
        ),
    ],
    key=[],
)
@triton.jit
def group_gemm(
    a, b_base_ptrs, c,
    m_cumsum,
    max_m,
    n,
    k,
    group_size,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    group_id = tl.program_id(2)
    n_start = tl.program_id(0) * BLOCK_N
    m_id = tl.program_id(1)
    
    group_m_start = tl.load(m_cumsum + group_id)
    group_m_end = tl.load(m_cumsum + group_id + 1)
    group_b_base_ptr = tl.load(b_base_ptrs + group_id)
    if group_m_start + m_id * BLOCK_M >= group_m_end:
        return
    a_block_ptr = tl.make_block_ptr(
        a + group_m_start * k,
        shape=(group_m_end - group_m_start, k),
        strides=(k, 1),
        offsets=(m_id * BLOCK_M, 0),
        block_shape=(BLOCK_M, BLOCK_K),
        order=(1, 0),
    )
    b_block_ptr = tl.make_block_ptr(
        base=group_b_base_ptr.to(tl.pointer_type(tl.float16)),
        shape=(k, n),
        strides=(n, 1),
        offsets=(0, n_start),
        block_shape=(BLOCK_K, BLOCK_N),
        order=(1, 0),
    )
    c_block_ptr = tl.make_block_ptr(
        base=c + group_m_start * n,
        shape=(group_m_end - group_m_start, n),
        strides=(n, 1),
        offsets=(m_id * BLOCK_M, n_start),
        block_shape=(BLOCK_M, BLOCK_N),
        order=(1, 0),
    )
    accum = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for i in range(0, k, BLOCK_K):
        a_block = tl.load(a_block_ptr, boundary_check=(0, 1))
        b_block = tl.load(b_block_ptr, boundary_check=(0, 1))
        accum += tl.dot(a_block, b_block) 
        a_block_ptr = tl.advance(a_block_ptr, offsets=(0, BLOCK_K))
        b_block_ptr = tl.advance(b_block_ptr, offsets=(BLOCK_K, 0))
    tl.store(c_block_ptr, accum.to(tl.float16), boundary_check=(0, 1), )
        
@triton.autotune(
    configs=[
        triton.Config(
            {
                "BLOCK_SIZE": 1024,
            },
            num_warps=4,
            num_stages=1,
        ),
    ],
    key=[],
)           
@triton.jit
def weighted_sum_gather(
    ffn1_out,
    token2index,
    topk_scores,
    shared_expert_ffn1,
    moe_out,
    d_hidden,
    batch_size,
    seq_len,
    n_experts_per_token,
    BLOCK_SIZE: tl.constexpr,
):
    
    segment_idx = tl.program_id(0)
    token_idx = tl.program_id(1)
    
    offsets = tl.arange(0, BLOCK_SIZE) + segment_idx * BLOCK_SIZE
    mask = offsets < d_hidden
    
    data = tl.load(
        shared_expert_ffn1 + token_idx * d_hidden + offsets,
        mask=mask,
    )
    for i in range(n_experts_per_token):
        weight = tl.load(topk_scores + token_idx * n_experts_per_token + i)
        index = tl.load(token2index + token_idx * n_experts_per_token + i)
        ffn1_out_block = tl.load(
            ffn1_out + index * d_hidden + offsets,
            mask=mask,
        )
        data += weight * ffn1_out_block
    tl.store(
        moe_out + token_idx * d_hidden + offsets,
        data,
        mask=mask,
    )
   
def moe_forward(
    input_tensor: torch.Tensor,
    weights: Dict[str, torch.Tensor],
    d_hidden: int,
    d_expert: int,
    n_routed_experts: int,
    n_shared_experts: int,
    n_experts_per_token: int,
    batch_size: int,
    seq_len: int
) -> torch.Tensor:
    input_tensor = input_tensor.reshape(-1, d_hidden)


    router_weight = weights['router.weight']
    logits = torch.matmul(input_tensor, router_weight.t())
    scores = F.softmax(logits, dim=-1) 
    topk_scores, topk_indices = torch.topk(scores, n_experts_per_token, dim=-1, sorted=False) 
    
    shared_expert_weight0 = weights['shared_experts.0.weight']
    shared_expert_weight1 = weights['shared_experts.1.weight']
    shared_expert_weight2 = weights['shared_experts.2.weight']
    shared_expert_ffn0 = F.silu(input_tensor @ shared_expert_weight0) * (input_tensor @ shared_expert_weight1) 
    
    shared_expert_ffn1 = shared_expert_ffn0 @ shared_expert_weight2

    BLOCK_SIZE = 512
    expert2token_count = torch.empty(size=(n_routed_experts,), device=input_tensor.device, dtype=torch.int32)
    within_expert_index = torch.empty(size=(batch_size * seq_len, n_experts_per_token), device=input_tensor.device, dtype=torch.int32)
    token2index = torch.empty(size=(batch_size * seq_len, n_experts_per_token), device=input_tensor.device, dtype=torch.int32)
    expert_token_count_cumsum = torch.zeros(size=(n_routed_experts + 1,), device=input_tensor.device, dtype=torch.int32)
    calculate_expert2token_count_and_within_expert_index[ (n_routed_experts, ) ](
        topk_indices,
        expert2token_count,
        within_expert_index,
        batch_size * seq_len,
        n_experts_per_token,
        n_routed_experts,
    )
    calculate_token2index[ lambda meta : (triton.cdiv(batch_size * seq_len * n_experts_per_token, meta["BLOCK_SIZE"]),) ](
        topk_indices,
        expert2token_count,
        within_expert_index,
        token2index,
        expert_token_count_cumsum,
        batch_size * seq_len,
        n_experts_per_token,
        n_routed_experts,
    )
    
    repeated_tokens = torch.empty(size=(batch_size * seq_len * n_experts_per_token, d_hidden), device=input_tensor.device, dtype=input_tensor.dtype)
    scatter[ (n_experts_per_token, seq_len, batch_size) ](
        input_tensor,
        token2index,
        repeated_tokens,
        batch_size,
        seq_len,
        n_experts_per_token,
        d_hidden,
    )
    
    b_ptrs0 = torch.tensor(
        [weights[f'experts.{i}.0.weight'].data_ptr()  for i in range(n_routed_experts)],
        device=repeated_tokens.device,
        dtype=torch.uint64,
    )
    b_ptrs1 = torch.tensor(
        [weights[f'experts.{i}.1.weight'].data_ptr()  for i in range(n_routed_experts)],
        device=repeated_tokens.device,
        dtype=torch.uint64,
    )
    b_ptrs2 = torch.tensor(
        [weights[f'experts.{i}.2.weight'].data_ptr()  for i in range(n_routed_experts)],
        device=repeated_tokens.device,
        dtype=torch.uint64,
    )
    ffn0_lhs = torch.empty(size=(batch_size * seq_len * n_experts_per_token, d_expert), device=input_tensor.device, dtype=input_tensor.dtype)
    BLOCK_M = 128
    BLOCK_N = 128
    BLOCK_K = 128
    max_m = batch_size * seq_len * n_experts_per_token
    group_gemm[ lambda meta: (triton.cdiv(d_expert, meta["BLOCK_N"]), triton.cdiv(max_m, meta["BLOCK_M"]), n_routed_experts) ](
        a = repeated_tokens,
        b_base_ptrs = b_ptrs0,
        c = ffn0_lhs,
        m_cumsum = expert_token_count_cumsum,
        max_m = batch_size * seq_len * n_experts_per_token,
        n = d_expert,
        k = d_hidden,
        group_size = n_routed_experts,
    )
    ffn0_rhs = torch.empty(size=(batch_size * seq_len * n_experts_per_token, d_expert), device=input_tensor.device, dtype=input_tensor.dtype)
    group_gemm[ lambda meta: (triton.cdiv(d_expert, meta["BLOCK_N"]), triton.cdiv(max_m, meta["BLOCK_M"]), n_routed_experts) ](
        a = repeated_tokens,
        b_base_ptrs = b_ptrs1,
        c = ffn0_rhs,
        m_cumsum = expert_token_count_cumsum,
        max_m = batch_size * seq_len * n_experts_per_token,
        n = d_expert,
        k = d_hidden,
        group_size = n_routed_experts,
    )
    
    ffn0_out = F.silu(ffn0_lhs) * ffn0_rhs 
    ffn1_out = torch.empty_like(repeated_tokens)
    group_gemm[ lambda meta: (triton.cdiv(d_hidden, meta["BLOCK_N"]), triton.cdiv(max_m, meta["BLOCK_M"]), n_routed_experts) ](
        a = ffn0_out,
        b_base_ptrs = b_ptrs2,
        c = ffn1_out,
        m_cumsum = expert_token_count_cumsum,
        max_m = batch_size * seq_len * n_experts_per_token,
        n = d_hidden,
        k = d_expert,
        group_size = n_routed_experts,
    )
    moe_out = torch.empty_like(input_tensor)
    BLOCK_SIZE = 1024
    weighted_sum_gather[ lambda meta: (triton.cdiv(d_hidden, meta["BLOCK_SIZE"]), batch_size * seq_len)](
        ffn1_out,
        token2index,
        topk_scores,
        shared_expert_ffn1,
        moe_out,
        d_hidden,
        batch_size,
        seq_len,
        n_experts_per_token,
    )
    # sleep 20 ms
    time.sleep(0.01)
    return moe_out.reshape(batch_size, seq_len, d_hidden)


def custom_kernel(data: input_t) -> output_t:
    input_tensor, weights, config = data
    output = moe_forward(input_tensor, weights, **config)
    return output

python eval.py leaderboard /tmp/tmpyggusfm8