# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import pytest
import torch
import os
import sys
import numpy as np
import copy

from aiter.ops.triton.triton_decode_attention import decode_attention_fwd

def pytorch_decode_attention_reference(
    q,
    k_buffer,
    v_buffer,
    o,
    req_to_page,
    b_seq_len,
    sm_scale,
    page_size=1,
):
    B, H_Q, D_QK = q.shape
    H_KV = k_buffer.shape[-2]
    D_V = v_buffer.shape[-1]
    kv_group_num = H_Q // H_KV
    
    # Always reshape buffers to [seq_len, H_KV, D] for unified handling
    if len(k_buffer.shape) == 4:
        raise ValueError("Stop here")
        k_buffer = k_buffer.view(-1, H_KV, D_QK)
        v_buffer = v_buffer.view(-1, H_KV, D_V)

    # Reshape to [B, num_pages]
    if len(req_to_page.shape) == 3:
        req_to_page = req_to_page.squeeze(2)

    # [B,]
    for b in range(B):
        page_idx = req_to_page[b][torch.arange(b_seq_len[b], device="cuda") // page_size]
        kv_loc = page_idx * page_size + torch.arange(b_seq_len[b], device="cuda") % page_size
        k = k_buffer[kv_loc, :, :]
        v = v_buffer[kv_loc, :, :]

        kk_re=[]
        vv_re=[]
        for i in range(H_KV):
            k_c = k[:,i:i+1,:]
            m_k_c = k_c.repeat(1, H_Q // H_KV, 1)
            kk_re.append(m_k_c)
            v_c = v[:,i:i+1,:]
            m_v_c = v_c.repeat(1, H_Q // H_KV, 1)
            vv_re.append(m_v_c)
        kk=torch.cat(kk_re,dim=1)
        vv=torch.cat(vv_re,dim=1)

        # [N, H_Q, D_QK] -> [H_Q, D_QK, N]
        kk = kk.permute(1, 2, 0)
        # [H_Q, D_QK] -> [H_Q, 1, D_QK]
        q_ = q[b][:, None, :]
        # [H_Q, 1, D_QK] @ [H_Q, D_QK, N] = [H_Q, 1, N]
        p = torch.matmul(q_, kk) * sm_scale
        p = torch.softmax(p, dim=-1)

        # [N, H_Q, D_V] -> [H_Q, N, D_V]
        vv = vv.permute(1, 0, 2)
        # [H_Q, 1, N] @ [H_Q, N, D_V] = [H_Q, 1, D_V]
        r = torch.matmul(p, vv).squeeze(1)
        o[b, :, :] = r


@pytest.mark.parametrize("B", [3, 5])
@pytest.mark.parametrize("L", [1027, 1025])
@pytest.mark.parametrize("H_Q", [32])
@pytest.mark.parametrize("H_KV", [32, 8])
@pytest.mark.parametrize("D_QK", [128, 192, 576])
@pytest.mark.parametrize("D_V", [128, 512])
@pytest.mark.parametrize("CACHE_SIZE", [16384])
@pytest.mark.parametrize("PAGE_SIZE", [1, 16])
def test_decode_attention(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE):
    assert CACHE_SIZE % PAGE_SIZE == 0
    dtype = torch.bfloat16
    seq_len = L  # This represents the number of tokens already in the sequence
    sm_scale = 1.0 / (D_QK**0.5)
    num_kv_splits = 8

    num_pages_per_batch = (seq_len + PAGE_SIZE - 1) // PAGE_SIZE
    req_to_page = torch.randint(0,
                                CACHE_SIZE // PAGE_SIZE,
                                (B, num_pages_per_batch, 1),
                                device="cuda")
    req_to_token = req_to_page * PAGE_SIZE
    req_to_token = req_to_token.expand(B, num_pages_per_batch, PAGE_SIZE)
    req_to_token = req_to_token + torch.arange(PAGE_SIZE, device="cuda").view(
        1, 1, -1)
    req_to_token = req_to_token.view(B, -1)
    req_to_token = req_to_token[:, :seq_len].contiguous()

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

    # k_buffer and v_buffer represent all previous tokens
    # Page size is 1.
    k_buffer = torch.randn(CACHE_SIZE, H_KV, D_QK, dtype=dtype, device="cuda")
    v_buffer = torch.randn(CACHE_SIZE, 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")

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

    attn_logits = torch.empty(
        (B, H_Q, num_kv_splits, D_V + 1),
        dtype=torch.float32,
        device="cuda",
    )

    # Call the original implementation.
    decode_attention_fwd(
        q,
        k_buffer,
        v_buffer,
        o,
        req_to_token,
        b_seq_len,
        attn_logits,
        num_kv_splits,
        sm_scale,
    )

    o_ref = torch.zeros_like(o)
    pytorch_decode_attention_reference(
        q,
        k_buffer,
        v_buffer,
        o_ref,
        req_to_token,
        b_seq_len,
        sm_scale,
    )
    torch.testing.assert_close(o, o_ref, atol=1e-2, rtol=0)

    # Page size can be larger than 1.
    k_buffer = k_buffer.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_QK)
    v_buffer = v_buffer.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_V)

    o1 = torch.zeros_like(o)

    decode_attention_fwd(
        q,
        k_buffer,
        v_buffer,
        o1,
        req_to_page,
        b_seq_len,
        attn_logits,
        num_kv_splits,
        sm_scale,
        PAGE_SIZE,
    )

    assert torch.allclose(o, o1)
    return True