# SPDX-License-Identifier: MIT

import pytest
import torch

from aiter.ops.triton.fla.fused_sigmoid_gating import (
    fused_sigmoid_gating_delta_rule_update,
)
from op_tests.triton_tests.utils.fused_sigmoid_gating_ref import (
    fused_sigmoid_gating_delta_rule_ref,
)


@pytest.mark.parametrize("tp_size", [1])
@pytest.mark.parametrize("num_reqs", [1, 2, 4])
@pytest.mark.parametrize("num_k_heads", [16])
@pytest.mark.parametrize("num_v_heads", [32])
@pytest.mark.parametrize("head_k_dim", [128])
@pytest.mark.parametrize("head_v_dim", [128])
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
def test_fused_sigmoid_gating_delta_rule_update_non_spec(
    tp_size: int,
    num_reqs: int,
    num_k_heads: int,
    num_v_heads: int,
    head_k_dim: int,
    head_v_dim: int,
    dtype: torch.dtype,
) -> None:
    if not torch.cuda.is_available():
        pytest.skip("CUDA is required")

    torch.manual_seed(0)
    key_dim = head_k_dim * num_k_heads
    value_dim = head_v_dim * num_v_heads
    mixed_qkv_dim = (key_dim * 2 + value_dim) // tp_size
    seq_len = 1
    num_tokens = num_reqs * seq_len
    total_entries = num_tokens * 2

    mixed_qkv = torch.rand(num_tokens, mixed_qkv_dim, dtype=dtype, device="cuda")
    query, key, value = torch.split(
        mixed_qkv,
        [key_dim // tp_size, key_dim // tp_size, value_dim // tp_size],
        dim=-1,
    )
    query = query.view(1, num_tokens, num_k_heads, head_k_dim)
    key = key.view(1, num_tokens, num_k_heads, head_k_dim)
    value = value.view(1, num_tokens, num_v_heads, head_v_dim)

    A_log = torch.rand(num_v_heads // tp_size, dtype=dtype, device="cuda")
    dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype, device="cuda")
    a = torch.rand(num_tokens, num_v_heads, dtype=dtype, device="cuda")
    b = torch.rand(num_tokens, num_v_heads, dtype=dtype, device="cuda")
    ssm_state = torch.rand(total_entries, num_v_heads, head_v_dim, head_k_dim, dtype=dtype, device="cuda")
    state_indices = torch.randperm(total_entries, device="cuda", dtype=torch.int64)[:num_tokens].to(torch.int32)
    cu_seqlens = torch.arange(0, num_tokens + 1, dtype=torch.int32, device="cuda")

    core_attn_out_ref, last_recurrent_state_ref = fused_sigmoid_gating_delta_rule_ref(
        A_log=A_log,
        a=a,
        b=b,
        dt_bias=dt_bias,
        q=query,
        k=key,
        v=value,
        initial_state=ssm_state.clone(),
        inplace_final_state=True,
        ssm_state_indices=state_indices,
        cu_seqlens=cu_seqlens,
        use_qk_l2norm_in_kernel=True,
    )

    core_attn_out, last_recurrent_state = fused_sigmoid_gating_delta_rule_update(
        A_log=A_log,
        a=a,
        b=b,
        dt_bias=dt_bias,
        q=query,
        k=key,
        v=value,
        initial_state=ssm_state,
        inplace_final_state=True,
        ssm_state_indices=state_indices,
        cu_seqlens=cu_seqlens,
        use_qk_l2norm_in_kernel=True,
    )

    torch.testing.assert_close(core_attn_out.float(), core_attn_out_ref.float(), atol=1e-2, rtol=1e-2)
    torch.testing.assert_close(
        last_recurrent_state.float(),
        last_recurrent_state_ref.float(),
        atol=1e-2,
        rtol=1e-2,
    )


@pytest.mark.parametrize("tp_size", [1])
@pytest.mark.parametrize("num_reqs", [1, 2, 4])
@pytest.mark.parametrize("num_k_heads", [16])
@pytest.mark.parametrize("num_v_heads", [32])
@pytest.mark.parametrize("head_k_dim", [128])
@pytest.mark.parametrize("head_v_dim", [128])
@pytest.mark.parametrize("num_speculative_tokens", [1, 3])
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
def test_fused_sigmoid_gating_delta_rule_update_spec(
    tp_size: int,
    num_reqs: int,
    num_k_heads: int,
    num_v_heads: int,
    head_k_dim: int,
    head_v_dim: int,
    num_speculative_tokens: int,
    dtype: torch.dtype,
) -> None:
    if not torch.cuda.is_available():
        pytest.skip("CUDA is required")

    torch.manual_seed(0)
    key_dim = head_k_dim * num_k_heads
    value_dim = head_v_dim * num_v_heads
    mixed_qkv_dim = (key_dim * 2 + value_dim) // tp_size
    num_tokens = num_reqs * (num_speculative_tokens + 1)
    total_entries = num_tokens * 2

    mixed_qkv = torch.rand(num_tokens, mixed_qkv_dim, dtype=dtype, device="cuda")
    query, key, value = torch.split(
        mixed_qkv,
        [key_dim // tp_size, key_dim // tp_size, value_dim // tp_size],
        dim=-1,
    )
    query = query.view(1, num_tokens, num_k_heads, head_k_dim)
    key = key.view(1, num_tokens, num_k_heads, head_k_dim)
    value = value.view(1, num_tokens, num_v_heads, head_v_dim)

    A_log = torch.rand(num_v_heads // tp_size, dtype=dtype, device="cuda")
    dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype, device="cuda")
    a = torch.rand(num_tokens, num_v_heads, dtype=dtype, device="cuda")
    b = torch.rand(num_tokens, num_v_heads, dtype=dtype, device="cuda")
    ssm_state = torch.rand(total_entries, num_v_heads, head_v_dim, head_k_dim, dtype=dtype, device="cuda")
    state_indices = torch.randperm(total_entries, device="cuda", dtype=torch.int64)[:num_tokens]
    state_indices = state_indices.to(torch.int32).view(num_reqs, num_speculative_tokens + 1)
    num_accepted_tokens = torch.randint(1, num_speculative_tokens + 1, (num_reqs,), dtype=torch.int32, device="cuda")
    cu_seqlens = torch.arange(0, num_tokens + 1, num_speculative_tokens + 1, dtype=torch.int32, device="cuda")

    core_attn_out_ref, last_recurrent_state_ref = fused_sigmoid_gating_delta_rule_ref(
        A_log=A_log,
        a=a,
        b=b,
        dt_bias=dt_bias,
        q=query,
        k=key,
        v=value,
        initial_state=ssm_state.clone(),
        inplace_final_state=True,
        ssm_state_indices=state_indices,
        cu_seqlens=cu_seqlens,
        num_accepted_tokens=num_accepted_tokens,
        use_qk_l2norm_in_kernel=True,
    )

    core_attn_out, last_recurrent_state = fused_sigmoid_gating_delta_rule_update(
        A_log=A_log,
        a=a,
        b=b,
        dt_bias=dt_bias,
        q=query,
        k=key,
        v=value,
        initial_state=ssm_state,
        inplace_final_state=True,
        ssm_state_indices=state_indices,
        cu_seqlens=cu_seqlens,
        num_accepted_tokens=num_accepted_tokens,
        use_qk_l2norm_in_kernel=True,
    )

    torch.testing.assert_close(core_attn_out.float(), core_attn_out_ref.float(), atol=1e-2, rtol=1e-2)
    torch.testing.assert_close(
        last_recurrent_state.float(),
        last_recurrent_state_ref.float(),
        atol=1e-2,
        rtol=1e-2,
    )


def test_fused_sigmoid_gating_delta_rule_update_spec_qwen() -> None:
    if not torch.cuda.is_available():
        pytest.skip("CUDA is required")

    torch.manual_seed(123)
    device = torch.device("cuda")

    # Target shape from runtime logs:
    # A_log:[16], a/b:[13628,16], q/k:[1,8,4,128], v:[1,8,16,128],
    # state:[1639,16,128,128], cu_seqlens:[3], ssm_state_indices:[2,4], num_accepted_tokens:[2]
    num_reqs = 2
    num_speculative_tokens = 3
    num_tokens = num_reqs * (num_speculative_tokens + 1)  # 8
    ab_tokens = 13628
    num_k_heads = 4
    num_v_heads = 16
    head_k_dim = 128
    head_v_dim = 128
    total_entries = 1639
    dtype = torch.bfloat16

    mixed_qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim
    mixed_qkv = torch.rand(num_tokens, mixed_qkv_dim, dtype=dtype, device=device)
    query, key, value = torch.split(
        mixed_qkv,
        [num_k_heads * head_k_dim, num_k_heads * head_k_dim, num_v_heads * head_v_dim],
        dim=-1,
    )
    query = query.view(1, num_tokens, num_k_heads, head_k_dim)
    key = key.view(1, num_tokens, num_k_heads, head_k_dim)
    value = value.view(1, num_tokens, num_v_heads, head_v_dim)

    A_log = torch.rand(num_v_heads, dtype=dtype, device=device)
    dt_bias = torch.rand(num_v_heads, dtype=dtype, device=device)
    a = torch.rand(ab_tokens, num_v_heads, dtype=dtype, device=device)
    b = torch.rand(ab_tokens, num_v_heads, dtype=dtype, device=device)
    ssm_state = torch.rand(total_entries, num_v_heads, head_v_dim, head_k_dim, dtype=dtype, device=device)

    state_indices = torch.randperm(total_entries, device=device, dtype=torch.int64)[:num_tokens]
    state_indices = state_indices.to(torch.int32).view(num_reqs, num_speculative_tokens + 1)
    num_accepted_tokens = torch.randint(1, num_speculative_tokens + 1, (num_reqs,), dtype=torch.int32, device=device)
    cu_seqlens = torch.tensor([0, 4, 8], dtype=torch.int32, device=device)

    core_attn_out_ref, last_recurrent_state_ref = fused_sigmoid_gating_delta_rule_ref(
        A_log=A_log,
        a=a,
        b=b,
        dt_bias=dt_bias,
        q=query,
        k=key,
        v=value,
        initial_state=ssm_state.clone(),
        inplace_final_state=True,
        ssm_state_indices=state_indices,
        cu_seqlens=cu_seqlens,
        num_accepted_tokens=num_accepted_tokens,
        use_qk_l2norm_in_kernel=True,
    )

    core_attn_out, last_recurrent_state = fused_sigmoid_gating_delta_rule_update(
        A_log=A_log,
        a=a,
        b=b,
        dt_bias=dt_bias,
        q=query,
        k=key,
        v=value,
        initial_state=ssm_state,
        inplace_final_state=True,
        ssm_state_indices=state_indices,
        cu_seqlens=cu_seqlens,
        num_accepted_tokens=num_accepted_tokens,
        use_qk_l2norm_in_kernel=True,
    )

    torch.testing.assert_close(core_attn_out.float(), core_attn_out_ref.float(), atol=1e-2, rtol=1e-2)
    torch.testing.assert_close(
        last_recurrent_state.float(),
        last_recurrent_state_ref.float(),
        atol=1e-2,
        rtol=1e-2,
    )
