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 = '''
#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,
    int num_tokens,
    int hidden_dim,
    int num_experts,
    float alpha
) {
    int token_idx = blockIdx.x;
    if (token_idx >= num_tokens) return;

    extern __shared__ float s_grounded[];
    
    const float* token = token_hidden + token_idx * hidden_dim;
    
    for (int e = threadIdx.x; e < num_experts; e += blockDim.x) {
        float dot = 0.0f;
        const float* expert = expert_ground + e * hidden_dim;
        for (int h = 0; h < hidden_dim; h++) {
            dot += token[h] * expert[h];
        }
        float router = router_logits[token_idx * num_experts + e];
        s_grounded[e] = router + alpha * dot;
    }
    __syncthreads();

    if (threadIdx.x == 0) {
        // Find top-2
        float max1 = -FLT_MAX, max2 = -FLT_MAX;
        int idx1 = -1, idx2 = -1;
        for (int e = 0; e < num_experts; e++) {
            float val = s_grounded[e];
            if (val > max1) {
                max2 = max1; idx2 = idx1;
                max1 = val; idx1 = e;
            } else if (val > max2) {
                max2 = val; idx2 = e;
            }
        }
        // Softmax on top2
        float m = fmaxf(max1, max2);
        float sum_exp = expf(max1 - m) + expf(max2 - m);
        float w1 = expf(max1 - m) / sum_exp;
        float w2 = expf(max2 - m) / sum_exp;
        
        // Output: [idx1, w1, idx2, w2] but stacked as [[idx1,w1],[idx2,w2]]
        float* out = output + token_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,
    float alpha
) {
    int num_tokens = token_hidden.size(0);
    int hidden_dim = token_hidden.size(1);
    int num_experts = router_logits.size(1);
    
    auto output = torch::empty({num_tokens, 2, 2}, token_hidden.options());
    
    int threads = 256;
    int smem = num_experts * sizeof(float);
    grounded_top2_kernel<<<num_tokens, threads, smem>>>(
        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 = '''
torch::Tensor grounded_top2_cuda(torch::Tensor token_hidden, torch::Tensor router_logits, torch::Tensor expert_ground, float 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, router_logits, expert_ground, alpha):
        return self.module.grounded_top2_cuda(token_hidden, router_logits, expert_ground, alpha)