# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang

import os

import pytest
import torch
import torch.nn.functional as F
from einops import rearrange, repeat

from aiter.ops.triton.gated_delta_net import (
    fused_recurrent_gated_delta_rule,
    chunk_gated_delta_rule,
)
from aiter.ops.triton._triton_kernels.gated_delta_rule.decode.fused_sigmoid_gating_recurrent import (
    fused_sigmoid_gating_delta_rule_update,
)
from aiter.ops.triton._triton_kernels.gated_delta_rule.gated_delta_rule_utils import (
    IS_INTEL_ALCHEMIST,
    assert_close,
    device,
)


def recurrent_gated_delta_rule_ref(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    beta: torch.Tensor,
    g: torch.Tensor,
    scale: float = None,
    initial_state: torch.Tensor = None,
    output_final_state: bool = False,
):
    q, k, v, beta, g = map(
        lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g]
    )
    B, H, T, K, V = *k.shape, v.shape[-1]
    o = torch.zeros(B, H, T, V).to(v)
    h = torch.zeros(B, H, K, V).to(v)
    if initial_state is not None:
        h = initial_state
    if scale is None:
        scale = 1 / (q.shape[-1] ** 0.5)
    q = q * scale
    for i in range(T):
        b_q = q[:, :, i]
        b_k = k[:, :, i]
        b_v = v[:, :, i].clone()
        h = h.clone() * g[:, :, i].exp()[..., None, None]
        b_beta = beta[:, :, i]
        b_v = b_v - (h.clone() * b_k[..., None]).sum(-2)
        b_v = b_v * b_beta[..., None]
        h = h.clone() + b_k.unsqueeze(-1) * b_v.unsqueeze(-2)
        o[:, :, i] = torch.einsum("bhd,bhdm->bhm", b_q, h)
    if not output_final_state:
        h = None
    o = o.transpose(1, 2).contiguous()
    return o, h


def chunk_gated_delta_rule_ref(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
    chunk_size: int = 64,
    scale: float = None,
    initial_state: torch.Tensor = None,
    output_final_state: bool = False,
):
    BT = chunk_size
    if scale is None:
        scale = 1 / (q.shape[-1] ** 0.5)
    # Calculate padding needed to make T a multiple of BT
    q, k, v, beta, g = map(
        lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g]
    )

    T = q.shape[-2]
    pad_len = (BT - (T % BT)) % BT
    if pad_len > 0:
        # Pad all tensors
        q = F.pad(q, (0, 0, 0, pad_len))
        k = F.pad(k, (0, 0, 0, pad_len))
        v = F.pad(v, (0, 0, 0, pad_len))
        beta = F.pad(beta, (0, pad_len))
        g = F.pad(g, (0, pad_len))
    q, k, v, beta, g = map(lambda x: x.to(torch.float32), [q, k, v, beta, g])
    decay = g
    chunk_size = BT
    b, h, seq_len, d_k = q.shape
    d_v = v.shape[-1]
    q = q * scale
    v = v * beta[..., None]
    k_beta = k * beta[..., None]
    assert seq_len % chunk_size == 0
    # note that diagonal is masked.
    mask = torch.triu(
        torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device),
        diagonal=0,
    )
    q, k, v, k_beta, decay = map(
        lambda x: rearrange(x, "b h (n c) d -> b h n c d", c=chunk_size),
        [q, k, v, k_beta, decay.unsqueeze(-1)],
    )
    decay = decay.squeeze(-1).cumsum(-1)
    decay_exp = decay.exp()[..., None]
    L_mask = ((decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril()
    attn = -((k_beta @ k.transpose(-1, -2)) * L_mask).masked_fill(mask, 0)
    for i in range(1, chunk_size):
        attn[..., i, :i] = attn[..., i, :i].clone() + (
            attn[..., i, :i, None].clone() * attn[..., :i, :i].clone()
        ).sum(-2)
    attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device)
    attn = attn
    k_cumsum = attn @ v
    k_cumdecay = attn @ (k_beta * decay_exp)
    v = k_cumsum
    S = k.new_zeros(b, h, d_k, d_v)
    if initial_state is not None:
        S = initial_state
    o = torch.zeros_like(v)
    mask = torch.triu(
        torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device),
        diagonal=1,
    )
    for i in range(0, seq_len // chunk_size):
        q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i]
        attn = (q_i @ k_i.transpose(-1, -2) * L_mask[:, :, i]).masked_fill_(mask, 0)
        v_prime = (k_cumdecay[:, :, i]) @ S
        v_new = v_i - v_prime
        o_inter = (q_i * decay[:, :, i, :, None].exp()) @ S
        o[:, :, i] = o_inter + attn @ v_new
        S = (
            S * decay[:, :, i, -1, None, None].exp()
            + (
                k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp()[..., None]
            ).transpose(-1, -2)
            @ v_new
        )
    if not output_final_state:
        S = None
    # unpad
    o = rearrange(o, "b h n c d -> b h (n c) d")
    o = o[:, :, :T]
    o = o.transpose(1, 2)
    return o, S


