import torch
import unittest
import torch_mlu_ops as ops
import random
from common_utils import *
import numpy as np

def quant2int8(input_fp: torch.Tensor,
               scale_fp: torch.Tensor):
    input_fp32 = input_fp.to(torch.float32)
    scaled_input = input_fp32 / scale_fp
    rounded = torch.round(scaled_input)
    clipped = torch.clip(rounded, -128, 127)
    return clipped.to(torch.int8)

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, True, True)
    cache_scale = args[10]
    slot_mapping = args[11]
    args = args[0:4]
    args.insert(2, cache_scale[0])
    args.insert(3, cache_scale[1])
    args.insert(4, slot_mapping)
    return args
class TestOfflineQuantToPagedCache(BtTestCase):
    def op_impl_base(self, *args):
        k, v, k_cache_scale, v_cache_scale, slot_mapping, k_cache, v_cache = args
        tokens_num = k.shape[0]
        block_size = k_cache.shape[2]
        for i in range(tokens_num):
            if slot_mapping[i] >= 0:
                key_i = k[i]
                block_id = torch.div(slot_mapping[i], block_size, rounding_mode='floor')
                block_offset = slot_mapping[i] % block_size
                key_cache_i = k_cache[block_id, :, block_offset, :]
                quant_key_i = quant2int8(key_i, k_cache_scale)
                key_cache_i[...] = quant_key_i
                if v is not None:
                    value_i = v[i]
                    value_cache_i = v_cache[block_id, :, block_offset, :]
                    quant_value_i = quant2int8(value_i, v_cache_scale)
                    value_cache_i[...] = quant_value_i
        return (k_cache, v_cache) if v is not None else k_cache

    @unittest.skipIf('MLU3' in torch.mlu.get_device_name(), "offline_quant_to_paged_cache not support MLU3XX device")
    def test_offline_quant_to_paged_cache(self):
        test_cases = 10
        token_list = torch.randint(low=1, high=512, size=(test_cases, ), dtype=torch.int32)
        head_num_list = torch.randint(low=1, high=64, size=(test_cases, ), dtype=torch.int32)
        head_size_list = torch.randint(low=1, high=64, size=(test_cases, ), dtype=torch.int32)
        block_size_list = torch.randint(low=1, high=64, size=(test_cases, ), dtype=torch.int32)
        dtype_list = torch.randint(low=0, high=3, size=(test_cases, ), dtype=torch.int32)
        dtype_map = [torch.bfloat16, torch.half, torch.float32]
        only_quant_key_list = [True, False]

        for i in range(test_cases):
            tokens_num = token_list[i].item()
            head_num = head_num_list[i].item()
            head_size = head_size_list[i].item()
            block_size = block_size_list[i].item()
            dtype = dtype_map[dtype_list[i]]
            print("tokens_num={}, head_num={}, head_size={}, block_size={}, dtype={} testing...".format(
            tokens_num, head_num, head_size, block_size, dtype))
            np.random.seed(1)
            only_quant_key = random.choice(only_quant_key_list)
            key_data = np.random.uniform(-1, 1, size=[tokens_num, head_num, head_size])
            key_cache_scale_data = np.random.uniform(-10, 10, size=[head_num, head_size])
            key = torch.tensor(key_data, dtype=dtype, device="mlu")
            key_cache_scale = torch.tensor(key_cache_scale_data, dtype=torch.float32, device="mlu")
            value_data = np.random.uniform(-0.25, 0.25, size=[tokens_num, head_num, head_size])
            value_cache_scale_data = np.random.uniform(-10, 10, size=[head_num, head_size])
            value = torch.tensor(value_data, dtype=dtype, device="mlu")
            value_cache_scale = torch.tensor(value_cache_scale_data, dtype=torch.float32, device="mlu")

            min_blocks = (int)((tokens_num + block_size - 1) / block_size)
            blocks_num = min(min_blocks + 10, 2 * min_blocks)
            num_slots = blocks_num * block_size
            slot_mapping = random.sample(range(num_slots), tokens_num)
            slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, device="mlu")
            slot_mapping[-1] = -1   # test mask

            torch.manual_seed(0)
            key_cache = torch.randint(-128, 127, (blocks_num, head_num, block_size, head_size), dtype=torch.int8).mlu()
            value_cache = torch.randint(-128, 127, (blocks_num, head_num, block_size, head_size), dtype=torch.int8).mlu()

            #python base result
            key_cache_base = key_cache.clone()
            value_cache_base = value_cache.clone()
            if only_quant_key:
                value, value_cache_scale, value_cache_base, value_cache = None, None, None, None
            self.op_impl_base(key, value, key_cache_scale, value_cache_scale,
                                 slot_mapping, key_cache_base, value_cache_base)
            #mlu result
            ops.offline_quant_to_paged_cache(key, value, key_cache_scale, value_cache_scale,
                                         slot_mapping, key_cache, value_cache)
            #compute diff
            baseline_key_cache = key_cache_base.cpu().flatten()
            mlu_key_cache = key_cache.cpu().flatten()
            key_cache_diff = (mlu_key_cache - baseline_key_cache).abs()
            assert torch.max(key_cache_diff) < 2, "key_cache_diff exceed threshold"
            if not only_quant_key:
                baseline_value_cache = value_cache_base.cpu().flatten()
                mlu_value_cache = value_cache.cpu().flatten()
                value_cache_diff = (mlu_value_cache - baseline_value_cache).abs()
                assert torch.max(value_cache_diff) < 2, "value_cache_diff exceed threshold"

    @unittest.skipIf('MLU3' in torch.mlu.get_device_name(), "offline_quant_to_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("offline_quant_to_paged_cache: test_large_tensor...")
            head_num = 16
            head_size = 128
            token_nums = 200
            block_size = 16
            block_nums = ((2**32 - 1) // 1 // head_num // head_size // block_size)
            num_slots = block_nums * block_size
            slot_mapping = torch.randperm(num_slots, dtype=torch.int32, device="mlu")[:token_nums]
            key = torch.randn(token_nums, head_num, head_size, dtype=dtype, device="mlu")
            value = torch.randn(token_nums, head_num, head_size, dtype=dtype, device="mlu")
            key_cache_scale = torch.randn(head_num, head_size, dtype=torch.float32, device="mlu")
            value_cache_scale = torch.randn(head_num, head_size, dtype=torch.float32, device="mlu")
            key_cache = torch.zeros(block_nums, head_num, block_size, head_size, dtype=torch.int8, device="mlu")
            value_cache = torch.zeros(block_nums, head_num, block_size, head_size, dtype=torch.int8, device="mlu")

            #python base result
            key_cache_base = key_cache.clone()
            value_cache_base = value_cache.clone()
            self.op_impl_base(key, value, key_cache_scale, value_cache_scale,
                                 slot_mapping, key_cache_base, value_cache_base)
            #mlu result
            ops.offline_quant_to_paged_cache(key, value, key_cache_scale, value_cache_scale,
                                         slot_mapping, key_cache, value_cache)
            #compute diff
            baseline_key_cache = key_cache_base.cpu().flatten()
            mlu_key_cache = key_cache.cpu().flatten()
            key_cache_diff = (mlu_key_cache - baseline_key_cache).abs()
            assert torch.max(key_cache_diff) < 2, "key_cache_diff exceed threshold"

            baseline_value_cache = value_cache_base.cpu().flatten()
            mlu_value_cache = value_cache.cpu().flatten()
            value_cache_diff = (mlu_value_cache - baseline_value_cache).abs()
            assert torch.max(value_cache_diff) < 2, "value_cache_diff exceed threshold"

            block_nums = (2**32 // 1 // head_num // head_size // block_size)
            num_slots = block_nums * block_size
            slot_mapping = torch.randperm(num_slots, dtype=torch.int32, device="mlu")[:token_nums]
            key_cache = torch.zeros(block_nums, head_num, block_size, head_size, dtype=torch.int8, device="mlu")
            value_cache = torch.zeros(block_nums, head_num, block_size, head_size, dtype=torch.int8, device="mlu")
            self.assertException("The addressing range of kv_cache cannot exceed 4G.", ops.offline_quant_to_paged_cache,
                                key, value, key_cache_scale, value_cache_scale, slot_mapping, key_cache, value_cache)

    @unittest.skipIf('MLU3' in torch.mlu.get_device_name(), "offline_quant_to_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):
        print("offline_quant_to_paged_cache: test prevent...")
        head_num = 16
        head_size = 128
        token_nums = 200
        block_size = 16
        block_nums = 20
        num_slots = block_nums * block_size
        dtype = torch.half
        slot_mapping = torch.randperm(num_slots, dtype=torch.int32, device="mlu")[:token_nums]
        key = torch.randn(token_nums, head_num, head_size, dtype=dtype, device="mlu")
        value = torch.randn(token_nums, head_num, head_size, dtype=dtype, device="mlu")
        key_cache_scale = torch.randn(head_num, head_size, dtype=torch.float32, device="mlu")
        value_cache_scale = torch.randn(head_num, head_size, dtype=torch.float32, device="mlu")
        key_cache = torch.zeros(block_nums, head_num, block_size, head_size, dtype=torch.int8, device="mlu")
        value_cache = None
        self.assertException("v.has_value() == v_cache.has_value() && v.has_value() == v_cache_scale.has_value().",
                             ops.offline_quant_to_paged_cache, key, value, key_cache_scale, value_cache_scale,
                             slot_mapping, key_cache, value_cache)
        value_cache = torch.zeros(block_nums, head_num, block_size, head_size, dtype=torch.int8, device="mlu")
        value = value.reshape(token_nums, head_num, head_size, 1)
        self.assertException("dim of v must be 3",
                             ops.offline_quant_to_paged_cache, key, value, key_cache_scale, value_cache_scale,
                             slot_mapping, key_cache, value_cache)
        value = value.squeeze()
        value_cache = value_cache.reshape(block_nums, head_num, block_size, head_size, 1)
        self.assertException("dim of v_cache must be 4",
                             ops.offline_quant_to_paged_cache, key, value, key_cache_scale, value_cache_scale,
                             slot_mapping, key_cache, value_cache)
        value_cache = value_cache.squeeze()
        value_cache_scale = value_cache_scale.reshape(head_num, head_size, 1)
        self.assertException("dim of v_cache_scale must be 2",
                             ops.offline_quant_to_paged_cache, key, value, key_cache_scale, value_cache_scale,
                             slot_mapping, key_cache, value_cache)
        value_cache_scale = value_cache_scale.squeeze()
        value = value.as_strided(size=(token_nums, head_num, head_size), stride=(head_num, 1, head_size))
        self.assertException("v last dim must be contiguous.",
                             ops.offline_quant_to_paged_cache, key, value, key_cache_scale, value_cache_scale,
                             slot_mapping, key_cache, value_cache)
        value = value.as_strided(size=(token_nums, head_num, head_size), stride=(head_num, token_nums, 1))
        self.assertException("v second dim must be contiguous.",
                             ops.offline_quant_to_paged_cache, key, value, key_cache_scale, value_cache_scale,
                             slot_mapping, key_cache, value_cache)

    @unittest.skipIf('MLU3' in torch.mlu.get_device_name(), "offline_quant_to_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.offline_quant_to_paged_cache, args)

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