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
class Expert(nn.Module):
    def __init__(self, config: Dict, d_expert: Optional[int] = None):
        super().__init__()
        self.config = config
        self.act_fn = nn.SiLU()
        self.d_hidden: int = config["d_hidden"]
        self.d_expert: int = config["d_expert"] if d_expert is None else d_expert

        self.W_gate = nn.Linear(self.d_hidden, self.d_expert, bias=False)
        self.W_up = nn.Linear(self.d_hidden, self.d_expert, bias=False)
        self.W_down = nn.Linear(self.d_expert, self.d_hidden, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate = self.act_fn(self.W_gate(x))
        out = self.W_down(gate * self.W_up(x))
        return out


class MoEGate(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.top_k: int = config["n_experts_per_token"]
        self.num_experts: int = config["n_routed_experts"]
        self.d_hidden: int = config["d_hidden"]

        self.W_g = nn.Linear(self.d_hidden, self.num_experts, bias=False)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        logits = self.W_g(x)
        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 MoE(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config
        self.experts = nn.ModuleList([
            Expert(config)
            for _ in range(config["n_routed_experts"])
        ])
        self.gating_network = MoEGate(config)
        shared_expert_dim = config["d_expert"] * config["n_shared_experts"]
        self.shared_expert = Expert(config=config, d_expert=shared_expert_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        shared_output = self.shared_expert(x)
        expert_indices, expert_scores = self.gating_network(x)
        batch_size, seq_len, hidden_dim = x.shape
        orig_shape = x.shape
        x_flat = x.view(-1, hidden_dim)
        flat_expert_indices = expert_indices.view(-1)
        flat_expert_weights = expert_scores.view(-1, 1)
        routed_output_flat = self.moe_infer(x_flat,
                                            flat_expert_indices,
                                            flat_expert_weights)

        routed_output = routed_output_flat.view(*orig_shape)
        return routed_output + shared_output

    @torch.no_grad()
    def moe_infer(self,
                  x: torch.Tensor,
                  flat_expert_indices: torch.Tensor,
                  flat_expert_weights: torch.Tensor
                 ) -> torch.Tensor:
        expert_cache = torch.zeros_like(x)
        idxs = flat_expert_indices.argsort()
        counts = flat_expert_indices.bincount().cpu().numpy()
        tokens_per_expert = counts.cumsum()
        num_per_tok = self.config["n_experts_per_token"]
        token_idxs = idxs // num_per_tok
        for expert_id, end_idx in enumerate(tokens_per_expert):
            start_idx = 0 if expert_id == 0 else tokens_per_expert[expert_id - 1]
            if start_idx == end_idx:
                continue

            expert = self.experts[expert_id]
            exp_token_idxs = token_idxs[start_idx:end_idx]
            expert_tokens = x[exp_token_idxs]
            expert_out    = expert(expert_tokens)
            expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
            expert_cache.scatter_reduce_(
                0,
                exp_token_idxs.view(-1, 1).repeat(1, x.shape[-1]),
                expert_out,
                reduce='sum'
            )

        return expert_cache


def my_ref_custom_kernel(data: input_t) -> output_t:
    """
    Submission template for DeepSeek-style Mixture of Experts using PyTorch.

    Args:
        data: Tuple of (input: torch.Tensor, weights: Dict[str, torch.Tensor], config: Dict)
            - input: Input tensor of shape [batch_size, seq_len, hidden_size]
            - weights: Dictionary containing model weights
            - config: Dictionary containing model configuration parameters

    Returns:
        Tuple containing:
            - output: Processed tensor [batch_size, seq_len, d_model]
            - aux_data: Dictionary with auxiliary data
    """
    input_tensor, weights, config = data
    num_experts = config["n_routed_experts"]
    moe = MoE(config)

    # Fill in the given weights of the model
    moe.gating_network.W_g.weight = nn.Parameter(weights['router.weight'])

    for i in range(num_experts):
        gate_proj_weight = weights[f'experts.{i}.0.weight']
        up_proj_weight = weights[f'experts.{i}.1.weight']
        down_proj_weight = weights[f'experts.{i}.2.weight']

        # Transpose weights to match expected shape for nn.Linear
        moe.experts[i].W_gate.weight = nn.Parameter(gate_proj_weight.t())
        moe.experts[i].W_up.weight = nn.Parameter(up_proj_weight.t())
        moe.experts[i].W_down.weight = nn.Parameter(down_proj_weight.t())

    moe.shared_expert.W_gate.weight = nn.Parameter(weights['shared_experts.0.weight'].t())
    moe.shared_expert.W_up.weight = nn.Parameter(weights['shared_experts.1.weight'].t())
    moe.shared_expert.W_down.weight = nn.Parameter(weights['shared_experts.2.weight'].t())

    # Run the model
    output = moe(input_tensor)

    return output


def my_generate_input(
    dhidden: int,
    dexpert: int,
    nroutedexperts: int,
    nsharedexperts: int,
    nexpertspertoken: int,
    bs: int,
    seqlen: int,
    seed: int
) -> input_t:

    # Really dumb but for now _ isn't parsing correctly.
    d_hidden = dhidden
    d_expert = dexpert
    n_routed_experts = nroutedexperts
    n_shared_experts = nsharedexperts
    n_experts_per_token = nexpertspertoken
    batch_size = bs
    seq_len = seqlen

    config = {
        "d_hidden": d_hidden,
        "d_expert": d_expert,
        "n_routed_experts": n_routed_experts,
        "n_shared_experts": n_shared_experts,
        "n_experts_per_token": n_experts_per_token,
        "batch_size": batch_size,
        "seq_len": seq_len,
    }

    gen = torch.Generator(device='cuda')
    gen.manual_seed(seed)

    num_experts = n_routed_experts
    expert_dim = d_expert
    weights = {}

    input_tensor = torch.randn(
        (batch_size, seq_len, d_hidden),
        device='cuda',
        dtype=torch.float16,
        generator=gen
    ).contiguous()

    # Initialize router weights
    weights['router.weight'] = torch.randn(
        (num_experts, d_hidden),
        device="cuda",
        dtype=torch.float16,
        generator=gen
    ) / math.sqrt(d_hidden)

    for i in range(num_experts):
        weights[f'experts.{i}.0.weight'] = torch.randn(
            (d_hidden, expert_dim),
            device='cuda',
            dtype=torch.float16,
            generator=gen
        ) / math.sqrt(expert_dim)

        weights[f'experts.{i}.1.weight'] = torch.randn(
            (d_hidden, expert_dim),
            device='cuda',
            dtype=torch.float16,
            generator=gen
        ) / math.sqrt(expert_dim)

        weights[f'experts.{i}.2.weight'] = torch.randn(
            (expert_dim, d_hidden),
            device='cuda',
            dtype=torch.float16,
            generator=gen
        ) / math.sqrt(d_hidden)

    weights['shared_experts.0.weight'] = torch.randn(
        (d_hidden, expert_dim * n_shared_experts),
        device='cuda',
        dtype=torch.float16,
        generator=gen
    ) / math.sqrt(expert_dim * n_shared_experts)
    weights['shared_experts.1.weight'] = torch.randn(
        (d_hidden, expert_dim * n_shared_experts),
        device='cuda',
        dtype=torch.float16,
        generator=gen
    ) / math.sqrt(expert_dim * n_shared_experts)
    weights['shared_experts.2.weight'] = torch.randn(
        (expert_dim * n_shared_experts, d_hidden),
        device='cuda',
        dtype=torch.float16,
        generator=gen
    ) / math.sqrt(d_hidden)

    return (input_tensor, weights, config)

@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.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.cat(
    #     tl.zeros((1, ), dtype=tl.int32),
    #     tl.cumsum(
    #         expert2token_count_block,
    #         axis=0,
    #     ),
    #     can_reorder=False,
    # ) # (BLOCK_SIZE + 1, )
    expert_token_count_cumsum_block = tl.cumsum(
        expert2token_count_block,
        axis=0,
    ) # (BLOCK_SIZE, )
    # tl.device_print("expert2token_count_block: ", expert2token_count_block)
    # tl.device_print("expert_token_count_cumsum_block: ", expert_token_count_cumsum_block)
    if tl.program_id(0) == 0:
        # tl.device_print("expert2token_count_block: ", expert2token_count_block)
        # tl.device_print("expert_token_count_cumsum_block: ", expert_token_count_cumsum_block)
        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,
    )
def torch_get_m_cumsum(
    topk_indices: torch.Tensor,
    expert_cnt: int,
):
    # topk_indices: (batch_size * seq_len, n_experts_per_token) : (n_experts_per_token, 1)
    # expert2token_count: (n_routed_experts,) : (1, )
    
    expert_count = torch.zeros((expert_cnt,), dtype=torch.int32)
    for i in range(topk_indices.shape[0]):
        for j in range(topk_indices.shape[1]):
            expert_count[topk_indices[i][j]] += 1
    expert_count_cumsum = torch.cat(
        (torch.zeros((1,), dtype=torch.int32), torch.cumsum(expert_count, dim=0)),
        dim=0,
    )
    return expert_count_cumsum.cuda()
def torch_get_token2index(
    topk_indices: torch.Tensor,
    expert_count: int,
):
    # topk_indices: (batch_size * seq_len, n_experts_per_token) : (n_experts_per_token, 1)
    # expert_token_count_cumsum: (n_routed_experts + 1,) : (1, )
    topk_indices = topk_indices.cpu()
    expert_token_count = torch.zeros((expert_count,), dtype=torch.int32)
    within_expert_index = torch.zeros((topk_indices.shape[0], topk_indices.shape[1]), dtype=torch.int32)
    for i in range(topk_indices.shape[0]):
        for j in range(topk_indices.shape[1]):
            expert_idx = topk_indices[i][j]
            within_expert_index[i][j] = expert_token_count[expert_idx]
            expert_token_count[expert_idx] += 1
    
    expert_token_count_cumsum = torch.cat(
        (torch.zeros((1,), dtype=torch.int32), torch.cumsum(expert_token_count, dim=0)),
        dim=0,
    )
    token2index = torch.zeros((topk_indices.shape[0], topk_indices.shape[1]), dtype=torch.int32)
    for i in range(topk_indices.shape[0]):
        for j in range(topk_indices.shape[1]):
            expert_idx = topk_indices[i][j]
            token2index[i][j] = expert_token_count_cumsum[expert_idx] + within_expert_index[i][j]
    return token2index.cuda(), expert_token_count_cumsum.cuda(), within_expert_index.cuda(), expert_token_count.cuda()
@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,
        )
def torch_scatter(
    input_tensor: torch.Tensor,
    token2index: torch.Tensor,
):
    # 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)
    repeated_tokens = torch.zeros(
        (input_tensor.shape[0] * token2index.shape[1], input_tensor.shape[1]),
        device=input_tensor.device,
        dtype=input_tensor.dtype,
    )
    for i in range(token2index.shape[0]):
        for j in range(token2index.shape[1]):
            dst_idx = token2index[i][j]
            repeated_tokens[dst_idx] = input_tensor[i]
    return repeated_tokens.cuda()
@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)
        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), )

