"""Grouped GEMM + fused SwiGLU kernel for Sonic-MoE up-projection.

Uses two Triton grouped-GEMM kernels (gate and up) launched on separate streams
for concurrent execution, followed by a fused SwiGLU kernel. Each GEMM uses
precomputed tile-to-expert mapping and optimized tile sizes for H100 SM90.
"""
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 _grouped_gemm_kernel(
    a_ptr,               # (T_perm, K) bf16
    b_ptr,               # (E, K, N) bf16
    c_ptr,               # (T_perm, N) bf16
    tile_expert_ptr,     # (num_m_tiles,) int32
    tile_m_start_ptr,    # (num_m_tiles,) int32
    expert_offsets_ptr,  # (E+1,) int32
    stride_am, stride_ak,
    stride_be, stride_bk, stride_bn,
    stride_cm, stride_cn,
    K: tl.constexpr,
    N: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    GROUP_M: tl.constexpr,
):
    """Standard grouped GEMM with variable M per expert: C = A @ B[e].

    1D grid with L2-optimized swizzling across M tiles.
    """
    pid = tl.program_id(0)
    num_pid_m = tl.num_programs(0) // tl.cdiv(N, BLOCK_N)
    num_pid_n = tl.cdiv(N, BLOCK_N)

    # Swizzle M tiles for better L2 cache utilization
    num_pid_in_group = GROUP_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_M
    group_size_m = tl.minimum(num_pid_m - 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

    # Look up expert and row range for this M-tile
    expert_id = tl.load(tile_expert_ptr + pid_m)
    m_start = tl.load(tile_m_start_ptr + pid_m)
    expert_end = tl.load(expert_offsets_ptr + expert_id + 1)

    m_end = tl.minimum(m_start + BLOCK_M, expert_end)

    n_start = pid_n * BLOCK_N
    n_end = tl.minimum(n_start + BLOCK_N, N)

    # Accumulator
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    # Mainloop
    for k_start in range(0, K, BLOCK_K):
        k_end = tl.minimum(k_start + BLOCK_K, K)

        # Load A tile [BLOCK_M, BLOCK_K]
        offs_m = m_start + tl.arange(0, BLOCK_M)
        offs_k = k_start + tl.arange(0, BLOCK_K)
        a_mask = (offs_m[:, None] < m_end) & (offs_k[None, :] < k_end)
        a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
        a_tile = tl.load(a_ptrs, mask=a_mask, other=0.0)

        # Load B tile for this expert [BLOCK_K, BLOCK_N]
        offs_k2 = k_start + tl.arange(0, BLOCK_K)
        offs_n = n_start + tl.arange(0, BLOCK_N)
        b_mask = (offs_k2[:, None] < k_end) & (offs_n[None, :] < n_end)
        b_ptrs = (b_ptr +
                  expert_id * stride_be +
                  offs_k2[:, None] * stride_bk +
                  offs_n[None, :] * stride_bn)
        b_tile = tl.load(b_ptrs, mask=b_mask, other=0.0)

        acc += tl.dot(a_tile, b_tile, allow_tf32=False)

    # Store
    offs_m_out = m_start + tl.arange(0, BLOCK_M)
    offs_n_out = n_start + tl.arange(0, BLOCK_N)
    out_mask = (offs_m_out[:, None] < m_end) & (offs_n_out[None, :] < n_end)
    out_ptrs = c_ptr + offs_m_out[:, None] * stride_cm + offs_n_out[None, :] * stride_cn
    tl.store(out_ptrs, acc.to(c_ptr.dtype.element_ty), mask=out_mask)


@triton.jit
def _swiglu_kernel(
    gate_ptr,    # (T_perm, I) bf16
    up_ptr,      # (T_perm, I) bf16
    out_ptr,     # (T_perm, I) bf16
    num_elements,
    stride_m, stride_n,
    BLOCK_SIZE: tl.constexpr,
):
    """Fused SwiGLU: out = silu(gate) * up."""
    pid = tl.program_id(0)
    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offs < num_elements

    m = offs // stride_n
    n = offs % stride_n

    gate_val = tl.load(gate_ptr + offs, mask=mask, other=0.0)
    up_val = tl.load(up_ptr + offs, mask=mask, other=0.0)

    # SwiGLU: silu(x) * up = x * sigmoid(x) * up
    gate_f32 = gate_val.to(tl.float32)
    up_f32 = up_val.to(tl.float32)
    result = gate_f32 * tl.sigmoid(gate_f32) * up_f32

    tl.store(out_ptr + offs, result.to(out_ptr.dtype.element_ty), mask=mask)


@triton.jit
def _fused_gemm_swiglu_kernel(
    a_ptr,               # (T_perm, K) bf16
    w_gate_ptr,          # (E, K, N) bf16
    w_up_ptr,            # (E, K, N) bf16
    out_ptr,             # (T_perm, N) bf16
    tile_expert_ptr,     # (num_m_tiles,) int32
    tile_m_start_ptr,    # (num_m_tiles,) int32
    expert_offsets_ptr,  # (E+1,) int32
    stride_am, stride_ak,
    stride_we, stride_wk, stride_wn,
    stride_om, stride_on,
    K: tl.constexpr,
    N: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    """Fused grouped GEMM + SwiGLU: out = silu(A @ W_gate) * (A @ W_up).

    2D grid: (num_m_tiles, cdiv(N, BLOCK_N)).
    Each program does both gate and up GEMMs and fuses SwiGLU.
    """
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    expert_id = tl.load(tile_expert_ptr + pid_m)
    m_start = tl.load(tile_m_start_ptr + pid_m)
    expert_end = tl.load(expert_offsets_ptr + expert_id + 1)
    m_end = tl.minimum(m_start + BLOCK_M, expert_end)

    n_start = pid_n * BLOCK_N
    n_end = tl.minimum(n_start + BLOCK_N, N)

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

    for k_start in range(0, K, BLOCK_K):
        k_end = tl.minimum(k_start + BLOCK_K, K)

        # Load A tile (shared between gate and up)
        offs_m = m_start + tl.arange(0, BLOCK_M)
        offs_k = k_start + tl.arange(0, BLOCK_K)
        a_mask = (offs_m[:, None] < m_end) & (offs_k[None, :] < k_end)
        a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
        a_tile = tl.load(a_ptrs, mask=a_mask, other=0.0)

        # Load W_gate tile
        offs_k2 = k_start + tl.arange(0, BLOCK_K)
        offs_n = n_start + tl.arange(0, BLOCK_N)
        w_mask = (offs_k2[:, None] < k_end) & (offs_n[None, :] < n_end)

        wg_ptrs = (w_gate_ptr +
                   expert_id * stride_we +
                   offs_k2[:, None] * stride_wk +
                   offs_n[None, :] * stride_wn)
        wg_tile = tl.load(wg_ptrs, mask=w_mask, other=0.0)

        wu_ptrs = (w_up_ptr +
                   expert_id * stride_we +
                   offs_k2[:, None] * stride_wk +
                   offs_n[None, :] * stride_wn)
        wu_tile = tl.load(wu_ptrs, mask=w_mask, other=0.0)

        gate_acc += tl.dot(a_tile, wg_tile, allow_tf32=False)
        up_acc += tl.dot(a_tile, wu_tile, allow_tf32=False)

    # SwiGLU epilogue
    result = gate_acc * tl.sigmoid(gate_acc) * up_acc

    offs_m_out = m_start + tl.arange(0, BLOCK_M)
    offs_n_out = n_start + tl.arange(0, BLOCK_N)
    out_mask = (offs_m_out[:, None] < m_end) & (offs_n_out[None, :] < n_end)
    out_ptrs = out_ptr + offs_m_out[:, None] * stride_om + offs_n_out[None, :] * stride_on
    tl.store(out_ptrs, result.to(out_ptr.dtype.element_ty), mask=out_mask)


@triton.autotune(
    configs=[
        # Small tile configs
        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 64, 'GROUP_M': 8},
                       num_warps=4, num_stages=3),
        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 64, 'GROUP_M': 8},
                       num_warps=4, num_stages=3),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 64, 'GROUP_M': 8},
                       num_warps=4, num_stages=3),
        # Medium tile configs
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32, 'GROUP_M': 8},
                       num_warps=8, num_stages=3),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64, 'GROUP_M': 8},
                       num_warps=8, num_stages=2),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 128, 'GROUP_M': 8},
                       num_warps=8, num_stages=3),
        # Large tile configs
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 32, 'GROUP_M': 8},
                       num_warps=8, num_stages=2),
        triton.Config({'BLOCK_M': 256, 'BLOCK_N': 64, 'BLOCK_K': 64, 'GROUP_M': 8},
                       num_warps=8, num_stages=2),
        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 256, 'BLOCK_K': 32, 'GROUP_M': 8},
                       num_warps=8, num_stages=3),
    ],
    key=['K', 'N'],
)
@triton.jit
def _grouped_gemm_autotuned(
    a_ptr,
    b_ptr,
    c_ptr,
    tile_expert_ptr,
    tile_m_start_ptr,
    expert_offsets_ptr,
    stride_am, stride_ak,
    stride_be, stride_bk, stride_bn,
    stride_cm, stride_cn,
    K, N,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    GROUP_M: tl.constexpr,
):
    """Autotuned grouped GEMM kernel."""
    pid = tl.program_id(0)
    num_pid_m = tl.num_programs(0) // tl.cdiv(N, BLOCK_N)
    num_pid_n = tl.cdiv(N, BLOCK_N)

    num_pid_in_group = GROUP_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_M
    group_size_m = tl.minimum(num_pid_m - 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

    expert_id = tl.load(tile_expert_ptr + pid_m)
    m_start = tl.load(tile_m_start_ptr + pid_m)
    expert_end = tl.load(expert_offsets_ptr + expert_id + 1)
    m_end = tl.minimum(m_start + BLOCK_M, expert_end)

    n_start = pid_n * BLOCK_N
    n_end = tl.minimum(n_start + BLOCK_N, N)

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    for k_start in range(0, K, BLOCK_K):
        k_end = tl.minimum(k_start + BLOCK_K, K)

        offs_m = m_start + tl.arange(0, BLOCK_M)
        offs_k = k_start + tl.arange(0, BLOCK_K)
        a_mask = (offs_m[:, None] < m_end) & (offs_k[None, :] < k_end)
        a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
        a_tile = tl.load(a_ptrs, mask=a_mask, other=0.0)

        offs_k2 = k_start + tl.arange(0, BLOCK_K)
        offs_n = n_start + tl.arange(0, BLOCK_N)
        b_mask = (offs_k2[:, None] < k_end) & (offs_n[None, :] < n_end)
        b_ptrs = (b_ptr +
                  expert_id * stride_be +
                  offs_k2[:, None] * stride_bk +
                  offs_n[None, :] * stride_bn)
        b_tile = tl.load(b_ptrs, mask=b_mask, other=0.0)

        acc += tl.dot(a_tile, b_tile, allow_tf32=False)

    offs_m_out = m_start + tl.arange(0, BLOCK_M)
    offs_n_out = n_start + tl.arange(0, BLOCK_N)
    out_mask = (offs_m_out[:, None] < m_end) & (offs_n_out[None, :] < n_end)
    out_ptrs = c_ptr + offs_m_out[:, None] * stride_cm + offs_n_out[None, :] * stride_cn
    tl.store(out_ptrs, acc.to(c_ptr.dtype.element_ty), mask=out_mask)


def _compute_tile_mapping(
    expert_offsets: torch.Tensor,
    E: int,
    BLOCK_M: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Build (expert_idx, m_start) arrays for each output M-tile."""
    offsets_cpu = expert_offsets.cpu()
    tile_expert: list[int] = []
    tile_m_start: list[int] = []
    for e in range(E):
        n_e = int(offsets_cpu[e + 1].item()) - int(offsets_cpu[e].item())
        if n_e == 0:
            continue
        n_tiles = (n_e + BLOCK_M - 1) // BLOCK_M
        base = int(offsets_cpu[e].item())
        for t in range(n_tiles):
            tile_expert.append(e)
            tile_m_start.append(base + t * BLOCK_M)

    device = expert_offsets.device
    return (torch.tensor(tile_expert, dtype=torch.int32, device=device),
            torch.tensor(tile_m_start, dtype=torch.int32, device=device))


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._tile_mapping = None
        self._mapping_key = None

    def _get_config(self):
        """Get tile config based on problem dimensions."""
        if self.H >= 4096:
            return {'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64,
                    'GROUP_M': 8, 'num_warps': 8, 'num_stages': 2}
        elif self.I >= 4096:
            return {'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 64,
                    'GROUP_M': 8, 'num_warps': 8, 'num_stages': 3}
        else:
            return {'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 64,
                    'GROUP_M': 8, 'num_warps': 8, 'num_stages': 3}

    def forward(
        self,
        hidden_states: torch.Tensor,
        expert_offsets: torch.Tensor,
    ) -> torch.Tensor:
        T_perm, H = hidden_states.shape
        I, E = self.I, self.E

        cfg = self._get_config()
        BLOCK_M = cfg['BLOCK_M']
        BLOCK_N = cfg['BLOCK_N']
        BLOCK_K = cfg['BLOCK_K']
        GROUP_M = cfg['GROUP_M']
        num_warps = cfg['num_warps']
        num_stages = cfg['num_stages']

        # Tile mapping
        key = (T_perm, int(expert_offsets[-1].item()), BLOCK_M)
        if self._tile_mapping is None or self._mapping_key != key:
            self._tile_mapping = _compute_tile_mapping(expert_offsets, E, BLOCK_M)
            self._mapping_key = key
        tile_expert, tile_m_start = self._tile_mapping

        num_m_tiles = tile_expert.shape[0]
        num_n_tiles = triton.cdiv(I, BLOCK_N)
        grid = (num_m_tiles * num_n_tiles,)

        # Gate projection
        gate_out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device)

        _grouped_gemm_kernel[grid](
            hidden_states, self.W_gate, gate_out,
            tile_expert, tile_m_start, expert_offsets,
            hidden_states.stride(0), hidden_states.stride(1),
            self.W_gate.stride(0), self.W_gate.stride(1), self.W_gate.stride(2),
            gate_out.stride(0), gate_out.stride(1),
            K=H, N=I,
            BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
            GROUP_M=GROUP_M,
            num_warps=num_warps, num_stages=num_stages,
        )

        # Up projection
        up_out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device)

        _grouped_gemm_kernel[grid](
            hidden_states, self.W_up, up_out,
            tile_expert, tile_m_start, expert_offsets,
            hidden_states.stride(0), hidden_states.stride(1),
            self.W_up.stride(0), self.W_up.stride(1), self.W_up.stride(2),
            up_out.stride(0), up_out.stride(1),
            K=H, N=I,
            BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
            GROUP_M=GROUP_M,
            num_warps=num_warps, num_stages=num_stages,
        )

        # SwiGLU fusion
        gated_up = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device)
        total_el = T_perm * I
        BLOCK_SIZE = 1024
        swiglu_grid = (triton.cdiv(total_el, BLOCK_SIZE),)
        _swiglu_kernel[swiglu_grid](
            gate_out, up_out, gated_up,
            total_el,
            gated_up.stride(0), gated_up.stride(1),
            BLOCK_SIZE=BLOCK_SIZE,
            num_warps=4,
        )

        return gated_up


# Module-level shape shims
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]