import math
import random
import torch
import torch_mlu
import unittest
import torch_mlu_ops as ops
from common_utils import *
from itertools import product

def gen_cache(batch, seq_q, num_kv_heads, head_size, is_pagedattn, seq_len, data_type, quant_bit, quant_mode, is_normal=True):
    int_max = float(2 ** (quant_bit - 1) - 1)
    int_min = -float(2 ** (quant_bit - 1))
    context_lens = torch.randint(seq_q, seq_len + 1, (batch, ), dtype=torch.int32).mlu()
    if is_normal is False and batch > 3: # replace some batch's context to 0
        num = batch // 3
        index = torch.randint(0, batch, (num,))
        context_lens[index] = 0
    max_context_len = context_lens.max().item()
    block_size = 16
    if is_pagedattn is False:
        block_size = seq_len
    num_blocks = batch * ((seq_len + block_size - 1) // block_size)
    max_num_blocks_per_seq = (max_context_len + block_size - 1) // block_size
    cache_shape = (num_blocks, num_kv_heads, block_size, head_size)
    if quant_mode == "per_token":
        scale_shape = (num_blocks, num_kv_heads, block_size, 1)
    else: # per channel
        scale_shape = (num_kv_heads, head_size)

    if quant_bit == 4:
        cache_shape_k_int4 = (num_blocks, num_kv_heads, block_size, head_size//2)
        cache_shape_v_int4 = (num_blocks, num_kv_heads, block_size//2, head_size)
        cache_shape_v_int4_tmp = (num_blocks, num_kv_heads, head_size, block_size//2)
        key_cache_int8 = torch.zeros(cache_shape).uniform_(int_min, int_max).to(torch.int8).mlu()
        value_cache_int8 = torch.zeros(cache_shape).uniform_(int_min, int_max).to(torch.int8).mlu()
        # pre_process int4_kv_cache
        key_cache_view = key_cache_int8.reshape(-1, head_size)
        value_cache_view = value_cache_int8.transpose(2, 3).reshape(-1, block_size)
        key_cache = PairlyPackInt8(key_cache_view).view(cache_shape_k_int4)
        value_cache = PairlyPackInt8(value_cache_view).view(cache_shape_v_int4_tmp).transpose(2,3).contiguous()
        key_scale = torch.randn(size=scale_shape, dtype=torch.float32).mlu()
        value_scale = torch.randn(size=scale_shape, dtype=torch.float32).mlu()
    elif quant_bit == 8:
        key_cache = torch.zeros(cache_shape).uniform_(int_min, int_max).to(torch.int8).mlu()
        value_cache = torch.zeros(cache_shape).uniform_(int_min, int_max).to(torch.int8).mlu()
        key_scale = torch.randn(size=scale_shape, dtype=torch.float32).mlu()
        value_scale = torch.randn(size=scale_shape, dtype=torch.float32).mlu()
    elif quant_bit == -1:
        key_cache = torch.randn(cache_shape, dtype=data_type).mlu()
        value_cache = torch.randn(cache_shape, dtype=data_type).mlu()
        key_scale = None
        value_scale = None
    else:
        print("gen case error, quant_bit_lp must be in {-1, 4, 8}")
    block_tables = random.sample(range(0, num_blocks), batch * max_num_blocks_per_seq)
    block_tables = torch.tensor(block_tables, dtype=torch.int32).mlu().view(batch, max_num_blocks_per_seq)
    output = [key_cache, value_cache, key_scale, value_scale, context_lens, block_tables]
    if quant_bit == 4:
        output.append(key_cache_int8)
        output.append(value_cache_int8)
    return output

def concate_cache_linear(cachek1, cachek2, cachev1, cachev2, context1, context2, block_tables1, block_tables2,
                          scalek1, scalek2, scalev1, scalev2):
    if scalek1 is not None:
        if scalek1.dim() == 2: # per_channel: [kv_head_num, head_size]
            scalek1 = scalek1.reshape(1, scalek1.shape[0], 1, scalek1.shape[1])
            scalev1 = scalev1.reshape(1, scalev1.shape[0], 1, scalev1.shape[1])
        elif scalek1.dim() == 3: # per_token: [num_blocks, k_head_num, block_size]
            scalek1 = scalek1.reshape(*scalek1.shape, 1)
            scalev1 = scalev1.reshape(*scalev1.shape, 1)
        cachek1 *= scalek1
        cachev1 *= scalev1
    if scalek2 is not None:
        if scalek2.dim() == 2: # per_channel: [kv_head_num, head_size]
            scalek2 = scalek2.reshape(1, scalek2.shape[0], 1, scalek2.shape[1])
            scalev2 = scalev2.reshape(1, scalev2.shape[0], 1, scalev2.shape[1])
        elif scalek2.dim() == 3: # per_token: [num_blocks, k_head_num, block_size]
            scalek2 = scalek2.reshape(*scalek2.shape, 1)
            scalev2 = scalev2.reshape(*scalev2.shape, 1)
        cachek2 *= scalek2
        cachev2 *= scalev2
    new_context = context1 + context2
    seq_len1 = cachek1.shape[2]
    seq_len2 = cachek2.shape[2]
    new_max_context = seq_len1 + seq_len2
    batch = cachek1.shape[0]
    num_head = cachek1.shape[1]
    head_size = cachek1.shape[3]
    new_cache_k = torch.randn(batch, num_head, new_max_context, head_size, dtype=torch.float32)
    new_cache_v = torch.randn(batch, num_head, new_max_context, head_size, dtype=torch.float32)
    new_block_table = torch.arange(0, batch)
    new_block_table = new_block_table.view(batch, 1)
    for i in range(batch):
        len1 = context1[i]
        len2 = context2[i]
        block_id1 = block_tables1[i]
        block_id2 = block_tables2[i]
        new_cache_k[i, :, :len1, :] = cachek1[block_id1, :, :len1, :]
        new_cache_v[i, :, :len1, :] = cachev1[block_id1, :, :len1, :]
        new_cache_k[i, :, len1:len1 + len2, :] = cachek2[block_id2, :, :len2, :]
        new_cache_v[i, :, len1:len1 + len2, :] = cachev2[block_id2, :, :len2, :]
    return new_cache_k.mlu(), new_cache_v.mlu(), new_block_table.mlu(), new_context.mlu()

# cache1 and cache2 are float
def concat_cache_paged(cachek1, cachek2, cachev1, cachev2, context1, context2, block_tables1, block_tables2,
                        scalek1, scalek2, scalev1, scalev2):
    batch = context1.shape[0]
    block_size = cachek1.shape[2]
    if scalek1 is not None:
        if scalek1.dim() == 2: # per_channel: [kv_head_num, head_size]
            scalek1 = scalek1.reshape(1, scalek1.shape[0], 1, scalek1.shape[1])
            scalev1 = scalev1.reshape(1, scalev1.shape[0], 1, scalev1.shape[1])
        elif scalek1.dim() == 3: # per_token: [num_blocks, k_head_num, block_size]
            scalek1 = scalek1.reshape(*scalek1.shape, 1)
            scalev1 = scalev1.reshape(*scalev1.shape, 1)
        cachek1 *= scalek1
        cachev1 *= scalev1
    if scalek2 is not None:
        if scalek2.dim() == 2: # per_channel: [kv_head_num, head_size]
            scalek2 = scalek2.reshape(1, scalek2.shape[0], 1, scalek2.shape[1])
            scalev2 = scalev2.reshape(1, scalev2.shape[0], 1, scalev2.shape[1])
        elif scalek2.dim() == 3: # per_token: [num_blocks, k_head_num, block_size]
            scalek2 = scalek2.reshape(*scalek2.shape, 1)
            scalev2 = scalev2.reshape(*scalev2.shape, 1)
        cachek2 *= scalek2
        cachev2 *= scalev2
    new_context = context1 + context2
    max_num_blocks_per_seq = block_tables1.shape[1] + block_tables1.shape[1]
    new_cache_k = torch.concat((cachek1, cachek2), dim = 0)
    new_cache_v = torch.concat((cachev1, cachev2), dim = 0)
    num_block1 = cachek1.shape[0]
    new_block_table = torch.zeros(batch, max_num_blocks_per_seq, dtype=torch.int32)
    for i in range(batch):
        len1 = context1[i]
        len2 = context2[i]
        block_num1 = (len1 + block_size - 1) // block_size
        block_num2 = (len2 + block_size - 1) // block_size
        len1_pad = block_num1 * block_size
        block1 = block_tables1[i]
        block2 = block_tables2[i]
        new_block_table[i, :block_num1] = block1[:block_num1]
        new_block_table[i, block_num1:block_num1 + block_num2] = block2[:block_num2] + num_block1
        if len1 != len1_pad:
            reg_block_id = new_block_table[i, block_num1 - 1] # last block of cache1
            cat_block_id = new_block_table[i, block_num1] # frist block of cache2
            reg_len = len1 % block_size
            pad_len = len1_pad - len1
            new_cache_k[reg_block_id, :, reg_len:, :] = new_cache_k[cat_block_id, :, :pad_len, :]
            new_cache_v[reg_block_id, :, reg_len:, :] = new_cache_v[cat_block_id, :, :pad_len, :]
            for j in range(block_num2-1):
                block_id1 = new_block_table[i, block_num1 + j] # current
                block_id2 = new_block_table[i, block_num1 + j + 1]  # next
                new_cache_k[block_id1, :, :reg_len, :] = new_cache_k[block_id1, :, pad_len:, :]
                new_cache_k[block_id1, :, reg_len:, :] = new_cache_k[block_id2, :, :pad_len, :]
                new_cache_v[block_id1, :, :reg_len, :] = new_cache_v[block_id1, :, pad_len:, :]
                new_cache_v[block_id1, :, reg_len:, :] = new_cache_v[block_id2, :, :pad_len, :]
            block_id = new_block_table[i, block_num1 + block_num2 - 1]
            new_cache_k[block_id, :, :reg_len, :] = new_cache_k[block_id, :, pad_len:, :]
            new_cache_v[block_id, :, :reg_len, :] = new_cache_v[block_id, :, pad_len:, :]
    return new_cache_k.mlu(), new_cache_v.mlu(), new_block_table.mlu(), new_context.mlu()

class TestSingleQueryMixedKVAttnOp(BtTestCase):
    def run_gen_case(self, dic):
        dump_data = dic.pop('dump_data')
        if dump_data:
            self.launch(*dic.values())
        else:
            q = create_tensor_from_dic(dic['q'])
            k_cache_lp = create_tensor_from_dic(dic['k_cache_lp'])
            v_cache_lp = create_tensor_from_dic(dic['v_cache_lp'])
            k_cache_hp = create_tensor_from_dic(dic['k_cache_hp'])
            v_cache_hp = create_tensor_from_dic(dic['v_cache_hp'])
            out = create_tensor_from_dic(dic['out'])
            block_tables_lp = dic['block_tables_lp']['data']
            block_tables_hp = dic['block_tables_hp']['data']
            context_lens_lp = dic['context_lens_lp']['data']
            context_lens_hp = dic['context_lens_hp']['data']
            k_cache_quant_scale_lp = create_tensor_from_dic(dic['k_cache_quant_scale_lp'])
            v_cache_quant_scale_lp = create_tensor_from_dic(dic['v_cache_quant_scale_lp'])
            k_cache_quant_scale_hp = create_tensor_from_dic(dic['k_cache_quant_scale_hp'])
            v_cache_quant_scale_hp = create_tensor_from_dic(dic['v_cache_quant_scale_hp'])
            alibi_slopes = create_tensor_from_dic(dic['alibi_slopes'])
            max_contxt_len_lp = dic['max_contxt_len_lp']['data']
            max_contxt_len_hp = dic['max_contxt_len_hp']['data']
            softmax_scale = dic['softmax_scale']['data']
            return_lse = dic['return_lse']['data']
            kv_cache_quant_bit_size_lp = dic['kv_cache_quant_bit_size_lp']['data']
            kv_cache_quant_bit_size_hp = dic['kv_cache_quant_bit_size_hp']['data']
            self.launch(q, k_cache_lp, v_cache_lp, k_cache_hp, v_cache_hp, out, block_tables_lp,
                        block_tables_hp, context_lens_lp, context_lens_hp, k_cache_quant_scale_lp,
                        v_cache_quant_scale_lp, k_cache_quant_scale_hp, v_cache_quant_scale_hp,
                        alibi_slopes, max_contxt_len_lp, max_contxt_len_hp, softmax_scale,
                        return_lse, kv_cache_quant_bit_size_lp, kv_cache_quant_bit_size_hp)

    def launch(self, *args):
        torch_output, torch_lse = self.op_impl_base(*args)
        tmo_output, tmo_lse = ops.single_query_mixed_cached_kv_attn(*args)
        self.assertTensorsEqual(torch_lse.cpu().float(), tmo_lse.cpu().float(),
                                0.0003, use_MSE=True)
        self.assertTensorsEqual(torch_output.cpu().float(), tmo_output.cpu().float(),
                                    0.003, use_MSE=True)

    def op_impl_base(self, *args):
        q, k_cache_lp, v_cache_lp, k_cache_hp, v_cache_hp, out, block_tables_lp, block_tables_hp, \
            context_lens_lp, context_lens_hp, k_cache_quant_scale_lp, v_cache_quant_scale_lp, k_cache_quant_scale_hp, \
                v_cache_quant_scale_hp, alibi_slopes, max_contxt_len_lp, max_contxt_len_hp, softmax_scale, return_lse, \
                    kv_cache_quant_bit_size_lp, kv_cache_quant_bit_size_hp = args

        if kv_cache_quant_bit_size_lp == 4:
            num_blocks, num_kv_heads, block_size_lp, head_size = v_cache_lp.size()
            block_size = block_size_lp * 2
            k_cache_lp = UnpackInt4(k_cache_lp).reshape(num_blocks, num_kv_heads, block_size, head_size)
            v_cache_lp = UnpackInt4(v_cache_lp.transpose(2,3)).reshape(num_blocks, num_kv_heads, head_size, block_size).transpose(2,3)

        torch_output_lp, torch_lse_lp = single_query_cached_kv_attn(q.contiguous().float(), k_cache_lp.float(), v_cache_lp.float(),
                block_tables_lp, context_lens_lp, k_cache_quant_scale_lp, v_cache_quant_scale_lp, alibi_slopes, -1, -1, softmax_scale, return_lse)
        torch_output_hp, torch_lse_hp = single_query_cached_kv_attn(q.contiguous().float(), k_cache_hp.float(), v_cache_hp.float(),
                block_tables_hp, context_lens_hp, k_cache_quant_scale_hp, v_cache_quant_scale_hp, alibi_slopes, -1, -1, softmax_scale, return_lse)
        torch_output, torch_lse = update_out_and_lse_torch(torch_output_lp, torch_lse_lp, torch_output_hp, torch_lse_hp, None, None, None)
        return (torch_output, torch_lse) if return_lse else torch_output

    def test_single_query_mixedkv_attention(self):
        head_num = 16
        batch = 12
        num_kv_heads = 4
        seq_q = 1
        head_size = 128
        seq_len_lp = 512
        seq_len_hp = 128
        is_pagedattn_list = [True, False]
        has_alibi_list = [True, False]
        is_normal_list = [True, False]  # if false, lp_k/v_len of some batch is 0
        quant_bit_list = [(-1, -1), (4, 8), (4, -1), (8, -1), (8, 8)]
        data_type_list = [torch.float, torch.float16]
        if torch_mlu.mlu.is_bf16_supported():
            data_type_list.append(torch.bfloat16)
        args = product(is_pagedattn_list, has_alibi_list, data_type_list, quant_bit_list, is_normal_list)
        for is_pagedattn, has_alibi, data_type, quant_bit, is_normal in args:
            quant_bit_lp, quant_bit_hp = quant_bit
            print("test separate:{} + {}: batch:{}, seq_len_lp:{}, seq_len_hp:{}, head_size:{}, is_pagedattn:{}, has_alibi:{}, data_type:{}, is_normal:{} ...".format(
              quant_bit_lp, quant_bit_hp, batch, seq_len_lp, seq_len_hp, head_size, is_pagedattn, has_alibi, data_type, is_normal))
            if is_pagedattn:
                mlu_name = torch.mlu.get_device_name()
                if "MLU3" in mlu_name:
                   print("pagedattn is not implement on mlu370, skip it")
                   continue
            torch.manual_seed(1)
            input_qkv = torch.randn((batch, 3 * seq_q, 3 * head_num, head_size), dtype=data_type).mlu()
            input_q = input_qkv[:, 0:seq_q, 0:head_num,:].view(batch, seq_q, head_num, head_size)
            #gen k/v cache
            params_lp = gen_cache(batch, seq_q, num_kv_heads, head_size, is_pagedattn, seq_len_lp, data_type, quant_bit_lp,
                                   "per_token", is_normal)
            params_hp = gen_cache(batch, seq_q, num_kv_heads, head_size, is_pagedattn, seq_len_hp, data_type, quant_bit_hp,
                                   "per_channel")
            if quant_bit_lp == 4:
               key_cache_lp, value_cache_lp, key_scale_lp, value_scale_lp, context_lens_lp, block_tables_lp, _, _ = params_lp
            else:
                key_cache_lp, value_cache_lp, key_scale_lp, value_scale_lp, context_lens_lp, block_tables_lp = params_lp
            key_cache_hp, value_cache_hp, key_scale_hp, value_scale_hp, context_lens_hp, block_tables_hp = params_hp
            max_context_len_lp = context_lens_lp.max().item()
            max_context_len_hp = context_lens_hp.max().item()

            alibi_slopes = None
            if has_alibi:
                alibi_slopes = torch.zeros((batch, head_num), dtype=torch.float32).mlu()
                alibi_slopes.uniform_(0, 0.125)
            softmax_scale = 1 / math.sqrt(head_size)
            torch_output, torch_lse = self.op_impl_base(input_q,
                                    key_cache_lp, value_cache_lp,
                                    key_cache_hp, value_cache_hp,
                                    None, #output
                                    block_tables_lp, block_tables_hp,
                                    context_lens_lp, context_lens_hp,
                                    key_scale_lp, value_scale_lp,
                                    key_scale_hp, value_scale_hp,
                                    alibi_slopes,
                                    max_context_len_lp, max_context_len_hp,
                                    softmax_scale, True,
                                    quant_bit_lp, quant_bit_hp)

            tmo_output, tmo_lse = ops.single_query_mixed_cached_kv_attn(input_q,
                                    key_cache_lp, value_cache_lp,
                                    key_cache_hp, value_cache_hp,
                                    None, #output
                                    block_tables_lp, block_tables_hp,
                                    context_lens_lp, context_lens_hp,
                                    key_scale_lp, value_scale_lp,
                                    key_scale_hp, value_scale_hp,
                                    alibi_slopes,
                                    max_context_len_lp, max_context_len_hp,
                                    softmax_scale, True,
                                    quant_bit_lp, quant_bit_hp)
            self.assertTensorsEqual(torch_lse.cpu().float(), tmo_lse.cpu().float(),
                                    0.0003, use_MSE=True)
            self.assertTensorsEqual(torch_output.cpu().float(), tmo_output.cpu().float(),
                                    0.003, use_MSE=True)

    def test_single_query_mixedkv_attention_concate(self):
        head_num = 16
        batch = 12
        num_kv_heads = 4
        seq_q = 1
        head_size = 128
        seq_len_lp = 512
        seq_len_hp = 128
        is_pagedattn_list = [True, False]
        has_alibi = False #only support without alibi
        is_normal_list = [True, False] # if false, lp_k/v_len of some batch is 0
        quant_bit_list = [(-1, -1), (4, 8), (4, -1), (8, -1), (8, 8)]
        data_type_list = [torch.float, torch.float16]
        if torch_mlu.mlu.is_bf16_supported():
            data_type_list.append(torch.bfloat16)
        args = product(is_pagedattn_list, data_type_list, quant_bit_list, is_normal_list)
        for is_pagedattn, data_type, quant_bit, is_normal in args:
            quant_bit_lp, quant_bit_hp = quant_bit
            print("test concate {} + {}: seq_len_lp: {}, seq_len_hp: {}, is_pagedattn: {}, data_type {}, is_normal {} ...".format(
              quant_bit_lp, quant_bit_hp, batch, seq_len_lp, seq_len_hp, head_size, is_pagedattn, data_type, is_normal))
            if is_pagedattn:
                mlu_name = torch.mlu.get_device_name()
                if "MLU3" in mlu_name:
                   print("pagedattn is not implement on mlu370, skip it")
                   continue
            torch.manual_seed(1)
            input_qkv = torch.randn((batch, 3 * seq_q, 3 * head_num, head_size), dtype=data_type).mlu()
            input_q = input_qkv[:, 0:seq_q, 0:head_num,:].view(batch, seq_q, head_num, head_size)
            #gen k/v cache
            params_lp = gen_cache(batch, seq_q, num_kv_heads, head_size, is_pagedattn, seq_len_lp, data_type, quant_bit_lp,
                                   "per_token", is_normal)
            params_hp = gen_cache(batch, seq_q, num_kv_heads, head_size, is_pagedattn, seq_len_hp, data_type, quant_bit_hp,
                                   "per_channel")
            if quant_bit_lp == 4:
               key_cache_lp, value_cache_lp, key_scale_lp, value_scale_lp, context_lens_lp, block_tables_lp, key_cache_lp_torch, value_cache_lp_torch = params_lp
            else:
                key_cache_lp, value_cache_lp, key_scale_lp, value_scale_lp, context_lens_lp, block_tables_lp = params_lp
                key_cache_lp_torch, value_cache_lp_torch = key_cache_lp, value_cache_lp
            key_cache_hp, value_cache_hp, key_scale_hp, value_scale_hp, context_lens_hp, block_tables_hp = params_hp
            max_context_len_lp = context_lens_lp.max().item()
            max_context_len_hp = context_lens_hp.max().item()

            alibi_slopes = None
            if has_alibi:
                alibi_slopes = torch.zeros((batch, head_num), dtype=torch.float32).mlu()
                alibi_slopes.uniform_(0, 0.125)
            softmax_scale = 1 / math.sqrt(head_size)
            # concate cache
            if is_pagedattn:
                concat_cache_k, concat_cache_v, concat_block_tables, concat_context = concat_cache_paged(
                    key_cache_lp_torch.float().cpu(), key_cache_hp.float().cpu(),
                    value_cache_lp_torch.float().cpu(), value_cache_hp.float().cpu(),
                    context_lens_lp.cpu(), context_lens_hp.cpu(),
                    block_tables_lp.cpu(), block_tables_hp.cpu(),
                    key_scale_lp.cpu() if key_scale_lp is not None else None,
                    key_scale_hp.cpu() if key_scale_hp is not None else None,
                    value_scale_lp.cpu() if value_scale_lp is not None else None,
                    value_scale_hp.cpu() if value_scale_hp is not None else None)
            else:
                concat_cache_k, concat_cache_v, concat_block_tables, concat_context = concate_cache_linear(
                    key_cache_lp_torch.float().cpu(), key_cache_hp.float().cpu(),
                    value_cache_lp_torch.float().cpu(), value_cache_hp.float().cpu(),
                    context_lens_lp.cpu(), context_lens_hp.cpu(),
                    block_tables_lp.cpu(), block_tables_hp.cpu(),
                    key_scale_lp.cpu() if key_scale_lp is not None else None,
                    key_scale_hp.cpu() if key_scale_hp is not None else None,
                    value_scale_lp.cpu() if value_scale_lp is not None else None,
                    value_scale_hp.cpu() if value_scale_hp is not None else None)

            torch_output_concat, torch_lse_concat = single_query_cached_kv_attn(input_q.contiguous().float(),
                                                 concat_cache_k.float(), concat_cache_v.float(), concat_block_tables,
                                                 concat_context, None, None, alibi_slopes, -1, -1, softmax_scale, True)
            tmo_output, tmo_lse = ops.single_query_mixed_cached_kv_attn(input_q,
                                    key_cache_lp, value_cache_lp,
                                    key_cache_hp, value_cache_hp,
                                    None, #output
                                    block_tables_lp, block_tables_hp,
                                    context_lens_lp, context_lens_hp,
                                    key_scale_lp, value_scale_lp,
                                    key_scale_hp, value_scale_hp,
                                    alibi_slopes,
                                    max_context_len_lp, max_context_len_hp,
                                    softmax_scale, True,
                                    quant_bit_lp, quant_bit_hp)
            if is_normal: # only compare lse when context_len_lp is normal, seq=0 case nan-value in lse
                self.assertTensorsEqual(torch_lse_concat.cpu().float(), tmo_lse.cpu().float(),
                                        0.0003, use_MSE=True)
            self.assertTensorsEqual(torch_output_concat.cpu().float(), tmo_output.cpu().float(),
                                    0.003, use_MSE=True)

    def test_inductor(self):
        return super().test_inductor()

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