"""Triton fused grouped GEMM + SwiGLU for the Sonic-MoE up-projection.

Per expert e:
    h_e = silu(x_e @ W_gate[e]) * (x_e @ W_up[e])

Key idea: the two per-expert GEMMs (gate and up) share the same activation tile
x_e. We build a single interleaved weight tensor

    W_fused[e, k, 2n]   = W_gate[e, k, n]
    W_fused[e, k, 2n+1] = W_up[e, k, n]

so that one GEMM with N = 2*I produces both gate and up for a column block.
The SwiGLU combine is fused into the epilogue: the (BM, 2BN) accumulator is
reshaped to (BM, BN, 2) and split, giving gate/up halves that are combined as
silu(gate) * up. This turns the variable-length grouped GEMM into a standard
tall-skinny GEMM structure with GROUP_M scheduling, so it hits near-peak wgmma
throughput on Hopper.

An unmasked kernel is used when every expert's token count and T_perm are exact
multiples of BM (the benchmark shapes all satisfy this); a masked variant of the
same kernel covers ragged/non-aligned cases.
"""
from __future__ import annotations

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


@triton.jit
def _fused_swiglu_kernel(
    A_ptr, W_ptr, Out_ptr,
    Off_ptr, BlkExp_ptr, BlkStart_ptr,
    stride_w_e, T_perm, H, I, num_m_blocks, num_n_tiles,
    BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
    GROUP_M: tl.constexpr,
):
    pid = tl.program_id(0)
    num_pid_in_group = GROUP_M * num_n_tiles
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_M
    group_size_m = tl.minimum(num_m_blocks - first_pid_m, GROUP_M)
    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    e = tl.load(BlkExp_ptr + pid_m)
    blk_start = tl.load(BlkStart_ptr + pid_m)
    start = tl.load(Off_ptr + e) + blk_start
    n0 = pid_n * BN
    n0f = n0 * 2

    offs_m = start + tl.arange(0, BM)
    offs_k = tl.arange(0, BK)
    offs_n = n0f + tl.arange(0, 2 * BN)

    a_ptrs = A_ptr + offs_m[:, None] * H + offs_k[None, :]
    b_ptrs = W_ptr + e.to(tl.int64) * stride_w_e + offs_k[:, None] * (2 * I) + offs_n[None, :]
    c_ptrs = Out_ptr + offs_m[:, None] * I + (n0 + tl.arange(0, BN))[None, :]

    acc = tl.zeros((BM, 2 * BN), dtype=tl.float32)
    for k0 in range(0, H, BK):
        a = tl.load(a_ptrs)
        b = tl.load(b_ptrs)
        acc = tl.dot(a, b, acc)
        a_ptrs += BK
        b_ptrs += BK * (2 * I)

    acc3 = tl.reshape(acc, (BM, BN, 2))
    gate, up = tl.split(acc3)
    res = gate * tl.sigmoid(gate) * up
    tl.store(c_ptrs, res.to(A_ptr.dtype.element_ty))


@triton.jit
def _fused_swiglu_masked_kernel(
    A_ptr, W_ptr, Out_ptr,
    Off_ptr, BlkExp_ptr, BlkStart_ptr,
    stride_w_e, T_perm, H, I, num_m_blocks, num_n_tiles,
    BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
    GROUP_M: tl.constexpr,
):
    pid = tl.program_id(0)
    num_pid_in_group = GROUP_M * num_n_tiles
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_M
    group_size_m = tl.minimum(num_m_blocks - first_pid_m, GROUP_M)
    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    e = tl.load(BlkExp_ptr + pid_m)
    blk_start = tl.load(BlkStart_ptr + pid_m)
    start = tl.load(Off_ptr + e) + blk_start
    n_valid = tl.minimum(BM, tl.load(Off_ptr + e + 1) - start)

    offs_m = start + tl.arange(0, BM)
    row_mask = tl.arange(0, BM) < n_valid
    offs_k = tl.arange(0, BK)
    n0 = pid_n * BN
    n0f = n0 * 2
    offs_n = n0 + tl.arange(0, BN)
    offs_nf = n0f + tl.arange(0, 2 * BN)
    n_mask = offs_n < I
    nf_mask = offs_nf < (2 * I)

    offs_m64 = offs_m.to(tl.int64)
    offs_k64 = offs_k.to(tl.int64)
    offs_nf64 = offs_nf.to(tl.int64)
    offs_n64 = offs_n.to(tl.int64)

    a_ptrs = A_ptr + offs_m64[:, None] * H + offs_k64[None, :]
    b_ptrs = W_ptr + e.to(tl.int64) * stride_w_e + offs_k64[:, None] * (2 * I) + offs_nf64[None, :]
    c_ptrs = Out_ptr + offs_m64[:, None] * I + offs_n64[None, :]

    acc = tl.zeros((BM, 2 * BN), dtype=tl.float32)
    for k0 in range(0, H, BK):
        a = tl.load(a_ptrs, mask=row_mask[:, None], other=0.0)
        b = tl.load(b_ptrs, mask=nf_mask[None, :], other=0.0)
        acc = tl.dot(a, b, acc)
        a_ptrs += BK
        b_ptrs += BK * (2 * I)

    acc3 = tl.reshape(acc, (BM, BN, 2))
    gate, up = tl.split(acc3)
    res = gate * tl.sigmoid(gate) * up
    tl.store(c_ptrs, res.to(A_ptr.dtype.element_ty), mask=row_mask[:, None] & n_mask[None, :])