def torch_group_gemm(
    a: torch.Tensor,
    b_tensors: List[torch.Tensor],
    m_cumsum: torch.Tensor,
    n: int,
):
    # a: (total_m, k) : (k, 1)
    # b_tensors: (group_size, ) : (1, )
    # b: (k, n) : (n, 1)
    # c: (total_m, n) : (n, 1)
    # m_cumsum: (group_size + 1, ) : (1, )
    
    group_size = len(b_tensors)
    c = torch.zeros(
        (m_cumsum[-1], n),
        device=a.device,
        dtype=a.dtype,
    )
    for i in range(group_size):
        
        mat_a = a[m_cumsum[i]:m_cumsum[i + 1]]
        mat_b = b_tensors[i]
        c[m_cumsum[i]:m_cumsum[i + 1]] = torch.matmul(
            mat_a,
            mat_b,
        )
    return c.cuda()
        
        
@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,
    )
       
def torch_weighted_sum_gather(
    ffn1_out: torch.Tensor,
    token2index: torch.Tensor,
    topk_scores: torch.Tensor,
    shared_expert_ffn1: torch.Tensor,
):
    # 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)
    token_count = token2index.shape[0]
    n_experts_per_token = token2index.shape[1]
    d_hidden = ffn1_out.shape[1]
    moe_out = torch.zeros(
        (token_count, d_hidden),
        device=ffn1_out.device,
        dtype=ffn1_out.dtype,
    )
    for i in range(token_count):
        for j in range(n_experts_per_token):
            index = token2index[i][j]
            weight = topk_scores[i][j]
            moe_out[i] += weight * ffn1_out[index]
        moe_out[i] += shared_expert_ffn1[i]
    return moe_out.cuda()
   
