import os
import torch
import triton
import pytest
import torch.nn.functional as F
from aiter.ops.triton.sage_attention import sageattn_qk_int8_pv_fp16


os.environ['AMDGPU_USE_BUFFER_OPS'] = "1"


def ref_sdp_attn(q, k, v, tensor_layout="HND", is_causal=False):
    if tensor_layout == "HND":
        q_sdp = q
        k_sdp = k
        v_sdp = v
    elif tensor_layout == "NHD":
        q_sdp = q.permute(0, 2, 1, 3)  # (B, L, H, D) -> (B, H, L, D)
        k_sdp = k.permute(0, 2, 1, 3)
        v_sdp = v.permute(0, 2, 1, 3)
    else:
        raise ValueError(f"Unsupported tensor_layout: {tensor_layout}")

    o = F.scaled_dot_product_attention(
        q_sdp,
        k_sdp,
        v_sdp,
        is_causal=is_causal,
    )

    return o.permute(0, 2, 1, 3) if tensor_layout == "NHD" else o


@pytest.mark.parametrize(
    "batch, num_qo_heads, qo_len, num_kv_heads, kv_len, head_dim, tensor_layout, is_causal",
    [
        (1, 14040, 40, 14040, 40, 128, "NHD", False),
        (1, 14040, 40, 512, 40, 128, "NHD", False),
        (2, 14040, 40, 14040, 40, 128, "NHD", False),
        (2, 14040, 40, 512, 40, 128, "NHD", False),
        (2, 24, 1280, 24, 1280, 128, "HND", False),
        (1, 24, 1280, 24, 1280, 128, "HND", False),
        (1, 24, 4352, 24, 4352, 128, "HND", False),
        (2, 24, 4352, 24, 4352, 128, "HND", False),
        (1, 19440, 1, 19440, 1, 128, "NHD", False),
        (1, 2430, 40, 512, 40, 128, "NHD", False),
        (1, 57600, 1, 57600, 1, 128, "NHD", False),
        (1, 7200, 40, 512, 40, 128, "NHD", False),
        (1, 14080, 1, 14080, 1, 128, "NHD", False),
        (1, 3520, 24, 512, 24, 128, "NHD", False),
    ]
)
@pytest.mark.parametrize("dtype", [torch.float16])
def test_qk_int8_pv_fp16(batch, num_qo_heads, qo_len, num_kv_heads, kv_len, head_dim, tensor_layout, is_causal, dtype):
    print(
        f"Test batch: {batch}, num_qo_heads: {num_qo_heads}, qo_len: {qo_len}, "
        f"num_kv_heads: {num_kv_heads}, kv_len: {kv_len}, head_dim: {head_dim}, "
        f"tensor_layout: {tensor_layout}, is_causal: {is_causal}"
    )

    device = "cuda"

    q = torch.randn((batch, num_qo_heads, qo_len, head_dim), dtype=dtype, device=device)
    k = torch.randn((batch, num_kv_heads, kv_len, head_dim), dtype=dtype, device=device)
    v = torch.randn((batch, num_kv_heads, kv_len, head_dim), dtype=dtype, device=device)

    o = sageattn_qk_int8_pv_fp16(q, k, v, tensor_layout=tensor_layout, is_causal=is_causal)

    # use F.scaled_dot_product_attention to validate result: o
    ref = ref_sdp_attn(q, k, v, tensor_layout=tensor_layout, is_causal=is_causal)
    ref = ref.to(o.dtype)

    diff = (ref - o).float()
    max_abs_err = diff.abs().max().item()
    mean_abs_err = diff.abs().mean().item()
    rel_err = (diff.abs() / (ref.abs() + 1e-6)).max().item()

    print(
        f"  validation: max_abs_err={max_abs_err:.4e}, "
        f"mean_abs_err={mean_abs_err:.4e}, "
        f"max_rel_err={rel_err:.4e}"
    )

    assert torch.allclose(o, ref, atol=2e-2, rtol=2e-2)