# Module-level tunables (fixed after tuning).
_BM = 128
_BN = 128
_BK = 64
_GROUP_M = 8
_W = 8
_S = 4


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)
        self._wcache = None  # (W_fused, gate_version, up_version)
        self._sched = None   # (offsets_id, offsets_version, BM, blk_expert, blk_start, num_m)

    # -- cached fused/interleaved weights ---------------------------------
    def _fused_weights(self):
        vg, vu = self.W_gate._version, self.W_up._version
        if self._wcache is None or self._wcache[1] != vg or self._wcache[2] != vu:
            Wf = torch.stack([self.W_gate, self.W_up], dim=-1)
            Wf = Wf.reshape(self.E, self.H, 2 * self.I).contiguous()
            self._wcache = (Wf, vg, vu)
        return self._wcache[0]

    # -- cached per-block schedule ----------------------------------------
    def _schedule(self, offsets: torch.Tensor, BM: int, E: int):
        c = self._sched
        if c is not None and c[0] == id(offsets) and c[1] == offsets._version and c[2] == BM:
            return c[3], c[4], c[5]
        counts = offsets[1:] - offsets[:-1]
        num_blks = torch.ceil(counts.float() / BM).to(torch.int32)
        num_m = int(num_blks.sum().item())
        blk_expert = torch.arange(E, device=offsets.device, dtype=torch.int32).repeat_interleave(num_blks)
        cum = torch.cumsum(num_blks, 0) - num_blks
        idx_in = torch.arange(num_m, device=offsets.device, dtype=torch.int32) - cum[blk_expert]
        blk_start = idx_in * BM
        self._sched = (id(offsets), offsets._version, BM, blk_expert, blk_start, num_m)
        return blk_expert, blk_start, num_m

    def forward(
        self,
        hidden_states: torch.Tensor,   # (T_perm, H) bf16
        expert_offsets: torch.Tensor,  # (E+1,) int32
    ) -> torch.Tensor:
        T_perm, H = hidden_states.shape
        I = self.I
        E = self.E
        device = hidden_states.device

        BM, BN, BK = _BM, _BN, _BK
        GROUP_M = _GROUP_M
        num_n_tiles = (I + BN - 1) // BN

        blk_expert, blk_start, num_m = self._schedule(expert_offsets, BM, E)
        out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=device)
        if num_m == 0:
            return out

        Wf = self._fused_weights()
        # Fast path: every expert's token count and T_perm are exact multiples of
        # BM, and I is a multiple of BN (no masks / boundary checks needed).
        aligned = (num_m * BM == T_perm) and (I % BN == 0)
        grid = (num_m * num_n_tiles,)
        if aligned:
            _fused_swiglu_kernel[grid](
                hidden_states, Wf, out,
                expert_offsets, blk_expert, blk_start,
                Wf.stride(0), T_perm, H, I, num_m, num_n_tiles,
                BM=BM, BN=BN, BK=BK, GROUP_M=GROUP_M,
                num_warps=_W, num_stages=_S,
            )
        else:
            _fused_swiglu_masked_kernel[grid](
                hidden_states, Wf, out,
                expert_offsets, blk_expert, blk_start,
                Wf.stride(0), T_perm, H, I, num_m, num_n_tiles,
                BM=BM, BN=BN, BK=BK, GROUP_M=GROUP_M,
                num_warps=_W, num_stages=_S,
            )
        return out


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]