@pytest.mark.parametrize(
    ("B", "T", "H", "HV", "D", "scale", "gate_logit_normalizer", "dtype"),
    [
        pytest.param(
            *test,
            id="B{}-T{}-H{}-HV{}-D{}-scale{}-gate_logit_normalizer{}-{}".format(*test),
        )
        for test in [
            (1, 63, 1, 1, 64, 1, 1, torch.float),
            (2, 500, 4, 4, 60, 1, 1, torch.float),
            (2, 1000, 2, 8, 128, 1, 0.1, torch.float),
            (3, 1024, 2, 2, 128, 0.1, 1, torch.float),
            (4, 1024, 3, 3, 128, 1, 10, torch.float),
            (4, 2048, 4, 4, 64, 0.1, 1, torch.float),
            (2, 1024, 4, 4, 128, 1, 0.1, torch.float16),
            (2, 1024, 4, 8, 128, 1, 10, torch.float16),
            (2, 1024, 4, 4, 128, 1, 0.1, torch.bfloat16),
            (2, 1024, 4, 8, 128, 1, 1, torch.bfloat16),
            (4, 2048, 4, 8, 64, 0.1, 1, torch.bfloat16),
        ]
    ],
)
def test_fused_recurrent(
    B: int,
    T: int,
    H: int,
    HV: int,
    D: int,
    scale: float,
    gate_logit_normalizer: float,
    dtype: torch.dtype,
):
    torch.manual_seed(42)
    q = torch.randn(B, T, H, D, dtype=torch.float32)
    k = torch.randn(B, T, H, D, dtype=torch.float32)
    v = torch.randn(B, T, HV, D, dtype=dtype)
    beta = torch.rand(B, T, HV, dtype=dtype).sigmoid()
    g = F.logsigmoid(torch.rand(B, T, HV, dtype=torch.float32))
    g = g / gate_logit_normalizer
    h0 = torch.randn(B, HV, D, D, dtype=torch.float32)
    q, k, v, beta, g, h0 = map(
        lambda x: x.to(device).requires_grad_(), (q, k, v, beta, g, h0)
    )
    ref, ref_ht = recurrent_gated_delta_rule_ref(
        q=F.normalize(
            repeat(q.clone(), "b t h d -> b t (h g) d", g=HV // H), p=2, dim=-1
        ).to(dtype),
        k=F.normalize(
            repeat(k.clone(), "b t h d -> b t (h g) d", g=HV // H), p=2, dim=-1
        ).to(dtype),
        v=v.clone(),
        beta=beta.clone(),
        g=g.clone(),
        scale=scale,
        initial_state=h0.clone(),
        output_final_state=True,
    )
    tri, tri_ht = fused_recurrent_gated_delta_rule(
        q=q.clone(),
        k=k.clone(),
        v=v.clone(),
        beta=beta.clone(),
        g=g.clone(),
        scale=scale,
        initial_state=h0.clone(),
        use_qk_l2norm_in_kernel=True,
        output_final_state=True,
    )
    # Use higher tolerance for bfloat16 due to lower precision
    tol = 0.005 if dtype == torch.bfloat16 else 0.002
    assert_close("o", ref, tri, tol)
    assert_close("ht", ref_ht, tri_ht, tol)


@pytest.mark.parametrize(
    (
        "B",
        "T",
        "H",
        "D",
        "scale",
        "gate_logit_normalizer",
        "mask_p",
        "use_qk_l2norm_in_kernel",
        "dtype",
    ),
    [
        pytest.param(
            *test,
            id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-mask_p{}-use_qk_l2norm_in_kernel{}-{}".format(
                *test
            ),
        )
        for test in [
            (1, 63, 1, 64, 1, 1, 0, False, torch.float16),
            (2, 500, 3, 60, 1, 1, 0, False, torch.float16),
            (2, 1000, 3, 64, 0.1, 1, 0.5, False, torch.float16),
            (3, 1024, 4, 100, 1, 0.1, 0, False, torch.float16),
            (4, 1024, 4, 128, 0.1, 1, 0, False, torch.float16),
            (4, 1024, 4, 128, 0.1, 1, 0, True, torch.float16),
            (2, 1500, 4, 128, 0.1, 10, 0, False, torch.float16),
            (4, 2048, 8, 64, 0.1, 1, 0, False, torch.float16),
            # bfloat16 tests
            (2, 500, 3, 60, 1, 1, 0, False, torch.bfloat16),
            (2, 1000, 3, 64, 0.1, 1, 0, False, torch.bfloat16),
            (4, 1024, 4, 128, 0.1, 1, 0, False, torch.bfloat16),
            (4, 1024, 4, 128, 0.1, 1, 0, True, torch.bfloat16),
            (4, 2048, 8, 64, 0.1, 1, 0, False, torch.bfloat16),
        ]
    ],
)
def test_chunk(
    B: int,
    T: int,
    H: int,
    D: int,
    scale: float,
    gate_logit_normalizer: float,
    mask_p: float,
    use_qk_l2norm_in_kernel: bool,
    dtype: torch.dtype,
):
    torch.manual_seed(42)
    if IS_INTEL_ALCHEMIST and D > 128:
        pytest.skip(
            reason="chunk_gated_delta_rule is not supported on alchemist for D>128"
        )

    q = torch.rand(B, T, H, D, dtype=dtype)
    k = torch.rand(B, T, H, D, dtype=dtype)
    v = torch.rand(B, T, H, D, dtype=dtype)
    beta = torch.rand(B, T, H, dtype=dtype).sigmoid()
    g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.float32))
    g = g / gate_logit_normalizer
    g = g * (torch.rand_like(g) > mask_p)
    h0 = torch.zeros(B, H, D, D, dtype=torch.float32)
    q, k, v, beta, g, h0 = map(
        lambda x: x.to(device).requires_grad_(True), (q, k, v, beta, g, h0)
    )

    tri, tri_ht = chunk_gated_delta_rule(
        q=(
            F.normalize(q.clone(), p=2, dim=-1)
            if not use_qk_l2norm_in_kernel
            else q.clone()
        ),
        k=(
            F.normalize(k.clone(), p=2, dim=-1)
            if not use_qk_l2norm_in_kernel
            else k.clone()
        ),
        v=v.clone(),
        g=g.clone(),
        beta=beta.clone(),
        scale=scale,
        initial_state=h0.clone(),
        output_final_state=True,
        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
    )
    # do = torch.randn_like(v)
    # dht = torch.randn_like(h0)

    ref, ref_ht = recurrent_gated_delta_rule_ref(
        q=F.normalize(q.clone(), p=2, dim=-1),
        k=F.normalize(k.clone(), p=2, dim=-1),
        v=v.clone(),
        beta=beta.clone(),
        g=g.clone(),
        scale=scale,
        output_final_state=True,
        initial_state=h0.clone(),
    )

    # ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True)
    # ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad
    assert_close("o", ref, tri, 0.005)
    assert_close("ht", ref_ht, tri_ht, 0.005)
    # assert_close('dq', ref_dq, tri_dq, 0.008)
    # assert_close('dk', ref_dk, tri_dk, 0.008)
    # assert_close('dv', ref_dv, tri_dv, 0.008)
    # assert_close('db', ref_dbeta, tri_dbeta, 0.02)
    # assert_close('dg', ref_dg, tri_dg, 0.02)
    # assert_close('dh0', ref_dh0, tri_dh0, 0.008)


