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'''
#include <torch/extension.h>
torch::Tensor moe_forward_cuda(torch::Tensor token_hidden, torch::Tensor router_logits, torch::Tensor expert_ground, double alpha);
torch::Tensor moe_forward(torch::Tensor token_hidden, torch::Tensor router_logits, torch::Tensor expert_ground, double alpha) {
    return moe_forward_cuda(token_hidden, router_logits, expert_ground, alpha);
}
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#define HIDDEN_DIM 128
#define NUM_EXPERTS 16
#define WARPS_PER_BLOCK 8
__device__ __forceinline__ float warp_reduce_sum(float v){
    v += __shfl_down_sync(0xffffffff, v, 16);
    v += __shfl_down_sync(0xffffffff, v, 8);
    v += __shfl_down_sync(0xffffffff, v, 4);
    v += __shfl_down_sync(0xffffffff, v, 2);
    v += __shfl_down_sync(0xffffffff, v, 1);
    return v;
}
__global__ __launch_bounds__(WARPS_PER_BLOCK*32, 3)
void moe_kernel(const float* __restrict__ token_hidden, const float* __restrict__ router_logits,
                const float* __restrict__ expert_ground, float alpha, float* __restrict__ out, int n_tokens){
    int tid = threadIdx.x;
    int lane = tid & 31;
    int warp = tid >> 5;
    int tok = blockIdx.x * WARPS_PER_BLOCK + warp;
    if (tok >= n_tokens) return;
    const float* th = token_hidden + tok * HIDDEN_DIM;
    float x0 = th[lane];
    float x1 = th[lane + 32];
    float x2 = th[lane + 64];
    float x3 = th[lane + 96];
    float best1=-3.402823466e38f, best2=-3.402823466e38f;
    int idx1=0, idx2=1;
    const float* rl = router_logits + tok * NUM_EXPERTS;
    #pragma unroll
    for(int e=0;e<NUM_EXPERTS;++e){
        const float* eg = expert_ground + e * HIDDEN_DIM;
        float s = 0.0f;
        s = fmaf(x0, eg[lane], s);
        s = fmaf(x1, eg[lane+32], s);
        s = fmaf(x2, eg[lane+64], s);
        s = fmaf(x3, eg[lane+96], s);
        s = warp_reduce_sum(s);
        if(lane == 0){
            float v = rl[e] + alpha * s;
            if(v > best1){ best2=best1; idx2=idx1; best1=v; idx1=e; }
            else if(v > best2){ best2=v; idx2=e; }
        }
    }
    if(lane == 0){
        float r = __expf(best2 - best1);
        float inv = 1.0f / (1.0f + r);
        int base = tok * 4;
        out[base] = (float)idx1;
        out[base+1] = inv;
        out[base+2] = (float)idx2;
        out[base+3] = r * inv;
    }
}
torch::Tensor moe_forward_cuda(torch::Tensor token_hidden, torch::Tensor router_logits, torch::Tensor expert_ground, double alpha){
    auto out = torch::empty({token_hidden.size(0), 2, 2}, token_hidden.options());
    int n_tokens = (int)token_hidden.size(0);
    int blocks = (n_tokens + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    moe_kernel<<<blocks, WARPS_PER_BLOCK*32, 0, stream>>>(token_hidden.data_ptr<float>(), router_logits.data_ptr<float>(), expert_ground.data_ptr<float>(), (float)alpha, out.data_ptr<float>(), n_tokens);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
'''
_ext = load_inline(name='grounded_top2_moe_ext_final', cpp_sources=[_cpp_src], cuda_sources=[_cuda_src], functions=['moe_forward'], extra_cuda_cflags=['-O3','--use_fast_math'], extra_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.moe_forward(token_hidden, router_logits, expert_ground, alpha)