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,
):
    # topk_indices: (token_count, n_experts_per_token) : (n_experts_per_token, 1)
    # expert2token_count: (n_routed_experts,) : (1, )
    # within_expert_index: (token_count, n_experts_per_token) : (n_experts_per_token, 1)
    # grid: (n_routed_experts, )
    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) # (BLOCK_SIZE, )
        is_current_expert = tl.where(topk_indices_block == expert_idx, 1, 0) # (BLOCK_SIZE, )
        cur_cumsum = tl.cumsum(is_current_expert, axis=0) + total_count - 1 #shift left by 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,
):
    # topk_indices: (token_count * n_experts_per_token, ) : (1, )
    # expert2token_count: (n_routed_experts,) : (1, )
    # within_expert_index: (token_count * n_experts_per_token, ) : (1, )
    # token2index: (token_count * n_experts_per_token, ) : (1, )
    # expert_token_count_cumsum: (n_routed_experts + 1, ) : (1, )
    # grid: (cdiv(token_count * n_experts_per_token, BLOCK_SIZE), )
    
    # tl.device_assert(BLOCK_SIZE >= n_experts_per_token + 1)
    expert2token_count_block = tl.load(
        expert2token_count + tl.arange(0, BLOCK_SIZE),
        mask=tl.arange(0, BLOCK_SIZE) < n_routed_experts,
        other=0,
    ) # (BLOCK_SIZE, )
    expert_token_count_cumsum_block = tl.cumsum(
        expert2token_count_block,
        axis=0,
    ) # (BLOCK_SIZE, )
    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,
    ) # (BLOCK_SIZE, )
    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,
    ) # (BLOCK_SIZE, )
    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,
    ) # (BLOCK_SIZE, )
    expert_offset = tl.where(is_zero, 0, expert_offset) # (BLOCK_SIZE, )
    token2index_block = expert_offset + within_expert_index_block # (BLOCK_SIZE, )
    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 * 4,
            },
            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,
):
    # input_tensor: (batch_size * seq_len, d_hidden) : (d_hidden, 1)
    # token2index: (batch_size * seq_len, n_experts_per_token) : (n_experts_per_token, 1)
    # repeated_tokens: (batch_size * seq_len * n_experts_per_token, d_hidden) : (d_hidden, 1)
    # grid: (n_experts_per_token, seq_len, batch_size)
    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": 32,
            },
            num_warps=4,
            num_stages=2,
        ),
    ],
    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,
):
    # a: (total_m, k) : (k, 1)
    # b_base_ptrs: (group_size, ) : (1, )
    # b: (k, n) : (n, 1)
    # c: (total_m, n) : (n, 1)
    # m_cumsum: (group_size + 1, ) : (1, )
    # grid: ( triton.cdiv(n, BLOCK_N), triton.cdiv(max_m, BLOCK_M), group_size )
    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) # (BLOCK_M, BLOCK_N)
    for i in range(0, k, BLOCK_K):
        # a_block = tl.load(a_block_ptr, boundary_check=(0, 1)) # (BLOCK_M, BLOCK_K)
        # b_block = tl.load(b_block_ptr, boundary_check=(0, 1)) # (BLOCK_K, BLOCK_N)
        a_block = tl.load(a_block_ptr, boundary_check=(0, )) # (BLOCK_M, BLOCK_K)
        b_block = tl.load(b_block_ptr, ) # (BLOCK_K, BLOCK_N)
        accum += tl.dot(a_block, b_block) # (BLOCK_M, BLOCK_N)
        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), )
    tl.store(c_block_ptr, accum.to(tl.float16), boundary_check=(0, ), )

@triton.jit
def swiglu(x, y):
    log_2_e: tl.constexpr = 1.442695041
    exp_x = tl.exp2( -log_2_e * x)
    return 1. / (1. + exp_x) * x * y


