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

# Compile CUDA kernel at module load time
cuda_source = r"""
#include <torch/extension.h>
#include <cublas_v2.h>
#include <cuda_runtime.h>
#include <cfloat>

static cublasHandle_t get_cublas_handle() {
    static cublasHandle_t handle = nullptr;
    if (handle == nullptr) {
        cublasCreate(&handle);
    }
    return handle;
}

__global__ void top2_softmax_kernel(
    const float* __restrict__ grounded,  // (num_tokens, num_experts)
    float* __restrict__ output,          // (num_tokens, 2, 2)
    int num_tokens,
    int num_experts
) {
    int token_idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (token_idx >= num_tokens) return;
    
    const float* row = grounded + token_idx * num_experts;
    
    // Find top-2 values and indices
    float max1 = -FLT_MAX, max2 = -FLT_MAX;
    int idx1 = 0, idx2 = 0;
    
    // First pass: find max1
    for (int e = 0; e < num_experts; e++) {
        float val = row[e];
        if (val > max1) {
            max2 = max1;
            idx2 = idx1;
            max1 = val;
            idx1 = e;
        } else if (val > max2) {
            max2 = val;
            idx2 = e;
        }
    }
    
    // Softmax on top-2 values
    float max_val = fmaxf(max1, max2);
    float e1 = expf(max1 - max_val);
    float e2 = expf(max2 - max_val);
    float sum = e1 + e2;
    float w1 = e1 / sum;
    float w2 = e2 / sum;
    
    // Write output: (num_tokens, 2, 2)
    // output[token_idx, :, :] = [[idx1, w1], [idx2, w2]]
    output[token_idx * 4 + 0] = static_cast<float>(idx1);  // [token_idx, 0, 0]
    output[token_idx * 4 + 1] = w1;                        // [token_idx, 0, 1]
    output[token_idx * 4 + 2] = static_cast<float>(idx2);  // [token_idx, 1, 0]
    output[token_idx * 4 + 3] = w2;                        // [token_idx, 1, 1]
}

// Fused kernel: computes grounded = router_logits + alpha * (H @ G^T), then top2+softmax
// Each thread handles one token, computes dot product with all experts
__global__ void fused_top2_kernel(
    const float* __restrict__ token_hidden,  // (num_tokens, hidden_dim)
    const float* __restrict__ router_logits,  // (num_tokens, num_experts)
    const float* __restrict__ expert_ground,  // (num_experts, hidden_dim)
    float* __restrict__ output,               // (num_tokens, 2, 2)
    int num_tokens,
    int hidden_dim,
    int num_experts,
    float alpha
) {
    int token_idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (token_idx >= num_tokens) return;
    
    const float* hidden_row = token_hidden + token_idx * hidden_dim;
    const float* logits_row = router_logits + token_idx * num_experts;
    
    // Find top-2 grounded values and indices
    // grounded[i,e] = router_logits[i,e] + alpha * dot(hidden[i], ground[e])
    float max1 = -FLT_MAX, max2 = -FLT_MAX;
    int idx1 = 0, idx2 = 0;
    
    for (int e = 0; e < num_experts; e++) {
        // Compute dot product with expert e
        float dot = 0.0f;
        const float* ground_row = expert_ground + e * hidden_dim;
        for (int h = 0; h < hidden_dim; h++) {
            dot += hidden_row[h] * ground_row[h];
        }
        
        float grounded_val = logits_row[e] + alpha * dot;
        
        if (grounded_val > max1) {
            max2 = max1;
            idx2 = idx1;
            max1 = grounded_val;
            idx1 = e;
        } else if (grounded_val > max2) {
            max2 = grounded_val;
            idx2 = e;
        }
    }
    
    // Softmax on top-2 values (numerically stable)
    float max_val = fmaxf(max1, max2);
    float e1 = expf(max1 - max_val);
    float e2 = expf(max2 - max_val);
    float sum = e1 + e2;
    float w1 = e1 / sum;
    float w2 = e2 / sum;
    
    // Write output
    output[token_idx * 4 + 0] = static_cast<float>(idx1);
    output[token_idx * 4 + 1] = w1;
    output[token_idx * 4 + 2] = static_cast<float>(idx2);
    output[token_idx * 4 + 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}, 
                               torch::dtype(torch::kFloat32)
                               .device(token_hidden.device()));
    
    int threads = 256;
    int blocks = (num_tokens + threads - 1) / threads;
    
    fused_top2_kernel<<<blocks, threads>>>(
        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;
}
"""

# Load the CUDA extension
grounded_top2_ext = load_inline(
    name='grounded_top2_ext',
    cpp_sources="torch::Tensor grounded_top2_cuda(torch::Tensor, torch::Tensor, torch::Tensor, float);",
    cuda_sources=cuda_source,
    functions=['grounded_top2_cuda'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    verbose=False
)


class ModelNew(nn.Module):
    """
    Grounded top-2 MoE router with custom CUDA kernel.
    """
    
    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 grounded_top2_ext.grounded_top2_cuda(
            token_hidden, router_logits, expert_ground, alpha
        )