import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Pack tokens into expert-major order for MoE execution.

    Each token is assigned to one expert and one slot within that expert's
    buffer. The kernel writes token activations into the packed dispatch layout.
    """

    def __init__(self):
        super().__init__()

    def forward(
        self,
        token_hidden: torch.Tensor,
        expert_idx: torch.Tensor,
        slot_idx: torch.Tensor,
        expert_offsets: torch.Tensor,
    ) -> torch.Tensor:
        num_rows = int(expert_offsets[-1].item())
        out = torch.zeros(num_rows, token_hidden.shape[1], dtype=token_hidden.dtype, device=token_hidden.device)
        for token in range(token_hidden.shape[0]):
            expert = int(expert_idx[token].item())
            row = int(expert_offsets[expert].item() + slot_idx[token].item())
            out[row] = token_hidden[token]
        return out


num_tokens = 2048
hidden_dim = 128
num_experts = 16


def get_inputs():
    return popcorn_pri.moe_dispatch_inputs(num_tokens, hidden_dim, num_experts)




def get_init_inputs():
    return []

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cuda_source = '''
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cstdint>

__global__ void pack_tokens_kernel(
    const float* __restrict__ token_hidden,
    const int32_t* __restrict__ expert_idx,
    const int32_t* __restrict__ slot_idx,
    const int32_t* __restrict__ expert_offsets,
    float* __restrict__ out,
    int num_tokens,
    int hidden_dim
) {
    int token = blockIdx.x * blockDim.x + threadIdx.x;
    if (token >= num_tokens) return;
    
    int32_t expert = expert_idx[token];
    int32_t slot = slot_idx[token];
    int32_t row = expert_offsets[expert] + slot;
    
    for (int d = 0; d < hidden_dim; d++) {
        out[row * hidden_dim + d] = token_hidden[token * hidden_dim + d];
    }
}

torch::Tensor pack_tokens_cuda(
    torch::Tensor token_hidden,
    torch::Tensor expert_idx,
    torch::Tensor slot_idx,
    torch::Tensor expert_offsets
) {
    int num_tokens = token_hidden.size(0);
    int hidden_dim = token_hidden.size(1);
    int num_experts_plus_one = expert_offsets.size(0);
    int32_t num_rows = expert_offsets.data_ptr<int32_t>()[num_experts_plus_one - 1];
    
    auto out = torch::zeros({num_rows, hidden_dim}, token_hidden.options());
    
    dim3 block(256);
    dim3 grid((num_tokens + 255) / 256);
    
    pack_tokens_kernel<<<grid, block>>>(
        token_hidden.data_ptr<float>(),
        expert_idx.data_ptr<int32_t>(),
        slot_idx.data_ptr<int32_t>(),
        expert_offsets.data_ptr<int32_t>(),
        out.data_ptr<float>(),
        num_tokens,
        hidden_dim
    );
    
    return out;
}
'''

cpp_source = '''
torch::Tensor pack_tokens_cuda(torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor);
'''

module = load_inline(
    name='pack_tokens',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['pack_tokens_cuda'],
    verbose=True
)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.module = module
    
    def forward(self, token_hidden, expert_idx, slot_idx, expert_offsets):
        return self.module.pack_tokens_cuda(token_hidden, expert_idx, slot_idx, expert_offsets)