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 test /tmp/tmpahkrptu9
/usr/local/lib/python3.10/dist-packages/torch/_inductor/lowering.py:7061: UserWarning: 
Online softmax is disabled on the fly since Inductor decides to
split the reduction. Cut an issue to PyTorch if this is an
important use case and you want to speed it up with online
softmax.

  warnings.warn(
W0512 19:50:42.391000 381 torch/fx/experimental/symbolic_shapes.py:6872] [2/0] failed during evaluate_expr(Eq(u0, 0), hint=None, size_oblivious=False, forcing_spec=False
