"""Paged-attention decode: Triton flash-decoding (split-K) baseline.

Two kernels per forward:
  1. _split_kernel: grid (S, B*Hkv). Each program streams a contiguous range of
     pages for one (batch, kv_head), computes online-softmax partials for the
     G query heads in that GQA group (padded to 16 rows for tl.dot).
  2. _reduce_kernel: grid (B*Hkv,). Merges the S partials into the final bf16 out.
"""
import math

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

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

BATCH = 8
NUM_HEADS = 32
NUM_KV_HEADS = 8
HEAD_DIM = 128
SEQ_LEN = 1024
PAGE_SIZE = 16

LOG2E = tl.constexpr(1.4426950408889634)


@triton.jit
def _split_kernel(
    Q, KV, BT, SL,
    Opart, Mpart, Lpart,
    scale,
    num_splits,
    stride_qb,            # H*D
    stride_bt,            # max_blocks per row of block table
    kv_page_stride,       # PAGE * Hkv * 2D
    kv_row_stride,        # Hkv * 2D
    HKV: tl.constexpr,
    G: tl.constexpr,
    D: tl.constexpr,
    PAGE: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    pid_s = tl.program_id(0)
    pid_bh = tl.program_id(1)
    b = pid_bh // HKV
    hkv = pid_bh % HKV

    len_b = tl.load(SL + b)
    np_b = (len_b + PAGE - 1) // PAGE
    chunk = (np_b + num_splits - 1) // num_splits
    p0 = pid_s * chunk
    p1 = tl.minimum(p0 + chunk, np_b)

    offs_m = tl.arange(0, 16)
    offs_d = tl.arange(0, D)
    offs_n = tl.arange(0, BLOCK_N)

    part_base = (pid_bh * num_splits + pid_s) * 16
    if p0 >= p1:
        # Empty split: mark l=0, m=-inf; reduce kernel masks these out.
        tl.store(Mpart + part_base + offs_m, tl.full((16,), float("-inf"), tl.float32))
        tl.store(Lpart + part_base + offs_m, tl.zeros((16,), tl.float32))
        return

    qmask = offs_m[:, None] < G
    q = tl.load(
        Q + b * stride_qb + (hkv * G + offs_m)[:, None] * D + offs_d[None, :],
        mask=qmask, other=0.0,
    )

    m_i = tl.full((16,), float("-inf"), tl.float32)
    l_i = tl.zeros((16,), tl.float32)
    acc = tl.zeros((16, D), tl.float32)
    qk_scale = scale * LOG2E

    for start in range(p0 * PAGE, p1 * PAGE, BLOCK_N):
        tok = start + offs_n
        page_idx = tok // PAGE
        in_page = tok % PAGE
        pg = tl.load(BT + b * stride_bt + page_idx, mask=page_idx < np_b, other=0)
        row = pg * kv_page_stride + in_page * kv_row_stride + hkv * (2 * D)
        kv_ptrs = KV + row[:, None] + offs_d[None, :]
        tmask = (tok < len_b)[:, None]
        k = tl.load(kv_ptrs, mask=tmask, other=0.0)
        v = tl.load(kv_ptrs + D, mask=tmask, other=0.0)

        s = tl.dot(q, tl.trans(k)) * qk_scale
        s = tl.where(tok[None, :] < len_b, s, float("-inf"))
        m_new = tl.maximum(m_i, tl.max(s, 1))
        alpha = tl.exp2(m_i - m_new)
        p = tl.exp2(s - m_new[:, None])
        l_i = l_i * alpha + tl.sum(p, 1)
        acc = acc * alpha[:, None] + tl.dot(p.to(tl.bfloat16), v)
        m_i = m_new

    tl.store(Mpart + part_base + offs_m, m_i)
    tl.store(Lpart + part_base + offs_m, l_i)
    optr = Opart + part_base * D + offs_m[:, None] * D + offs_d[None, :]
    tl.store(optr, acc)


@triton.jit
def _reduce_kernel(
    Opart, Mpart, Lpart, Out,
    num_splits,
    stride_ob,            # H*D
    HKV: tl.constexpr,
    G: tl.constexpr,
    D: tl.constexpr,
    S_POW2: tl.constexpr,
):
    pid = tl.program_id(0)
    b = pid // HKV
    hkv = pid % HKV

    offs_m = tl.arange(0, 16)
    offs_d = tl.arange(0, D)
    offs_s = tl.arange(0, S_POW2)
    smask = offs_s < num_splits

    base = pid * num_splits * 16
    m = tl.load(Mpart + base + offs_s[:, None] * 16 + offs_m[None, :],
                mask=smask[:, None], other=float("-inf"))          # (S~, 16)
    l = tl.load(Lpart + base + offs_s[:, None] * 16 + offs_m[None, :],
                mask=smask[:, None], other=0.0)
    M = tl.max(m, 0)                                               # (16,)
    w = tl.exp2(m - M[None, :])                                    # (S~, 16)
    L = tl.sum(l * w, 0)                                           # (16,)

    acc = tl.zeros((16, D), tl.float32)
    for s in range(0, num_splits):
        m_s = tl.load(Mpart + base + s * 16 + offs_m)
        l_s = tl.load(Lpart + base + s * 16 + offs_m)
        w_s = tl.exp2(m_s - M)
        o_s = tl.load(Opart + (base + s * 16) * D + offs_m[:, None] * D + offs_d[None, :],
                      mask=(l_s > 0)[:, None], other=0.0)
        acc += (w_s * tl.where(l_s > 0, 1.0, 0.0))[:, None] * o_s

    out = acc / tl.maximum(L, 1e-30)[:, None]
    optr = Out + b * stride_ob + (hkv * G + offs_m)[:, None] * D + offs_d[None, :]
    tl.store(optr, out.to(Out.dtype.element_ty), mask=offs_m[:, None] < G)


def _pick_splits(batch: int, num_kv_heads: int, num_pages: int) -> int:
    target = 512
    pairs = batch * num_kv_heads
    s = max(1, (target + pairs - 1) // pairs)
    # keep at least 4 pages per split
    s = min(s, max(1, num_pages // 4))
    return min(s, 32)


class Model(nn.Module):
    def __init__(self, batch, num_heads, num_kv_heads, head_dim, seq_len, page_size):
        super().__init__()
        assert num_heads % num_kv_heads == 0
        self.batch = batch
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = head_dim
        self.seq_len = seq_len
        self.page_size = page_size
        self.group_size = num_heads // num_kv_heads
        self.scale = 1.0 / math.sqrt(head_dim)
        self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16), persistent=False)

        num_pages = (seq_len + page_size - 1) // page_size
        self.num_splits = _pick_splits(batch, num_kv_heads, num_pages)
        self._scratch = None

    def _get_scratch(self, device):
        if self._scratch is None or self._scratch[0].device != device:
            B, Hkv, S, D = self.batch, self.num_kv_heads, self.num_splits, self.head_dim
            opart = torch.empty(B * Hkv * S * 16 * D, dtype=torch.float32, device=device)
            mpart = torch.empty(B * Hkv * S * 16, dtype=torch.float32, device=device)
            lpart = torch.empty(B * Hkv * S * 16, dtype=torch.float32, device=device)
            self._scratch = (opart, mpart, lpart)
        return self._scratch

    def forward(self, query, kv_cache, block_table, seq_lens):
        B, H, D = query.shape
        Hkv = self.num_kv_heads
        G = self.group_size
        P = self.page_size
        S = self.num_splits
        device = query.device

        opart, mpart, lpart = self._get_scratch(device)
        out = torch.empty_like(query)

        kv_row_stride = kv_cache.stride(1)      # Hkv * 2D
        kv_page_stride = kv_cache.stride(0)     # P * Hkv * 2D

        BLOCK_N = 128 if D == 128 else 128
        grid = (S, B * Hkv)
        _split_kernel[grid](
            query, kv_cache, block_table, seq_lens,
            opart, mpart, lpart,
            self.scale, S,
            H * D,
            block_table.stride(0),
            kv_page_stride, kv_row_stride,
            HKV=Hkv, G=G, D=D, PAGE=P, BLOCK_N=BLOCK_N,
            num_warps=4, num_stages=3,
        )
        s_pow2 = triton.next_power_of_2(S)
        _reduce_kernel[(B * Hkv,)](
            opart, mpart, lpart, out,
            S,
            H * D,
            HKV=Hkv, G=G, D=D, S_POW2=s_pow2,
            num_warps=4, num_stages=2,
        )
        return out


def get_inputs():
    B, H, Hkv, D, L, P = BATCH, NUM_HEADS, NUM_KV_HEADS, HEAD_DIM, SEQ_LEN, PAGE_SIZE
    pages_per_seq = (L + P - 1) // P
    total_pages = max(B * pages_per_seq + 8, 64)
    query = torch.randn(B, H, D, dtype=torch.bfloat16) * 0.1
    kv_cache = torch.randn(total_pages, P, Hkv, 2 * D, dtype=torch.bfloat16) * 0.1
    perm = torch.randperm(total_pages)[: B * pages_per_seq].reshape(B, pages_per_seq).int()
    block_table = perm.contiguous()
    seq_lens = torch.full((B,), L, dtype=torch.int32)
    return [query, kv_cache, block_table, seq_lens]


def get_init_inputs():
    return [BATCH, NUM_HEADS, NUM_KV_HEADS, HEAD_DIM, SEQ_LEN, PAGE_SIZE]