# 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)
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: (batch_size, seq_len, d_hidden)
    input_tensor = input_tensor.reshape(-1, d_hidden) # (batch_size * seq_len, d_hidden)


    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)    
    # topk_scores: token_id -> topk scores (batch_size * seq_len, n_experts_per_token)
    # topk_indices: token_id -> topk expert_id (batch_size * seq_len, n_experts_per_token)

    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
    shared_expert_ffn0 = F.silu(input_tensor @ shared_expert_weight0) * (input_tensor @ shared_expert_weight1) # (batch_size * seq_len, d_expert * n_shared_experts)
    # TODO: implement matmul? or use torch.matmul?
    shared_expert_ffn1 = shared_expert_ffn0 @ shared_expert_weight2 # (batch_size * seq_len, d_hidden)

    # TODO: calculate expert2token_count (n_routed_experts + 1,)
    # TODO: calculate token2within_expert_index (batch_size * seq_len, n_experts_per_token)
    # TODO: calculate token2index (batch_size * seq_len, n_experts_per_token)
    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,
        BLOCK_SIZE=BLOCK_SIZE,
    )
    # print(f"within_expert_index: {within_expert_index}")
    calculate_token2index[ (triton.cdiv(batch_size * seq_len * n_experts_per_token, 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,
        BLOCK_SIZE=BLOCK_SIZE,
    )
    token2index_ref, expert_token_count_cumsum_ref, within_expert_index_ref, expert_token_count_ref = torch_get_token2index(
        topk_indices,
        n_routed_experts,
    )
    assert torch.all(token2index == token2index_ref), f"token2index: {token2index}, token2index_ref: {token2index_ref}"
    assert torch.all(expert_token_count_cumsum == expert_token_count_cumsum_ref), f"expert_token_count_cumsum: {expert_token_count_cumsum}, expert_token_count_cumsum_ref: {expert_token_count_cumsum_ref}"
    assert torch.all(within_expert_index == within_expert_index_ref), f"within_expert_index: {within_expert_index}, within_expert_index_ref: {within_expert_index_ref}"
    assert torch.all(expert2token_count == expert_token_count_ref), f"expert2token_count: {expert2token_count}, expert_token_count_ref: {expert_token_count_ref}"
    assert torch.all(token2index>=0) and torch.all(token2index < batch_size * seq_len * n_experts_per_token), f"out of bound"
    assert expert_token_count_cumsum[0] == 0 and expert_token_count_cumsum[-1] == batch_size * seq_len * n_experts_per_token
    assert torch.all(expert_token_count_cumsum[1:] - expert_token_count_cumsum[:-1] >= 0) and torch.all(expert_token_count_cumsum[1:] - expert_token_count_cumsum[:-1] <= batch_size * seq_len)
    # exit()  
    # print(f"topk_indices: {topk_indices}")
    # print(f"expert2token_count : {expert2token_count}")
    # print(f"expert_token_count_ref : {expert_token_count_ref}")
    # exit()
    # print(f"expert_token_count_cumsum: {expert_token_count_cumsum}")
    # print(f"expert_token_count_cumsum_ref: {expert_token_count_cumsum_ref}")
    # exit()
    # token2index = token2index_ref
    # expert_token_count_cumsum = expert_token_count_cumsum_ref

    # for token_idx in range(batch_size * seq_len):
    #     diff = torch.abs(token2index[token_idx] - token2index_ref[token_idx])
    #     if diff.sum() != 0:
    #         print(f"token_idx: {token_idx}, diff: {diff}")
    #         print(f"token2index: {token2index[token_idx]}")
    #         print(f"token2index_ref: {token2index_ref[token_idx]}")
    #         exit()
    # print(f"within_expert_index_diff: {within_expert_index_diff}")
    
    # print(f"token2index: {token2index}")
    
    # TODO: scatter input_tensor into repeated_tokens according to token2index
    # TODO: pre allocate memory for repeated_tokens
    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,
        BLOCK_SIZE=1024,
    )
    # repeated_tokens_ref = torch_scatter(
    #     input_tensor,
    #     token2index,
    # )
    # diff = torch.abs(repeated_tokens - repeated_tokens_ref)
    # print(f"repeated_tokens_diff: {diff}")
    # print(f"repeated_tokens: {repeated_tokens}")
    # print(f"before group gemm")
    # exit()
    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,
    )
    # TODO: group gemm0
    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[ (triton.cdiv(d_expert, BLOCK_N), triton.cdiv(max_m, 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,
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        BLOCK_K=BLOCK_K,
    )
    # ffn0_lhs_ref = torch_group_gemm(
    #     repeated_tokens,
    #     [weights[f'experts.{i}.0.weight'] for i in range(n_routed_experts)],
    #     expert_token_count_cumsum,
    #     d_expert,
    # )
    # diff = torch.abs(ffn0_lhs - ffn0_lhs_ref)
    # print(f"ffn0_lhs: {ffn0_lhs.shape}")
    # print(f"ffn0_lhs_ref: {ffn0_lhs_ref.shape}")
    # print(f"ffn0_lhs_diff: {diff}")
    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[ (triton.cdiv(d_expert, BLOCK_N), triton.cdiv(max_m, 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,
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        BLOCK_K=BLOCK_K,
    )
    
    # # TODO: SiLU activation
    ffn0_out = F.silu(ffn0_lhs) * ffn0_rhs # (batch_size * seq_len * n_experts_per_token, d_expert)
    # # TODO: group gemm1
    ffn1_out = torch.empty_like(repeated_tokens)
    group_gemm[ (triton.cdiv(d_hidden, BLOCK_N), triton.cdiv(max_m, 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,
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        BLOCK_K=BLOCK_K,
    )
    # # TODO: weighted sum + gather
    moe_out = torch.empty_like(input_tensor)
    BLOCK_SIZE = 1024
    weighted_sum_gather[ (triton.cdiv(d_hidden, 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,
    )
    # moe_out_ref = torch_weighted_sum_gather(
    #     ffn1_out,
    #     token2index,
    #     topk_scores,
    #     shared_expert_ffn1,
    # )
    # diff = torch.abs(moe_out - moe_out_ref)
    # print(f"moe_out_diff: {diff}")
    return moe_out.reshape(batch_size, seq_len, d_hidden) # (batch_size, seq_len, d_hidden)


def custom_kernel(data: input_t) -> output_t:
    """
    Submission template for DeepSeek-style Mixture of Experts using PyTorch.
    
    Args:
        data: Tuple of (input: torch.Tensor, weights: Dict[str, torch.Tensor], config: Dict)
            - input: Input tensor of shape [batch_size, seq_len, hidden_size]
            - weights: Dictionary containing model weights
            - config: Dictionary containing model configuration parameters
            
    Returns:
        Tuple containing:
            - output: Processed tensor [batch_size, seq_len, d_model]
            - aux_data: Dictionary with auxiliary data
    """
    input_tensor, weights, config = data
    output = moe_forward(input_tensor, weights, **config)
    return output


# if __name__ == "__main__":
# torch.cuda.set_device(1)
configs = [
    (7168, 2048, 4, 1, 4, 1, 512, 9371),
    (7168, 2048, 8, 1, 4, 2, 512, 2291),
    (7168, 2048, 8, 1, 4, 1, 8192, 81934),
]
# for config_dict in configs:
for dhidden, dexpert, nroutedexperts, nsharedexperts, nexpertspertoken, bs, seqlen, seed in configs:

    input_tensor, weights, config = my_generate_input(
        dhidden,
        dexpert,
        nroutedexperts,
        nsharedexperts,
        nexpertspertoken,
        bs,
        seqlen,
        seed
    )
    print("Config dictionary:", config)
    my_input = (input_tensor, weights, config)
    ref_input = deepcopy(my_input)
    my_moe_out = custom_kernel(my_input)
    ref_moe_out = my_ref_custom_kernel(ref_input)
    diff = torch.abs(my_moe_out - ref_moe_out)
    print("Reference MoE output shape:", ref_moe_out.shape)
    print("My MoE output shape:", my_moe_out.shape)
    print("Final difference:", diff)

python eval.py benchmark /tmp/tmpyxhjnhe2
/usr/local/lib/python3.10/dist-packages/triton/language/semantic.py:1687: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int32
  warnings.warn(
/usr/local/lib/python3.10/dist-packages/triton/language/semantic.py:1687: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int32
  warnings.warn(

Config dictionary: {'d_hidden': 7168, 'd_expert': 2048, 'n_routed_experts': 4, 'n_shared_experts': 1, 'n_experts_per_token': 4, 'batch_size': 1, 'seq_len': 512}
Reference MoE output shape: torch.Size([1, 512, 7168])
My MoE output shape: torch.Size([1, 512, 7168])
Final difference: tensor([[[0.0000, 0.0003, 0.0000,  ..., 0.0005, 0.0000, 0.0005],
         [0.0000, 0.0000, 0.0000,  ..., 0.0004, 0.0002, 0.0010],
         [0.0010, 0.0000, 0.0000,  ..., 0.0010, 0.0010, 0.0000],
         ...,
         [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0020, 0.0005],
         [0.0000, 0.0010, 0.0000,  ..., 0.0005, 0.0000, 0.0010],
         [0.0020, 0.0020, 0.0005,  ..., 0.0020, 0.0020, 0.0001]]],
       device='cuda:0', dtype=torch.float16, grad_fn=<AbsBackward0>)
Config dictionary: {'d_hidden': 7168, 'd_expert': 2048, 'n_routed_experts': 8, 'n_shared_experts': 1, 'n_experts_per_token': 4, 'batch_size': 2, 'seq_len': 512}
Reference MoE output shape: torch.Size([2, 512, 7168])
My MoE output shape: torch.Size([2, 512, 7168])
Final difference: tensor([[[0.0000e+00, 0.0000e+00, 1.9531e-03,  ..., 9.7656e-04,
          0.0000e+00, 2.4414e-04],
         [2.4414e-04, 0.0000e+00, 4.8828e-04,  ..., 4.8828e-04,
          0.0000e+00, 1.6403e-04],
         [0.0000e+00, 1.2207e-04, 2.4414e-04,  ..., 0.0000e+00,
          4.8828e-04, 1.9531e-03],
         ...,
         [0.0000e+00, 2.4414e-04, 0.0000e+00,  ..., 9.7656e-04,
          7.6294e-05, 2.4414e-04],
         [9.7656e-04, 0.0000e+00, 0.0000e+00,  ..., 1.9531e-03,
          4.8828e-04, 0.0000e+00],
         [0.0000e+00, 4.8828e-04, 9.7656e-04,  ..., 2.4414e-04,
          0.0000e+00, 2.4414e-04]],

        [[0.0000e+00, 0.0000e+00, 9.7656e-04,  ..., 0.0000e+00,
          1.9531e-03, 0.0000e+00],
         [3.6621e-04, 2.4414e-04, 0.0000e+00,  ..., 2.4414e-04,
          9.7656e-04, 4.8828e-04],
         [4.8828e-04, 0.0000e+00, 2.4414e-04,  ..., 0.0000e+00,
          9.7656e-04, 7.6294e-06],
         ...,
         [0.0000e+00, 0.0000e+00, 6.1035e-05,  ..., 0.0000e+00,
          0.0000e+00, 0.0000e+00],
         [0.0000e+00, 4.8828e-04, 0.0000e+00,  ..., 0.0000e+00,
          0.0000e+00, 0.0000e+00],
         [4.8828e-04, 0.0000e+00, 9.7656e-04,  ..., 0.0000e+00,
          0.0000e+00, 0.0000e+00]]], device='cuda:0', dtype=torch.float16,
       grad_fn=<AbsBackward0>)
Config dictionary: {'d_hidden': 7168, 'd_expert': 2048, 'n_routed_experts': 8, 'n_shared_experts': 1, 'n_experts_per_token': 4, 'batch_size': 1, 'seq_len': 8192}
Reference MoE output shape: torch.Size([1, 8192, 7168])
My MoE output shape: torch.Size([1, 8192, 7168])
Final difference: tensor([[[0.0005, 0.0005, 0.0010,  ..., 0.0000, 0.0001, 0.0020],
         [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0010],
         [0.0010, 0.0002, 0.0000,  ..., 0.0000, 0.0000, 0.0005],
         ...,
         [0.0005, 0.0005, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
         [0.0002, 0.0000, 0.0005,  ..., 0.0020, 0.0000, 0.0000],
         [0.0002, 0.0000, 0.0000,  ..., 0.0000, 0.0005, 0.0002]]],
       device='cuda:0', dtype=torch.float16, grad_fn=<AbsBackward0>)