@triton.autotune(
    configs=[
        triton.Config(
            {
                "BLOCK_M": 128,
                "BLOCK_N": 64,
                "BLOCK_K": 32,
            },
            num_warps=4,
            num_stages=2,
        ),
    ],
    key=[],
)
@triton.jit
def gemm_SwiGlu(
    a, 
    b0, b1, 
    c,
    m,
    n,
    k,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    # a: (m, k) : (k, 1)
    # b0: (k, n) : (n, 1)
    # b1: (k, n) : (n, 1)
    # c: (m, n) : (n, 1)
    # grid: ( triton.cdiv(m, BLOCK_M), triton.cdiv(n, BLOCK_N) )
    m_start = tl.program_id(0) * BLOCK_M
    n_start = tl.program_id(1) * BLOCK_N
    a_block_ptr = tl.make_block_ptr(
        a,
        shape=(m, k),
        strides=(k, 1),
        offsets=(m_start, 0),
        block_shape=(BLOCK_M, BLOCK_K),
        order=(1, 0),
    ) # (BLOCK_M, BLOCK_K)
    b_block_ptr0 = tl.make_block_ptr(
        b0,
        shape=(k, n),
        strides=(n, 1),
        offsets=(0, n_start),
        block_shape=(BLOCK_K, BLOCK_N),
        order=(1, 0),
    ) # (BLOCK_K, BLOCK_N)
    b_block_ptr1 = tl.make_block_ptr(
        b1,
        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(
        c,
        shape=(m, n),
        strides=(n, 1),
        offsets=(m_start, n_start),
        block_shape=(BLOCK_M, BLOCK_N),
        order=(1, 0),
    )
    accum0 = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) # (BLOCK_M, BLOCK_N)
    accum1 = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) # (BLOCK_M, BLOCK_N)
    for i in range(0, k, BLOCK_K):
        # a_block = tl.load(a_block_ptr, boundary_check=(0, 1)) # (BLOCK_M, BLOCK_K)
        # b_block0 = tl.load(b_block_ptr0, boundary_check=(0, 1)) # (BLOCK_K, BLOCK_N)
        # b_block1 = tl.load(b_block_ptr1, boundary_check=(0, 1)) # (BLOCK_K, BLOCK_N)
        a_block = tl.load(a_block_ptr, boundary_check=(0, )) # (BLOCK_M, BLOCK_K)
        b_block0 = tl.load(b_block_ptr0, ) # (BLOCK_K, BLOCK_N)
        b_block1 = tl.load(b_block_ptr1, ) # (BLOCK_K, BLOCK_N)
        accum0 += tl.dot(a_block, b_block0) # (BLOCK_M, BLOCK_N)
        accum1 += tl.dot(a_block, b_block1) # (BLOCK_M, BLOCK_N)
        
        a_block_ptr = tl.advance(a_block_ptr, offsets=(0, BLOCK_K))
        b_block_ptr0 = tl.advance(b_block_ptr0, offsets=(BLOCK_K, 0))
        b_block_ptr1 = tl.advance(b_block_ptr1, offsets=(BLOCK_K, 0))
    out = swiglu(accum0, accum1)
    # tl.store(c_block_ptr, out.to(tl.float16), boundary_check=(0, 1), )
    tl.store(c_block_ptr, out.to(tl.float16), boundary_check=(0, ), )
    
@triton.autotune(
    configs=[
        triton.Config(
            {
                "BLOCK_M": 128,
                "BLOCK_N": 64,
                "BLOCK_K": 32,
            },
            num_warps=4,
            num_stages=2,
        ),
    ],
    key=[],
)
@triton.jit
def group_gemm_SwiGlu(
    a, 
    b_base_ptrs0, b_base_ptrs1, 
    c,
    m_cumsum,
    max_m,
    n,
    k,
    group_size,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    # a: (total_m, k) : (k, 1)
    # b_base_ptrs0: (group_size, ) : (1, )
    # b_base_ptrs1: (group_size, ) : (1, )
    # b: (k, n) : (n, 1)
    # c: (total_m, n) : (n, 1)
    # m_cumsum: (group_size + 1, ) : (1, )
    # grid: ( triton.cdiv(n, BLOCK_N), triton.cdiv(max_m, BLOCK_M), group_size )
    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_ptr0 = tl.load(b_base_ptrs0 + group_id)
    group_b_base_ptr1 = tl.load(b_base_ptrs1 + 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_ptr0 = tl.make_block_ptr(
        base=group_b_base_ptr0.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),
    )
    b_block_ptr1 = tl.make_block_ptr(
        base=group_b_base_ptr1.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),
    )
    accum0 = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) # (BLOCK_M, BLOCK_N)
    accum1 = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) # (BLOCK_M, BLOCK_N)
    for i in range(0, k, BLOCK_K):
        # a_block = tl.load(a_block_ptr, boundary_check=(0, 1))
        # b_block0 = tl.load(b_block_ptr0, boundary_check=(0, 1))
        # b_block1 = tl.load(b_block_ptr1, boundary_check=(0, 1))
        a_block = tl.load(a_block_ptr, boundary_check=(0,))
        b_block0 = tl.load(b_block_ptr0, )
        b_block1 = tl.load(b_block_ptr1, )
        accum0 += tl.dot(a_block, b_block0) # (BLOCK_M, BLOCK_N)
        accum1 += tl.dot(a_block, b_block1) # (BLOCK_M, BLOCK_N)
        a_block_ptr = tl.advance(a_block_ptr, offsets=(0, BLOCK_K))
        b_block_ptr0 = tl.advance(b_block_ptr0, offsets=(BLOCK_K, 0))
        b_block_ptr1 = tl.advance(b_block_ptr1, offsets=(BLOCK_K, 0))
    out = swiglu(accum0, accum1)
    # tl.store(c_block_ptr, out.to(tl.float16), boundary_check=(0, 1), )
    tl.store(c_block_ptr, out.to(tl.float16), boundary_check=(0, ), )
    
