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

_cpp_src = r'''
#include <torch/extension.h>

torch::Tensor combine_cuda(torch::Tensor expert_hidden, torch::Tensor token_idx, torch::Tensor gates, int64_t num_tokens);

torch::Tensor combine_forward(torch::Tensor expert_hidden, torch::Tensor token_idx, torch::Tensor gates, int64_t num_tokens) {
    return combine_cuda(expert_hidden, token_idx, gates, num_tokens);
}
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void combine_h128_f2_kernel(const float* __restrict__ expert_hidden,
                                       const float* __restrict__ gates,
                                       float* __restrict__ out,
                                       int num_tokens) {
    int warp = threadIdx.x >> 5;          // 0..3
    int lane = threadIdx.x & 31;
    int token = (blockIdx.x << 2) + warp;
    if (token >= num_tokens) return;

    int row0 = token << 1;
    int in0 = row0 << 7;                 // row0 * 128
    int in1 = in0 + 128;
    int out_base = token << 7;
    float g0 = gates[row0];
    float g1 = gates[row0 + 1];

    int h = lane;
    out[out_base + h] = g0 * expert_hidden[in0 + h] + g1 * expert_hidden[in1 + h];
    h += 32;
    out[out_base + h] = g0 * expert_hidden[in0 + h] + g1 * expert_hidden[in1 + h];
    h += 32;
    out[out_base + h] = g0 * expert_hidden[in0 + h] + g1 * expert_hidden[in1 + h];
    h += 32;
    out[out_base + h] = g0 * expert_hidden[in0 + h] + g1 * expert_hidden[in1 + h];
}

__global__ void combine_generic_sorted_kernel(const float* __restrict__ expert_hidden,
                                      const float* __restrict__ gates,
                                      float* __restrict__ out,
                                      int hidden_dim, int num_tokens, int fanout) {
    int token = blockIdx.x;
    int h = threadIdx.x;
    if (h >= hidden_dim || token >= num_tokens) return;
    int base = token * fanout;
    float acc = 0.0f;
    for (int f = 0; f < fanout; ++f) {
        int row = base + f;
        acc += gates[row] * expert_hidden[row * hidden_dim + h];
    }
    out[token * hidden_dim + h] = acc;
}

torch::Tensor combine_cuda(torch::Tensor expert_hidden, torch::Tensor token_idx, torch::Tensor gates, int64_t num_tokens64) {
    const int rows = (int)expert_hidden.size(0);
    const int hidden_dim = (int)expert_hidden.size(1);
    const int num_tokens = (int)num_tokens64;
    const int fanout = rows / num_tokens;
    auto out = torch::empty({num_tokens64, (int64_t)hidden_dim}, expert_hidden.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (hidden_dim == 128 && fanout == 2) {
        int blocks = (num_tokens + 3) >> 2;
        combine_h128_f2_kernel<<<blocks, 128, 0, stream>>>(expert_hidden.data_ptr<float>(), gates.data_ptr<float>(), out.data_ptr<float>(), num_tokens);
    } else {
        combine_generic_sorted_kernel<<<num_tokens, 128, 0, stream>>>(expert_hidden.data_ptr<float>(), gates.data_ptr<float>(), out.data_ptr<float>(), hidden_dim, num_tokens, fanout);
    }
    return out;
}
'''

_ext = load_inline(name='weighted_expert_combine_ext_v4',
                   cpp_sources=_cpp_src,
                   cuda_sources=_cuda_src,
                   functions=['combine_forward'],
                   extra_cuda_cflags=['-O3'],
                   extra_cflags=['-O3'],
                   verbose=False)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, expert_hidden: torch.Tensor, token_idx: torch.Tensor, gates: torch.Tensor, num_tokens: int) -> torch.Tensor:
        return _ext.combine_forward(expert_hidden, token_idx, gates, num_tokens)