import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Weighted expert-output combine back into token-major order.

    Expert outputs are scattered back to token positions and scaled by the
    routing weight associated with each expert contribution.
    """

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

    def forward(
        self,
        expert_hidden: torch.Tensor,
        token_idx: torch.Tensor,
        gates: torch.Tensor,
        num_tokens: int,
    ) -> torch.Tensor:
        out = torch.zeros(num_tokens, expert_hidden.shape[1], dtype=expert_hidden.dtype, device=expert_hidden.device)
        for row in range(expert_hidden.shape[0]):
            token = int(token_idx[row].item())
            out[token] += gates[row] * expert_hidden[row]
        return out


num_tokens = 1536
hidden_dim = 128
fanout = 2


def get_inputs():
    return popcorn_pri.moe_combine_inputs(num_tokens, hidden_dim, fanout)




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.h>
#include <cuda_runtime.h>
#include <cstdint>

__global__ void scatter_add_kernel(
    const float* __restrict__ expert_hidden,
    const int32_t* __restrict__ token_idx,
    const float* __restrict__ gates,
    float* __restrict__ out,
    int num_expert_rows,
    int hidden_dim
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = num_expert_rows * hidden_dim;
    if (idx >= total) return;

    int row = idx / hidden_dim;
    int d = idx % hidden_dim;

    int32_t token = token_idx[row];
    float gate = gates[row];
    float val = gate * expert_hidden[idx];
    atomicAdd(&out[token * hidden_dim + d], val);
}

torch::Tensor scatter_add_cuda(
    torch::Tensor expert_hidden,
    torch::Tensor token_idx,
    torch::Tensor gates,
    int64_t num_tokens
) {
    int num_expert_rows = expert_hidden.size(0);
    int hidden_dim = expert_hidden.size(1);

    auto out = torch::zeros({num_tokens, hidden_dim}, expert_hidden.options());

    token_idx = token_idx.to(torch::kInt32).contiguous();

    int total = num_expert_rows * hidden_dim;
    const int threads = 256;
    const int blocks = (total + threads - 1) / threads;

    scatter_add_kernel<<<blocks, threads>>>(
        expert_hidden.data_ptr<float>(),
        token_idx.data_ptr<int32_t>(),
        gates.data_ptr<float>(),
        out.data_ptr<float>(),
        num_expert_rows,
        hidden_dim
    );

    return out;
}
"""

cpp_source = """
torch::Tensor scatter_add_cuda(torch::Tensor expert_hidden, torch::Tensor token_idx, torch::Tensor gates, int64_t num_tokens);
"""

module = load_inline(
    name="scatter_add_module",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["scatter_add_cuda"],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.module = module

    def forward(self, expert_hidden, token_idx, gates, num_tokens):
        return self.module.scatter_add_cuda(expert_hidden, token_idx, gates, num_tokens)
```