import torch

from aiter.ops.triton._triton_kernels.attention.fp8_mqa_logits import (
    _fp8_mqa_logits_kernel,
    _fp8_mqa_clear_logits_kernel,
)


def fp8_mqa_logits(
    Q,
    KV,
    kv_scales,
    weights,
    cu_starts,
    cu_ends,
):
    """
    This function computes the logits to be used by a topk function for sparse attention.

    Q:           [seq_len, NUM_HEADS, HEAD_SIZE], dtype float8
    KV:          [seq_len_kv, HEAD_SIZE], dtype float8
    kv_scales:   [seq_len_kv], dtype float32
    weights:     [seq_len, NUM_HEADS], dtype float32
    cu_starts:   [seq_len], dtype int32, start indices
    cu_ends:     [seq_len], dtype int32, end indices

    Returns:
    logits:      [seq_len, seq_len_kv], dtype float32 (must be initialized to -inf, because of causal masking)
    """
    #BLOCK_KV = 128
    BLOCK_KV = 64
    seq_len, num_heads, head_size = Q.shape
    seq_len_kv = KV.shape[0]
    logits_masking = (seq_len_kv % BLOCK_KV != 0)
    # TODO: Currently assuming num_heads and head_size is power of 2.
    assert num_heads & (num_heads - 1) == 0, "num q. heads should be power of 2."
    assert head_size & (head_size - 1) == 0, "head size should be power of 2."
    # Initialize with -inf because of causal masking
    if logits_masking:
        logits = torch.full(
            (seq_len, seq_len_kv),
            fill_value=-float("inf"),
            dtype=torch.float32,
            device=Q.device,
        )
    else:
        logits = torch.empty(
            (seq_len, seq_len_kv),
            dtype=torch.float32,
            device=Q.device,
        )

    stride_q_s, stride_q_h, stride_q_d = Q.stride()
    stride_kv_s, stride_kv_d = KV.stride()
    stride_w_s, stride_w_h = weights.stride()
    stride_logits_s, stride_logits_k = logits.stride()

    # heuristic for MFMA instruction shape
    #matrix_instr_nonkdim = 32
    #if seq_len <= 1024:
    #    matrix_instr_nonkdim = 16
    matrix_instr_nonkdim = 16

    _fp8_mqa_logits_kernel[(seq_len,)](
        Q_ptr=Q,
        KV_ptr=KV,
        kv_scales_ptr=kv_scales,
        weights_ptr=weights,
        cu_start_ptr=cu_starts,
        cu_end_ptr=cu_ends,
        logits_ptr=logits,
        seq_len=seq_len,
        seq_len_kv=seq_len_kv,
        NUM_HEADS=num_heads,
        HEAD_SIZE=head_size,
        stride_q_s=stride_q_s,
        stride_q_h=stride_q_h,
        stride_q_d=stride_q_d,
        stride_kv_s=stride_kv_s,
        stride_kv_d=stride_kv_d,
        stride_w_s=stride_w_s,
        stride_w_h=stride_w_h,
        stride_logits_s=stride_logits_s,
        stride_logits_k=stride_logits_k,
        BLOCK_KV=BLOCK_KV,
        num_warps=4,
        ##num_stages=2,
        num_stages=1,
        waves_per_eu=2,
        matrix_instr_nonkdim=matrix_instr_nonkdim,
        LOGITS_MASKING=logits_masking,
    )

    # Clear per-row invalid KV positions (outside [cu_start, cu_end)) to -inf.
    if not logits_masking:
        _fp8_mqa_clear_logits_kernel[(seq_len,)](
            logits_ptr=logits,
            cu_start_ptr=cu_starts,
            cu_end_ptr=cu_ends,
            seq_len_kv=seq_len_kv,
            stride_logits_s=stride_logits_s,
            stride_logits_k=stride_logits_k,
            BLOCK_KV=256,
            waves_per_eu=2,
            num_warps=4,
            num_stages=1,
        )

    return logits
