import random
import unittest

import numpy as np
import math
import torch
import torch_mlu_ops as ops

from common_utils import *

def gen_args(batch_size,
             max_context_len,
             head_num_q,
             head_num_kv,
             cache_mem_len,
             head_size,
             group_size,
             block_size,
             use_seq_offset,
             dtype,
             quant_mode,
             quant_bit,
             has_value = True):
    # Preprocess arguments
    assert cache_mem_len >= max_context_len, \
        "cache_mem_len should greater then or equal to max_context_len."
    assert head_size % group_size == 0, \
        "head_size should be a multiply of groupwise."
    total_heads = head_num_q + head_num_kv * 2
    max_seq_offset = cache_mem_len - max_context_len
    max_block_num = int(math.ceil(max_context_len / block_size))
    total_blocks = int(math.ceil(cache_mem_len / block_size)) * batch_size
    block_tables = random.sample(range(0, total_blocks), batch_size * max_block_num)
    block_tables = torch.tensor(block_tables, dtype=torch.int32).mlu().view(batch_size,
                                                                            max_block_num)
    # Generates key and cache from context
    context_lens = torch.randint(size=[batch_size], low=1, high=max_context_len + 1,
                                 dtype=torch.int32, device="mlu")
    if use_seq_offset:
        context_paddings = torch.randint(size=[batch_size], low=0, high=max_seq_offset,
                                         dtype=torch.int32, device="mlu")
    else:
        context_paddings = torch.zeros_like(context_lens)
    cu_context_lens = torch.cumsum(context_lens + context_paddings, dim=-1)
    total_seqlen = cu_context_lens[-1]
    context_seq_offset = torch.zeros([batch_size], dtype=torch.int32, device="mlu")
    context_seq_offset[1:] = cu_context_lens[:-1]
    context = torch.randn([total_seqlen, total_heads, head_size], dtype=dtype, device="mlu")
    key = context[..., head_num_q:head_num_q + head_num_kv, :]
    value = context[..., head_num_q + head_num_kv:head_num_q + 2 * head_num_kv, :]
    # Generates key_cache and value_cache
    cache = torch.randint(size=(2, total_blocks, head_num_kv, block_size, head_size),
                          low=-128, high=127, dtype=torch.int32, device="mlu")
    cache = cache.to(torch.int8)
    key_cache, value_cache = cache[[0, 1]]

    # Generates key_cache_scale and value_cache_scale
    if quant_mode == 0:  # quant_mode == 0 is per channel
        cache_scale = torch.randn((2, head_num_kv, head_size), dtype=torch.float, device="mlu")
    else:  # quant_mode != 1 (== 1 for extend) is per head
        cache_scale = torch.randn((2, total_blocks, head_num_kv, block_size),
                                  dtype=torch.float, device="mlu")
    key_cache_scale, value_cache_scale = cache_scale[[0, 1]]

    # Prepare arguments
    if has_value == False:
        value = None
        value_cache = None
        value_cache_scale = None
    args = [key, value, key_cache, value_cache, key_cache_scale, value_cache_scale]
    args += [context_lens, max_context_len, context_seq_offset if use_seq_offset else None,
             block_tables]
    args += [quant_mode, quant_bit]
    return args

