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>

__global__ void scatter_add_gated_kernel(
    const float* __restrict__ expert_hidden,
    const int* __restrict__ token_idx,
    const float* __restrict__ gates,
    float* __restrict__ out,
    int num_experts,
    int hidden_dim
) {
    int row = blockIdx.x;
    int col = threadIdx.x;
    
    if (row < num_experts && col < hidden_dim) {
        int token = token_idx[row];
        float gate = gates[row];
        float val = expert_hidden[row * hidden_dim + col];
        atomicAdd(&out[token * hidden_dim + col], gate * val);
    }
}

torch::Tensor scatter_add_gated_cuda(
    torch::Tensor expert_hidden,
    torch::Tensor token_idx,
    torch::Tensor gates
) {
    int num_experts = expert_hidden.size(0);
    int hidden_dim = expert_hidden.size(1);
    
    // Compute num_tokens as max token_idx + 1
    auto token_idx_accessor = token_idx.accessor<int, 1>();
    int n_tokens = 0;
    for (int i = 0; i < num_experts; i++) {
        n_tokens = std::max(n_tokens, token_idx_accessor[i] + 1);
    }
    
    auto out = torch::zeros({n_tokens, hidden_dim}, expert_hidden.options());
    
    int threads = (hidden_dim + 31) / 32 * 32;
    int blocks = num_experts;
    
    scatter_add_gated_kernel<<<blocks, threads>>>(
        expert_hidden.data_ptr<float>(),
        token_idx.data_ptr<int>(),
        gates.data_ptr<float>(),
        out.data_ptr<float>(),
        num_experts,
        hidden_dim
    );
    
    return out;
}
"""

cpp_source = "torch::Tensor scatter_add_gated_cuda(torch::Tensor expert_hidden, torch::Tensor token_idx, torch::Tensor gates);"

scatter_ext = load_inline(
    name='scatter_add_gated',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['scatter_add_gated_cuda'],
    extra_cuda_cflags=['-O3'],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, expert_hidden, token_idx, gates, num_tokens):
        return scatter_ext.scatter_add_gated_cuda(expert_hidden, token_idx, gates)