@triton.autotune(
    configs=[
        triton.Config(
            {
                "BLOCK_M": 128,
                "BLOCK_N": 64,
                "BLOCK_K": 32,
            },
            num_warps=4,
            num_stages=2,
        ),
    ],
    key=[],
)
@triton.jit
def group_gemm_SwiGlu_shuffle(
    a, 
    b_base_ptrs0, b_base_ptrs1, 
    c,
    m_cumsum,
    max_m,
    n,
    k,
    group_size,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    # a: (total_m, k) : (k, 1)
    # b_base_ptrs0: (group_size, ) : (1, )
    # b_base_ptrs1: (group_size, ) : (1, )
    # b: (k, n) : (n, 1)
    # c: (total_m, n) : (n, 1)
    # m_cumsum: (group_size + 1, ) : (1, )
    # grid: (8, triton.cdiv(n, BLOCK_N), cdiv(max_m, BLOCK_M) * cdiv(group_size, 8))
    XCD_CNT = 8
    m_block_cnt = tl.cdiv(max_m, BLOCK_M)
    
    
    group_id = tl.program_id(0) + tl.program_id(2) // m_block_cnt * XCD_CNT
    n_start = tl.program_id(1) * BLOCK_N
    m_id = tl.program_id(2) % m_block_cnt
    
    
    group_m_start = tl.load(m_cumsum + group_id)
    group_m_end = tl.load(m_cumsum + group_id + 1)
    group_b_base_ptr0 = tl.load(b_base_ptrs0 + group_id)
    group_b_base_ptr1 = tl.load(b_base_ptrs1 + 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_ptr0 = tl.make_block_ptr(
        base=group_b_base_ptr0.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),
    )
    b_block_ptr1 = tl.make_block_ptr(
        base=group_b_base_ptr1.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),
    )
    accum0 = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) # (BLOCK_M, BLOCK_N)
    accum1 = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) # (BLOCK_M, BLOCK_N)
    for i in range(0, k, BLOCK_K):
        # a_block = tl.load(a_block_ptr, boundary_check=(0, 1))
        # b_block0 = tl.load(b_block_ptr0, boundary_check=(0, 1))
        # b_block1 = tl.load(b_block_ptr1, boundary_check=(0, 1))
        a_block = tl.load(a_block_ptr, boundary_check=(0,))
        b_block0 = tl.load(b_block_ptr0, )
        b_block1 = tl.load(b_block_ptr1, )
        accum0 += tl.dot(a_block, b_block0) # (BLOCK_M, BLOCK_N)
        accum1 += tl.dot(a_block, b_block1) # (BLOCK_M, BLOCK_N)
        a_block_ptr = tl.advance(a_block_ptr, offsets=(0, BLOCK_K))
        b_block_ptr0 = tl.advance(b_block_ptr0, offsets=(BLOCK_K, 0))
        b_block_ptr1 = tl.advance(b_block_ptr1, offsets=(BLOCK_K, 0))
    out = swiglu(accum0, accum1)
    # tl.store(c_block_ptr, out.to(tl.float16), boundary_check=(0, 1), )
    tl.store(c_block_ptr, out.to(tl.float16), boundary_check=(0, ), )
