import math
import random
import torch
import torch_mlu
import unittest
import torch_mlu_ops as ops
from common_utils import *
from itertools import product
from typing import List, Tuple
import os
import copy

class TestCopyBlocksOp(BtTestCase):
    def create_kv_caches(
        self,
        num_blocks: int,
        block_size: int,
        num_layers: int,
        num_heads: int,
        head_size: int,
        dtype: torch.dtype,
        seed: int
        ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:

        torch.random.manual_seed(seed)
        torch.cuda.manual_seed(seed)

        scale = head_size**-0.5
        # vllm scale
        # x = 16 // torch.tensor([], dtype=dtype).element_size()
        # key_cache_shape = (num_blocks, num_heads, head_size // x, block_size, x)
        key_cache_shape = (num_blocks, num_heads, block_size, head_size)
        key_caches = []
        for _ in range(num_layers):
            if dtype in {torch.int8, torch.uint8, torch.int16, torch.int32, torch.int64}:
                info = torch.iinfo(dtype)
                key_cache = torch.randint(info.min, info.max, size=key_cache_shape, dtype=dtype).mlu()
            else:
                key_cache = torch.empty(size=key_cache_shape, dtype=dtype).mlu()
                key_cache.uniform_(-scale, scale)
            key_caches.append(key_cache)

        value_cache_shape = (num_blocks, num_heads, head_size, block_size)
        value_caches = []
        for _ in range(num_layers):
            if dtype in {torch.int8, torch.uint8, torch.int16, torch.int32, torch.int64}:
                info = torch.iinfo(dtype)
                value_cache = torch.randint(info.min, info.max, size=value_cache_shape, dtype=dtype).mlu()
            else:
                value_cache = torch.empty(size=value_cache_shape, dtype=dtype).mlu()
                value_cache.uniform_(-scale, scale)
            value_caches.append(value_cache)

        return key_caches, value_caches

    def create_block_mapping(self, num_blocks, num_mappings, seed = 0):
        random.seed(seed)
        torch.random.manual_seed(seed)
        assert 3 * num_mappings <= num_blocks
        block_mapping = {}
        src_blocks = random.sample(range(num_blocks), num_mappings)
        remainig_blocks = list(set(range(num_blocks)) - set(src_blocks))
        dst_blocks = random.sample(remainig_blocks, 2 * num_mappings)
        for i in range(num_mappings):
            src = src_blocks[i]
            dst1 = dst_blocks[2 * i]
            dst2 = dst_blocks[2 * i + 1]
            block_mapping[src] = [dst1, dst2]
        return block_mapping

    def op_impl_base(self, *args):
        k_caches, v_caches, block_mapping = args
        for src, dsts in block_mapping.items():
            srcs = [src for i in range(len(dsts))]
            srcs_ind = torch.tensor(srcs, dtype=torch.int64)
            dsts_ind = torch.tensor(dsts, dtype=torch.int64)
            for key_cache in k_caches:
                key_cache[dsts_ind] = key_cache[srcs_ind]
            if v_caches is not None:
                for value_cache in v_caches:
                    value_cache[dsts_ind] = value_cache[srcs_ind]
        return (k_caches, v_caches) if v_caches is not None else k_caches

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_copy_blocks due to ASan issues")
    def test_copy_blocks(self):
        dtype_list = [torch.int8, torch.uint8, torch.int16, torch.int32, torch.int64, torch.half, torch.float]
        if torch_mlu.mlu.is_bf16_supported():
            dtype_list.append(torch.bfloat16)
        num_tokens_list = [83]
        num_heads_list = [8]
        head_size_list = [64, 512]
        num_blocks_list = [3600]
        block_size_list = [8]
        num_layers_list = [1, 6]
        num_mappings_list = [128, 600]
        seeds_list = [0]
        only_key_cache_list = [True, False]

        args = product(num_tokens_list, num_heads_list, head_size_list, num_blocks_list, block_size_list, dtype_list,
        num_layers_list, num_mappings_list, seeds_list)
        for num_tokens, num_heads, head_size, num_blocks, block_size, dtype, num_layers, num_mappings, seed in args:
            print("num_tokens: {}, num_heads: {}, head_size: {}, num_blocks: {}, block_size: {}, dtype: {}, num_layers: {}, \
num_mappings: {}, seed: {}, testing...".format(
                num_tokens, num_heads, head_size, num_blocks, block_size, dtype, num_layers, num_mappings, seed), flush=True)
            block_mapping = self.create_block_mapping(num_blocks, num_mappings, seed)
            only_key_cache = random.choice(only_key_cache_list)
            # Create the KV caches.
            key_caches, value_caches = self.create_kv_caches(num_blocks, block_size,
                                                num_layers, num_heads,
                                                head_size, dtype, seed)
            # Clone the KV caches.
            cloned_key_caches = [key_cache.cpu().clone() for key_cache in key_caches]
            cloned_value_caches = [value_cache.cpu().clone() for value_cache in value_caches]

            # Call the copy blocks kernel.
            if only_key_cache:
                value_caches = None
                cloned_value_caches = None
            ops.copy_blocks(key_caches, value_caches, block_mapping)

            self.op_impl_base(cloned_key_caches, cloned_value_caches, block_mapping)

            # Compare the results.
            for key_cache, cloned_key_cache in zip(key_caches, cloned_key_caches):
                self.assertTensorsEqual(key_cache.cpu().float(), cloned_key_cache.cpu().float(), 0, use_MSE=True, use_RAE=True, use_RMA=True)
            if not only_key_cache:
                for value_cache, cloned_value_cache in zip(value_caches, cloned_value_caches):
                    self.assertTensorsEqual(value_cache.cpu().float(), cloned_value_cache.cpu().float(), 0, use_MSE=True, use_RAE=True, use_RMA=True)

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_inductor due to ASan issues")
    def test_prevent(self):
        print("test_copy_block: test_prevent...")
        num_blocks, block_size, head_num, head_size, block_mapping = 384, 6, 32, 128, 128
        k_cache = torch.randn(num_blocks * head_num, block_size, head_size, dtype=torch.half, device="mlu")
        v_cache = torch.randn(num_blocks * head_num, head_size, block_size + 1, dtype=torch.float, device="mlu")
        key_caches = [k_cache,]
        value_caches = None
        block_mapping = self.create_block_mapping(num_blocks, block_mapping)
        self.assertException("every layer k_cache must be 4d.", ops.copy_blocks,
                             key_caches, value_caches, block_mapping)
        k_cache = k_cache.reshape(num_blocks, head_num, block_size, head_size)
        key_caches = [k_cache, k_cache,]
        value_caches = [v_cache]
        self.assertException("k_caches size must equal to v_caches size if v_caches is not none.",
                             ops.copy_blocks, key_caches, value_caches, block_mapping)
        key_caches = [k_cache]
        value_caches = [v_cache]
        self.assertException("the data type of k_caches and v_caches are not the same.",
                             ops.copy_blocks, key_caches, value_caches, block_mapping)
        value_caches[0] = value_caches[0].to(torch.half)
        self.assertException("every layer k_cache dim must equal to v_cache dim.",
                             ops.copy_blocks, key_caches, value_caches, block_mapping)
        value_caches[0] = value_caches[0].reshape(num_blocks, head_num, head_size, block_size + 1)
        self.assertException("the block_size of k_caches and v_caches are not the same.",
                             ops.copy_blocks, key_caches, value_caches, block_mapping)

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_inductor due to ASan issues")
    def test_inductor(self):
        num_heads, head_size, num_blocks, block_size, num_layers, num_mappings = 8, 64, 384, 8, 1, 128
        key_caches, value_caches = self.create_kv_caches(num_blocks, block_size,
                                                num_layers, num_heads,
                                                head_size, torch.float16, 0)
        block_mapping = self.create_block_mapping(num_blocks, num_mappings)
        self.base_opcheck(torch.ops.torch_mlu_ops.copy_blocks, (key_caches, value_caches, block_mapping))

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