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

_cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>

__inline__ __device__ float warp_sum(float v) {
    unsigned mask = 0xffffffffu;
    v += __shfl_down_sync(mask, v, 16);
    v += __shfl_down_sync(mask, v, 8);
    v += __shfl_down_sync(mask, v, 4);
    v += __shfl_down_sync(mask, v, 2);
    v += __shfl_down_sync(mask, v, 1);
    return v;
}

__global__ void grounded_router_kernel(const float* __restrict__ token_hidden,
                                      const float* __restrict__ router_logits,
                                      const float* __restrict__ expert_ground,
                                      float* __restrict__ out,
                                      int N, float alpha) {
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    int token = blockIdx.x * 16 + warp;
    if (token >= N) return;

    const float* t = token_hidden + ((long long)token) * 128;
    float t0 = t[lane];
    float t1 = t[lane + 32];
    float t2 = t[lane + 64];
    float t3 = t[lane + 96];

    float best1 = -FLT_MAX;
    float best2 = -FLT_MAX;
    int idx1 = 0;
    int idx2 = 1;

    #pragma unroll
    for (int ex = 0; ex < 16; ++ex) {
        const float* g = expert_ground + ex * 128;
        float s = 0.0f;
        s = fmaf(t0, g[lane], s);
        s = fmaf(t1, g[lane + 32], s);
        s = fmaf(t2, g[lane + 64], s);
        s = fmaf(t3, g[lane + 96], s);
        s = warp_sum(s);
        if (lane == 0) {
            float v = fmaf(alpha, s, router_logits[((long long)token) * 16 + ex]);
            if (v > best1) {
                best2 = best1; idx2 = idx1;
                best1 = v; idx1 = ex;
            } else if (v > best2) {
                best2 = v; idx2 = ex;
            }
        }
    }

    if (lane == 0) {
        float ee = __expf(best2 - best1);
        float w1 = 1.0f / (1.0f + ee);
        long long base = ((long long)token) * 4;
        out[base + 0] = (float)idx1;
        out[base + 1] = w1;
        out[base + 2] = (float)idx2;
        out[base + 3] = ee * w1;
    }
}

torch::Tensor grounded_router_forward(torch::Tensor token_hidden, torch::Tensor router_logits, torch::Tensor expert_ground, double alpha) {
    int N = (int)token_hidden.size(0);
    auto out = torch::empty({N, 2, 2}, token_hidden.options());
    grounded_router_kernel<<<(N + 15) / 16, 512>>>(token_hidden.data_ptr<float>(), router_logits.data_ptr<float>(), expert_ground.data_ptr<float>(), out.data_ptr<float>(), N, (float)alpha);
    return out;
}
"""

_ext = load_inline(name="grounded_router_ext_final", cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=["grounded_router_forward"], extra_cflags=["-O3"], extra_cuda_cflags=["-O3"], 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_router_forward(token_hidden, router_logits, expert_ground, alpha)