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_forward_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_forward_cuda(expert_hidden, token_idx, gates, num_tokens);
}
'''

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

__global__ void combine_pair128_kernel(const float* __restrict__ expert,
                                       const float* __restrict__ gates,
                                       float* __restrict__ out,
                                       int num_tokens) {
    int warp = threadIdx.x >> 5;
    int lane = threadIdx.x & 31;
    int t = blockIdx.x * 4 + warp;
    if (t >= num_tokens) return;

    int row0 = t * 2;
    int row1 = row0 + 1;
    float g0 = gates[row0];
    float g1 = gates[row1];

    int col = lane * 4;
    const float4 a = *reinterpret_cast<const float4*>(expert + row0 * 128 + col);
    const float4 b = *reinterpret_cast<const float4*>(expert + row1 * 128 + col);

    float4 r;
    r.x = g0 * a.x + g1 * b.x;
    r.y = g0 * a.y + g1 * b.y;
    r.z = g0 * a.z + g1 * b.z;
    r.w = g0 * a.w + g1 * b.w;

    *reinterpret_cast<float4*>(out + t * 128 + col) = r;
}

__global__ void combine_atomic_kernel(const float* __restrict__ expert,
                                      const int* __restrict__ token_idx,
                                      const float* __restrict__ gates,
                                      float* __restrict__ out,
                                      int rows, int H) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int n = rows * H;
    if (idx >= n) return;
    int h = idx % H;
    int r = idx / H;
    int tok = token_idx[r];
    atomicAdd(out + tok * H + h, gates[r] * expert[idx]);
}

torch::Tensor combine_forward_cuda(torch::Tensor expert_hidden, torch::Tensor token_idx, torch::Tensor gates, int64_t num_tokens64) {
    int64_t rows64 = expert_hidden.size(0);
    int64_t H64 = expert_hidden.size(1);
    auto out = torch::empty({num_tokens64, H64}, expert_hidden.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const float* expert = expert_hidden.data_ptr<float>();
    const float* g = gates.data_ptr<float>();
    float* o = out.data_ptr<float>();

    int num_tokens = (int)num_tokens64;
    if (H64 == 128 && rows64 == num_tokens64 * 2) {
        combine_pair128_kernel<<<(num_tokens + 3) / 4, 128, 0, stream>>>(expert, g, o, num_tokens);
    } else {
        out.zero_();
        const int* tidx = token_idx.data_ptr<int>();
        int rows = (int)rows64;
        int H = (int)H64;
        int total = rows * H;
        combine_atomic_kernel<<<(total + 255) / 256, 256, 0, stream>>>(expert, tidx, g, o, rows, H);
    }
    return out;
}
'''

_ext = load_inline(
    name='expert_combine_ext_final',
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=['combine_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    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)