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

cuda_source = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>

__global__ void grounded_top2_kernel(
    const float* __restrict__ token_hidden,
    const float* __restrict__ router_logits,
    const float* __restrict__ expert_ground,
    float* __restrict__ output,
    const int N,
    const int H,
    const int E,
    const float alpha
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N) return;

    const float* token = token_hidden + idx * H;
    const float* router = router_logits + idx * E;

    float scores[16];
    #pragma unroll
    for (int e = 0; e < E; ++e) {
        float sum = 0.0f;
        const float* expert = expert_ground + e * H;
        #pragma unroll
        for (int h = 0; h < H; ++h) {
            sum += token[h] * expert[h];
        }
        scores[e] = router[e] + alpha * sum;
    }

    // Find top-2
    float max1 = -FLT_MAX, max2 = -FLT_MAX;
    int idx1 = -1, idx2 = -1;
    #pragma unroll
    for (int e = 0; e < E; ++e) {
        float v = scores[e];
        if (v > max1) {
            max2 = max1; idx2 = idx1;
            max1 = v; idx1 = e;
        } else if (v > max2) {
            max2 = v; idx2 = e;
        }
    }

    float v1 = max1, v2 = max2;
    float exp1 = expf(v1 - v1);
    float exp2 = expf(v2 - v1);
    float sum_exp = exp1 + exp2;
    float w1 = exp1 / sum_exp;
    float w2 = exp2 / sum_exp;

    // Output layout: [N, 2, 2] -> for each token: [[idx1, w1], [idx2, w2]]
    float* out = output + idx * 4;
    out[0] = (float)idx1;
    out[1] = w1;
    out[2] = (float)idx2;
    out[3] = w2;
}

torch::Tensor grounded_top2_cuda(
    torch::Tensor token_hidden,
    torch::Tensor router_logits,
    torch::Tensor expert_ground,
    double alpha
) {
    int N = token_hidden.size(0);
    int H = token_hidden.size(1);
    int E = router_logits.size(1);

    auto options = torch::TensorOptions().dtype(torch::kFloat32).device(token_hidden.device());
    auto output = torch::empty({N, 2, 2}, options);

    dim3 block(256);
    dim3 grid((N + 255) / 256);

    grounded_top2_kernel<<<grid, block>>>(
        token_hidden.data_ptr<float>(),
        router_logits.data_ptr<float>(),
        expert_ground.data_ptr<float>(),
        output.data_ptr<float>(),
        N, H, E, (float)alpha
    );

    return output;
}
'''

cpp_source = r'''
torch::Tensor grounded_top2_cuda(torch::Tensor token_hidden, torch::Tensor router_logits, torch::Tensor expert_ground, double alpha);
'''

module = load_inline(
    name='grounded_top2',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['grounded_top2_cuda'],
    verbose=False
)

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

    def forward(
        self,
        token_hidden: torch.Tensor,
        router_logits: torch.Tensor,
        expert_ground: torch.Tensor,
        alpha: float,
    ) -> torch.Tensor:
        return self.module.grounded_top2_cuda(token_hidden, router_logits, expert_ground, alpha)
```