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

def gen_args(num_blocks, num_heads, block_size, head_size, dtype, cpy, num_pair):
    shape = (num_blocks, num_heads, block_size, head_size)
    if dtype in {torch.int8, torch.uint8, torch.int16, torch.int32, torch.int64}:
        info = torch.iinfo(dtype)
        if cpy == "mlu to mlu":
            src = torch.randint(info.min, info.max, size=shape, dtype=dtype).mlu()
            dst = torch.randint(info.min, info.max, size=shape, dtype=dtype).mlu()
        elif cpy == "mlu to cpu":
            src = torch.randint(info.min, info.max, size=shape, dtype=dtype).mlu()
            dst = torch.randint(info.min, info.max, size=shape, dtype=dtype).cpu()
        elif cpy == "cpu to mlu":
            src = torch.randint(info.min, info.max, size=shape, dtype=dtype).cpu()
            dst = torch.randint(info.min, info.max, size=shape, dtype=dtype).mlu()
        else:
            print("unkown copy direction.", flush=True)
            exit(1)
    else:
        if cpy == "mlu to mlu":
            src = torch.randn(size=shape, dtype=dtype).mlu()
            dst = torch.randn(size=shape, dtype=dtype).mlu()
        elif cpy == "mlu to cpu":
            src = torch.randn(size=shape, dtype=dtype).mlu()
            dst = torch.randn(size=shape, dtype=dtype).cpu()
        elif cpy == "cpu to mlu":
            src = torch.randn(size=shape, dtype=dtype).cpu()
            dst = torch.randn(size=shape, dtype=dtype).mlu()
        else:
            print("unkown copy direction.", flush=True)
            exit(1)

    values = list(range(num_pair))
    random.shuffle(values)
    src_to_dst = {key: value for key, value in zip(range(num_pair), values)}
    return dst, src, src_to_dst

class TestSwapBlocksOp(BtTestCase):
    def op_impl_base(self, *args):
        dst, src, block_mapping = args
        for key, value in block_mapping.items():
            dst[value] = src[key]
        return dst

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test swap blocks due to ASan issues")
    def test_swap_blocks(self):
        num_blocks_list = [3600]
        num_heads_list = [8]
        head_size_list = [64,128]
        block_size_list = [16]
        num_pairs_list = [6,512]

        types = [torch.int8, torch.uint8, torch.int16, torch.int32, torch.int64, torch.half, torch.float]
        if torch_mlu.mlu.is_bf16_supported():
            types.append(torch.bfloat16)
        cpys = ["mlu to mlu", "mlu to cpu", "cpu to mlu"]
        args = product(num_blocks_list, num_heads_list, block_size_list, head_size_list, types, cpys, num_pairs_list)
        for num_blocks, num_heads, block_size, head_size, dtype, cpy, num_pair in args:
            print("num_blocks: {}, num_heads: {}, block_size: {}, head_size: {}, dtype: {}, dir: {}, num_pairs: {} testing..."
                .format(num_blocks, num_heads, block_size, head_size, dtype, cpy, num_pair), flush=True)

            dst, src, src_to_dst = gen_args(num_blocks, num_heads, block_size, head_size, dtype, cpy, num_pair)

            ref_src, ref_dst = src.clone(), dst.clone()
            # cpu
            self.op_impl_base(ref_dst, ref_src, src_to_dst)
            # mlu
            ops.swap_blocks(dst, src, src_to_dst)
            # diff
            self.assertTensorsEqual(src.cpu().float(), ref_src.cpu().float(), 0, use_MSE=True, use_RAE=True, use_RMA=True)
            self.assertTensorsEqual(dst.cpu().float(), ref_dst.cpu().float(), 0, use_MSE=True, use_RAE=True, use_RMA=True)

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_prevent due to ASan issues")
    def test_inductor(self):
        num_blocks, num_heads, block_size, head_size, dtype, cpy, num_pair = 3600, 8, 16, 64, torch.half, "mlu to mlu", 512
        dst, src, block_mapping = gen_args(num_blocks, num_heads, block_size, head_size, dtype, cpy, num_pair)
        self.base_opcheck(torch.ops.torch_mlu_ops.swap_blocks, (dst, src, block_mapping))


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