import triton
import triton.language as tl


def _get_autotune_configs():
    configs = []
#    for BLOCK_Q in [1, 2, 4]:
    for BLOCK_KV in [64, 128, 256]:
        for num_stages in [1, 2]:
            for num_warps in [4, 8]:
                for waves_per_eu in [0, 2]:
                    configs.append(
                        triton.Config(
                            {
                                #"BLOCK_Q": BLOCK_Q,
                                "BLOCK_KV": BLOCK_KV,
                                "waves_per_eu": waves_per_eu,
                            },
                            num_stages=num_stages,
                            num_warps=num_warps,
                        )
                    )
    return configs


def _get_clear_autotune_configs():
    configs = []
    for BLOCK_KV in [64, 128, 256, 512, 1024, 2048]:
        for num_warps in [2, 4, 8]:
            for num_stages in [1, 2]:
                for waves_per_eu in [0, 2]:
                    configs.append(
                        triton.Config(
                            {"BLOCK_KV": BLOCK_KV, "waves_per_eu": waves_per_eu},
                            num_warps=num_warps,
                            num_stages=num_stages,
                        )
                    )
    return configs
#@triton.autotune(
#    configs=_get_autotune_configs(),
#    #configs=[
#    #    triton.Config(
#    #        {"BLOCK_Q": 1, "BLOCK_KV": 64, "waves_per_eu": 2}, num_warps=4, num_stages=1)
#    #],
#    key=["NUM_HEADS", "HEAD_SIZE", "seq_len_kv"],
#)
@triton.jit
def _fp8_mqa_logits_kernel(
    Q_ptr,  # fp8e4m3 [seq_len, H, D]
    KV_ptr,  # fp8e4m3 [seq_len_kv, D]
    kv_scales_ptr,  # fp32 [seq_len_kv]
    weights_ptr,  # fp32 [seq_len, H]
    cu_start_ptr,  # int32 [seq_len]
    cu_end_ptr,  # int32 [seq_len]
    logits_ptr,  # fp32 [seq_len, seq_len_kv]
    seq_len,
    seq_len_kv,
    NUM_HEADS: tl.constexpr,
    HEAD_SIZE: tl.constexpr,
    # strides
    stride_q_s: tl.int64,
    stride_q_h: tl.constexpr,
    stride_q_d: tl.constexpr,
    stride_kv_s: tl.int64,
    stride_kv_d: tl.constexpr,
    stride_w_s: tl.int64,
    stride_w_h: tl.constexpr,
    stride_logits_s: tl.int64,
    stride_logits_k: tl.int64,
    # block sizes
    BLOCK_KV: tl.constexpr,
    LOGITS_MASKING: tl.constexpr,
):
    row_id = tl.program_id(0)
    # go from larger to smaller in terms of work
    # to reduce the tail effect
    row_id = tl.num_programs(0) - row_id - 1
    tl.assume(row_id >= 0)
    tl.assume(stride_q_s > 0)
    tl.assume(stride_q_h > 0)
    tl.assume(stride_q_d > 0)
    tl.assume(stride_kv_s > 0)
    tl.assume(stride_kv_d > 0)
    tl.assume(stride_w_s > 0)
    tl.assume(stride_w_h > 0)

    logits_row_ptrs = logits_ptr + row_id * stride_logits_s

    h_inds = tl.arange(0, NUM_HEADS)[:, None]
    d_inds = tl.arange(0, HEAD_SIZE)

    # load Q[BLOCK_Q, NUM_HEADS, HEAD_SIZE]
    q_ptrs = (
        Q_ptr + row_id * stride_q_s + h_inds * stride_q_h + d_inds[None, :] * stride_q_d
    )

    q_block = tl.load(q_ptrs, cache_modifier=".cg")
    w_ptrs = weights_ptr + row_id * stride_w_s + h_inds * stride_w_h
    w_block = tl.load(w_ptrs, cache_modifier=".cg").to(tl.float32)

    # Load start/end for each row in this block
    start_ind = tl.load(cu_start_ptr + row_id)
    end_ind = tl.load(cu_end_ptr + row_id)

    if LOGITS_MASKING:
        start_ind = tl.maximum(start_ind, 0)
        end_ind = tl.minimum(end_ind, seq_len_kv)
    else:
        start_ind = tl.maximum(start_ind, 0) // BLOCK_KV * BLOCK_KV
        end_ind = tl.cdiv(tl.minimum(end_ind, seq_len_kv), BLOCK_KV) * BLOCK_KV
    shifted_end = end_ind - start_ind
    shifted_unmasked_end = shifted_end // BLOCK_KV * BLOCK_KV

    kv_col_offsets = tl.arange(0, BLOCK_KV) + start_ind
    kv_ptrs = (
        KV_ptr + kv_col_offsets[None, :] * stride_kv_s + d_inds[:, None] * stride_kv_d
    )

    kv_scales_ptrs = kv_scales_ptr + kv_col_offsets

    logits_ptrs = logits_row_ptrs + kv_col_offsets * stride_logits_k

    # Loop over KV tiles
    for _ in tl.range(0, shifted_unmasked_end, BLOCK_KV):
        kv_block = tl.load(kv_ptrs)
        kv_scales = tl.load(kv_scales_ptrs)

        # [NUM_HEADS, BLOCK_KV] = [NUM_HEADS, HEAD_SIZE] x [HEAD_SIZE, BLOCK_KV]
        scores = tl.dot(q_block, kv_block, input_precision="ieee")
        # Multiply by kv_scales (broadcast along rows)
        scores = scores * kv_scales[None, :]
        # ReLU
        scores = tl.maximum(scores, 0.0)
        scores = scores * w_block
        # [NUM_HEADS, BLOCK_KV] -> [BLOCK_KV, ]
        scores = tl.sum(scores, axis=0)
        tl.store(logits_ptrs, scores)

        kv_ptrs += BLOCK_KV * stride_kv_s
        kv_scales_ptrs += BLOCK_KV
        logits_ptrs += BLOCK_KV * stride_logits_k
        kv_col_offsets += BLOCK_KV

    # masked load
    if LOGITS_MASKING:
        kv_col_mask = kv_col_offsets < end_ind
        kv_block = tl.load(kv_ptrs, mask=kv_col_mask[None, :], other=0.0)
        kv_scales = tl.load(kv_scales_ptrs, mask=kv_col_mask, other=0.0)

        # [NUM_HEADS, BLOCK_KV] = [NUM_HEADS, HEAD_SIZE] x [HEAD_SIZE, BLOCK_KV]
        scores = tl.dot(q_block, kv_block, input_precision="ieee")
        # Multiply by kv_scales (broadcast along rows)
        scores = scores * kv_scales[None, :]
        # ReLU
        scores = tl.maximum(scores, 0.0)
        scores = scores * w_block
        # [NUM_HEADS, BLOCK_KV] -> [BLOCK_KV, ]
        scores = tl.sum(scores, axis=0)
        # masked store
        in_window = (kv_col_offsets >= start_ind) & (kv_col_offsets < end_ind)
        tl.store(logits_ptrs, scores, mask=in_window)


