import torch
import torch_mlu
import unittest
import torch_mlu_ops
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(__file__)))
from common_utils import *
import random
from itertools import product
import time
import numpy as np
import os

class TestSoftmaxTopkOp(BtTestCase):
    def op_impl_base(self, *args):
        input, topk, normalize, num_expert_group, topk_group, origin_mask, normed_by = args
        softmax = torch.softmax(input.float(), dim=-1)
        if num_expert_group <= 1:
            if origin_mask is not None:
                softmax = softmax * origin_mask
            reduce_weight, expert_id = torch.topk(softmax, k=topk, dim=-1)
            if normalize:
                if normed_by == "topk_logit":
                    reduce_weight = reduce_weight / reduce_weight.sum(dim=-1, keepdim=True)
                if normed_by == "softmax_logit":
                    reduce_weight = reduce_weight / softmax.sum(dim=-1, keepdim=True)
            return reduce_weight, expert_id
        else:
            group_size = softmax.shape[-1] // num_expert_group
            new_shape = softmax.shape[:-1] + (num_expert_group, group_size)
            group_data = softmax.view(new_shape)
            group_max_value = group_data.max(dim=-1).values
            group_idx = torch.topk(group_max_value, k=topk_group, dim=-1)[1]
            mask_shape = softmax.shape[:-1] + (num_expert_group,)
            mask = torch.zeros((mask_shape), dtype = torch.bool, device = group_idx.device)
            mask.scatter_(-1, group_idx, True)
            mask = mask.unsqueeze(-1).expand(new_shape)
            masked_data = group_data.masked_fill(~mask, 0.0)
            masked_data = masked_data.reshape(softmax.shape)
            reduce_weight, expert_id = torch.topk(masked_data, k=topk, dim=-1)
            if normalize:
                if normed_by == "topk_logit":
                    reduce_weight = reduce_weight / reduce_weight.sum(dim=-1, keepdim=True)
                if normed_by == "softmax_logit":
                    reduce_weight = reduce_weight / softmax.sum(dim=-1, keepdim=True)
            return reduce_weight, expert_id

    # 接口测试
    def test_interface(self):
        num_token, num_expert, topk, num_expert_group, topk_group, normalize = 1024, 160, 6, 10, 5, False
        input = torch.randn(num_token, num_expert, dtype=torch.half, device='mlu')
        mask = None
        normed_by = "topk_logit"
        base_reduce_weight, base_expert_id = self.op_impl_base(input, topk, normalize, num_expert_group, topk_group, mask, normed_by)
        reduce_weight, expert_id = torch_mlu_ops.moe_softmax_topk(input, topk, normalize, num_expert_group, topk_group)
        base_expert_id, _ = base_expert_id.sort()
        expert_id, _ = expert_id.sort()
        self.assertTensorsEqual(base_reduce_weight.cpu().float(), reduce_weight.cpu().float(),
                                0.003, use_MSE=True, use_RAE=True)
        self.assertTensorsEqual(base_expert_id.cpu().float(), expert_id.cpu().float(),
                                0.003, use_MSE=True, use_RAE=True)


   # 随机遍历测试
    def test_kernel_random(self):
        dtype_list = [torch.half, torch.float]
        if torch_mlu.mlu.is_bf16_supported():
            dtype_list.append(torch.bfloat16)
        for i in range(1000):
            num_batch = random.randint(1, 64)
            num_mask = random.randint(1, 1024)
            num_expert = random.randint(1, 512)
            factors = [i for i in range(1, num_expert + 1) if num_expert % i == 0]
            num_expert_group = random.choice(factors)
            topk_group = random.randint(1, num_expert_group)
            topk = random.randint(1, num_expert / num_expert_group * topk_group)
            normalize = random.sample([True, False], 1)[0]
            dtype = random.sample(dtype_list, 1)[0]
            group_invalid = random.sample([True, False], 1)[0]
            if group_invalid:
                num_expert_group = -1
            normed_by = random.choice(["topk_logit", "softmax_logit"])
            print("===============================================================================")
            print(f"[{i}]: num_batch: {num_batch}, num_mask: {num_mask}, num_expert: {num_expert}, num_expert_group: {num_expert_group}")
            print(f"       topk_group: {topk_group}, topk: {topk}, normalize: {normalize}, dtype: {dtype}, normed_by: {normed_by}")
            input = torch.randn(num_batch, num_mask, num_expert, dtype=dtype).mlu()
            mask = torch.randint(0, 2, (1, num_mask, num_expert), dtype = dtype).mlu()
            if num_expert_group > 1:
                mask = None
            base_reduce_weight, base_expert_id = self.op_impl_base(input, topk, normalize, num_expert_group, topk_group, mask, normed_by)
            reduce_weight, expert_id = torch.ops.torch_mlu_ops.moe_softmax_topk(input,
                                                                                topk,
                                                                                num_expert_group,
                                                                                topk_group,
                                                                                normalize,
                                                                                mask,
                                                                                normed_by)
            # softmax后的值可能因数值差异较小，造成topk的值存在顺序上的差异，例如：
            # base_reduce_weight[N, 10:12] = [0.012, 0.011]
            # reduce_weight[N, 10:12] = [0.011, 0.012]
            # base_expert_id[N, 10:12] = [7, 8]
            # expert_id[N, 10:12] = [8, 7]
            # 这种情况产生的顺序上的差异不是错误的，但这样会造成结果对比错误，因此需要对结果先排序再对比结果
            base_expert_id, _ = base_expert_id.sort()
            expert_id, _ = expert_id.sort()
            self.assertTensorsEqual(base_reduce_weight.cpu().float(), reduce_weight.cpu().float(),
                                    0.003, use_MSE=True, use_RAE=True)
            self.assertTensorsEqual(base_expert_id.cpu().float(), expert_id.cpu().float(),
                                    0.003, use_MSE=True, use_RAE=True)
   # 防呆测试
    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_prevent due to ASan issues")
    def test_prevent(self):
        func = torch.ops.torch_mlu_ops.moe_softmax_topk
        num_batch, num_mask, num_expert, topk, num_expert_group, topk_group, normalize, normed_by = 2, 1024, 160, 0, 1, 1, True, "abc_logit"
        input = torch.randn(num_batch, num_mask, num_expert, dtype = torch.half, device='mlu')
        input_permute = input.permute(0, 2, 1)
        mask = torch.randint(0, 2, (num_mask, num_expert), dtype = torch.half, device='mlu')
        mask_permute = mask.permute(1, 0)
        self.assertException("input must be contiguous.",
                             func, input_permute, topk, num_expert_group, topk_group, normalize, mask, normed_by)
        self.assertException("mask must be contiguous.",
                             func, input, topk, num_expert_group, topk_group, normalize, mask_permute, normed_by)
        self.assertException("normed_by must be 'topk_logit' or 'softmax_logit'",
                             func, input, topk, num_expert_group, topk_group, normalize, mask, normed_by)
        input = torch.randn(num_expert, dtype = torch.half, device = 'mlu')
        normed_by = "softmax_logit"
        self.assertException("input.dim() >= 2",
                             func, input, topk, num_expert_group, topk_group, normalize, mask, normed_by)
        input = torch.randn(num_batch, num_mask, num_expert, dtype = torch.half, device='mlu')
        self.assertException("topk > 0 && topk <= num_expert",
                             func, input, topk, num_expert_group, topk_group, normalize, mask, normed_by)
        topk = 5
        self.assertException("the dim of mask should be the same as input",
                             func, input, topk, num_expert_group, topk_group, normalize, mask, normed_by)
        mask = torch.randint(0, 2, (1, num_mask, num_expert - 1), dtype = torch.half, device='mlu')
        self.assertException("the last dim of mask should be the same as the last dim of input",
                             func, input, topk, num_expert_group, topk_group, normalize, mask, normed_by)
        mask = torch.randint(0, 2, (1, num_mask - 1, num_expert), dtype = torch.half, device='mlu')
        self.assertException("the penultimate dim of mask should be the same as the penultimate dim of input",
                             func, input, topk, num_expert_group, topk_group, normalize, mask, normed_by)
        mask = torch.randint(0, 2, (2, num_mask, num_expert), dtype = torch.half, device='mlu')
        self.assertException("the product of all but the lower two dimensions of mask is 1",
                             func, input, topk, num_expert_group, topk_group, normalize, mask, normed_by)
        mask = torch.randint(0, 2, (1, num_mask, num_expert), dtype = torch.float, device='mlu')
        self.assertException("the dtype of mask should be the same as input",
                             func, input, topk, num_expert_group, topk_group, normalize, mask, normed_by)
        mask = torch.randint(0, 2, (1, num_mask, num_expert), dtype = torch.half, device='mlu')
        num_expert_group = 8
        self.assertException("if num_expert_group > 1, mask should be None",
                             func, input, topk, num_expert_group, topk_group, normalize, mask, normed_by)
        mask = None
        topk = 160
        self.assertException("topk <= (num_expert / num_expert_group) * topk_group",
                             func, input, topk, num_expert_group, topk_group, normalize, mask, normed_by)

        num_expert_group = 11
        self.assertException("num_expert % num_expert_group == 0",
                             func, input, topk, num_expert_group, topk_group, normalize, mask, normed_by)

        num_expert_group = 8
        topk_group = 9
        self.assertException("topk_group > 0 && topk_group <= num_expert_group",
                             func, input, topk, num_expert_group, topk_group, normalize, mask, normed_by)

    # 单条测例
    def test_single(self):
        num_batch, num_mask, num_expert, topk, num_expert_group, topk_group, normalize, normed_by = 32, 16, 128, 34, 1, 1, True, "softmax_logit"
        input = torch.randn(2, num_batch, num_mask, num_expert, dtype = torch.float32, device='mlu')
        mask = torch.randint(0, 2, (1, 1, num_mask, num_expert), dtype = torch.float32, device='mlu')
