import numpy
from typing import List, Optional
from itertools import accumulate

import pytest
import torch
from vllm.model_executor.layers.rotary_embedding import get_rope, _ROPE_DICT, LinearScalingRotaryEmbedding
from vllm_mlu.model_executor.layers.rotary_embedding import MLURotaryEmbedding

ROPE_THRESHOLD_DIFF1 = 5e-3
ROPE_THRESHOLD_DIFF2 = 5e-3

def compute_diff(baseline: numpy.ndarray, compare: numpy.ndarray):
    '''add diff1  diff2 accuracy method'''
    error = numpy.abs(baseline - compare)
    diff1 = numpy.sum(error) / numpy.sum(numpy.abs(baseline))
    diff2 = numpy.sqrt(numpy.sum(error**2)/numpy.sum(baseline**2))
    return diff1, diff2

IS_NEOX_STYLE = [True, False]
DTYPES = [torch.half, torch.bfloat16, torch.float]
HEAD_SIZES = [64, 80, 96, 112, 128, 256]
ROTARY_DIMS = [32]  # None means rotary dim == head size
NUM_HEADS = [9, 17]  # Arbitrary values for testing
BATCH_SIZES = [1, 5]  # Arbitrary values for testing
SEQ_LENS = [11, 8192]

@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("is_neox_style", IS_NEOX_STYLE)
@pytest.mark.parametrize("batch_size", BATCH_SIZES)
@pytest.mark.parametrize("seq_len", SEQ_LENS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
@pytest.mark.parametrize("rotary_dim", ROTARY_DIMS)
@torch.inference_mode()
def test_rotary_embedding(
    is_neox_style: bool,
    batch_size: int,
    seq_len: int,
    num_heads: int,
    head_size: int,
    rotary_dim: Optional[int],
    dtype: torch.dtype,
    max_position: int = 128,
    base: int = 10000,
) -> None:
    if rotary_dim is None:
        rotary_dim = head_size
    if rotary_dim is None:
        rotary_dim = head_size

    total_seq_len = batch_size * seq_len

    MLURotaryEmbedding.max_seq_len = max_position
    rope = MLURotaryEmbedding(head_size, rotary_dim, max_position, base, is_neox_style, dtype)
    rope = rope.to(dtype=dtype, device=0)

    positions = torch.randint(0,
                              max_position, ([batch_size*seq_len]),
                              device=0).to(dtype=torch.int32)

    context_shape = (total_seq_len, num_heads, head_size)
    context = torch.randn(size=context_shape, dtype=dtype).mlu()
    qk = context[..., 0 : num_heads, :]
    ref_qk = qk.clone()
    
    cu_seq_lens = torch.arange(0, batch_size + 1, dtype=torch.int32).mlu() * seq_len
    MLURotaryEmbedding.set_cos_sin = False
    MLURotaryEmbedding.cu_seq_lens = cu_seq_lens
    MLURotaryEmbedding.is_prompt = False
    MLURotaryEmbedding.is_chunked = False
    qk_out = rope.forward(positions, qk)
    rope_base = MLURotaryEmbedding(head_size, rotary_dim, max_position, base, is_neox_style, dtype)
    # for simular CPU re_init_cos_sin_cache
    if "cos_sin_cache" in rope_base._buffers:
        del rope_base._buffers["cos_sin_cache"]
    cache = rope_base._compute_cos_sin_cache()
    cache = cache.to(dtype).mlu()
    rope_base.cos_sin_cache: torch.Tensor
    rope_base.register_buffer("cos_sin_cache", cache, persistent=False)

    num_q_heads = num_heads - 1
    ref_q = ref_qk[:, :num_q_heads, :].reshape(-1,head_size)
    ref_k = ref_qk[:, num_q_heads:, :].reshape(-1,head_size)
    ref_q_o, ref_k_o = rope_base.forward_native(positions, ref_q, ref_k)
    ref_q_o_reshape = ref_q_o.reshape(-1, num_q_heads, head_size)
    ref_k_o_reshape = ref_k_o.reshape(-1, 1, head_size)
    ref_qk_out = torch.cat((ref_q_o_reshape, ref_k_o_reshape), dim=1).cpu()
    qk_out_cpu = qk_out.cpu()
    MLURotaryEmbedding.unset_mlu_var()
    diff1, diff2 = compute_diff(baseline=ref_qk_out.float().detach().numpy(),
                                compare=qk_out_cpu.float().detach().numpy())
    assert diff1 <= ROPE_THRESHOLD_DIFF1 and diff2 <= ROPE_THRESHOLD_DIFF2


@pytest.mark.parametrize("is_neox_style", IS_NEOX_STYLE)
@pytest.mark.parametrize("batch_size", BATCH_SIZES)
@pytest.mark.parametrize("seq_len", [1, 11, 1024])
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_size", [64, 80, 128])
@pytest.mark.parametrize("rotary_dim", ROTARY_DIMS)
@pytest.mark.parametrize("dtype", DTYPES)
@torch.inference_mode()
def test_batched_rotary_embedding_multi_lora(
    is_neox_style: bool,
    batch_size: int,
    seq_len: int,
    num_heads: int,
    head_size: int,
    rotary_dim: Optional[int],
    dtype: torch.dtype,
    device: str = "mlu",
    seed: int = 0,
    max_position: int = 4096,
    base: int = 10000,
) -> None:
    """test linear scaling rope kernel"""
    assert device == "mlu"
    assert torch.mlu.is_available()

    torch.random.manual_seed(seed)
    torch.mlu.manual_seed(seed)
    torch.set_default_device(device)

    if rotary_dim is None:
        rotary_dim = head_size
    is_prompt = seq_len == 1
    scaling_factors: List[int] = [1, 2, 4]

    MLURotaryEmbedding.max_seq_len = max_position
    MLURotaryEmbedding.set_cos_sin = False
    MLURotaryEmbedding.is_prompt = is_prompt
    MLURotaryEmbedding.is_chunked = False
    MLURotaryEmbedding.positions_ = None
    MLURotaryEmbedding.cu_seq_lens = seq_len * torch.arange(
        0, batch_size + 1, dtype=torch.int32)

    rope = get_rope(head_size, rotary_dim, max_position, base, is_neox_style, {
        "rope_type": "linear",
        "factor": tuple(scaling_factors)
    })
    rope = rope.to(dtype=dtype)

    positions = torch.randint(0, max_position, (batch_size * seq_len, ),
                              dtype=torch.int32)
    query = torch.randn(batch_size,
                        seq_len,
                        num_heads * head_size,
                        dtype=dtype)
    key = torch.randn_like(query)

    offset_map = torch.tensor(
        list(
            accumulate([0] + [
                max_position * scaling_factor * 2
                for scaling_factor in scaling_factors[:-1]
            ])), dtype=torch.int32)
    query_types = torch.randint(0,
                                len(scaling_factors), (batch_size * seq_len, ))
    query_offsets = offset_map[query_types]

    qk = torch.cat([query, key], dim=-1)
    qk = qk.view(batch_size * seq_len, num_heads + num_heads, head_size)

    out_qk = rope.forward(positions, qk, query_offsets)

    scaling_factor = tuple(scaling_factors)
    rope_base = LinearScalingRotaryEmbedding(head_size, rotary_dim,
                                                      max_position, base,
                                                      is_neox_style,
                                                      scaling_factor,
                                                      torch.get_default_dtype())

    ref_query, ref_key = rope_base.forward_native(positions, query, key,
                                             query_offsets)
    ref_qk = torch.cat([ref_query, ref_key], dim=-1)
    ref_qk = ref_qk.view(batch_size * seq_len, num_heads + num_heads, head_size)

    # delete rope cache to init rope instance every time
    _ROPE_DICT.clear()

    # compare the results
    diff1, diff2 = compute_diff(baseline=ref_qk.cpu().float().detach().numpy(),
                                compare=out_qk.cpu().float().detach().numpy())
    assert diff1 <= ROPE_THRESHOLD_DIFF1 and diff2 <= ROPE_THRESHOLD_DIFF2
