import torch
import unittest
import torch_mlu_ops as ops
import random
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, True, True)
    cache_scale = args[10]
    args = args[0:10]
    args.insert(4, cache_scale[0])
    args.insert(5, cache_scale[1])
    args.insert(8, False)
    return args

class TestOfflineQuantToLinearCache(BtTestCase):
    def op_impl_base(self, *args):
        def quant2Int8(input_fp: torch.Tensor,
                       scale_fp: torch.tensor,
                       quant_mode: torch.int64):
            head_num = input_fp.size(0)
            seq = input_fp.size(1)
            head_size = input_fp.size(2)
            input_fp32 = input_fp.to(torch.float32)
            if quant_mode == 0:  # per_channel
                scale = scale_fp.reshape((head_num, 1, head_size))
            else:
                scale = scale_fp.reshape((head_num, seq, 1))
            scaled_context = input_fp32 / scale
            rounded = torch.round(scaled_context)
            clipped = torch.clip(rounded, -128, 127)
            return clipped.to(torch.int8)

        key, value, key_cache, value_cache, key_cache_quant_scale, value_cache_quant_scale, \
            context_lengths, max_context_len, quant_mode, packed, context_seq_offset, cache_bs_id, \
                cache_seqlen_offset = args
        batch_size = context_lengths.size(0) - 1 if packed else context_lengths.size(0)
        for i in range(batch_size):
            if packed:
                key_i = key[context_lengths[i]:context_lengths[i+1]].transpose(1, 0)
                value_i = value[context_lengths[i]:context_lengths[i+1]].transpose(1, 0) if value is not None else None
                context_len_i = context_lengths[i+1] - context_lengths[i]
                context_seq_offset_i = 0
            else:
                key_i = key[i].transpose(1, 0)
                value_i = value[i].transpose(1, 0) if value is not None else None
                context_len_i = context_lengths[i]
                context_seq_offset_i = context_seq_offset[i] if context_seq_offset is not None else 0
            cache_bs_id_i = cache_bs_id[i] if cache_bs_id is not None else i
            cache_seq_begin = cache_seqlen_offset[i] if cache_seqlen_offset is not None else 0
            if cache_bs_id_i < 0 or cache_seq_begin < 0:
                continue
            cache_seq_end = cache_seq_begin + context_len_i

            # quant key to int8
            key_i = key_i[:, context_seq_offset_i:context_seq_offset_i + context_len_i]
            if quant_mode == 0:
                key_cache_scale_i = key_cache_quant_scale
            else:
                key_cache_scale_i = key_cache_quant_scale[:, cache_seq_begin:cache_seq_end]
            quant_key_i = quant2Int8(key_i, key_cache_scale_i, quant_mode)
            key_cache_i = key_cache[cache_bs_id_i, :, cache_seq_begin:cache_seq_end]
            key_cache_i[...] = quant_key_i

            # quant value to int8
            if value_cache is not None and value is not None and value_cache_quant_scale is not None:
                value_i = value_i[:, context_seq_offset_i:context_seq_offset_i + context_len_i]
                if quant_mode == 0:
                    value_cache_scale_i = value_cache_quant_scale
                else:
                    value_cache_scale_i = value_cache_quant_scale[:, cache_seq_begin:cache_seq_end]
                quant_value_i = quant2Int8(value_i, value_cache_scale_i, quant_mode)
                value_cache_i = value_cache[cache_bs_id_i, :, cache_seq_begin:cache_seq_end]
                value_cache_i[...] = quant_value_i
        return (key_cache, value_cache, key_cache_quant_scale, value_cache_quant_scale) if value_cache is not None else (key_cache, key_cache_quant_scale)

    def test_offline_quant_to_linear_cache(self):
        test_cases = 100
        bs_list = torch.randint(low=1, high=64, 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
        cache_memory_len_list = torch.randint(low=2, high=1024, size=(test_cases, ), dtype=torch.int32)
        packed_list = torch.randint(low=0, high=2, size=(test_cases, ), dtype=torch.int32)
        mode_list= torch.randint(low=0, high=2, size=(test_cases, ), dtype=torch.int32)
        dtype_list = torch.randint(low=0, high=3, size=(test_cases, ), dtype=torch.int32)
        dtype_map = [torch.half, torch.bfloat16, torch.float32]

        for i in range(test_cases):
            q_heads = 1
            batch_size       = bs_list[i].item()
            invalid_batch = batch_size // 10
            num_heads        = num_heads_list[i].item()
            head_size        = head_size_list[i].item()
            cache_memory_len = cache_memory_len_list[i].item()
            packed           = packed_list[i].item()
            quant_mode       = mode_list[i].item()
            total_heads      = q_heads + num_heads * 2
            dtype            = dtype_map[dtype_list[i]]
            mlu_name = torch.mlu.get_device_name()
            if "MLU3" in mlu_name and dtype == torch.bfloat16:
                dtype = torch.half
                print("BFLOAT16 is not support on {}, use half instead".format(mlu_name))

            print("batch_size={}, num_heads={}, head_size={}, cache_memory_len={}, packed={}, mode ={}, dtype={} testing...".format(
                   batch_size, num_heads, head_size, cache_memory_len, packed >
                   0, quant_mode, dtype))

            torch.manual_seed(1)
            max_bs = batch_size + 1
            context_lens = torch.randint(size=(batch_size, ), low=1,
                                         high=cache_memory_len // 2,
                                         dtype=torch.int32, device='mlu')
            max_context_len = context_lens.max().item()
            max_seq_offset = max_context_len // 3 + 1
            context_seq_offsets = torch.randint(size=(batch_size, ),
                                                low=0, high=max_seq_offset,
                                                dtype=torch.int32, device='mlu')
            cache_seq_offsets = torch.randint(size=(batch_size, ), low=0,
                                              high=(cache_memory_len - max_context_len) // 3 + 1,
                                              dtype=torch.int32, device='mlu')
            cache_seq_offsets[random.sample([*range(0, batch_size)], invalid_batch)] = -1
            cu_context_lens = torch.cumsum(context_lens, dim=-1)
            cu_context_lens = torch.nn.functional.pad(cu_context_lens, (1,0), "constant", 0).to(torch.int32)
            total_seqlen = cu_context_lens[-1]
            if packed > 0:
                context = torch.randn((total_seqlen, total_heads, head_size),
                                      dtype=torch.float, device='mlu')
            else:
                context = torch.randn((batch_size, max_context_len + max_seq_offset, total_heads, head_size),
                                      dtype=torch.float, device='mlu')
            context = context.to(dtype)
            key = context[..., q_heads:q_heads + num_heads, :]
            value = context[..., q_heads + num_heads:q_heads + 2 * num_heads, :]

            # prepare key_scale and value_scale
            if quant_mode == 0 :   # per_channel
                cache_scale = torch.randn((2, num_heads, head_size), dtype=torch.float, device='mlu')
            else:
                cache_scale = torch.randn((2, num_heads, cache_memory_len), dtype=torch.float, device='mlu')
            key_cache_scale = cache_scale[0]
            value_cache_scale = cache_scale[1]

            cache = torch.randn((2, max_bs, num_heads, cache_memory_len, head_size), dtype=torch.float, device='mlu')
            cache = (cache - 0.5) * 256
            cache = cache.to(torch.int8)


            key_cache = cache[0]
            value_cache = cache[1]

            ref_cache = cache.clone()
            ref_key_cache = ref_cache[0]
            ref_value_cache = ref_cache[1]

            cache_bs_id = random.sample([*range(0, max_bs)], batch_size)
            cache_bs_id = torch.IntTensor(cache_bs_id).mlu()
            cache_bs_id[random.sample([*range(0, batch_size)], invalid_batch)] = -1 # set invalid batch

            if packed > 0:
                ops.offline_quant_to_linear_cache(key, value, key_cache, value_cache, key_cache_scale,
                                                  value_cache_scale, cu_context_lens, max_context_len,
                                                  quant_mode,  packed > 0, None, cache_bs_id,
                                                  cache_seq_offsets)
                self.op_impl_base(key, value, ref_key_cache, ref_value_cache, key_cache_scale,
                                  value_cache_scale, cu_context_lens, max_context_len,
                                  quant_mode,  packed > 0, None, cache_bs_id,
                                  cache_seq_offsets)
            else:
                ops.offline_quant_to_linear_cache(key, value, key_cache, value_cache, key_cache_scale,
                                                  value_cache_scale, context_lens, max_context_len,
                                                  quant_mode, packed > 0, context_seq_offsets, cache_bs_id,
                                                  cache_seq_offsets)
                self.op_impl_base(key, value, ref_key_cache, ref_value_cache, key_cache_scale,
                                  value_cache_scale, context_lens, max_context_len,
                                  quant_mode, packed > 0, context_seq_offsets, cache_bs_id,
                                  cache_seq_offsets)
            # for debug
            cache = cache.cpu().flatten()
            ref_cache = ref_cache.cpu().flatten()
            diff = cache - ref_cache
            diff = diff.abs()
            assert torch.max(diff) < 2, "ref_cache must equal cache or absolute values differ by 1 due to round_mode!"

    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_linear_cache, args)

        args = gen_args(batch_size, num_heads, head_size, cache_memory_len, 0, dtype)
        self.base_opcheck(torch.ops.torch_mlu_ops.offline_quant_to_linear_cache, args)

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