import random
import torch
import torch_mlu
import unittest
import torch_mlu_ops as ops
from common_utils import *

def gen_args(batch_size, num_heads, head_size, cache_memory_len, packed, dtype):
    args = generate_cache_func_args(batch_size, num_heads, head_size, cache_memory_len, packed, dtype)
    slot_mapping = args[11]
    args = args[0:4] + [slot_mapping]
    return args

class TestReshapePagedCacheOp(BtTestCase):
    def op_impl_base(self, *args):
        k, v, k_cache, v_cache, slot_mapping = args
        num_tokens = k.shape[0]
        block_size = k_cache.shape[2]
        for i in range(num_tokens):
            if slot_mapping[i] >= 0:
                block_id = torch.div(slot_mapping[i], block_size, rounding_mode='floor')
                block_offset = slot_mapping[i] % block_size
                k_cache[block_id, :, block_offset, :] = k[i]
                if v is not None:
                    v_cache[block_id, :, block_offset, :] = v[i]
        return (k_cache, v_cache) if v is not None else k_cache

    @unittest.skipIf('MLU3' in torch.mlu.get_device_name(), "reshape_paged_cache not support MLU3XX device")
    def test_reshape_paged_cache(self):
        test_cases = 100
        num_tokens_list = torch.randint(low=1, high=1024, size=(test_cases, ), dtype=torch.int32)
        num_heads_list = torch.randint(low=1, high=64, size=(test_cases, ), dtype=torch.int32)
        head_size_list = torch.randint(low=1, high=16, size=(test_cases, ), dtype=torch.int32)
        head_size_list *= 16
        block_size_list = torch.randint(low=1, high=4, size=(test_cases, ), dtype=torch.int32)
        block_size_list *= 16
        only_reshape_key_list = [True, False]

        for i in range(test_cases):
            num_tokens = num_tokens_list[i]
            num_heads = num_heads_list[i]
            head_size = head_size_list[i]
            block_size = block_size_list[i]
            min_blocks = (num_tokens + block_size - 1) // block_size
            num_blocks = min(min_blocks + 10, 2 * min_blocks)
            dtype_list = [torch.half, torch.float]
            if torch_mlu.mlu.is_bf16_supported():
                dtype_list.append(torch.bfloat16)
            only_reshape_key = random.choice(only_reshape_key_list)
            for dtype in dtype_list:
                print("num_tokens: {}, num_heads: {}, head_size: {}, num_blocks: {}, block_size: {}, testing...".format(
                    num_tokens, num_heads, head_size, num_blocks, block_size), flush=True)
                qkv = torch.randn(num_tokens, 3, num_heads, head_size, dtype=dtype, device="mlu")
                _, key, value = qkv.unbind(dim=1)

                key_cache   = torch.randn(num_blocks, num_heads, block_size, head_size, dtype=dtype, device="mlu")
                value_cache = torch.randn(num_blocks, num_heads, block_size, head_size, dtype=dtype, device="mlu")

                num_slots = num_blocks * block_size
                slot_mapping = random.sample(range(num_slots.item()), num_tokens.item())
                slot_mapping = torch.tensor(slot_mapping, dtype=torch.int, device="mlu")
                slot_mapping[-1] = -1

                ref_key_cache, ref_value_cache = key_cache.clone(), value_cache.clone()
                if only_reshape_key:
                    value, ref_value_cache, value_cache = None, None, None
                self.op_impl_base(key, value, ref_key_cache, ref_value_cache, slot_mapping)
                ops.reshape_paged_cache(key, value, key_cache, value_cache, slot_mapping)
                self.assertTensorsEqual(key_cache.cpu().float(), ref_key_cache.cpu().float(), 0, use_MSE=True, use_RAE=True, use_RMA=True)
                if not only_reshape_key:
                    self.assertTensorsEqual(value_cache.cpu().float(), ref_value_cache.cpu().float(), 0, use_MSE=True, use_RAE=True, use_RMA=True)

    @unittest.skipIf('MLU3' in torch.mlu.get_device_name(), "reshape_paged_cache not support MLU3XX device")
    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_prevent due to ASan issues")
    def test_large_tensor(self):
        dtype_list = [torch.half, torch.float]
        if torch_mlu.mlu.is_bf16_supported():
            dtype_list.append(torch.bfloat16)
        for dtype in dtype_list:
            print("reshape_paged_cache: test_large_tensor")
            if dtype == torch.float:
                dtype_size = 4
            elif dtype == torch.half or dtype == torch.bfloat16:
                dtype_size = 2
            head_num = 16
            head_size = 128
            token_num = 20
            block_size = 16
            block_num = ((2**32 - 1) // dtype_size // head_num // head_size // block_size)
            k = torch.randn(token_num, head_num, head_size, dtype=dtype, device="mlu")
            v = torch.randn(token_num, head_num, head_size, dtype=dtype, device="mlu")
            k_cache = torch.randn(block_num, head_num, block_size, head_size, dtype=dtype, device="mlu")
            v_cache = torch.randn(block_num, head_num, block_size, head_size, dtype=dtype, device="mlu")
            num_slots = block_num * block_size
            slot_mapping = torch.randperm(num_slots, dtype=torch.int32, device="mlu")[:token_num]
            ref_key_cache, ref_value_cache = k_cache.clone(), v_cache.clone()
            self.op_impl_base(k, v, ref_key_cache, ref_value_cache, slot_mapping)
            ops.reshape_paged_cache(k, v, k_cache, v_cache, slot_mapping)
            for i in range(block_size):
                self.assertTensorsEqual(k_cache[:, :, i, :].cpu().float(), ref_key_cache[:, :, i, :].cpu().float(), 0, use_MSE=True, use_RAE=True, use_RMA=True)
                self.assertTensorsEqual(v_cache[:, :, i, :].cpu().float(), ref_value_cache[:, :, i, :].cpu().float(), 0, use_MSE=True, use_RAE=True, use_RMA=True)
            block_num = (2**32 // dtype_size // head_num // head_size // block_size)
            k_cache = torch.randn(block_num, head_num, block_size, head_size, dtype=dtype, device="mlu")
            v_cache = torch.randn(block_num, head_num, block_size, head_size, dtype=dtype, device="mlu")
            num_slots = block_num * block_size
            slot_mapping = torch.randperm(num_slots, dtype=torch.int32, device="mlu")[:token_num]
            self.assertException("The addressing range of kv_cache cannot exceed 4G.", ops.reshape_paged_cache,
                                k, v, k_cache, v_cache, slot_mapping)

    @unittest.skipIf('MLU3' in torch.mlu.get_device_name(), "reshape_paged_cache not support MLU3XX device")
    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_prevent due to ASan issues")
    def test_prevent(self):
        k = torch.randn(1024, 8, 128, dtype=torch.half, device="mlu")
        v = torch.randn(1024, 8, 128, dtype=torch.half, device="mlu")
        k_cache = torch.randn(1024, 8, 4, 128, dtype=torch.half, device="mlu")
        v_cache = None
        slot_mapping = random.sample(range(1024 * 4), 1024)
        slot_mapping = torch.tensor(slot_mapping, dtype=torch.int, device="mlu")
        self.assertException("v.has_value() == v_cache.has_value().", ops.reshape_paged_cache,
                             k, v, k_cache, v_cache, slot_mapping)
        v_cache = torch.randn(1024, 8, 4, 128, dtype=torch.half, device="mlu")
        v_cache = v_cache.as_strided(size=(1024, 8, 4, 128), stride=(8*4*128, 4* 128, 127, 1))
        self.assertException("v_cache need be contiguous.", ops.reshape_paged_cache,
                             k, v, k_cache, v_cache, slot_mapping)
        v_cache = v_cache.contiguous()
        v = v.as_strided(size=(1024, 8, 128), stride=(1, 1024, 1024 * 8))
        self.assertException("v last dim must be contiguous.", ops.reshape_paged_cache,
                             k, v, k_cache, v_cache, slot_mapping)
        v = v.as_strided(size=(1024, 8, 128), stride=(1024, 8, 1))
        self.assertException("v second dim must be contiguous.", ops.reshape_paged_cache,
                             k, v, k_cache, v_cache, slot_mapping)


    @unittest.skipIf('MLU3' in torch.mlu.get_device_name(), "reshape_paged_cache not support MLU3XX device")
    def test_inductor(self):
        batch_size, num_heads, head_size, cache_memory_len, dtype = 4, 8, 64, 128, torch.float16
        args = gen_args(batch_size, num_heads, head_size, cache_memory_len, 1, dtype)
        self.base_opcheck(torch.ops.torch_mlu_ops.reshape_paged_cache, args)

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