#        mask = None
        base_reduce_weight, base_expert_id = self.op_impl_base(input, topk, normalize, num_expert_group, topk_group, mask, normed_by)
        reduce_weight, expert_id = torch.ops.torch_mlu_ops.moe_softmax_topk(input,
                                                                            topk,
                                                                            num_expert_group,
                                                                            topk_group,
                                                                            normalize,
                                                                            mask,
                                                                            normed_by)
        base_expert_id, _ = base_expert_id.sort()
        expert_id, _ = expert_id.sort()
        self.assertTensorsEqual(base_reduce_weight.cpu().float(), reduce_weight.cpu().float(),
                                0.003, use_MSE=True, use_RAE=True)
        self.assertTensorsEqual(base_expert_id.cpu().float(), expert_id.cpu().float(),
                                0.003, use_MSE=True, use_RAE=True)

    def test_inductor(self):
        num_batch, num_mask, num_expert, topk, num_expert_group, topk_group, normalize, normed_by = 32, 16, 128, 34, 4, 3, True, "softmax_logit"
        input = torch.randn(num_batch, num_mask, num_expert, dtype=torch.half, device='mlu')
        mask = torch.randint(0, 2, (1, num_mask, num_expert), dtype = torch.half, device='mlu')
        normed_by = "softmax_logit"
        if num_expert_group > 1:
            mask = None
        args = (input, topk, num_expert_group, topk_group, normalize, mask, normed_by)
        self.base_opcheck(torch.ops.torch_mlu_ops.moe_softmax_topk, args)

if __name__ == '__main__':
    random.seed(0)
    torch.manual_seed(0)
    exit(run_unittest(TestSoftmaxTopkOp))
