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

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)
    return args[0:10]

class TestReshapeLinearCache(BtTestCase):
    def op_impl_base(self, *args):
        key, value, key_cache, value_cache, context_lengths, max_context_len, packed, \
            context_seq_offset, cache_bs_id, cache_seqlen_offset = args

        batch_size = context_lengths.shape[0] - 1 if packed else context_lengths.shape[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]
            cache_bs_id_i = cache_bs_id[i]
            cache_seqlen_offset_i = cache_seqlen_offset[i]
            if cache_seqlen_offset_i < 0 or cache_bs_id_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_i[...] = key_i[:, context_seq_offset_i:context_seq_offset_i + context_len_i]
            if value_cache is not None and value is not None:
                value_cache_i = \
                    value_cache[cache_bs_id_i, :, \
                            cache_seqlen_offset_i:cache_seqlen_offset_i + context_len_i]
                value_cache_i[...] = value_i[:, context_seq_offset_i:context_seq_offset_i + context_len_i]
        return (key_cache, value_cache, ) if value_cache is not None else key_cache

    def test_reshape_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=16, high=1024, size=(test_cases, ), dtype=torch.int32)
        packed_list = torch.randint(low=0, high=2, size=(test_cases, ), dtype=torch.int32)
        dtype_list = torch.randint(low=0, high=4, size=(test_cases, ), dtype=torch.int32)
        dtype_map = [torch.int8, torch.half, torch.bfloat16, torch.float]

        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()
            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={}, dtype={} testing...".format(
                batch_size, num_heads, head_size, cache_memory_len, packed > 0, dtype))

            max_bs = batch_size + 1
            context_lens = torch.randint(size=(batch_size, ), low=1,
                                         high=cache_memory_len // 2,
                                         dtype=torch.int32, device='mlu')
            # print(context_lens)
            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=-1,
                                              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')
            cache = torch.randn((2, max_bs, num_heads, cache_memory_len, head_size), dtype=torch.float, device='mlu')
            context = context.to(dtype)
            cache = cache.to(dtype)
            ref_cache = cache.clone()
            key = context[..., q_heads:q_heads + num_heads, :]
            value = context[..., q_heads + num_heads:q_heads + 2 * num_heads, :]
            key_cache = cache[0]
            value_cache = cache[1]
            ref_key_cache = ref_cache[0]
            ref_value_cache = ref_cache[1]

            cache_bs_id = None
            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, cu_context_lens,
                                        max_context_len, packed > 0, None,
                                        cache_bs_id, cache_seq_offsets)
                ops.reshape_linear_cache(key, value, key_cache, value_cache, cu_context_lens,
                                        max_context_len, packed > 0, None,
                                        cache_bs_id, cache_seq_offsets)
            else:
                self.op_impl_base(key, value, ref_key_cache, ref_value_cache, context_lens,
                                        max_context_len, packed > 0, context_seq_offsets,
                                        cache_bs_id, cache_seq_offsets)
                ops.reshape_linear_cache(key, value, key_cache, value_cache, context_lens,
                                        max_context_len, packed > 0, context_seq_offsets,
                                        cache_bs_id, cache_seq_offsets)
            self.assertTensorsEqual(cache.cpu().float(), ref_cache.cpu().float(), 0, "ref_cache must equal cache",
                True, True, True, True)

    def test_reshape_linear_key_cache(self):
        batch_size, num_heads, head_size, cache_memory_len = 2, 2, 16, 128
        print("[test_reshape_linear_key_cache] batch_size={}, num_heads={}, head_size={}, cache_memory_len={} testing...".format(
            batch_size, num_heads, head_size, cache_memory_len))

        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=-1,
                                            high=(cache_memory_len - max_context_len) // 3 + 1,
                                            dtype=torch.int32, device='mlu')
        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)

        context = torch.randn((batch_size, max_context_len + max_seq_offset, num_heads, head_size),
                                dtype=torch.float, device='mlu')
        key_cache = torch.randn((max_bs, num_heads, cache_memory_len, head_size), dtype=torch.float, device='mlu')
        ref_key_cache = key_cache.clone()
        key = context[..., :]

        cache_bs_id = None
        cache_bs_id = random.sample([*range(0, max_bs)], batch_size)
        cache_bs_id = torch.IntTensor(cache_bs_id).mlu()

        self.op_impl_base(key, None, ref_key_cache, None, context_lens,
                                    max_context_len, False, context_seq_offsets,
                                    cache_bs_id, cache_seq_offsets)
        ops.reshape_linear_cache(key, None, key_cache, None, context_lens,
                                    max_context_len, False, context_seq_offsets,
                                    cache_bs_id, cache_seq_offsets)
        self.assertTensorsEqual(key_cache.cpu().float(), ref_key_cache.cpu().float(), 0, "ref_cache must equal cache",
            True, True, True, True)


    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.reshape_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.reshape_linear_cache, args)

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