import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Grounded top-2 MoE router.

    Router logits are biased by token-expert similarity against learned
    grounding embeddings, then normalized over the selected top-2 experts.
    """

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

    def forward(
        self,
        token_hidden: torch.Tensor,
        router_logits: torch.Tensor,
        expert_ground: torch.Tensor,
        alpha: float,
    ) -> torch.Tensor:
        grounded = router_logits + alpha * (token_hidden @ expert_ground.t())
        top_vals, top_idx = torch.topk(grounded, k=2, dim=-1)
        top_weights = torch.softmax(top_vals, dim=-1)
        return torch.stack((top_idx.to(torch.float32), top_weights), dim=-1)


num_tokens = 4096
hidden_dim = 128
num_experts = 16
alpha = 0.35


def get_inputs():
    p = popcorn_pri
    nt = p.jitter_int(num_tokens)
    hd = p.jitter_int(hidden_dim, align=16)
    ne = p.jitter_int(num_experts, minimum=2)
    token_hidden = torch.randn(nt, hd, dtype=torch.float32)
    router_logits = torch.randn(nt, ne, dtype=torch.float32)
    expert_ground = torch.randn(ne, hd, dtype=torch.float32)
    return [token_hidden, router_logits, expert_ground, alpha]



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 grounded_top2_forward(torch::Tensor token_hidden,
                                    torch::Tensor router_logits,
                                    torch::Tensor expert_ground,
                                    double alpha);
"""

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

__global__ __launch_bounds__(256, 2)
void grounded_top2_kernel(const float* __restrict__ token_hidden,
                          const float* __restrict__ router_logits,
                          const float* __restrict__ expert_ground,
                          float* __restrict__ out,
                          int T,
                          float alpha) {
    const int tok = blockIdx.x;
    if (tok >= T) return;

    const int tid = threadIdx.x;       // 0..255
    const int expert = tid >> 4;       // 0..15
    const int lane = tid & 15;         // 0..15 within expert group

    __shared__ float vals[16];

    const float* th = token_hidden + ((long long)tok) * 128;
    const float* eg = expert_ground + expert * 128;

    float sum = 0.0f;
#pragma unroll
    for (int k = 0; k < 8; ++k) {
        int h = lane + (k << 4);
        sum = fmaf(th[h], eg[h], sum);
    }

    // Reduce the 16 lanes assigned to each expert.  Masks are warp-local.
    const unsigned mask = (threadIdx.x & 16) ? 0xffff0000u : 0x0000ffffu;
#pragma unroll
    for (int off = 8; off > 0; off >>= 1) {
        sum += __shfl_down_sync(mask, sum, off, 16);
    }

    if (lane == 0) {
        vals[expert] = router_logits[((long long)tok) * 16 + expert] + alpha * sum;
    }
    __syncthreads();

    if (tid == 0) {
        float best1 = -CUDART_INF_F;
        float best2 = -CUDART_INF_F;
        int idx1 = 0;
        int idx2 = 0;

#pragma unroll
        for (int e = 0; e < 16; ++e) {
            float v = vals[e];
            if (v > best1) {
                best2 = best1;
                idx2 = idx1;
                best1 = v;
                idx1 = e;
            } else if (v > best2) {
                best2 = v;
                idx2 = e;
            }
        }

        float r = expf(best2 - best1);
        float w1 = 1.0f / (1.0f + r);
        float w2 = r * w1;

        long long base = ((long long)tok) * 4;
        out[base + 0] = (float)idx1;
        out[base + 1] = w1;
        out[base + 2] = (float)idx2;
        out[base + 3] = w2;
    }
}

torch::Tensor grounded_top2_forward(torch::Tensor token_hidden,
                                    torch::Tensor router_logits,
                                    torch::Tensor expert_ground,
                                    double alpha) {
    const int T = (int)token_hidden.size(0);
    auto out = torch::empty({T, 2, 2}, token_hidden.options());

    const dim3 block(256);
    const dim3 grid(T);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    grounded_top2_kernel<<<grid, block, 0, stream>>>(
        token_hidden.data_ptr<float>(),
        router_logits.data_ptr<float>(),
        expert_ground.data_ptr<float>(),
        out.data_ptr<float>(),
        T,
        (float)alpha);
    return out;
}
"""

_ext = load_inline(
    name="grounded_top2_moe_router_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["grounded_top2_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    with_cuda=True,
    verbose=False,
)


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

    def forward(self, token_hidden: torch.Tensor, router_logits: torch.Tensor, expert_ground: torch.Tensor, alpha: float) -> torch.Tensor:
        return _ext.grounded_top2_forward(token_hidden, router_logits, expert_ground, float(alpha))