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

def gen_args(batch_size, num_heads, head_size, cache_memory_len, packed, dtype, quant_bit):
    args = generate_cache_func_args(batch_size, num_heads, head_size, cache_memory_len, packed, dtype, True)
    cache_scale = args[10]
    args = args[0:10]
    args.insert(4, cache_scale[0])
    args.insert(5, cache_scale[1])
    args.insert(12, quant_bit)
    return args

class TestQuantToLinearCache(BtTestCase):
    def op_impl_base(self, *args):
        key, value, key_cache, value_cache, key_cache_quant_scale, value_cache_quant_scale, \
            context_lengths, max_context_len, packed, context_seq_offset, cache_bs_id, cache_seqlen_offset, \
                quant_bit = args
        def quant(context: torch.Tensor,
                  quant_bit: int,
                  group_size: int):
            # context:[head_num, seq, head_size]
            head_num = context.shape[0]
            head_size = context.shape[-1]
            if group_size != head_size:
                context = context.reshape(head_num, -1, group_size)
            context_fp32 = context.to(torch.float32)
            max_value, _ = torch.max(context_fp32.abs(), dim=-1, keepdim=True)
            int_max = float(2 ** (quant_bit - 1) - 1)
            scale = max_value / int_max
            scaled_context = context_fp32 / scale
            return scaled_context.reshape(head_num, -1, head_size), scale[..., 0]

        batch_size = context_lengths.shape[0] - 1 if packed else context_lengths.shape[0]
        head_num = key.shape[-2]
        head_size = key.shape[-1]

        if key_cache_quant_scale.dim() == 3:
            group_size = head_size
        else:
            group_size = head_size // key_cache_quant_scale.size(-1)
        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]
            cache_bs_id_i = cache_bs_id[i]
            cache_seqlen_offset_i = cache_seqlen_offset[i] if cache_seqlen_offset is not None else 0
            if cache_bs_id_i < 0 or cache_seqlen_offset_i < 0:
                continue
            key_cache_i = \
                key_cache[cache_bs_id_i, :, \
                          cache_seqlen_offset_i:cache_seqlen_offset_i + context_len_i]
            key_cache_scale_i = \
                key_cache_quant_scale[cache_bs_id_i, :, \
                          cache_seqlen_offset_i:cache_seqlen_offset_i + context_len_i]

            # key_i[head_num, context_len[i], head_size]
            key_i = key_i[:, context_seq_offset_i:context_seq_offset_i + context_len_i]
            float_key_i, key_scale_i = quant(key_i, quant_bit, group_size)
            key_cache_scale_i = key_cache_scale_i.reshape(key_cache_scale_i.shape[0], -1)
            key_cache_scale_i[...] = key_scale_i
            rounded = torch.round(float_key_i)
            clipped = torch.clip(rounded, -2 ** (quant_bit - 1), 2 ** (quant_bit - 1) - 1)
            quant_key_i = clipped.to(torch.int8)
            if quant_bit == 4:
                quant_key_flat = quant_key_i.flatten()
                d0 = quant_key_flat[0::2].to(torch.uint8)
                d1 = quant_key_flat[1::2].to(torch.uint8)
                dp = (d1 << 4) + (d0 & 0x0F)
                quant_key_i = dp.to(torch.int8).reshape(head_num, -1, head_size // 2)
            key_cache_i[...] = quant_key_i

            if value_cache is not None and value is not None:
                value_cache_scale_i = \
                  value_cache_quant_scale[cache_bs_id_i, :, \
                            cache_seqlen_offset_i:cache_seqlen_offset_i + context_len_i]
                value_i = value_i[:, context_seq_offset_i:context_seq_offset_i + context_len_i]
                float_value_i, value_scale_i = quant(value_i, quant_bit, group_size)
                value_cache_scale_i = value_cache_scale_i.reshape(value_cache_scale_i.shape[0], -1)
                value_cache_scale_i[...] = value_scale_i
                rounded = torch.round(float_value_i)
                clipped = torch.clip(rounded, -2 ** (quant_bit - 1), 2 ** (quant_bit - 1) - 1)
                quant_value_i = clipped.to(torch.int8)
                if quant_bit == 8:
                    value_cache_i = \
                      value_cache[cache_bs_id_i, :, \
                                cache_seqlen_offset_i:cache_seqlen_offset_i + context_len_i]
                else:
                    value_cache_i = \
                      value_cache[cache_bs_id_i, :, \
                                cache_seqlen_offset_i // 2:math.ceil((cache_seqlen_offset_i + context_len_i) / 2)]
                    if cache_seqlen_offset_i % 2 == 1:
                        front_vec = value_cache[cache_bs_id_i, :, cache_seqlen_offset_i // 2, :]
                        front_low_bits = front_vec & (0x0F)
                        front_low_bits_expand = front_low_bits.unsqueeze(1) # [head_num, 1, head_size]
                        quant_value_i = torch.cat((front_low_bits_expand, quant_value_i), dim=1)
                    if (cache_seqlen_offset_i + context_len_i) % 2 == 1:
                        back_vec = value_cache[cache_bs_id_i, :, math.ceil((cache_seqlen_offset_i + context_len_i) / 2) - 1, :]
                        back_high_bits = (back_vec >> 4) & (0x0F)
                        back_high_bits_expand = back_high_bits.unsqueeze(1) # [head_num, 1, head_size]
                        quant_value_i = torch.cat((quant_value_i, back_high_bits_expand), dim=1)
                    value_temp = quant_value_i.reshape(head_num, -1, 2, head_size)
                    quant_value_flat = value_temp.permute(0, 1, 3, 2).flatten()
                    v0 = quant_value_flat[0::2].to(torch.uint8)
                    v1 = quant_value_flat[1::2].to(torch.uint8)
                    vp = (v1 << 4) + (v0 & 0x0F)
                    quant_value_i = vp.to(torch.int8).reshape(head_num, -1, head_size)
                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 int8_to_int4(self,
                     input):
        input_flat = input.flatten()
        size = input_flat.size(0)
        output = torch.zeros(size * 2, dtype=torch.int8, device=input.device)
        high = input_flat >> 4
        low = input_flat << 4
        low = low >> 4
        output[0::2] = low
        output[1::2] = high
        return output

    def test_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=32, 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=8, high=512, size=(test_cases, ), dtype=torch.int32)
        cache_memory_len_list = cache_memory_len_list * 2
        packed_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.float]
        quant_bit_list = torch.randint(low=0, high=2, size=(test_cases, ), dtype=torch.int32)
        quant_bit_map = [4, 8]

        for i in range(test_cases):
            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()
            group_size_factors = [i for i in range(4, head_size + 1) if head_size % i == 0] # group_size should > 1
            group_size = random.choice(group_size_factors)
            cache_memory_len = cache_memory_len_list[i].item()
            packed = packed_list[i].item()
            dtype = dtype_map[dtype_list[i]]
            quant_bit = quant_bit_map[quant_bit_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("case num={}, batch_size={}, num_heads={}, head_size={}, cache_memory_len={}, packed={}, dtype={}, quant_bit={}, group_size={} testing...".format(
                i, batch_size, num_heads, head_size, cache_memory_len, packed > 0, dtype, quant_bit, group_size))

            max_bs = batch_size + 1
            context_lens = torch.randint(size=(batch_size, ), low=1,
                                         high=cache_memory_len // 4,
                                         dtype=torch.int32, device='mlu')
            context_lens = context_lens * 2
            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 # set invalid

            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:
                key = torch.randn((total_seqlen, num_heads, head_size),
                                      dtype=torch.float, device='mlu')
                value = torch.randn((total_seqlen, num_heads, head_size),
                                      dtype=torch.float, device='mlu')
            else:
                key = torch.randn((batch_size, max_context_len + max_seq_offset, num_heads, head_size),
                                      dtype=torch.float, device='mlu')
                value = torch.randn((batch_size, max_context_len + max_seq_offset, num_heads, head_size),
                                      dtype=torch.float, device='mlu')
            key = key.to(dtype)
            value = value.to(dtype)
            if quant_bit == 8 and group_size == head_size:
                key_cache = torch.randint(-128, 127, (max_bs, num_heads, cache_memory_len, head_size), dtype=torch.int8).mlu()
                value_cache = torch.randint(-128, 127, (max_bs, num_heads, cache_memory_len, head_size), dtype=torch.int8).mlu()
                key_cache_scale = torch.randn((max_bs, num_heads, cache_memory_len), dtype=torch.float32).mlu()
                value_cache_scale = torch.randn((max_bs, num_heads, cache_memory_len), dtype=torch.float32).mlu()
            if quant_bit == 8 and group_size != head_size:
                key_cache = torch.randint(-128, 127, (max_bs, num_heads, cache_memory_len, head_size), dtype=torch.int8).mlu()
                value_cache = torch.randint(-128, 127, (max_bs, num_heads, cache_memory_len, head_size), dtype=torch.int8).mlu()
                key_cache_scale = torch.randn((max_bs, num_heads, cache_memory_len, head_size // group_size), dtype=torch.float32).mlu()
                value_cache_scale = torch.randn((max_bs, num_heads, cache_memory_len, head_size // group_size), dtype=torch.float32).mlu()
            if quant_bit == 4 and group_size == head_size:
                key_cache = torch.randint(-128, 127, (max_bs, num_heads, cache_memory_len, head_size // 2), dtype=torch.int8).mlu()
                value_cache = torch.randint(-128, 127, (max_bs, num_heads, cache_memory_len // 2, head_size), dtype=torch.int8).mlu()
                key_cache_scale = torch.randn((max_bs, num_heads, cache_memory_len), dtype=torch.float32).mlu()
                value_cache_scale = torch.randn((max_bs, num_heads, cache_memory_len), dtype=torch.float32).mlu()
            if quant_bit == 4 and group_size != head_size:
                key_cache = torch.randint(-128, 127, (max_bs, num_heads, cache_memory_len, head_size // 2), dtype=torch.int8).mlu()
                value_cache = torch.randint(-128, 127, (max_bs, num_heads, cache_memory_len // 2, head_size), dtype=torch.int8).mlu()
                key_cache_scale = torch.randn((max_bs, num_heads, cache_memory_len, head_size // group_size), dtype=torch.float32).mlu()
                value_cache_scale = torch.randn((max_bs, num_heads, cache_memory_len, head_size // group_size), dtype=torch.float32).mlu()
            ref_key_cache = key_cache.clone()
            ref_value_cache = value_cache.clone()
            ref_key_cache_scale = key_cache_scale.clone()
            ref_value_cache_scale = value_cache_scale.clone()

            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:
                self.op_impl_base(key, value, ref_key_cache, ref_value_cache, ref_key_cache_scale,
                                  ref_value_cache_scale, cu_context_lens, max_context_len,
                                  packed > 0, None, cache_bs_id, cache_seq_offsets,
                                  quant_bit)
                ops.quant_to_linear_cache(key, value, key_cache, value_cache, key_cache_scale,
                                          value_cache_scale, cu_context_lens, max_context_len,
                                          packed > 0, None, cache_bs_id, cache_seq_offsets,
                                          quant_bit)
            else:
                self.op_impl_base(key, value, ref_key_cache, ref_value_cache, ref_key_cache_scale,
                                  ref_value_cache_scale, context_lens, max_context_len,
                                  packed > 0, context_seq_offsets, cache_bs_id,
                                  cache_seq_offsets, quant_bit)
                ops.quant_to_linear_cache(key, value, key_cache, value_cache, key_cache_scale,
                                          value_cache_scale, context_lens, max_context_len,
                                          packed > 0, context_seq_offsets, cache_bs_id,
                                          cache_seq_offsets, quant_bit)
            if quant_bit == 8:
                self.assertTensorsEqual(key_cache.cpu().float(), ref_key_cache.cpu().float(), 0.003,
                                        "key_cache must equal ref_key_cache", True, True, True, True)
                self.assertTensorsEqual(value_cache.cpu().float(), ref_value_cache.cpu().float(), 0.003,
                                        "value_cache must equal ref_value_cache", True, True, True, True)
            else:
                key_cache_int8 = self.int8_to_int4(key_cache)
                ref_key_cache_int8 = self.int8_to_int4(ref_key_cache)
                value_cache_int8 = self.int8_to_int4(value_cache)
                ref_value_cache_int8 = self.int8_to_int4(ref_value_cache)

                key_cache_diff = (key_cache_int8.cpu() - ref_key_cache_int8.cpu()).abs()
                assert torch.max(key_cache_diff) < 2, "ref_key_cache must equal key_cache or absolute values differ by 1 due to round_mode!"
                value_cache_diff = (value_cache_int8.cpu() - ref_value_cache_int8.cpu()).abs()
                assert torch.max(value_cache_diff) < 2, "ref_value_cache must equal value_cache or absolute values differ by 1 due to round_mode!"

            self.assertTensorsEqual(key_cache_scale.cpu().float(), ref_key_cache_scale.cpu().float(), 0.003,
                                    "key_cache_scale must equal ref_key_cache_scale", True, True, True, True)
            self.assertTensorsEqual(value_cache_scale.cpu().float(), ref_value_cache_scale.cpu().float(), 0.003,
                                    "value_cache_scale must equal ref_value_cache_scale", True, True, True, True)

    def test_inductor(self):
        batch_size, num_heads, head_size, cache_memory_len, dtype, quant_bit = 4, 8, 64, 128, torch.float16, 8
        args = gen_args(batch_size, num_heads, head_size, cache_memory_len, 1, dtype, quant_bit)
        self.base_opcheck(torch.ops.torch_mlu_ops.quant_to_linear_cache, args)

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

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