"""
AMD-optimised Mixture-of-Experts submission for the DeepSeek-style benchmark.
*   Heavy kernels use **aiter** Triton paths (MI300 tuned) but silently fall back to PyTorch
    when unavailable.
*   Bug-fixed routing path – duplicates are aggregated with `scatter_add_`, then the
    flat buffer is reshaped directly without a second permutation.
*   Public interface (`custom_kernel`) and the weight-dict layout are unchanged.
"""

from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Optional

# -----------------------------------------------------------------------------
# Simple cache so we only build the MoE structure once per device/dtype/config.
# -----------------------------------------------------------------------------
_cached: Dict[str, Optional[object]] = {
    "moe": None,
    "cfg": None,
    "device": None,
    "dtype": None,
}

# -----------------------------------------------------------------------------
# Optional fast kernels (skips to eager PyTorch on import failure)
# -----------------------------------------------------------------------------
try:
    from aiter.ops.triton.gemm_a16w16 import gemm_a16w16  # noqa: F401
    from aiter.ops.triton.silu_mul import silu_mul  # noqa: F401

    _AITER_AVAILABLE = True
except (ModuleNotFoundError, ImportError):
    _AITER_AVAILABLE = False


# -----------------------------------------------------------------------------
# Helpers: transparent dispatch to Triton vs. PyTorch
# -----------------------------------------------------------------------------
def _matmul(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
    # x: [B,K], w: [K,N]
    if _AITER_AVAILABLE and x.is_cuda and w.is_cuda and x.dtype in (torch.float16, torch.bfloat16):
        return gemm_a16w16(x, w)  # type: ignore[name-defined]
    return x @ w


def _silu_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    # x,y: same shape
    if _AITER_AVAILABLE and x.is_cuda and x.dtype in (torch.float16, torch.bfloat16):
        return silu_mul(x, y)  # type: ignore[name-defined]
    return F.silu(x) * y


# -----------------------------------------------------------------------------
# Expert: FFN block with SiLU gating
# -----------------------------------------------------------------------------
class Expert(nn.Module):
    def __init__(self, config: Dict, *, d_expert: Optional[int] = None):
        super().__init__()
        self.d_hid = config["d_hidden"]
        self.d_exp = config["d_expert"] if d_expert is None else d_expert

        self.W_gate = nn.Parameter(torch.empty(self.d_hid, self.d_exp))
        self.W_up = nn.Parameter(torch.empty(self.d_hid, self.d_exp))
        self.W_down = nn.Parameter(torch.empty(self.d_exp, self.d_hid))

    @torch.no_grad()
    def set_weights(
        self, *, gate_w: torch.Tensor, up_w: torch.Tensor, down_w: torch.Tensor
    ) -> None:
        # Copy directly; weights already in [in,out] layout
        self.W_gate.copy_(gate_w.contiguous())
        self.W_up.copy_(up_w.contiguous())
        self.W_down.copy_(down_w.contiguous())

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [N, d_hid]
        g = _matmul(x, self.W_gate)  # [N, d_exp]
        u = _matmul(x, self.W_up)  # [N, d_exp]
        f = _silu_mul(g, u)  # [N, d_exp]
        return _matmul(f, self.W_down)  # [N, d_hid]


# -----------------------------------------------------------------------------
# Router: top-k gating
# -----------------------------------------------------------------------------
class Router(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.k = config["n_experts_per_token"]
        self.n_experts = config["n_routed_experts"]
        d_h = config["d_hidden"]
        self.W_g = nn.Parameter(torch.empty(self.n_experts, d_h))

    def forward(self, x: torch.Tensor):
        # x: [B,T,H]
        B, T, H = x.shape
        logits = _matmul(x.view(-1, H), self.W_g.t()).view(B, T, self.n_experts)
        scores = logits.softmax(-1)
        topk_s, topk_i = torch.topk(scores, self.k, dim=-1)
        return topk_i, topk_s


# -----------------------------------------------------------------------------
# MoE block
# -----------------------------------------------------------------------------
class MoE(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.cfg = config
        self.k = config["n_experts_per_token"]
        self.E = config["n_routed_experts"]

        self.experts = nn.ModuleList([Expert(config) for _ in range(self.E)])
        self.router = Router(config)
        shared_dim = config["d_expert"] * config["n_shared_experts"]
        self.shared_expert = Expert(config, d_expert=shared_dim)

    @torch.no_grad()
    def load_from_weightdict(self, wd: Dict[str, torch.Tensor]):
        # Router weights
        self.router.W_g.copy_(wd["router.weight"].contiguous())
        # Per-expert weights
        for i, expert in enumerate(self.experts):
            expert.set_weights(
                gate_w=wd[f"experts.{i}.0.weight"],
                up_w=wd[f"experts.{i}.1.weight"],
                down_w=wd[f"experts.{i}.2.weight"],
            )
        # Shared expert
        self.shared_expert.set_weights(
            gate_w=wd["shared_experts.0.weight"],
            up_w=wd["shared_experts.1.weight"],
            down_w=wd["shared_experts.2.weight"],
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [B,T,H]
        B, T, H = x.shape

        # Shared path
        shared_out = self.shared_expert(x.view(-1, H)).view(B, T, H)

        # Routing
        topk_i, topk_s = self.router(x)  # [B,T,k]
        flat_i = topk_i.view(-1)  # [B*T*k]
        flat_s = topk_s.view(-1, 1)  # [B*T*k,1]
        flat_x = x.view(-1, H)  # [B*T, H]
        routed_buf = torch.zeros_like(flat_x)  # [B*T, H]

        # Sort by expert
        order = torch.argsort(flat_i)
        sorted_i = flat_i[order]
        tokens = order // self.k

        counts = torch.bincount(sorted_i, minlength=self.E).cpu()
        ptrs = counts.cumsum(0)

        start = 0
        for e, end in enumerate(ptrs.tolist()):
            if start == end:
                continue
            sel = tokens[start:end]
            inp = flat_x[sel]
            out = self.experts[e](inp)
            out.mul_(flat_s[order[start:end]])
            routed_buf.scatter_add_(0, sel.unsqueeze(1).expand(-1, H), out)
            start = end

        routed = routed_buf.view(B, T, H)
        return routed + shared_out


# -----------------------------------------------------------------------------
# Public entry point
# -----------------------------------------------------------------------------
def custom_kernel(data):  # type: ignore[override]
    """
    Signature expected by the benchmark:
      (input_tensor, weight_dict, config) -> output_tensor
    """
    input_tensor, weight_dict, config = data
    device, dtype = input_tensor.device, input_tensor.dtype

    # Build/rebuild MoE only when structure, device or dtype changes
    if (
        _cached["moe"] is None
        or _cached["cfg"] != config
        or _cached["device"] != device
        or _cached["dtype"] != dtype
    ):
        if _cached["moe"] is not None:
            # free old model memory
            del _cached["moe"]
            torch.cuda.empty_cache()
        # instantiate & cache
        moe = MoE(config).to(device=device, dtype=dtype)
        _cached.update(moe=moe, cfg=config, device=device, dtype=dtype)
    else:
        moe = _cached["moe"]

    # Always reload new weights (ops are in-place)
    moe.load_from_weightdict(weight_dict)

    # Forward without grad, then release any freed cache
    with torch.no_grad():
        out = moe(input_tensor)
    torch.cuda.synchronize()
    torch.cuda.empty_cache()

    return out

python eval.py benchmark /tmp/tmp4fc96azu