@pytest.mark.parametrize(
    ("B", "T", "H", "HV", "D", "scale", "gate_logit_normalizer", "dtype"),
    [
        pytest.param(
            *test,
            id="B{}-T{}-H{}-HV{}-D{}-scale{}-gate_logit_normalizer{}-{}".format(*test),
        )
        for test in [
            (1, 63, 1, 1, 64, 1, 1, torch.float),
            (2, 500, 4, 4, 60, 1, 1, torch.float),
            (2, 1000, 2, 8, 128, 1, 0.1, torch.float),
            (3, 1024, 2, 2, 128, 0.1, 1, torch.float),
            (4, 1024, 3, 3, 128, 1, 10, torch.float),
            (4, 2048, 4, 4, 64, 0.1, 1, torch.float),
            (2, 1024, 4, 4, 128, 1, 0.1, torch.float16),
            (2, 1024, 4, 8, 128, 1, 10, torch.float16),
            (2, 1024, 4, 4, 128, 1, 0.1, torch.bfloat16),
            (2, 1024, 4, 8, 128, 1, 1, torch.bfloat16),
            (4, 2048, 4, 8, 64, 0.1, 1, torch.bfloat16),
        ]
    ],
)
def test_fused_sigmoid_gating_delta_rule_update(
    B: int,
    T: int,
    H: int,
    HV: int,
    D: int,
    scale: float,
    gate_logit_normalizer: float,
    dtype: torch.dtype,
):
    """Test fused sigmoid gating delta rule update kernel."""
    torch.manual_seed(42)

    # Create input tensors
    q = torch.randn(B, T, H, D, dtype=torch.float32)
    k = torch.randn(B, T, H, D, dtype=torch.float32)
    v = torch.randn(B, T, HV, D, dtype=dtype)
    b = torch.randn(B, T, HV, dtype=dtype)  # beta logits

    # Gating parameters
    A_log = torch.randn(HV, dtype=torch.float32)
    a = torch.randn(B, T, HV, dtype=torch.float32)
    dt_bias = torch.randn(HV, dtype=torch.float32)
    softplus_beta = 1.0
    softplus_threshold = 20.0

    # Initial state
    h0 = torch.randn(B, HV, D, D, dtype=torch.float32)
    initial_state_indices = torch.arange(B, dtype=torch.long)

    # Move to device
    q, k, v, b, A_log, a, dt_bias, h0, initial_state_indices = map(
        lambda x: x.to(device),
        (q, k, v, b, A_log, a, dt_bias, h0, initial_state_indices),
    )

    # Compute reference using recurrent implementation
    # Compute g = -exp(A_log) * softplus(a + dt_bias)
    # This is already in log-space (negative values for decay)
    x = a + dt_bias[None, None, :]
    softplus_x = F.softplus(x, beta=softplus_beta, threshold=softplus_threshold)
    g = -torch.exp(A_log[None, None, :]) * softplus_x
    beta = torch.sigmoid(b)

    # Expand q and k to match HV heads for reference implementation
    q_expanded = F.normalize(
        repeat(q.clone(), "b t h d -> b t (h g) d", g=HV // H), p=2, dim=-1
    ).to(dtype)
    k_expanded = F.normalize(
        repeat(k.clone(), "b t h d -> b t (h g) d", g=HV // H), p=2, dim=-1
    ).to(dtype)

    ref, _ = recurrent_gated_delta_rule_ref(
        q=q_expanded,
        k=k_expanded,
        v=v.clone(),
        beta=beta.clone(),
        g=g.clone(),
        scale=scale,
        initial_state=h0.clone(),
        output_final_state=True,
    )

    # Compute using fused kernel
    tri = fused_sigmoid_gating_delta_rule_update(
        A_log=A_log,
        a=a,
        dt_bias=dt_bias,
        softplus_beta=softplus_beta,
        softplus_threshold=softplus_threshold,
        q=q.clone(),
        k=k.clone(),
        v=v.clone(),
        b=b.clone(),
        initial_state_source=h0.clone(),
        initial_state_indices=initial_state_indices,
        scale=scale,
        use_qk_l2norm_in_kernel=True,
    )

    # Use higher tolerance for bfloat16 due to lower precision
    tol = 0.005 if dtype == torch.bfloat16 else 0.002
    assert_close("o", ref, tri, tol)


@pytest.mark.parametrize(
    ("H", "D", "mask_p", "cu_seqlens", "dtype"),
    [
        pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}".format(*test))
        for test in [
            (4, 60, 0, [0, 15], torch.float16),
            (4, 64, 0, [0, 256, 500, 1000], torch.float16),
            (4, 64, 0.5, [0, 256, 500, 1000], torch.float16),
            (4, 100, 0, [0, 15, 100, 300, 1200, 2000], torch.float16),
            # bfloat16 tests
            (4, 60, 0, [0, 15], torch.bfloat16),
            (4, 64, 0, [0, 256, 500, 1000], torch.bfloat16),
            (4, 100, 0, [0, 15, 100, 300, 1200, 2000], torch.bfloat16),
        ]
    ],
)
@pytest.mark.skipif(
    os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1",
    reason="Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set",
)
def test_chunk_varlen(
    H: int,
    D: int,
    mask_p: float,
    cu_seqlens: list[int],
    dtype: torch.dtype,
):
    if IS_INTEL_ALCHEMIST and D > 128:
        pytest.skip(
            reason="chunk_gated_delta_rule is not supported on alchemist for D>128"
        )
    torch.manual_seed(42)
    os.environ["TRITON_F32_DEFAULT"] = "ieee"
    # randomly split the sequence into N segments
    cu_seqlens = torch.LongTensor(cu_seqlens).to(device)
    T = cu_seqlens[-1]
    N = len(cu_seqlens) - 1

    # seq-first required for inputs with variable lengths
    q = torch.randn((1, T, H, D), dtype=dtype)
    k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype)
    v = torch.randn((1, T, H, D), dtype=dtype)
    g = F.logsigmoid(torch.rand(1, T, H, dtype=dtype))
    g = g * (torch.rand_like(g) > mask_p)
    beta = torch.rand(1, T, H, dtype=dtype).sigmoid()
    h0 = torch.randn((N, H, D, D), dtype=dtype)

    q, k, v, beta, g, h0 = map(
        lambda x: x.to(device).requires_grad_(False), (q, k, v, beta, g, h0)
    )
    # do = torch.randn_like(v)
    # dht = torch.rand_like(h0)

    tri, tri_ht = chunk_gated_delta_rule(
        q=q.clone(),
        k=k.clone(),
        v=v.clone(),
        beta=beta.clone(),
        g=g.clone(),
        initial_state=h0.clone(),
        output_final_state=True,
        cu_seqlens=cu_seqlens,
    )
    # ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True)
    # tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dg, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad
    # q.grad = k.grad = v.grad = beta.grad = g.grad = h0.grad = None

    ref = []
    ref_ht = []
    for i in range(N):
        ref_i, ref_ht_i = recurrent_gated_delta_rule_ref(
            q=q[:, cu_seqlens[i] : cu_seqlens[i + 1]],
            k=k[:, cu_seqlens[i] : cu_seqlens[i + 1]],
            v=v[:, cu_seqlens[i] : cu_seqlens[i + 1]],
            beta=beta[:, cu_seqlens[i] : cu_seqlens[i + 1]],
            g=g[:, cu_seqlens[i] : cu_seqlens[i + 1]],
            initial_state=h0[i],
            output_final_state=True,
        )
        ref.append(ref_i)
        ref_ht.append(ref_ht_i)
    ref = torch.cat(ref, 1)
    ref_ht = torch.cat(ref_ht, 0)

    # ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True)
    # ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad

    assert_close("o", ref, tri, 0.005)
    assert_close("ht", ref_ht, tri_ht, 0.005)
    # assert_close('dq', ref_dq, tri_dq, 0.007)
    # assert_close('dk', ref_dk, tri_dk, 0.008)
    # assert_close('dv', ref_dv, tri_dv, 0.007)
    # assert_close('db', ref_dbeta, tri_dbeta, 0.015)
    # assert_close('dg', ref_dg, tri_dg, 0.015)
    # assert_close('dh0', ref_dh0, tri_dh0, 0.007)
