# /// script
# dependencies = [
#   "kernels @ git+https://github.com/huggingface/kernels.git@32ec496c5a3169ef02e06a21fb864e0412",
#   "numpy",
#   "torch",
# ]
# ///
import torch
from torch import nn

from kernels import use_kernel_forward_from_hub
from kernels import (
    kernelize,
    register_kernel_mapping,
    use_kernel_mapping,
    LayerRepository,
)

# Make reproducible
torch.manual_seed(42)
torch.cuda.manual_seed(42)


@use_kernel_forward_from_hub("MoE")
class MoE(nn.Module):

    def forward(
        self,
        hidden_states: torch.Tensor,
        router_idx: torch.Tensor,
        router_wt: torch.Tensor,
        alpha: float,
        gate_up_weights: torch.Tensor,
        gate_up_bias: torch.Tensor,
        down_weights: torch.Tensor,
        down_bias: torch.Tensor,
    ) -> torch.Tensor:
        return 0.0


kernel_layer_mapping = {
    "MoE": {
        "cuda": LayerRepository(
            repo_id="kernels-community/triton-moe",
            layer_name="MoE",
            revision="main",
        )
    }
}

layer = MoE()

register_kernel_mapping(kernel_layer_mapping)

with use_kernel_mapping(kernel_layer_mapping):
    layer = kernelize(
        layer,
        needs_backward=True,
        # needs_torch_compile=True,
        use_fallback=False,
        device="cuda",
    )

num_experts = 4
batch_tokens = 8
hidden_size = 16
expert_dim = hidden_size // num_experts
topk = 2


# Example usage
hidden_states = torch.randn(
    batch_tokens,
    hidden_size,
    device="cuda",
    dtype=torch.float32,
    requires_grad=True,
)
router_idx = torch.randint(0, num_experts, (batch_tokens, topk), device="cuda")
router_wt = torch.rand(batch_tokens, topk, device="cuda")
router_wt = router_wt / router_wt.sum(dim=-1, keepdim=True)  # Normalize

# Initialize parameters (shared between both implementations)
gate_up_proj = nn.Parameter(
    torch.randn(num_experts, hidden_size, 2 * expert_dim, device="cuda")
)
gate_up_proj_bias = nn.Parameter(
    torch.randn(num_experts, 2 * expert_dim, device="cuda")
)
down_proj = nn.Parameter(
    torch.randn(num_experts, expert_dim, hidden_size, device="cuda")
)
down_proj_bias = nn.Parameter(torch.randn(num_experts, hidden_size, device="cuda"))
alpha = 1.702


# Run the MoE layer
output = layer(
    hidden_states,
    router_idx,
    router_wt,
    alpha,
    gate_up_proj,
    gate_up_proj_bias,
    down_proj,
    down_proj_bias,
)

print("Output sum:", output.sum().item())
print("Output shape:", output.shape)

loss = output.sum()
loss.backward()
print("Loss:", loss.item())

print("Gate up weights grad shape:", gate_up_proj.grad.shape)
print("Gate up weights grad shape:", gate_up_proj_bias.grad)