#@triton.autotune(
#    configs=_get_autotune_configs(),
#    #configs=[
#    #    triton.Config(
#    #        {"BLOCK_Q": 1, "BLOCK_KV": 64, "waves_per_eu": 2}, num_warps=4, num_stages=1)
#    #],
#    key=["NUM_HEADS", "HEAD_SIZE", "seq_len_kv"],
#)
@triton.jit
def _fp8_mqa_logits_kernel_grouped(
    Q_ptr,  # fp8e4m3 [seq_len, H, D]
    KV_ptr,  # fp8e4m3 [seq_len_kv, D]
    kv_scales_ptr,  # fp32 [seq_len_kv]
    weights_ptr,  # fp32 [seq_len, H]
    cu_start_ptr,  # int32 [seq_len]
    cu_end_ptr,  # int32 [seq_len]
    logits_ptr,  # fp32 [seq_len, seq_len_kv]
    seq_len,
    seq_len_kv,
    NUM_HEADS: tl.constexpr,
    HEAD_SIZE: tl.constexpr,
    # strides
    stride_q_s: tl.int64,
    stride_q_h: tl.constexpr,
    stride_q_d: tl.constexpr,
    stride_kv_s: tl.int64,
    stride_kv_d: tl.constexpr,
    stride_w_s: tl.int64,
    stride_w_h: tl.constexpr,
    stride_logits_s: tl.int64,
    stride_logits_k: tl.int64,
    # block sizes
    BLOCK_Q: tl.constexpr,
    BLOCK_KV: tl.constexpr,
):
    block_id = tl.program_id(0)
    # go from larger to smaller in terms of work
    # to reduce the tail effect
    block_id = tl.num_programs(0) - block_id - 1
    tl.assume(block_id >= 0)
    tl.assume(stride_q_s > 0)
    tl.assume(stride_q_h > 0)
    tl.assume(stride_q_d > 0)
    tl.assume(stride_kv_s > 0)
    tl.assume(stride_kv_d > 0)
    tl.assume(stride_w_s > 0)
    tl.assume(stride_w_h > 0)

    seq_start = block_id * BLOCK_Q

    # Scalar-load cu_start/cu_end for each row in the block and reduce to the
    # union range [cu_k_s_min, cu_k_e_max).  BLOCK_Q is a compile-time constant
    # (typically 1/2/4), so tl.static_range unrolls this completely.
    cu_k_s_min = seq_len_kv
    cu_k_e_max = 0
    for bq_i in tl.static_range(BLOCK_Q):
        s_i = tl.load(cu_start_ptr + seq_start + bq_i)
        e_i = tl.load(cu_end_ptr + seq_start + bq_i)
        s_i = tl.maximum(s_i, 0)
        e_i = tl.minimum(e_i, seq_len_kv)
        cu_k_s_min = tl.minimum(cu_k_s_min, s_i)
        cu_k_e_max = tl.maximum(cu_k_e_max, e_i)
    cu_k_s_min = tl.maximum(cu_k_s_min, 0)
    cu_k_e_max = tl.minimum(cu_k_e_max, seq_len_kv)

    # Round up to full BLOCK_KV blocks — the clear kernel handles out-of-window positions.
    total_blocks = tl.cdiv(tl.maximum(cu_k_e_max - cu_k_s_min, 0), BLOCK_KV)

    h_inds = tl.arange(0, NUM_HEADS)[:, None]
    d_inds = tl.arange(0, HEAD_SIZE)

    kv_col_offsets = tl.arange(0, BLOCK_KV) + cu_k_s_min
    kv_ptrs = (
        KV_ptr + kv_col_offsets[None, :] * stride_kv_s + d_inds[:, None] * stride_kv_d
    )
    kv_scales_ptrs = kv_scales_ptr + kv_col_offsets

    # Loop over all KV tiles (including partial tail) — no masking needed,
    # _fp8_mqa_clear_logits_kernel will set invalid positions to -inf.
    for _ in tl.range(0, total_blocks):
        kv_block = tl.load(kv_ptrs)
        kv_scales = tl.load(kv_scales_ptrs)

        for bq_i in tl.static_range(BLOCK_Q):
            q_ptrs = (
                Q_ptr + (seq_start + bq_i) * stride_q_s
                + h_inds * stride_q_h + d_inds[None, :] * stride_q_d
            )
            q_block = tl.load(q_ptrs, cache_modifier=".cg")

            w_ptrs = (
                weights_ptr + (seq_start + bq_i) * stride_w_s
                + h_inds * stride_w_h
            )
            w_block = tl.load(w_ptrs, cache_modifier=".cg").to(tl.float32)

            # [NUM_HEADS, BLOCK_KV] = [NUM_HEADS, HEAD_SIZE] x [HEAD_SIZE, BLOCK_KV]
            scores = tl.dot(q_block, kv_block, input_precision="ieee")
            scores = scores * kv_scales[None, :]
            # ReLU
            scores = tl.maximum(scores, 0.0)
            scores = scores * w_block
            # [NUM_HEADS, BLOCK_KV] -> [BLOCK_KV]
            scores = tl.sum(scores, axis=0)

            logits_ptrs = (
                logits_ptr + (seq_start + bq_i) * stride_logits_s
                + kv_col_offsets * stride_logits_k
            )
            tl.store(logits_ptrs, scores)

        kv_ptrs += BLOCK_KV * stride_kv_s
        kv_scales_ptrs += BLOCK_KV
        kv_col_offsets += BLOCK_KV