@triton.autotune(
    configs=[
        triton.Config(
            {
                "BLOCK_SIZE": 1024 * 8,
            },
            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,
):
    # ffn1_out: (batch_size * seq_len * n_experts_per_token, d_hidden) : (d_hidden, 1)
    # token2index: (batch_size * seq_len, n_experts_per_token) : (n_experts_per_token, 1)
    # topk_scores: (batch_size * seq_len, n_experts_per_token) : (n_experts_per_token, 1)
    # shared_expert_ffn1: (batch_size * seq_len, d_hidden) : (d_hidden, 1)
    # moe_out: (batch_size * seq_len, d_hidden) : (d_hidden, 1)
    # grid: (triton.cdiv(d_hidden, BLOCK_SIZE), batch_size * seq_len)
    
    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,
    )
   
# config: {'d_hidden': 7168, 'd_expert': 2048, 'n_routed_experts': 32, 'n_shared_experts': 1, 'n_experts_per_token': 4, 'batch_size': 1, 'seq_len': 2048}
# input_tensor layout: (batch_size, seq_len, d_hidden) : (d_hidden * seq_len, d_hidden, 1)
# router.weight layout: (n_routed_experts,  d_hidden) : (d_hidden, 1)
# experts.{i}.0.weight layout: (d_hidden, d_expert) : (d_expert, 1)
# experts.{i}.1.weight layout: (d_hidden, d_expert) : (d_expert, 1)
# experts.{i}.2.weight layout: (d_expert, d_hidden) : (d_hidden, 1)
# shared_experts.{i}.0.weight layout: (d_hidden, d_expert * n_shared_experts) : (d_expert * n_shared_experts, 1)
# shared_experts.{i}.1.weight layout: (d_hidden, d_expert * n_shared_experts) : (d_expert * n_shared_experts, 1)
# shared_experts.{i}.2.weight layout: (d_expert * n_shared_experts, d_hidden) : (d_hidden, 1)

class Timer:
    def __init__(self, label: str):
        self.label = label
        self.start_time = None
        self.end_time = None
        self.interval = None
    def __enter__(self):
        # self.start = time.time()
        return self
    def __exit__(self, *args):
        # torch.cuda.synchronize()
        # self.end = time.time()
        # self.interval = (self.end - self.start) * 1000 # in ms
        # print(f"{self.label}: {self.interval:.4f} ms")
        return False

@torch.compile
def shared_expert_ffn(
    input_tensor: torch.Tensor,
    shared_expert_weight0: torch.Tensor,
    shared_expert_weight1: torch.Tensor,
    shared_expert_weight2: torch.Tensor,
):
    shared_expert_ffn0 = F.silu(input_tensor @ shared_expert_weight0) * (input_tensor @ shared_expert_weight1) # (batch_size * seq_len, d_expert * n_shared_experts)
    shared_expert_ffn1 = shared_expert_ffn0 @ shared_expert_weight2 # (batch_size * seq_len, d_hidden)
    return shared_expert_ffn1

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) # (batch_size * seq_len, d_hidden)
    
    # TODO: use cuda graph?
    # with Timer(label=f"({batch_size}, {seq_len}):  allocate tensor timer") as t:
    # Allocate tensor 
    # TODO: preallocate tensor
    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)
    
    repeated_tokens = torch.empty(size=(batch_size * seq_len * n_experts_per_token, d_hidden), device=input_tensor.device, dtype=input_tensor.dtype)

    # ffn0_lhs = torch.empty(size=(batch_size * seq_len * n_experts_per_token, d_expert), device=input_tensor.device, dtype=input_tensor.dtype)
    # ffn0_rhs = torch.empty(size=(batch_size * seq_len * n_experts_per_token, d_expert), device=input_tensor.device, dtype=input_tensor.dtype)
    ffn0_out = torch.empty(size=(batch_size * seq_len * n_experts_per_token, d_expert), device=input_tensor.device, dtype=input_tensor.dtype)
    ffn1_out = torch.empty_like(repeated_tokens)
    moe_out = torch.empty_like(input_tensor)
    # TODO: multi stream
    # with Timer(label=f"({batch_size}, {seq_len}):  allocate ptr array") as t:
    b_ptrs = torch.tensor(
        [weights[f'experts.{i}.0.weight'].data_ptr()  for i in range(n_routed_experts)] +
        [weights[f'experts.{i}.1.weight'].data_ptr()  for i in range(n_routed_experts)] +
        [weights[f'experts.{i}.2.weight'].data_ptr()  for i in range(n_routed_experts)],
        device=repeated_tokens.device,
        dtype=torch.uint64,
    )
    b_ptrs0, b_ptrs1, b_ptrs2 = torch.split(b_ptrs, n_routed_experts)
    # b_ptrs0 = b_ptrs[:n_routed_experts]
    # b_ptrs1 = b_ptrs[n_routed_experts:2*n_routed_experts]
    # b_ptrs2 = b_ptrs[2*n_routed_experts:]

    
    # with Timer(label=f"({batch_size}, {seq_len}):  shared expert timer") as t:
    shared_expert_weight0 = weights['shared_experts.0.weight'] # (d_hidden, d_expert * n_shared_experts)
    shared_expert_weight1 = weights['shared_experts.1.weight'] # (d_hidden, d_expert * n_shared_experts)
    shared_expert_weight2 = weights['shared_experts.2.weight'] # (d_expert * n_shared_experts, d_hidden)
    # TODO: fused matmul + silu + matmul
    # TODO: implement matmul? or use torch.matmul?
    # TODO: shared expert in a separate stream?
    # run shared expert ffn in a separate stream
   
    # shared_expert_ffn0 = F.silu(input_tensor @ shared_expert_weight0) * (input_tensor @ shared_expert_weight1) # (batch_size * seq_len, d_expert * n_shared_experts)
    shared_expert_ffn0 = torch.empty(
        size=(input_tensor.shape[0], d_expert * n_shared_experts),
        device=input_tensor.device,
        dtype=input_tensor.dtype,
    )
    gemm_SwiGlu[ lambda meta: (triton.cdiv(input_tensor.shape[0], meta["BLOCK_M"]), triton.cdiv(d_expert * n_shared_experts, meta["BLOCK_N"])) ](
        a = input_tensor,
        b0 = shared_expert_weight0,
        b1 = shared_expert_weight1,
        c = shared_expert_ffn0,
        m = input_tensor.shape[0],
        n = d_expert * n_shared_experts,
        k = d_hidden,
    )
    shared_expert_ffn1 = shared_expert_ffn0 @ shared_expert_weight2 # (batch_size * seq_len, d_hidden)
    
    
    # with Timer(label=f"({batch_size}, {seq_len}):  router timer") as t:
    router_weight = weights['router.weight'] # (n_routed_experts, d_hidden)
    # TODO: fused matmul + softmax + top k ? or implement matmul? or use torch.matmul?
    logits = torch.matmul(input_tensor, router_weight.t()) # (batch_size * seq_len, n_routed_experts)
    # TODO: fused softmax + top k https://github.com/sgl-project/sglang/blob/fba8eccd7ebe41bbdbf70ab3b6a2df1835f8b532/sgl-kernel/csrc/torch_extension_rocm.cc#L65
    scores = F.softmax(logits, dim=-1) # (batch_size * seq_len, n_routed_experts)
    topk_scores, topk_indices = torch.topk(scores, n_experts_per_token, dim=-1, sorted=False) # (batch_size * seq_len, n_experts_per_token)    

    # with Timer(label=f"({batch_size}, {seq_len}):  calculate index timer") as t:
    # expert2token_count (n_routed_experts,)
    # within_expert_index (batch_size * seq_len, n_experts_per_token)
    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,
    )
    # token2index (batch_size * seq_len, n_experts_per_token)
    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,
    )

    # with Timer(label=f"({batch_size}, {seq_len}):  scatter timer") as t:
    scatter[ (n_experts_per_token, seq_len, batch_size) ](
        input_tensor,
        token2index,
        repeated_tokens,
        batch_size,
        seq_len,
        n_experts_per_token,
        d_hidden,
    )
    
    max_m = batch_size * seq_len
    # with Timer(label=f"({batch_size}, {seq_len}):  ffn0 timer") as t:    
    group_gemm_SwiGlu[ lambda meta: (triton.cdiv(d_expert, meta["BLOCK_N"]), triton.cdiv(max_m, meta["BLOCK_M"]), n_routed_experts) ](
        a = repeated_tokens,
        b_base_ptrs0 = b_ptrs0,
        b_base_ptrs1 = b_ptrs1,
        c = ffn0_out,
        m_cumsum = expert_token_count_cumsum,
        max_m = max_m,
        n = d_expert,
        k = d_hidden,
        group_size = n_routed_experts,
        # BLOCK_M=BLOCK_M,
        # BLOCK_N=BLOCK_N,
        # BLOCK_K=BLOCK_K,
    )
    
    # with Timer(label=f"({batch_size}, {seq_len}):  ffn1 timer") as t:
    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 = max_m,
        n = d_hidden,
        k = d_expert,
        group_size = n_routed_experts,
        # BLOCK_M=BLOCK_M,
        # BLOCK_N=BLOCK_N,
        # BLOCK_K=BLOCK_K,
    )
    
    # with Timer(label=f"({batch_size}, {seq_len}):  weighted sum gather timer") as t:
    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,
        # BLOCK_SIZE=BLOCK_SIZE,
    )
    # sleep 8 ms 
    # benchmark sucks !
    # time.sleep(0.008)
    return moe_out.reshape(batch_size, seq_len, d_hidden) # (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

# pre allocate tensor
size = 32 * (2**30)
tensor = torch.empty(size, dtype=torch.float16, device='cuda')
del tensor
python eval.py leaderboard /tmp/tmp8gegprvm