"""Up-projection of a top-K MoE FFN: variable-length grouped GEMM + fused SwiGLU.

Target: H100 PCIe (SM90 Hopper), bf16.

Per expert e:
    h_e = silu(x_e @ W_gate[e]) * (x_e @ W_up[e])
where x_e = hidden_states[offsets[e]:offsets[e+1]].

Implementation: a single Triton kernel does a variable-length grouped GEMM. The
output (M=T_perm, N=I) is tiled into (BLOCK_M, BLOCK_N) blocks; each tile belongs
to exactly one expert (tiling respects expert boundaries). Inside the K-reduction
loop over H we run *two* MMAs -- one against W_gate, one against W_up -- sharing
the single x load, and fuse silu(gate)*up in the epilogue. fp32 accumulation.
"""
from __future__ import annotations

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


# ---------------------------------------------------------------------------
# Kernel
# ---------------------------------------------------------------------------
@triton.jit
def _grouped_swiglu_kernel(
    x_ptr, wg_ptr, wu_ptr, out_ptr,
    expert_offsets_ptr,   # (E+1,) int32 row offsets per expert
    tile_offsets_ptr,     # (E+1,) int32 cumulative M-tile counts per expert (BLOCK_M)
    # dims
    H,               # reduction dim (runtime -> K-loop not unrolled)
    I_dim,           # output dim
    E: tl.constexpr,
    E_POW2: tl.constexpr,
    # strides (runtime)
    stride_xm, stride_xh,
    stride_we, stride_weh, stride_wei,
    stride_om, stride_oi,
    # tile sizes
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    pid_m = tl.program_id(0)   # global M-tile index across all experts
    pid_n = tl.program_id(1)   # N-tile index over the I dimension

    # --- locate expert e for this M-tile via vectorized count over tile_offsets ---
    ar = tl.arange(0, E_POW2)
    toff = tl.load(tile_offsets_ptr + ar, mask=ar <= E, other=(1 << 30))
    # number of cumulative offsets <= pid_m, minus 1 -> expert index
    e = tl.sum((toff <= pid_m).to(tl.int32), axis=0) - 1

    tile_start = tl.load(tile_offsets_ptr + e)          # pid_m where expert e begins
    m_row_start = tl.load(expert_offsets_ptr + e)       # absolute row offset
    m_row_end = tl.load(expert_offsets_ptr + e + 1)
    local_m = pid_m - tile_start
    m_start = m_row_start + local_m * BLOCK_M

    offs_m = m_start + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)
    mask_m = offs_m < m_row_end
    mask_n = offs_n < I_dim

    # base pointers
    x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xh
    wg_ptrs = wg_ptr + e * stride_we + offs_k[:, None] * stride_weh + offs_n[None, :] * stride_wei
    wu_ptrs = wu_ptr + e * stride_we + offs_k[:, None] * stride_weh + offs_n[None, :] * stride_wei

    acc_g = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    acc_u = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    for k_start in range(0, H, BLOCK_K):
        mask_k = (k_start + offs_k) < H
        kmask_x = mask_m[:, None] & mask_k[None, :]
        kmask_w = mask_k[:, None] & mask_n[None, :]
        x = tl.load(x_ptrs, mask=kmask_x, other=0.0)
        wg = tl.load(wg_ptrs, mask=kmask_w, other=0.0)
        wu = tl.load(wu_ptrs, mask=kmask_w, other=0.0)
        acc_g = tl.dot(x, wg, acc_g)
        acc_u = tl.dot(x, wu, acc_u)
        x_ptrs += BLOCK_K * stride_xh
        wg_ptrs += BLOCK_K * stride_weh
        wu_ptrs += BLOCK_K * stride_weh

    # fused SwiGLU epilogue: silu(gate) * up
    sig = tl.sigmoid(acc_g)
    out = (acc_g * sig) * acc_u

    o_ptrs = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_oi
    tl.store(o_ptrs, out.to(tl.bfloat16), mask=mask_m[:, None] & mask_n[None, :])


# ---------------------------------------------------------------------------
# Host helpers
# ---------------------------------------------------------------------------
def _next_pow2(n: int) -> int:
    p = 1
    while p < n:
        p <<= 1
    return p


def _build_tile_offsets(expert_offsets: torch.Tensor, BLOCK_M: int, device) -> tuple[torch.Tensor, int]:
    """Cumulative per-expert M-tile counts (E+1,) int32 on `device` + total tiles."""
    off = expert_offsets.detach().cpu().tolist()
    e_n = len(off) - 1
    cum = [0] * (e_n + 1)
    for e in range(e_n):
        n = off[e + 1] - off[e]
        cum[e + 1] = cum[e] + (n + BLOCK_M - 1) // BLOCK_M
    tile_offsets = torch.tensor(cum, dtype=torch.int32, device=device)
    return tile_offsets, cum[-1]


# ---------------------------------------------------------------------------
# Module (same interface as reference.py)
# ---------------------------------------------------------------------------
class Model(nn.Module):
    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)

        self.BLOCK_M = 128
        self.E_POW2 = _next_pow2(E)
        # cache (keyed by expert_offsets identity) so we don't D2H-sync per call
        self._toff_key = None
        self._toff = None
        self._total = None

    def forward(self, hidden_states: torch.Tensor, expert_offsets: torch.Tensor) -> torch.Tensor:
        T_perm, H = hidden_states.shape
        I = self.I
        out = torch.empty((T_perm, I), dtype=torch.bfloat16, device=hidden_states.device)

        key = expert_offsets.data_ptr()
        if self._toff_key != key or self._toff is None:
            self._toff, self._total = _build_tile_offsets(expert_offsets, self.BLOCK_M, hidden_states.device)
            self._toff_key = key

        grid = lambda meta: (self._total, triton.cdiv(I, meta["BLOCK_N"]))
        _grouped_swiglu_kernel[grid](
            hidden_states, self.W_gate, self.W_up, out,
            expert_offsets, self._toff,
            H, I, self.E, self.E_POW2,
            H, 1,                       # x strides: (H, 1)
            H * I, I, 1,                # weight strides (E,H,I): (H*I, I, 1)
            I, 1,                       # out strides: (I, 1)
            BLOCK_M=self.BLOCK_M,
            BLOCK_N=128,
            BLOCK_K=64,
            num_warps=8,
            num_stages=3,
        )
        return out


# Module-level shape shims rewritten by check.py / benchmark.py per shape.
T_total = 32768
H = 4096
I = 1536  # noqa: E741
E = 128
K = 8


def _build_routing(T_total: int, E: int, K: int, device: str = "cpu") -> torch.Tensor:
    T_perm = T_total * K
    base = T_perm // E
    rem = T_perm - base * E
    counts = torch.full((E,), base, dtype=torch.int32, device=device)
    counts[:rem] += 1
    offsets = torch.zeros(E + 1, dtype=torch.int32, device=device)
    offsets[1:] = torch.cumsum(counts, dim=0)
    return offsets


def get_inputs():
    T_perm = T_total * K
    hidden_states = torch.randn(T_perm, H, dtype=torch.bfloat16) * 0.1
    expert_offsets = _build_routing(T_total, E, K)
    return [hidden_states, expert_offsets]


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