#@triton.autotune(
#    configs=_get_clear_autotune_configs(),
#    key=["seq_len_kv"],
#)
@triton.jit
def _fp8_mqa_clear_logits_kernel(
    logits_ptr,  # fp32 [seq_len, seq_len_kv]
    cu_start_ptr,  # int32 [seq_len]
    cu_end_ptr,  # int32 [seq_len]
    seq_len_kv,
    stride_logits_s: tl.int64,
    stride_logits_k: tl.int64,
    BLOCK_KV: tl.constexpr,
):
    """Standalone kernel to set logits outside each row's valid
    [cu_start, cu_end) window to -inf.  Grid: (seq_len,)."""
    row_id = tl.program_id(0)

    cu_start = tl.load(cu_start_ptr + row_id)
    cu_end = tl.load(cu_end_ptr + row_id)

    kv_offsets = tl.arange(0, BLOCK_KV)

    for _ in tl.range(0, tl.cdiv(seq_len_kv, BLOCK_KV)):
        invalid = (kv_offsets < cu_start) | (kv_offsets >= cu_end)
        in_bounds = kv_offsets < seq_len_kv
        store_mask = invalid & in_bounds

        logits_ptrs = (
            logits_ptr + row_id * stride_logits_s + kv_offsets * stride_logits_k
        )
        tl.store(
            logits_ptrs,
            tl.full([BLOCK_KV], float("-inf"), dtype=tl.float32),
            mask=store_mask,
        )

        kv_offsets += BLOCK_KV
