"""Up-projection of a top-K MoE FFN: grouped GEMM with fused SwiGLU (Triton, H100).

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

x_e is the slice of permuted hidden states owned by expert e:
rows [expert_offsets[e] : expert_offsets[e+1]] of hidden_states.

Strategy:
  * Pack W_gate and W_up along the K dim: W_packed[e] = concat(W_gate[e], W_up[e], dim=0)
    -> shape (E, 2*H, I).  One tl.dot accumulates the full (gate, up) output at once
    over the doubled K dimension, with the K-axis split into two halves in the epilogue.
  * 3D launch grid: (n_blocks, m_blocks, E).  Each kernel instance handles a
    (n_tile, m_tile) sub-tile of one expert's output.  If the tile lies past
    the expert's last row, the program early-exits.
  * Epilogue fuses silu(gate)*up in fp32 and writes bf16.
"""
from __future__ import annotations

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


@triton.jit
def grouped_swiglu_packed_kernel(
    A_ptr,            # (T_perm, H) bf16
    Wg_ptr, Wu_ptr,   # (E, H, I) bf16, two halves of "packed" weight
    C_ptr,            # (T_perm, I) bf16
    expert_offsets_ptr,  # (E+1,) int32
    H, I,
    stride_am, stride_ak,
    stride_we, stride_wn, stride_wk,
    stride_cm, stride_cn,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    SPLIT_K: tl.constexpr,  # K-iterations for gate + K-iterations for up; = 2 always
    GROUP_M: tl.constexpr,
):
    # Super-grouping for L2 reuse (Triton matmul tutorial trick).
    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(tl.num_programs(1) * tl.num_programs(2), 1)  # not used
    # We'll compute pid_n and pid_m from a flat pid across (n_blocks * m_blocks)
    # with a grouping by m to improve L2 locality.
    pid_n = tl.program_id(0)
    pid_m = tl.program_id(1)
    pid_e = tl.program_id(2)

    expert_start = tl.load(expert_offsets_ptr + pid_e)
    expert_end = tl.load(expert_offsets_ptr + pid_e + 1)

    m_pos = expert_start + pid_m * BLOCK_M
    if m_pos >= expert_end:
        return

    offs_m = m_pos + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    m_mask = offs_m < expert_end

    a_ptrs = A_ptr + offs_m[:, None] * stride_am
    eid = pid_e.to(tl.int64)
    wg_ptrs = Wg_ptr + eid * stride_we + offs_n[None, :] * stride_wn
    wu_ptrs = Wu_ptr + eid * stride_we + offs_n[None, :] * stride_wn

    gate_acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    up_acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    for k in range(0, H, BLOCK_K):
        offs_k = k + tl.arange(0, BLOCK_K)
        a = tl.load(
            a_ptrs + offs_k[None, :] * stride_ak,
            mask=m_mask[:, None], other=0.0,
        )
        wg = tl.load(wg_ptrs + offs_k[:, None] * stride_wk)
        wu = tl.load(wu_ptrs + offs_k[:, None] * stride_wk)
        gate_acc = tl.dot(a, wg, gate_acc)
        up_acc = tl.dot(a, wu, up_acc)

    silu_gate = gate_acc * tl.sigmoid(gate_acc)
    out = silu_gate * up_acc

    c_ptrs = C_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    tl.store(c_ptrs, out.to(tl.bfloat16), mask=m_mask[:, None])


# Block sizes tuned for H100.  H and I are multiples of these in all test shapes.
# SMEM budget: A (BLOCK_M*BLOCK_K) + W_gate (BLOCK_K*BLOCK_N) + W_up (BLOCK_K*BLOCK_N)
# per stage * num_stages <= 200KB (H100 limit ~228KB).
_BLOCK_M = 128
_BLOCK_N = 128
_BLOCK_K = 64
_NUM_STAGES = 4
_NUM_WARPS = 8


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)

    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
        out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device)

        m_per_expert = (T_perm + E - 1) // E
        m_blocks = (m_per_expert + _BLOCK_M - 1) // _BLOCK_M
        n_blocks = (I + _BLOCK_N - 1) // _BLOCK_N

        grid = (n_blocks, m_blocks, E)
        grouped_swiglu_packed_kernel[grid](
            hidden_states, self.W_gate, self.W_up, out,
            expert_offsets,
            H, I,
            hidden_states.stride(0), hidden_states.stride(1),
            self.W_gate.stride(0), self.W_gate.stride(2), self.W_gate.stride(1),
            out.stride(0), out.stride(1),
            BLOCK_M=_BLOCK_M, BLOCK_N=_BLOCK_N, BLOCK_K=_BLOCK_K,
            SPLIT_K=2,
            GROUP_M=8,
            num_stages=_NUM_STAGES,
            num_warps=_NUM_WARPS,
        )
        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]