#!POPCORN leaderboard amd-mixture-of-experts
#!POPCORN gpus MI300

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple
from task import input_t, output_t


@torch.jit.script
def fused_expert_forward(
    x: torch.Tensor,
    W_gate: torch.Tensor,
    W_up: torch.Tensor,
    W_down: torch.Tensor,
    weight: torch.Tensor
) -> torch.Tensor:
    x = x.contiguous()
    gate = F.silu(x @ W_gate)
    up = x @ W_up
    return (gate * up * weight) @ W_down


class Expert(nn.Module):
    def __init__(self):
        super().__init__()
        self.W_gate = None
        self.W_up = None
        self.W_down = None


class MoEGate(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.top_k = config["n_experts_per_token"]
        self.num_experts = config["n_routed_experts"]
        self.d_hidden = config["d_hidden"]
        self.W_g = None

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        logits = x @ self.W_g
        probs = F.softmax(logits, dim=-1)
        topk_scores, topk_indices = torch.topk(probs, k=self.top_k, dim=-1)
        return topk_indices, topk_scores


class MoE(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config
        self.num_experts = config["n_routed_experts"]
        self.top_k = config["n_experts_per_token"]
        self.experts = [Expert() for _ in range(self.num_experts)]
        self.gating_network = MoEGate(config)

        self.W_shared_gate = None
        self.W_shared_up = None
        self.W_shared_down = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, D = x.shape
        N = B * T
        x_flat = x.view(N, D)

        shared_out = (
            F.silu(x @ self.W_shared_gate)
            * (x @ self.W_shared_up)
        ) @ self.W_shared_down

        expert_indices, expert_scores = self.gating_network(x_flat)
        expert_indices = expert_indices.reshape(-1)
        expert_scores = expert_scores.reshape(-1, 1)
        token_indices = torch.arange(N, device=x.device).repeat_interleave(self.top_k)

        sorted_idx = expert_indices.argsort()
        expert_indices = expert_indices[sorted_idx]
        token_indices = token_indices[sorted_idx]
        expert_scores = expert_scores[sorted_idx]

        counts = torch.bincount(expert_indices, minlength=self.num_experts)
        cumsum = counts.cumsum(0).cpu().numpy()

        output = torch.zeros_like(x_flat)

        for i in range(self.num_experts):
            start = 0 if i == 0 else cumsum[i - 1]
            end = cumsum[i]
            if start == end:
                continue

            tokens = token_indices[start:end]
            weights = expert_scores[start:end]
            inputs = x_flat[tokens]

            expert = self.experts[i]
            out = fused_expert_forward(
                inputs,
                expert.W_gate,
                expert.W_up,
                expert.W_down,
                weights
            )

            output.scatter_add_(0, tokens.unsqueeze(1).expand(-1, D), out)

        return (shared_out.view(N, D) + output).view(B, T, D)


def custom_kernel(data: input_t) -> output_t:
    x, weights, config = data
    moe = MoE(config)

    moe.gating_network.W_g = weights["router.weight"].t()
    moe.W_shared_gate = weights["shared_experts.0.weight"]
    moe.W_shared_up = weights["shared_experts.1.weight"]
    moe.W_shared_down = weights["shared_experts.2.weight"]

    for i in range(config["n_routed_experts"]):
        moe.experts[i].W_gate = weights[f"experts.{i}.0.weight"]
        moe.experts[i].W_up = weights[f"experts.{i}.1.weight"]
        moe.experts[i].W_down = weights[f"experts.{i}.2.weight"]

    return moe(x)

python eval.py benchmark /tmp/tmpy9z8r5im