# SPDX-License-Identifier: MIT

import pytest
import torch

from aiter.ops.triton.fla.fused_recurrent import (
    fused_recurrent_gated_delta_rule_packed_decode,
)
from op_tests.triton_tests.utils.fused_recurrent_ref import (
    fused_recurrent_gated_delta_rule_packed_decode_ref,
)


@pytest.mark.parametrize("use_l2norm", [False, True])
def test_fused_recurrent_packed_decode_against_ref(use_l2norm: bool):
    if not torch.cuda.is_available():
        pytest.skip("CUDA is required")

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

    bsz, h_dim, hv, k_dim, v_dim = 8, 4, 16, 128, 128
    pool = 16

    qkv_dim = 2 * h_dim * k_dim + hv * v_dim
    mixed_qkv = torch.randn((bsz, qkv_dim), device=device, dtype=torch.float16) * 0.2
    a = torch.randn((bsz, hv), device=device, dtype=torch.float16) * 0.2
    b = torch.randn((bsz, hv), device=device, dtype=torch.float16) * 0.2
    A_log = torch.randn((hv,), device=device, dtype=torch.float32) * 0.2
    dt_bias = torch.randn((hv,), device=device, dtype=torch.float16) * 0.2
    ssm_state_indices = torch.randperm(pool, device=device, dtype=torch.int64)[:bsz].to(torch.int32)
    scale = k_dim ** -0.5

    out_cur = torch.empty((bsz, 1, hv, v_dim), device=device, dtype=torch.float16)
    out_ref = torch.empty_like(out_cur)
    state_init = torch.randn((pool, hv, v_dim, k_dim), device=device, dtype=torch.float32) * 0.05

    out_cur, state_cur = fused_recurrent_gated_delta_rule_packed_decode(
        mixed_qkv=mixed_qkv,
        a=a,
        b=b,
        A_log=A_log,
        dt_bias=dt_bias,
        scale=scale,
        initial_state=state_init.clone(),
        out=out_cur,
        ssm_state_indices=ssm_state_indices,
        use_qk_l2norm_in_kernel=use_l2norm,
    )

    out_ref, state_ref = fused_recurrent_gated_delta_rule_packed_decode_ref(
        mixed_qkv=mixed_qkv,
        a=a,
        b=b,
        A_log=A_log,
        dt_bias=dt_bias,
        scale=scale,
        initial_state=state_init.clone(),
        out=out_ref,
        ssm_state_indices=ssm_state_indices,
        use_qk_l2norm_in_kernel=use_l2norm,
    )

    torch.testing.assert_close(out_cur.float(), out_ref.float(), atol=3e-2, rtol=2e-2)
    torch.testing.assert_close(state_cur.float(), state_ref.float(), atol=3e-2, rtol=2e-2)


def test_fused_recurrent_packed_decode_pad_slot_id_against_ref():
    if not torch.cuda.is_available():
        pytest.skip("CUDA is required")

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

    bsz, h_dim, hv, k_dim, v_dim = 10, 4, 16, 128, 128
    pool = 20

    qkv_dim = 2 * h_dim * k_dim + hv * v_dim
    mixed_qkv = torch.randn((bsz, qkv_dim), device=device, dtype=torch.float16) * 0.2
    a = torch.randn((bsz, hv), device=device, dtype=torch.float16) * 0.2
    b = torch.randn((bsz, hv), device=device, dtype=torch.float16) * 0.2
    A_log = torch.randn((hv,), device=device, dtype=torch.float32) * 0.2
    dt_bias = torch.randn((hv,), device=device, dtype=torch.float16) * 0.2
    valid_unique = torch.randperm(pool, device=device, dtype=torch.int64)[:bsz].to(torch.int32)
    ssm_state_indices = valid_unique.clone()
    ssm_state_indices[0] = -1
    ssm_state_indices[3] = -1
    scale = k_dim ** -0.5

    out_cur = torch.empty((bsz, 1, hv, v_dim), device=device, dtype=torch.float16)
    out_ref = torch.empty_like(out_cur)
    state_init = torch.randn((pool, hv, v_dim, k_dim), device=device, dtype=torch.float32) * 0.05

    out_cur, state_cur = fused_recurrent_gated_delta_rule_packed_decode(
        mixed_qkv=mixed_qkv,
        a=a,
        b=b,
        A_log=A_log,
        dt_bias=dt_bias,
        scale=scale,
        initial_state=state_init.clone(),
        out=out_cur,
        ssm_state_indices=ssm_state_indices,
        use_qk_l2norm_in_kernel=False,
    )

    out_ref, state_ref = fused_recurrent_gated_delta_rule_packed_decode_ref(
        mixed_qkv=mixed_qkv,
        a=a,
        b=b,
        A_log=A_log,
        dt_bias=dt_bias,
        scale=scale,
        initial_state=state_init.clone(),
        out=out_ref,
        ssm_state_indices=ssm_state_indices,
        use_qk_l2norm_in_kernel=False,
    )

    torch.testing.assert_close(out_cur.float(), out_ref.float(), atol=3e-2, rtol=2e-2)
    torch.testing.assert_close(state_cur.float(), state_ref.float(), atol=3e-2, rtol=2e-2)
