# SPDX-License-Identifier: MIT

import pytest
import torch

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


def _make_tensors(N, T, H, HV, K, V, device="cuda", seed=2025):
    torch.manual_seed(seed)
    A_log = torch.randn(HV, dtype=torch.float32, device=device)
    dt_bias = torch.randn(HV, dtype=torch.bfloat16, device=device)
    a = torch.randn(1, N * T, HV, dtype=torch.bfloat16, device=device)
    b = torch.randn(1, N * T, HV, dtype=torch.bfloat16, device=device)
    q = torch.randn(1, N * T, H, K, dtype=torch.bfloat16, device=device)
    k = torch.randn(1, N * T, H, K, dtype=torch.bfloat16, device=device)
    v = torch.randn(1, N * T, HV, V, dtype=torch.bfloat16, device=device)
    indices = torch.arange(N, dtype=torch.int32, device=device)
    # Match runtime pool style in logs: keep one extra slot.
    initial_state = torch.randn(N + 1, HV, K, V, dtype=torch.float, device=device)
    cu_seqlens = torch.arange(0, N * T + 1, T, dtype=torch.int32, device=device)
    return A_log, dt_bias, a, b, q, k, v, initial_state, indices, cu_seqlens


@pytest.mark.parametrize("N", [1, 8, 16])
@pytest.mark.parametrize("T", [1, 4, 8])
def test_fused_gdn_mtp_precision(N: int, T: int):
    if not torch.cuda.is_available():
        pytest.skip("CUDA is required")
    H, HV, K, V = 4, 16, 128, 128
    A_log, dt_bias, a, b, q, k, v, state, indices, cu_seqlens = _make_tensors(
        N, T, H, HV, K, V
    )
    indices_ref = indices[:, None].expand(N, T).contiguous()

    out_ref, _ = fused_sigmoid_gating_delta_rule_ref(
        A_log=A_log,
        a=a.view(-1, HV),
        b=b.view(-1, HV),
        dt_bias=dt_bias,
        q=q,
        k=k,
        v=v,
        initial_state=state.clone(),
        inplace_final_state=True,
        ssm_state_indices=indices_ref,
        cu_seqlens=cu_seqlens,
        use_qk_l2norm_in_kernel=True,
    )
    state_before = state.clone()
    out_fused = fused_sigmoid_gating_delta_rule_update(
        A_log=A_log,
        dt_bias=dt_bias,
        q=q,
        k=k,
        v=v,
        a=a,
        b=b,
        initial_state_source=state,
        initial_state_indices=indices,
        cu_seqlens=cu_seqlens,
        use_qk_l2norm_in_kernel=True,
        softplus_beta=1.0,
        softplus_threshold=20.0,
        is_kda=False,
        disable_state_update=True,
        intermediate_states_buffer=torch.empty(
            N, T, HV, K, V, dtype=state.dtype, device=state.device
        ),
        intermediate_state_indices=torch.arange(
            N, dtype=torch.int32, device=state.device
        ),
        cache_steps=T,
    )

    torch.testing.assert_close(out_ref, out_fused, rtol=1e-2, atol=1e-2)
    torch.testing.assert_close(state, state_before, rtol=0, atol=0)


@pytest.mark.parametrize("N", [1, 16, 64])
def test_mtp_single_step_decode(N: int):
    if not torch.cuda.is_available():
        pytest.skip("CUDA is required")
    T = 1
    H, HV, K, V = 4, 16, 128, 128
    A_log, dt_bias, a, b, q, k, v, state, indices, cu_seqlens = _make_tensors(
        N, T, H, HV, K, V
    )

    out_ref, state_ref = fused_sigmoid_gating_delta_rule_ref(
        A_log=A_log,
        a=a.view(-1, HV),
        b=b.view(-1, HV),
        dt_bias=dt_bias,
        q=q,
        k=k,
        v=v,
        initial_state=state.clone(),
        inplace_final_state=True,
        ssm_state_indices=indices,
        cu_seqlens=cu_seqlens,
        use_qk_l2norm_in_kernel=True,
    )
    out_fused = fused_sigmoid_gating_delta_rule_update(
        A_log=A_log,
        dt_bias=dt_bias,
        q=q,
        k=k,
        v=v,
        a=a,
        b=b,
        initial_state_source=state,
        initial_state_indices=indices,
        cu_seqlens=cu_seqlens,
        use_qk_l2norm_in_kernel=True,
        softplus_beta=1.0,
        softplus_threshold=20.0,
        is_kda=False,
        disable_state_update=False,
    )
    torch.testing.assert_close(out_ref, out_fused, rtol=1e-2, atol=1e-2)
    torch.testing.assert_close(state, state_ref, rtol=1e-2, atol=1e-2)


def test_fused_gdn_mtp_precision_runtime_shape():
    if not torch.cuda.is_available():
        pytest.skip("CUDA is required")
    N, T = 48, 4
    H, HV, K, V = 4, 16, 128, 128
    A_log, dt_bias, a, b, q, k, v, state, indices, cu_seqlens = _make_tensors(
        N, T, H, HV, K, V
    )
    indices_ref = indices[:, None].expand(N, T).contiguous()
    out_ref, _ = fused_sigmoid_gating_delta_rule_ref(
        A_log=A_log,
        a=a.view(-1, HV),
        b=b.view(-1, HV),
        dt_bias=dt_bias,
        q=q,
        k=k,
        v=v,
        initial_state=state.clone(),
        inplace_final_state=True,
        ssm_state_indices=indices_ref,
        cu_seqlens=cu_seqlens,
        use_qk_l2norm_in_kernel=True,
    )
    out_fused = fused_sigmoid_gating_delta_rule_update(
        A_log=A_log,
        dt_bias=dt_bias,
        q=q,
        k=k,
        v=v,
        a=a,
        b=b,
        initial_state_source=state,
        initial_state_indices=indices,
        cu_seqlens=cu_seqlens,
        use_qk_l2norm_in_kernel=True,
        softplus_beta=1.0,
        softplus_threshold=20.0,
        disable_state_update=True,
        intermediate_states_buffer=torch.empty(
            N + 1, T, HV, K, V, dtype=state.dtype, device=state.device
        ),
        intermediate_state_indices=torch.arange(
            N, dtype=torch.int32, device=state.device
        ),
        cache_steps=T,
    )
    torch.testing.assert_close(out_ref, out_fused, rtol=1e-2, atol=1e-2)
