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

class TestQuantToPagedCache(BtTestCase):
  def run_gen_case(self, dic):
    dump_data = dic.pop('dump_data')
    if dump_data:
      self.launch(*dic.values())
    else:
      k = create_tensor_from_dic(dic['k'], is_uniform=True, low=-1, high=1)
      v = create_tensor_from_dic(dic['v'], is_uniform=True, low=-0.25, high=0.25)
      k_cache = create_tensor_from_dic(dic['k_cache'])
      v_cache = create_tensor_from_dic(dic['v_cache'])
      k_cache_quant_scale = create_tensor_from_dic(dic['k_cache_quant_scale'])
      v_cache_quant_scale = create_tensor_from_dic(dic['v_cache_quant_scale'])
      slot_mapping = dic['slot_mapping']['data']
      self.launch(k, v, k_cache, v_cache, k_cache_quant_scale, v_cache_quant_scale, slot_mapping)

  def launch(self, *args):
    v = args[1]
    args_bak = copy.deepcopy(args)
    torch_out = self.op_impl_base(*args_bak)
    tmo_out = ops.quant_to_paged_cache(*args)
    self.assertTensorsEqual(torch_out[0].cpu().float(), tmo_out[0].cpu().float(), 9e-3,
                              use_MSE=True, use_RAE=True)
    self.assertTensorsEqual(torch_out[1].cpu().float(), tmo_out[1].cpu().float(), 9e-3,
                              use_MSE=True, use_RAE=True)
    if v is not None:
      self.assertTensorsEqual(torch_out[2].cpu().float(), tmo_out[2].cpu().float(), 3e-3,
                                use_MSE=True, use_RAE=True)
      self.assertTensorsEqual(torch_out[3].cpu().float(), tmo_out[3].cpu().float(), 3e-3,
                                    use_MSE=True, use_RAE=True)

  def op_impl_base(self, *args):
    k, v, k_cache, v_cache, k_cache_quant_scale, v_cache_quant_scale, slot_mapping = args
    # input_fp[head_num, head_size]
    def quant(input_fp: torch.Tensor):
        input_fp32 = input_fp.to(torch.float32)
        max_value, _ = torch.max(input_fp32.abs(), dim=-1, keepdim=True)
        scale = max_value / 127.0
        scaled_input = input_fp32 / scale
        return scaled_input.to(torch.int8), scale[..., 0]

    tokens_num = k.shape[0] # [token_num, head_num, head_size]
    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, :]
            key_cache_scale_i = k_cache_quant_scale[block_id, :, block_offset]
            quant_key_i, key_scale_i = quant(key_i)
            key_cache_i[...] = quant_key_i
            key_cache_scale_i[...] = key_scale_i
            if v is not None:
              value_i = v[i]
              value_cache_i = v_cache[block_id, :, block_offset, :]
              value_cache_scale_i = v_cache_quant_scale[block_id, :, block_offset]
              quant_value_i, value_scale_i = quant(value_i)
              value_cache_i[...] = quant_value_i
              value_cache_scale_i[...] = value_scale_i
    return (k_cache, v_cache, k_cache_quant_scale, v_cache_quant_scale) if v is not None else (k_cache, k_cache_quant_scale)

  @unittest.skipIf('MLU3' in torch.mlu.get_device_name(), "quant_to_paged_cache not support MLU3XX device")
  def test_quant_to_paged_cache(self):
    token_nums = random.randint(1, 2048)
    head_num_kv = random.randint(1, 128)
    head_size = random.randint(1, 1024)
    block_size = random.randint(1, 50)
    min_blocks = (int)((token_nums + block_size - 1) / block_size)
    block_nums = min(min_blocks + 10, 2 * min_blocks)
    num_slots = block_nums * block_size
    slot_mapping = random.sample(range(num_slots), token_nums)
    slot_mapping = torch.tensor(slot_mapping, dtype=torch.int).mlu()
    slot_mapping[-1] = -1   # test mask
    dtype_list = [torch.half, torch.float]
    if torch_mlu.mlu.is_bf16_supported():
        dtype_list.append(torch.bfloat16)
    only_quant_key_list = [True, False]

    for _ in range(100):
       print("test_quant_to_paged_cache...")
       dtype = random.choice(dtype_list)
       only_quant_key = random.choice(only_quant_key_list)
       key = torch.randn(token_nums, head_num_kv, head_size, dtype=dtype, device="mlu").uniform_(-1, 1)
       key_cache = torch.zeros(block_nums, head_num_kv, block_size, head_size, dtype=torch.int8, device="mlu")
       key_cache_scale = torch.zeros(block_nums, head_num_kv, block_size, dtype=torch.float, device="mlu")
       if only_quant_key:
           value, value_cache, value_cache_scale = None, None, None
       else:
           value = torch.randn(token_nums, head_num_kv, head_size, dtype=dtype, device="mlu").uniform_(-0.25, 0.25)
           value_cache = torch.zeros_like(key_cache)
           value_cache_scale = torch.zeros_like(key_cache_scale)
       self.launch(key, value, key_cache, value_cache, key_cache_scale, value_cache_scale, slot_mapping)

  @unittest.skipIf('MLU3' in torch.mlu.get_device_name(), "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('quant_to_paged_cache: test_large_tensor')
          head_num_kv = 16
          head_size = 128
          token_nums = 20
          block_size = 16
          block_nums = ((2**32 - 1) // 1 // head_num_kv // head_size // block_size)
          num_slots = block_nums * block_size
          key = torch.randn(token_nums, head_num_kv, head_size, dtype=dtype, device="mlu").uniform_(-1, 1)
          value = torch.randn(token_nums, head_num_kv, head_size, dtype=dtype, device="mlu").uniform_(-0.25, 0.25)
          key_cache_torch = torch.zeros(block_nums, head_num_kv, block_size, head_size, dtype=torch.int8, device="mlu")
          value_cache_torch = torch.zeros_like(key_cache_torch)
          key_cache_scale_torch = torch.zeros(block_nums, head_num_kv, block_size, dtype=torch.float, device="mlu")
          value_cache_scale_torch = torch.zeros_like(key_cache_scale_torch)
          slot_mapping = torch.randperm(num_slots, dtype=torch.int32, device="mlu")[:token_nums]
          key_cache_tmo = torch.zeros_like(key_cache_torch)
          value_cache_tmo = torch.zeros_like(value_cache_torch)
          key_cache_scale_tmo = torch.zeros_like(key_cache_scale_torch)
          value_cache_scale_tmo = torch.zeros_like(value_cache_scale_torch)
          self.op_impl_base(key, value, key_cache_torch, value_cache_torch, key_cache_scale_torch, value_cache_scale_torch, slot_mapping)
          ops.quant_to_paged_cache(key, value, key_cache_tmo, value_cache_tmo, key_cache_scale_tmo, value_cache_scale_tmo, slot_mapping)
          self.assertTensorsEqual(key_cache_torch.cpu(), key_cache_tmo.cpu(), 1)
          self.assertTensorsEqual(value_cache_torch.cpu(), value_cache_tmo.cpu(), 1)
          self.assertTensorsEqual(key_cache_scale_torch.cpu().float(), key_cache_scale_tmo.cpu().float(), 3e-3,
                                      use_MSE=True, use_RAE=True)
          self.assertTensorsEqual(value_cache_scale_torch.cpu().float(), value_cache_scale_tmo.cpu().float(), 3e-3,
                                      use_MSE=True, use_RAE=True)

          block_nums = (2**32 // 1 // head_num_kv // 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 = torch.zeros(block_nums, head_num_kv, block_size, head_size, dtype=torch.int8, device="mlu")
          value_cache_torch = torch.zeros_like(key_cache_torch)
          key_cache_scale_torch = torch.zeros(block_nums, head_num_kv, block_size, dtype=torch.float, device="mlu")
          value_cache_scale_torch = torch.zeros_like(key_cache_scale_torch)
          key_cache_tmo = torch.zeros_like(key_cache_torch)
          value_cache_tmo = torch.zeros_like(value_cache_torch)
          key_cache_scale_tmo = torch.zeros_like(key_cache_scale_torch)
          value_cache_scale_tmo = torch.zeros_like(value_cache_scale_torch)
          self.assertException("The addressing range of kv_cache cannot exceed 4G.", ops.quant_to_paged_cache,
                                key, value, key_cache_tmo, value_cache_tmo, key_cache_scale_tmo, value_cache_scale_tmo, slot_mapping)

  @unittest.skipIf('MLU3' in torch.mlu.get_device_name(), "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):
    token_nums = random.randint(1, 2048)
    head_num_kv = random.randint(1, 128)
    head_size = random.randint(1, 1024)
    block_size = random.randint(1, 50)
    min_blocks = (int)((token_nums + block_size - 1) / block_size)
    block_nums = min(min_blocks + 10, 2 * min_blocks)
    num_slots = block_nums * block_size
    slot_mapping = random.sample(range(num_slots), token_nums)
    slot_mapping = torch.tensor(slot_mapping, dtype=torch.int).mlu()
    slot_mapping[-1] = -1   # test mask
    dtype_list = [torch.half, torch.float]
    if torch_mlu.mlu.is_bf16_supported():
        dtype_list.append(torch.bfloat16)
    dtype = random.choice(dtype_list)

    print("quant_to_paged_cache: test_prevent...")
    key = torch.randn(token_nums, head_num_kv, head_size, dtype=dtype, device="mlu").uniform_(-1, 1)
    value = None
    key_cache = torch.zeros(block_nums, head_num_kv, block_size, head_size, dtype=torch.int8, device="mlu")
    value_cache = torch.zeros_like(key_cache)
    key_cache_scale = torch.zeros(block_nums, head_num_kv, block_size, dtype=torch.float, device="mlu")
    value_cache_scale = torch.zeros_like(key_cache_scale)
    self.assertException("v.has_value() == v_cache.has_value() && v.has_value() == v_cache_scale.has_value().",
                         ops.quant_to_paged_cache, key, value, key_cache, value_cache, key_cache_scale,
                         value_cache_scale, slot_mapping)
    value = torch.randn(token_nums, head_num_kv, head_size, dtype=dtype, device="mlu").uniform_(-0.25, 0.25)
    value = value.as_strided(size=(token_nums, head_num_kv, head_size), stride=(1, token_nums, token_nums * head_num_kv))
    self.assertException("v last dim must be contiguous.",
                         ops.quant_to_paged_cache, key, value, key_cache, value_cache, key_cache_scale,
                         value_cache_scale, slot_mapping)
    value = value.as_strided(size=(token_nums, head_num_kv, head_size), stride=(head_size, head_num_kv, 1))
    self.assertException("v second dim must be contiguous.",
                         ops.quant_to_paged_cache, key, value, key_cache, value_cache, key_cache_scale,
                         value_cache_scale, slot_mapping)


  @unittest.skipIf('MLU3' in torch.mlu.get_device_name(), "quant_to_paged_cache not support MLU3XX device")
  def test_inductor(self):
    token_nums = 20
    head_num_kv = 30
    head_size = 20
    block_size = 10
    min_blocks = (int)((token_nums + block_size - 1) / block_size)
    block_nums = min(min_blocks + 10, 2 * min_blocks)
    num_slots = block_nums * block_size
    slot_mapping = random.sample(range(num_slots), token_nums)
    slot_mapping = torch.tensor(slot_mapping, dtype=torch.int).mlu()
    slot_mapping[-1] = -1   # test mask

    dtype = torch.half
    key = torch.randn(token_nums, head_num_kv, head_size, dtype=dtype, device="mlu").uniform_(-1, 1)
    value = torch.randn(token_nums, head_num_kv, head_size, dtype=dtype, device="mlu").uniform_(-0.25, 0.25)
    key_cache = torch.zeros(block_nums, head_num_kv, block_size, head_size, dtype=torch.int8, device="mlu")
    value_cache = torch.zeros_like(key_cache)
    key_cache_scale = torch.zeros(block_nums, head_num_kv, block_size, dtype=torch.float, device="mlu")
    value_cache_scale = torch.zeros_like(key_cache_scale)
    args = (key, value, key_cache, value_cache, key_cache_scale, value_cache_scale, slot_mapping)
    self.base_opcheck(torch.ops.torch_mlu_ops.quant_to_paged_cache, args)

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