import torch
import torch_mlu
from itertools import product
import torch_mlu_ops as ops
import random
import torch.testing._internal.optests as optests
from torch._subclasses.fake_tensor import FakeTensorMode
from torch.testing._internal.common_utils import (
    run_tests,
    TestCase,
)

dtype_dict = {
  torch.half: "half",
  torch.bfloat16: "bfloat16",
  torch.float: "float",
}

fake_tensor_mode = FakeTensorMode(allow_fallback_kernels=False)

class FakeTensorTest(TestCase):
    def test_matmul(self):
        with fake_tensor_mode:
            mat_m, mat_n, mat_k, alpha, beta, act_mode, fast_act, approximate = 32, 256, 128, 0.8, 0.3, 'silu', True, True
            trans_a_list = [True, False]
            trans_b_list = [True, False]
            dtype_list = [torch.half, torch.float]
            args = product(trans_a_list, trans_b_list, dtype_list)
            for trans_a, trans_b, dtype in args:
                print(f"matmul... trans_a: {trans_a}, trans_b: {trans_b}, dtype: {dtype}")
                shape_a, shape_b = (mat_m, mat_k), (mat_k, mat_n)
                if trans_a:
                    shape_a = (mat_k, mat_m)
                if trans_b:
                    shape_b = (mat_n, mat_k)
                a = torch.randn(shape_a, dtype=dtype, device='mlu')
                b = torch.randn(shape_b, dtype=dtype, device='mlu')
                bias = torch.randn((mat_n), dtype=dtype, device='mlu')
                c = torch.randn((mat_m, mat_n), dtype=dtype, device='mlu')
                output = torch.ops.torch_mlu_ops.matmul(a, b, None, bias, c, None, act_mode, alpha, beta, fast_act, approximate, 1.0, 1.0, trans_a, trans_b)
                self.assertEqual(output.shape, (mat_m, mat_n))

                a8 = torch.randint(-128, 127, shape_a).to(torch.int8).mlu()
                b8 = torch.randint(-128, 127, shape_b).to(torch.int8).mlu()
                a_scale = 280.8
                b_scale = 190.8
                str_dtype = dtype_dict[dtype]
                output = torch.ops.torch_mlu_ops.matmul(a8, b8, None, bias, c, str_dtype, act_mode, alpha, beta, fast_act, approximate, a_scale, b_scale, trans_a, trans_b)
                self.assertEqual(output.shape, (mat_m, mat_n))

    def test_weight_only_quant_matmul(self):
        with fake_tensor_mode:
            M, K, N, group_num = 2, 256, 32, 4
            quant_bit_size, act_mode, use_hp_active, act_coef, alpha, beta = 8, 'none', True, 1., 0.8, 0.3
            group_quant_list = [True, False]
            trans_a_list = [True, False]
            trans_b_list = [True, False]
            dtype_list = [torch.half, torch.bfloat16] if torch_mlu.mlu.is_bf16_supported() else [torch.half]
            args = product(group_quant_list, dtype_list, trans_a_list, trans_b_list)
            for group_quant, dtype, trans_a, trans_b in args:
                print(f"weight_only_quant_matmul... group_quant: {group_quant}, dtype: {dtype}, trans_a: {trans_a}, trans_b: {trans_b}")
                a = torch.randn((M, K), dtype=dtype).mlu()
                b = torch.randint(-128, 127, (N, K), dtype=torch.int8).mlu()
                c = torch.randn(M, N, device="mlu", dtype=dtype)
                bias = torch.randn(N, device="mlu", dtype=dtype)
                group_wise_scale = torch.randn((N, group_num), device="mlu", dtype=dtype)
                b_quant_layout = "quantize_group_wise" if group_quant else "quantize_per_channel"
                b_scale = group_wise_scale if group_quant else None
                gemm_output_scale =  None if group_quant else torch.randn(N, device="mlu", dtype=torch.float)
                a_scale, a_zero, b_zero, c_zero = None, None, None, None
                c_scale, gemm_output_zero = None, None
                quant_algo, a_quant_layout = "weight_only", "quantize_none"
                str_dtype = dtype_dict[dtype]
                output = torch.ops.torch_mlu_ops.quant_matmul(a, a_scale, a_zero,
                        b, b_scale, b_zero,
                        bias, c, c_scale, c_zero,
                        gemm_output_scale, gemm_output_zero,
                        str_dtype, None, quant_algo,
                        a_quant_layout, b_quant_layout,
                        quant_bit_size, act_mode, use_hp_active, act_coef,
                        alpha, beta, trans_a, trans_b,)
                self.assertEqual(output.shape, (M, N))

    def test_smooth_quant_matmul(self):
        with fake_tensor_mode:
            act_mode_list = ["none", "silu", "gelu"]
            dtype_list = [torch.half, torch.bfloat16] if torch_mlu.mlu.is_bf16_supported() else [torch.half]
            trans_a_list = [True, False]
            trans_b_list = [True, False]
            M, K, N = 2, 16, 32
            quant_bit_size, use_hp_active, act_coef, alpha, beta = 8, True, 1., 0.8, 0.3
            arg = product(act_mode_list, dtype_list, trans_a_list, trans_b_list)
            for act_mode, dtype, trans_a, trans_b in arg:
                print(f"smooth_quant_matmul...  act_mode: {act_mode}, dtype: {dtype}, trans_a: {trans_a}, trans_b: {trans_b}")
                a = torch.randint(-128, 127, (M, K), dtype=torch.int8).mlu()
                b = torch.randint(-128, 127, (N, K), dtype=torch.int8).mlu()
                c = None
                bias = torch.randn(N, device="mlu", dtype=dtype)
                a_scale = torch.randn(M, device="mlu", dtype=torch.float)
                b_scale = torch.randn(N, device="mlu", dtype=torch.float)
                a_zero, b_zero, c_zero = None, None, None
                c_scale, gemm_output_scale, gemm_output_zero = None, None, None
                quant_algo, a_quant_layout, b_quant_layout = "smooth_quant", "quantize_per_token", "quantize_per_channel"
                str_dtype = dtype_dict[dtype]
                output = torch.ops.torch_mlu_ops.quant_matmul(a, a_scale, a_zero,
                        b, b_scale, b_zero,
                        bias, c, c_scale, c_zero,
                        gemm_output_scale, gemm_output_zero,
                        str_dtype, None, quant_algo,
                        a_quant_layout, b_quant_layout,
                        quant_bit_size, act_mode, use_hp_active, act_coef,
                        alpha, beta, trans_a, trans_b)
                self.assertEqual(output.shape, (M, N))

    def test_group_gemm(self):
        with fake_tensor_mode:
            batch, seq, k, n, experts_num, topk = 1, 8, 1024, 2048, 8, 5
            dtype_list = [torch.half, torch.bfloat16] if torch_mlu.mlu.is_bf16_supported() else [torch.half]
            idx_list =  [False, True]
            has_bias_list = [True, False]
            args = product( dtype_list, idx_list, has_bias_list)
            for data_type, idx_mode, has_bias in args:
                print(f"group_gemm... has_bias: {has_bias}, idx_mode: {idx_mode}, dtype: {data_type}")
                bs = batch * seq
                token_topk = bs * topk
                expert_id = torch.randint(experts_num, (token_topk,), device="mlu")
                sorted_expert_id, indices = expert_id.sort()
                gather_idx = indices // topk
                gather_idx = gather_idx.to(torch.int32)
                token_count = torch.randint(0, token_topk, (experts_num,)).to(torch.int32)
                a = torch.randn(bs, k, device="mlu", dtype=data_type)
                if not idx_mode:
                    a = a[gather_idx]
                b = torch.randn(experts_num, n, k, device="mlu", dtype=data_type)
                c = torch.randn(token_topk, n, device="mlu", dtype=data_type)
                alpha = torch.randn(experts_num, device="mlu", dtype=torch.float32)
                beta = torch.randn(experts_num, device="mlu", dtype=torch.float32)
                a_scale = None
                b_scale = None
                bias = None
                if has_bias:
                    bias = torch.randn(experts_num, n, device="mlu", dtype=data_type)
                gather_idx_ = gather_idx if idx_mode else None
                output = torch.ops.torch_mlu_ops.group_gemm(a, b, token_count, gather_idx_, c, alpha, beta, a_scale, b_scale, bias, None, None, None, bs)
                self.assertEqual(output.shape, (token_topk, n))

    def test_smoothquant_group_gemm(self):
        with fake_tensor_mode:
            batch, seq, k, n, experts_num, topk = 1, 8, 1024, 2048, 8, 5
            dtype_list = [torch.half, torch.bfloat16] if torch_mlu.mlu.is_bf16_supported() else [torch.half]
            idx_list =  [False, True]
            has_bias_list = [True, False]
            args = product( dtype_list, idx_list, has_bias_list)
            for data_type, idx_mode, has_bias in args:
                print(f"smoothquant_group_gemm... has_bias: {has_bias}, idx_mode: {idx_mode}, dtype: {data_type}")
                bs = batch * seq
                token_topk = bs * topk
                expert_id = torch.randint(experts_num, (token_topk,), device="mlu")
                sorted_expert_id, indices = expert_id.sort()
                gather_idx = indices // topk
                gather_idx = gather_idx.to(torch.int32)
                token_count = torch.randint(0, token_topk, (experts_num,)).to(torch.int32)
                a8 = torch.randint(-128, 127, (bs, k)).to(torch.int8).mlu()
                b8 = torch.randint(-128, 127, (experts_num, n, k)).to(torch.int8).mlu()
                a_scale = torch.randn(token_topk, dtype=torch.float32).mlu()
                b_scale = torch.randn(experts_num, n, dtype=torch.float32).mlu()
                c = torch.randn(token_topk, n, device="mlu", dtype=data_type)
                if not idx_mode:
                    a8 = a8[gather_idx]
                alpha = torch.randn(experts_num, device="mlu", dtype=torch.float32)
                beta = torch.randn(experts_num, device="mlu", dtype=torch.float32)
                bias = None
                if has_bias:
                    bias = torch.randn(experts_num, n, device="mlu", dtype=data_type)
                gather_idx_ = gather_idx if idx_mode else None
                str_dtype = dtype_dict[data_type]
                output = torch.ops.torch_mlu_ops.group_gemm(a8, b8, token_count, gather_idx_, c, alpha, beta, a_scale, b_scale, bias, str_dtype, None, None, bs)
                self.assertEqual(output.shape, (token_topk, n))

    def test_moe_expand_input(self):
        with fake_tensor_mode:
            token_num, hidden_size, expert_num, topk, start_expert_id, expert_size = 2048, 4096, 32, 8, 3, 20
            dtype_list = [torch.half, torch.float, torch.int8, torch.bfloat16] if torch_mlu.mlu.is_bf16_supported() else [torch.half, torch.float, torch.int8]
            for dtype in dtype_list:
                print(f"moe_expand_input... token_num: {token_num}, expert_num: {expert_num}, topk: {topk}, dtype: {dtype}")
                input = torch.randn(token_num, hidden_size, device='mlu').to(dtype)
                gather_idx = torch.randint(low=0, high=token_num, size=(token_num * topk,), dtype=torch.int32, device='mlu')
                cusum_token_count = torch.zeros(expert_num + 1, dtype=torch.int32).mlu()
                output=torch.ops.torch_mlu_ops.moe_expand_input(input, gather_idx, cusum_token_count, start_expert_id, expert_size)
                self.assertEqual(output.shape, (token_num*topk, hidden_size))

    def test_moe_gen_idx(self):
        with fake_tensor_mode:
            token_num, expert_num, topk = 2048, 32, 8
            print(f"moe_gen_idx... token_num: {token_num}, expert_num: {expert_num}, topk: {topk}")
            expert_id = torch.randint(low=0, high=expert_num, size=(token_num, topk)).to(torch.int32).to('mlu')
            expand_idx, combine_idx, token_count, cusum_token_count =torch.ops.torch_mlu_ops.moe_gen_idx(expert_id, expert_num)
            self.assertEqual(expand_idx.shape, (token_num * topk,))
            self.assertEqual(combine_idx.shape, (token_num * topk,))
            self.assertEqual(token_count.shape, (expert_num,))
            self.assertEqual(cusum_token_count.shape, (expert_num + 1,))

    def test_moe_combine_result(self):
        with fake_tensor_mode:
            has_bias_list = [True, False]
            num_tokens, hidden_size, num_expert, topk, start_expert_id = 1, 2048, 8, 2, 0
            expert_size_list = [5, 8]
            dtype_list = [torch.half, torch.bfloat16] if torch_mlu.mlu.is_bf16_supported() else [torch.half]
            args = product(has_bias_list, expert_size_list, dtype_list)
            for has_bias, expert_size, dtype in args:
                print(f"moe_combine_result... has_bias: {has_bias}, expert_size: {expert_size}, dtype: {dtype}")
                input = torch.randn((num_tokens * topk, hidden_size), dtype=dtype, device='mlu')
                reduce_weight = torch.randn((num_tokens, topk), dtype=torch.float32, device='mlu')
                gather_ids = torch.randperm(num_tokens * topk, dtype=torch.int32, device='mlu')
                bias = None
                residual = None
                cusum_token_count = None
                if has_bias:
                    bias = torch.randn((num_expert, hidden_size), dtype=dtype, device='mlu')
                residual = torch.randn((num_tokens, hidden_size), dtype=dtype, device='mlu')
                if has_bias or expert_size < num_expert:
                    cusum_token_count =torch.zeros(num_expert + 1, dtype=torch.int32).mlu()
                output=torch.ops.torch_mlu_ops.moe_combine_result(input, reduce_weight, gather_ids, residual, cusum_token_count, start_expert_id, expert_size, bias)
                self.assertEqual(output.shape, (num_tokens, hidden_size))

    def test_moe(self):
        with fake_tensor_mode:
            act_mode = 'gelu'
            case_list = set()
            while (len(case_list) < 100):
                batch = random.randint(1, 10)
                seq = random.randint(1, 10)
                hidden_size = random.randrange(1024, 3072, 512)
                inner_size = random.randrange(1024, 3072, 512)
                expert_num = random.randint(1, 40)
                topk = random.randint(1, expert_num)
                gated = random.choice([True, False])
                renormalize = random.choice([True, False])
                quant_mode = random.choice(["no_quant", "w4", "w8", "w4w8"])
                quant_wise = random.choice([128, 256, 512])
                data_type = random.choice([torch.bfloat16, torch.float16])
                if not torch_mlu.mlu.is_bf16_supported():
                    data_type = torch.float16
                case = (batch, seq, hidden_size, inner_size, expert_num, topk, gated, renormalize, data_type, quant_mode, quant_wise, act_mode)
                if case in case_list:
                    continue
                case_list.add(case)
                print(f"bs: {batch}, seq_len: {seq}, hidden_size: {hidden_size}, inner_size: {inner_size}, experts_num: {expert_num}, \
topk: {topk}, gated: {gated}, renormalize: {renormalize}, data_type: {data_type}, quant_mode: {quant_mode}, quant_wise: {quant_wise}, act_mode: {act_mode} testing...", flush=True)
                # if expert_size == -1:
                #     expert_size = expert_num
                if quant_mode == "no_quant":
                    w1 = torch.randn((expert_num, inner_size*(1+gated), hidden_size), device="mlu", dtype=data_type)
                    w2 = torch.randn((expert_num, hidden_size, inner_size), device="mlu", dtype=data_type)
                    w1_quant_flag = None
                    w2_quant_flag = None
                elif quant_mode == "w4":
                    w1_quant_group = hidden_size // quant_wise
                    w2_quant_group = inner_size // quant_wise
                    w1_quant_flag = None
                    w2_quant_flag = None
                    w1 = torch.randint(-128, 127, (expert_num, inner_size*(1+gated), hidden_size // 2), device="mlu", dtype=torch.int32).to(torch.int8)
                    w2 = torch.randint(-128, 127, (expert_num, hidden_size, inner_size // 2), device="mlu", dtype=torch.int32).to(torch.int8)
                    w1_scale = torch.empty((w1_quant_group, expert_num, inner_size*(1+gated)), device="mlu", dtype=torch.float32).uniform_(-0.05, 0.05)
                    w2_scale = torch.empty((w2_quant_group, expert_num, hidden_size), device="mlu", dtype=torch.float32).uniform_(-0.05, 0.05)
                elif quant_mode == "w8":
                    w1_quant_flag = None
                    w2_quant_flag = None
                    w1 = torch.randint(-128, 127, (expert_num, inner_size*(1+gated), hidden_size), device="mlu", dtype=torch.int32).to(torch.int8)
                    w2 = torch.randint(-128, 127, (expert_num, hidden_size, inner_size), device="mlu", dtype=torch.int32).to(torch.int8)
                    w1_scale = torch.empty((expert_num, inner_size*(1+gated)), device="mlu", dtype=torch.float32).uniform_(-0.05, 0.05)
                    w2_scale = torch.empty((expert_num, hidden_size), device="mlu", dtype=torch.float32).uniform_(-0.05, 0.05)
                elif quant_mode == "w4w8":
                    w1_quant_group = hidden_size // quant_wise
                    w2_quant_group = inner_size // quant_wise
                    w1_quant_flag = random.choices([4,8], k=expert_num * w1_quant_group)
                    w2_quant_flag = random.choices([4,8], k=expert_num * w2_quant_group)
                    w1_count = (sum(w1_quant_flag) // 4) * (quant_wise // 2) * inner_size*(1+gated)
                    w2_count = (sum(w2_quant_flag) // 4) * (quant_wise // 2) * hidden_size
                    w1 = torch.randint(-128, 127, (w1_count,), device="mlu", dtype=torch.int32).to(torch.int8)
                    w2 = torch.randint(-128, 127, (w2_count,), device="mlu", dtype=torch.int32).to(torch.int8)
                    w1_scale = torch.empty((w1_quant_group, expert_num, inner_size*(1+gated)), device="mlu", dtype=torch.float32).uniform_(-0.05, 0.05)
                    w2_scale = torch.empty((w2_quant_group, expert_num, hidden_size), device="mlu", dtype=torch.float32).uniform_(-0.05, 0.05)
                hidden_states = torch.randn(batch, seq, hidden_size, device="mlu", dtype=data_type)
                residual = None  # torch.randn(batch, seq, hidden_size, device="mlu", dtype=data_type)
                router_logit = torch.randn(batch, seq, expert_num, device="mlu", dtype=torch.float32)
                input_smooth = None if quant_mode == "no_quant" else torch.empty(expert_num, hidden_size, device="mlu", dtype=torch.float32).uniform_(0.01, 0.05)
                act_smooth = None if quant_mode == "no_quant" else torch.empty(expert_num, inner_size, device="mlu", dtype=torch.float32).uniform_(0.01, 0.05)
                bias1, bias2 = None, None
                output = torch.ops.torch_mlu_ops.fused_moe(hidden_states, router_logit, w1, w2, bias1, bias2, residual,
                                                           input_smooth, act_smooth, w1_scale, w2_scale, w1_quant_flag,
                                                           w2_quant_flag, topk, renormalize, gated, act_mode, 0,
                                                           0, 0)
                self.assertEqual(output.shape, hidden_states.shape)


if __name__ == "__main__":
    if torch.__version__ >= '2.3.0':
        run_tests()
