import random
import torch
import pytest

from aiter.ops.triton.grouped_decode_attention import decode_attention_fwd_grouped, decode_attention_fwd_normal


def _set_all_seeds(seed):
    """Set all random seeds for reproducibility."""
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


def setUp():
    # Set seeds before each test method
    _set_all_seeds(42)


@pytest.mark.parametrize(
    "B, S, H_Q, H_KV, D, D_V",
    [
        (1, 8, 16, 1, 576, 512),
        (1, 13, 16, 1, 576, 512),
        (1, 135, 16, 1, 576, 512),
        (4, 152, 16, 1, 576, 512),
        (4, 256, 16, 1, 576, 512),
        (4, 593, 16, 1, 576, 512),
    ]
)
@pytest.mark.parametrize("dtype", [torch.float16])
@pytest.mark.parametrize("max_kv_splits", [16])
def test_grouped_decode_attention(B, S, H_Q, H_KV, D, D_V, max_kv_splits, dtype):
    seq_len = S  # This represents the number of tokens already in the sequence
    total_tokens = B * seq_len
    sm_scale = 1.0 / (D**0.5)
    num_kv_splits = torch.full((B,), 4, dtype=torch.int32, device="cuda")

    # q represents the new token being generated, one per batch
    q = torch.randn(B, H_Q, D, dtype=dtype, device="cuda")

    # k_buffer and v_buffer represent all previous tokens
    k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
    v_buffer = torch.randn(total_tokens, H_KV, D_V, dtype=dtype, device="cuda")

    # o will have the same shape as q
    o = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")
    o_grouped = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")

    b_seq_len = torch.full((B,), seq_len, device="cuda")

    kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device="cuda")
    kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len[:B], dim=0)
    kv_indices = torch.arange(total_tokens, dtype=torch.int32, device="cuda")

    attn_logits = torch.empty(
        (B, H_Q, max_kv_splits, D_V),
        dtype=torch.float32,
        device="cuda",
    )
    attn_lse = torch.empty(
        (B, H_Q, max_kv_splits),
        dtype=torch.float32,
        device="cuda",
    )

    decode_attention_fwd_normal(
        q,
        k_buffer,
        v_buffer,
        o,
        kv_indptr,
        kv_indices,
        attn_logits,
        attn_lse,
        num_kv_splits,
        max_kv_splits,
        sm_scale,
    )

    attn_logits1 = torch.empty(
        (B, H_Q, max_kv_splits, D_V),
        dtype=torch.float32,
        device="cuda",
    )
    attn_lse1 = torch.empty(
        (B, H_Q, max_kv_splits, D_V),
        dtype=torch.float32,
        device="cuda",
    )

    decode_attention_fwd_grouped(
        q,
        k_buffer,
        v_buffer,
        o_grouped,
        kv_indptr,
        kv_indices,
        attn_logits1,
        attn_lse1,
        num_kv_splits,
        max_kv_splits,
        sm_scale,
    )

    cos_sim = torch.nn.functional.cosine_similarity(
        o.flatten(), o_grouped.flatten(), dim=0
    )
    assert cos_sim.item() > 0.99
    assert torch.allclose(o, o_grouped, atol=2e-2)


if __name__ == "__main__":
    test_grouped_decode_attention(1, 8, 16, 1, 576, 512, 16, torch.float16)