class TestDequantFromPagedCache(BtTestCase):
    def run_gen_case(self, dic):
        dump_data = dic.pop('dump_data')
        if dump_data:
          self.launch(*dic.values())
        else:
          key = create_tensor_from_dic(dic['key'])
          value = create_tensor_from_dic(dic['value'])
          key_cache = create_tensor_from_dic(dic['key_cache'])
          value_cache = create_tensor_from_dic(dic['value_cache'])
          key_cache_quant_scale = create_tensor_from_dic(dic['key_cache_quant_scale'])
          value_cache_quant_scale = create_tensor_from_dic(dic['value_cache_quant_scale'])
          context_lengths = dic['context_lengths']['data']
          max_context_len = dic['max_context_len']['data']
          context_seq_offset = dic['context_seq_offset']['data']
          block_tables = dic['block_tables']['data']
          quant_mode = dic['quant_mode']['data']
          quant_bit = dic['quant_bit']['data']

          self.launch(key, value, key_cache, value_cache, key_cache_quant_scale,
                      value_cache_quant_scale, context_lengths, max_context_len,
                      context_seq_offset, block_tables, quant_mode, quant_bit)

    def launch(self, *args):
        key, value, key_cache, value_cache, key_cache_scale, value_cache_scale, \
            context_lengths, max_context_len, context_seq_offset, block_tables, \
                quant_mode, quant_bit = args
        if value is None or value_cache is None or value_cache_scale is None:
            has_value = False
        else:
            has_value = True
        ops.dequant_from_paged_cache(key, value, key_cache, value_cache, key_cache_scale,
                                     value_cache_scale, context_lengths, max_context_len,
                                     context_seq_offset, block_tables, quant_mode, quant_bit)
        key_clone = key.clone()
        if has_value:
            value_clone = value.clone()
        self.op_impl_base(key, value, key_cache, value_cache, key_cache_scale,
                          value_cache_scale, context_lengths, max_context_len,
                          context_seq_offset, block_tables, quant_mode, quant_bit)
        self.assertTensorsEqual(key_clone.cpu().float(), key.cpu().float(), 0.001, use_MSE=True)
        if has_value:
            self.assertTensorsEqual(value_clone.cpu().float(), value.cpu().float(), 0.001,
                                    use_MSE=True)

    def op_impl_base(self, *args):
        def dequant_from_cache(quant_data: torch.Tensor,
                               scale_data: torch.Tensor,
                               quant_mode: int):
            quant_data_fp32 = quant_data.clone().to(torch.float)
            scale_data_fp32 = scale_data.clone().to(torch.float)
            if quant_mode == 0:  # per channel [head_num, 1, head_size]
                scale_data_fp32 = scale_data[..., None, :]
            else:  # per head/token [head_num, context_len, 1]
                scale_data_fp32 = scale_data[..., None]
            dequant_data_fp32 = quant_data_fp32 * scale_data_fp32

            return dequant_data_fp32

        key, value, key_cache, value_cache, key_cache_scale, value_cache_scale, \
            context_lengths, max_context_len, context_seq_offset, block_tables, \
                quant_mode, quant_bit = args
        batch_size = context_lengths.size(0)
        if context_seq_offset is None:
            cu_seq_offset = torch.cumsum(context_lengths, dim=-1)
            context_seq_offset = torch.zeros_like(cu_seq_offset)
            context_seq_offset[1:] = cu_seq_offset[:-1]

        total_seqlen = 0
        block_size = key_cache.size(2)
        for i in range(batch_size):
            context_len = context_lengths[i].item()
            seq_begin = context_seq_offset[i].item()
            seq_end = seq_begin + context_len
            total_seqlen += context_len
            full_block_num = context_len // block_size
            rem_token_num = context_len % block_size
            key_i = key[seq_begin:seq_end].transpose(1, 0)
            # [head_num, seq_num, head_size]
            key_cache_i = torch.concat(
                [key_cache[block_tables[i, j], ...] for j in range(full_block_num)] +
                ([key_cache[block_tables[i, full_block_num], :, :rem_token_num, :]] \
                if rem_token_num > 0 else []), dim=-2
            )

            # dequant key from cache
            if quant_mode == 0:
                # [head_num, head_size]
                key_cache_scale_i = key_cache_scale
            else:
                # [head_num, seq_num]
                key_cache_scale_i = torch.concat(
                    [key_cache_scale[block_tables[i, j],...] for j in range(full_block_num)] +
                    ([key_cache_scale[block_tables[i, full_block_num], :, :rem_token_num]] \
                    if rem_token_num > 0 else []), dim=-1
                )

            dequant_key_i = dequant_from_cache(key_cache_i, key_cache_scale_i, quant_mode)
            key_i[...] = dequant_key_i.to(key_i.dtype)
            # dequant value from cache
            if not (value_cache is None or value is None or value_cache_scale is None):
                value_i = value[seq_begin:seq_end].transpose(1, 0)
                # [head_num, seq_num, head_size]
                value_cache_i = torch.concat(
                    [value_cache[block_tables[i, j], ...] for j in range(full_block_num)] +
                    ([value_cache[block_tables[i, full_block_num], :, :rem_token_num, :]] \
                    if rem_token_num > 0 else []), dim=-2
                )

                # dequant value from cache
                if quant_mode == 0:
                    # [head_num, head_size]
                    value_cache_scale_i = value_cache_scale
                else:
                    # [head_num, seq_num]
                    value_cache_scale_i = torch.concat(
                        [value_cache_scale[block_tables[i, j],...] for j in range(full_block_num)] +
                        ([value_cache_scale[block_tables[i, full_block_num], :, :rem_token_num]] \
                        if rem_token_num > 0 else []), dim=-1
                    )
                dequant_value_i = dequant_from_cache(value_cache_i, value_cache_scale_i, quant_mode)
                value_i[...] = dequant_value_i.to(value_i.dtype)

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ or "MLU3" in torch.mlu.get_device_name(),
                     "Skipping test_prevent due to ASan issues or in MLU300 series.")
    def test_prevent(self):
        batch_size, max_context_len, head_num_q, head_num_kv, cache_mem_len, \
            head_size, block_size = 8, 1024, 16, 32, 2048, 128, 16
        dtype, quant_mode, quant_bit = torch.float16, 0, 8
        default_args = gen_args(batch_size, max_context_len, head_num_q, head_num_kv, cache_mem_len,
                                head_size, head_size, block_size, 1, dtype, quant_mode, quant_bit)
        args = gen_args(batch_size, max_context_len, head_num_q, head_num_kv, cache_mem_len,
                        head_size, head_size, block_size, 1, torch.float32, quant_mode, quant_bit)
        self.assertException("Tensor key type should be half or bfloat16.",
                             torch.ops.torch_mlu_ops.dequant_from_paged_cache, *args)
        args = gen_args(batch_size, max_context_len, head_num_q, head_num_kv, cache_mem_len,
                        head_size, head_size, block_size, 1, dtype, 2, quant_bit)
        self.assertException("quantization mode support 0 and 1.",
                             torch.ops.torch_mlu_ops.dequant_from_paged_cache, *args)
        args = gen_args(batch_size, max_context_len, head_num_q, head_num_kv, cache_mem_len,
                        head_size, head_size, block_size, 1, dtype, quant_mode, 4)
        self.assertException("quantization bit width only supports 8.",
                             torch.ops.torch_mlu_ops.dequant_from_paged_cache, *args)
        default_args[-5] = 10240
        self.assertException("max_context_len should smaller than or equal to block_size * max_block_num.",
                             torch.ops.torch_mlu_ops.dequant_from_paged_cache, *default_args)
        default_args[-5] = max_context_len
        default_args[1] = default_args[1].transpose(-1, -2)
        self.assertException("Tensor value last dim must be contiguous.",
                             torch.ops.torch_mlu_ops.dequant_from_paged_cache, *default_args)

    @unittest.skipIf("MLU3" in torch.mlu.get_device_name(),
                     "Skipping test_dequant_from_paged_cache in MLU300 series")
    def test_dequant_from_paged_cache(self):
        test_cases           = 100
        head_size_times      = 16
        mlu_name             = torch.mlu.get_device_name()
        batch_size_list      = torch.randint(low=1, high=32, size=[test_cases], dtype=torch.int32)
        max_context_len_list = torch.randint(low=2, high=1024, size=[test_cases],
                                             dtype=torch.int32)
        head_num_q_list     = torch.randint(low=1, high=64, size=[test_cases], dtype=torch.int32)
        head_num_kv_list    = torch.randint(low=1, high=8, size=[test_cases], dtype=torch.int32)
        head_size_list      = torch.randint(low=1, high=16, size=[test_cases],
                                            dtype=torch.int32) * head_size_times
        block_size_list     = torch.randint(low=1, high=32, size=[test_cases], dtype=torch.int32)
        cache_mem_len_list  = torch.randint(low=512, high=1024, size=[test_cases],
                                            dtype=torch.int32) * 2
        quant_mode_list     = np.random.choice([0, 1], test_cases)
        quant_bit_list      = np.random.choice([8], test_cases)
        use_offset_list     = np.random.choice([False, True], test_cases)
        has_value_list      = np.random.choice([False, True], test_cases)
        dtype_list          = [torch.half, torch.bfloat16]
        dtype_list          = np.random.choice(dtype_list, test_cases)
        for i in range(test_cases):
            batch_size      = batch_size_list[i].item()
            head_num_q      = head_num_q_list[i].item()
            max_context_len = max_context_len_list[i].item()
            head_num_kv     = head_num_kv_list[i].item()
            head_size       = head_size_list[i].item()
            block_size      = block_size_list[i].item()
            cache_mem_len   = cache_mem_len_list[i].item()
            quant_mode      = quant_mode_list[i]
            quant_bit       = quant_bit_list[i]
            use_seq_offset  = use_offset_list[i]
            has_value       = has_value_list[i]
            dtype           = dtype_list[i]
            print("batch_size={}, head_num={}, head_size={}, max_context_len={}, quant_mode={}, "
                  "quant_bit={}, dtype={} testing...".format(batch_size, head_num_kv, head_size,
                      max_context_len, quant_mode, quant_bit, dtype))

            torch.manual_seed(2766)
            args = gen_args(batch_size, max_context_len, head_num_q, head_num_kv, cache_mem_len,
                            head_size, head_size, block_size, use_seq_offset, dtype, quant_mode,
                            quant_bit, has_value)
            self.launch(*args)

    @unittest.skipIf("MLU3" in torch.mlu.get_device_name(),
                     "Skipping test_dequant_from_paged_cache in MLU300 series")
    def test_inductor(self):
        batch_size, max_context_len, head_num_q, head_num_kv, cache_mem_len, \
            head_size, block_size = 8, 1024, 16, 32, 2048, 128, 16
        dtype, quant_mode, quant_bit = torch.float16, 0, 8
        args = gen_args(batch_size, max_context_len, head_num_q, head_num_kv, cache_mem_len,
                        head_size, head_size, block_size, 1, dtype, quant_mode, quant_bit)
        self.base_opcheck(torch.ops.torch_mlu_ops.dequant_from_paged_cache, args)
        args = gen_args(batch_size, max_context_len, head_num_q, head_num_kv, cache_mem_len,
                        head_size, head_size, block_size, 0, dtype, quant_mode, quant_bit)
        self.base_opcheck(torch.ops.torch_mlu_ops.dequant_from_paged_cache, args)

if __name__ == '__main__':
    exit(run_unittest(TestDequantFromPagedCache))
