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


class Expert(nn.Module):
    """Single expert with fused gate & up projection."""

    def __init__(self, config: Dict, d_expert: Optional[int] = None):
        super().__init__()
        self.d_hidden: int = config["d_hidden"]
        self.d_expert: int = config["d_expert"] if d_expert is None else d_expert
        self.act_fn = nn.SiLU()

        # Fuse W_gate and W_up -> one GEMM instead of two
        self.W_gate_up = nn.Linear(self.d_hidden, 2 * self.d_expert, bias=False)
        self.W_down = nn.Linear(self.d_expert, self.d_hidden, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate_up = self.W_gate_up(x)
        gate, up = gate_up.chunk(2, dim=-1)
        gate = self.act_fn(gate)
        return self.W_down(gate * up)


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

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        scores = self.W_g(x).softmax(dim=-1)
        topk_scores, topk_indices = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
        return topk_indices, topk_scores


class MoE(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config
        self.experts = nn.ModuleList([Expert(config) for _ in range(config["n_routed_experts"])])
        self.gating_network = MoEGate(config)
        shared_dim = config["d_expert"] * config["n_shared_experts"]
        self.shared_expert = Expert(config=config, d_expert=shared_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        shared_out = self.shared_expert(x)
        expert_idxs, expert_scores = self.gating_network(x)

        b, s, h = x.shape
        flat_x = x.view(-1, h)
        flat_idxs = expert_idxs.view(-1)
        flat_scores = expert_scores.view(-1, 1)
        routed = self._dispatch(flat_x, flat_idxs, flat_scores)
        return shared_out + routed.view(b, s, h)

    @torch.no_grad()
    def _dispatch(self, x: torch.Tensor, e_idx: torch.Tensor, e_w: torch.Tensor) -> torch.Tensor:
        """Per‑expert dispatch with index_add_ (fast & memory‑friendly)."""
        out = torch.zeros_like(x)
        order = e_idx.argsort()
        counts = e_idx.bincount().cpu().tolist()
        cum = 0
        tokens_per_exp = []
        for c in counts:
            cum += c
            tokens_per_exp.append(cum)
        tok_per_tok = self.config["n_experts_per_token"]
        token_indices = order // tok_per_tok
        for exp_id, end in enumerate(tokens_per_exp):
            start = 0 if exp_id == 0 else tokens_per_exp[exp_id - 1]
            if start == end:
                continue
            sel_tok = token_indices[start:end]
            expert_tokens = x[sel_tok]
            result = self.experts[exp_id](expert_tokens)
            result.mul_(e_w[order[start:end]])
            out.index_add_(0, sel_tok, result)
        return out


def custom_kernel(data: input_t) -> output_t:
    """Entry point expected by the evaluator."""
    input_tensor, weights, config = data
    num_experts = config["n_routed_experts"]

    moe = MoE(config)

    # --- load weights ---
    moe.gating_network.W_g.weight = nn.Parameter(weights["router.weight"])

    for i in range(num_experts):
        gate_w = weights[f"experts.{i}.0.weight"].t()
        up_w = weights[f"experts.{i}.1.weight"].t()
        down_w = weights[f"experts.{i}.2.weight"].t()
        moe.experts[i].W_gate_up.weight = nn.Parameter(torch.cat([gate_w, up_w], dim=0))
        moe.experts[i].W_down.weight = nn.Parameter(down_w)

    sg = weights["shared_experts.0.weight"].t()
    sup = weights["shared_experts.1.weight"].t()
    sd = weights["shared_experts.2.weight"].t()
    moe.shared_expert.W_gate_up.weight = nn.Parameter(torch.cat([sg, sup], dim=0))
    moe.shared_expert.W_down.weight = nn.Parameter(sd)

    # --- compile for fused kernels & scheduling ---
    moe = torch.compile(moe, mode="default", fullgraph=False)

    return moe(input_tensor)

python eval.py leaderboard /tmp/tmp9qu6sazr
W0512 19:53:53.259000 3142 torch/fx/experimental/symbolic_shapes.py:6872] [2/0] failed during evaluate_expr(Eq(u0, 0), hint=None, size_oblivious=False, forcing_spec=False
