"""Optimized Triton GEMM + fused SwiGLU kernel for MoE up-projection.

Uses per-expert kernel launches with optimal block sizes for each expert's M dimension.
Fuses both GEMMs (gate and up) in a single kernel to maximize data reuse.
"""
from __future__ import annotations

import torch
import torch.nn as nn
import triton
import triton.language as tl

OP_TYPE = "grouped_gemm_swiglu"
SUPPORTED_PRECISIONS = ["bf16"]
HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]


@triton.jit
def _gemm_swiglu_kernel(
    x_ptr, wg_ptr, wu_ptr, out_ptr,
    M: tl.constexpr,
    N: tl.constexpr,
    K_dim: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    """Single-expert GEMM + SwiGLU.

    x: (M, K_dim), W_gate: (K_dim, N), W_up: (K_dim, N), out: (M, N)
    """
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)

    mask_m = rm < M
    mask_n = rn < N

    acc_gate = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    acc_up = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    for ki in range(tl.cdiv(K_dim, BLOCK_K)):
        rk = ki * BLOCK_K + tl.arange(0, BLOCK_K)
        mask_k = rk < K_dim

        x_ptrs = x_ptr + rm[:, None] * K_dim + rk[None, :]
        x_tile = tl.load(x_ptrs, mask=mask_m[:, None] & mask_k[None, :], other=0.0)

        wg_ptrs = wg_ptr + rk[:, None] * N + rn[None, :]
        wg_tile = tl.load(wg_ptrs, mask=mask_k[:, None] & mask_n[None, :], other=0.0)

        wu_ptrs = wu_ptr + rk[:, None] * N + rn[None, :]
        wu_tile = tl.load(wu_ptrs, mask=mask_k[:, None] & mask_n[None, :], other=0.0)

        acc_gate += tl.dot(x_tile, wg_tile)
        acc_up += tl.dot(x_tile, wu_tile)

    # SwiGLU
    neg_gate = -acc_gate
    exp_neg = tl.exp(neg_gate)
    sigmoid_val = 1.0 / (1.0 + exp_neg)
    result = acc_gate * sigmoid_val * acc_up

    result_bf16 = result.to(tl.bfloat16)
    out_ptrs = out_ptr + rm[:, None] * N + rn[None, :]
    tl.store(out_ptrs, result_bf16, mask=mask_m[:, None] & mask_n[None, :])


class Model(nn.Module):
    """Up-projection of a top-K MoE FFN with fused SwiGLU."""

    def __init__(self, T_total: int, H: int, I: int, E: int, K: int):  # noqa: E741
        super().__init__()
        self.T_total = T_total
        self.H = H
        self.I = I
        self.E = E
        self.K = K
        self.W_gate = nn.Parameter(torch.empty(E, H, I, dtype=torch.bfloat16))
        self.W_up = nn.Parameter(torch.empty(E, H, I, dtype=torch.bfloat16))
        nn.init.normal_(self.W_gate, std=0.02)
        nn.init.normal_(self.W_up, std=0.02)

    def forward(
        self,
        hidden_states: torch.Tensor,
        expert_offsets: torch.Tensor,
    ) -> torch.Tensor:
        T_perm = hidden_states.shape[0]
        H = self.H
        I = self.I  # noqa: E741
        E = self.E

        out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device)

        x = hidden_states.contiguous()
        W_gate = self.W_gate.contiguous()
        W_up = self.W_up.contiguous()

        offsets_cpu = expert_offsets.cpu().numpy()

        # Block sizes tuned for tensor cores (bf16 MMA 16x16x16)
        # Larger blocks = better memory coalescing, but need to fit per-expert M
        BLOCK_M = 64
        BLOCK_N = 64
        BLOCK_K = 64

        for e in range(E):
            start = int(offsets_cpu[e])
            end = int(offsets_cpu[e + 1])
            m_e = end - start
            if m_e == 0:
                continue

            grid = (triton.cdiv(m_e, BLOCK_M), triton.cdiv(I, BLOCK_N))

            _gemm_swiglu_kernel[grid](
                x[start:end],
                W_gate[e],
                W_up[e],
                out[start:end],
                M=m_e, N=I, K_dim=H,
                BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
            )

        return out


T_total = 32768
H = 4096
I = 1536  # noqa: E741
E = 128
K = 8


def get_inputs():
    T_perm = T_total * K
    hidden_states = torch.randn(T_perm, H, dtype=torch.bfloat16) * 0.1
    base = T_perm // E
    rem = T_perm - base * E
    counts = torch.full((E,), base, dtype=torch.int32)
    counts[:rem] += 1
    expert_offsets = torch.zeros(E + 1, dtype=torch.int32)
    expert_offsets[1:] = torch.cumsum(counts, dim=0)
    return [hidden_states, expert_offsets]


def get_init_inputs():
    return [T_total, H, I, E, K]