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
from torch.utils.cpp_extension import load_inline

cuda_source = r'''
#include <torch/extension.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 num_tokens,
    const int hidden_dim,
    const int num_experts,
    const float alpha
) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= num_tokens) return;

    float scores[16];  // assume num_experts <=16
    const float* token = token_hidden + tid * hidden_dim;
    for (int e = 0; e < num_experts; ++e) {
        float dot = 0.0f;
        const float* exp_g = expert_ground + e * hidden_dim;
        for (int h = 0; h < hidden_dim; ++h) {
            dot += token[h] * exp_g[h];
        }
        scores[e] = router_logits[tid * num_experts + e] + alpha * dot;
    }

    // Find top-2 indices and values
    int idx0 = 0, idx1 = 1;
    float val0 = scores[0], val1 = scores[1];
    if (val1 > val0) {
        float tmp = val0; val0 = val1; val1 = tmp;
        int itmp = idx0; idx0 = idx1; idx1 = itmp;
    }
    for (int i = 2; i < num_experts; ++i) {
        if (scores[i] > val0) {
            val1 = val0; idx1 = idx0;
            val0 = scores[i]; idx0 = i;
        } else if (scores[i] > val1) {
            val1 = scores[i]; idx1 = i;
        }
    }

    // softmax on top_vals
    float maxv = fmaxf(val0, val1);
    float sum = expf(val0 - maxv) + expf(val1 - maxv);
    float w0 = expf(val0 - maxv) / sum;
    float w1 = expf(val1 - maxv) / sum;

    // output: [tid, 0, 0] = idx0, [tid, 0, 1] = w0, output[tid,1,0]=idx1, output[tid,1,1]=w1
    float* out = output + tid * 4;
    out[0] = (float)idx0;
    out[1] = w0;
    out[2] = (float)idx1;
    out[3] = w1;
}

at::Tensor grounded_top2_forward(
    at::Tensor token_hidden,
    at::Tensor router_logits,
    at::Tensor expert_ground,
    float alpha
) {
    TORCH_CHECK(token_hidden.dim() == 2, "token_hidden must be 2D");
    int num_tokens = token_hidden.size(0);
    int hidden_dim = token_hidden.size(1);
    int num_experts = router_logits.size(1);

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

    dim3 block(256);
    dim3 grid((num_tokens + 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>(),
        num_tokens,
        hidden_dim,
        num_experts,
        alpha
    );

    return output;
}
'''

cpp_source = '''
at::Tensor grounded_top2_forward(at::Tensor token_hidden, at::Tensor router_logits, at::Tensor expert_ground, float alpha);
'''

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

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

    def forward(self, token_hidden, router_logits, expert_ground, alpha):
        return self.module.grounded_top2_forward(token_hidden, router_logits, expert_ground, alpha)