import torch
import torch_mlu
import unittest
import torch_mlu_ops as 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 os

def gen_data(num_expert,
             total_tokens,
             inner_size,
             output_stride,
             dtype,
             is_gated,
             has_bias,
             is_ep):
    ci = inner_size * (1 + is_gated)
    input = torch.randn(total_tokens, ci, dtype=dtype, device='mlu')
    cusum_token_count, token_count = generate_token_count(num_expert, total_tokens)
    output = torch.empty((total_tokens, inner_size), dtype=dtype, device='mlu')
    output.as_strided(output.size(), (output_stride, 1))
    start_expert_id = random.randint(0, num_expert - 1) if is_ep else 0
    expert_size = random.randint(1, num_expert - start_expert_id) if is_ep else num_expert
    bias = torch.randn(num_expert, ci, dtype=dtype, device='mlu') if has_bias else None
    return input, bias, token_count, cusum_token_count, output, start_expert_id, expert_size

class TestMoeActiveKernel(BtTestCase):
    def op_impl_base(self, *args):
        input, act_mode, is_gated, output, bias, cusum_token_count, start_expert_id, expert_size = args
        act_fun = torch.nn.functional.gelu if act_mode == 'gelu' else torch.nn.functional.silu
        total_token_num = input.size(0)
        inner_size = input.size(1) // 2 if is_gated else input.size(1)
        input_ = input.clone()
        if bias is not None:
            deal_token_num = cusum_token_count[start_expert_id+expert_size] - cusum_token_count[start_expert_id]
            input_ = input_[:deal_token_num, :]
            token_count = cusum_token_count[1:] - cusum_token_count[:-1]
            token_count_ = token_count[start_expert_id:start_expert_id+expert_size].tolist()
            input_list = list(input_.split(token_count_))
            for i in range(expert_size):
                input_list[i] += bias[i]
            input_ = torch.cat(input_list, dim=0)
            if cusum_token_count.size(0) - 1 != expert_size:
                pad = torch.zeros(total_token_num-deal_token_num, input_.size(-1)).to(input_.dtype).mlu()
                input_ = torch.cat((input_, pad), dim=0)
        acted = act_fun(input_[:, :inner_size])
        acted = acted * input_[:, inner_size:] if is_gated else acted
        if output is None:
            return acted
        else:
            return output.copy_(acted)

    # 功能测试
    def test_functional(self):
        num_expert_list = [5, 32]
        total_tokens_list = [64, 1024]
        inner_size_list = [1024, 8192]
        dtype_list = [torch.half, torch.float32]
        if torch_mlu.mlu.is_bf16_supported():
            dtype_list.append(torch.bfloat16)
        is_gated_list = [True, False]
        is_ep_list = [True, False]
        # change True when test
        has_bias_list = [False, False, True]
        act_mode_list = ['silu', 'gelu']
        args = product(num_expert_list, total_tokens_list, inner_size_list, dtype_list,
                       is_gated_list, is_ep_list, has_bias_list, act_mode_list)

        for num_expert, total_tokens, inner_size, dtype, is_gated, is_ep, has_bias, act_mode in args:
            print("===============================================================================")
            print(f"num_expert: {num_expert}, total_tokens: {total_tokens}")
            print(f"inner_size: {inner_size}, dtype: {dtype}, is_gated: {is_gated}")
            print(f"is_ep: {is_ep}, has_bias: {has_bias}, act_mode: {act_mode}")

            input, bias, token_count, cusum_token_count, output, start_expert_id, expert_size = \
                gen_data(num_expert, total_tokens, inner_size, inner_size, dtype, is_gated, has_bias, is_ep)
            bias_real = bias[start_expert_id:start_expert_id + expert_size] if has_bias else None
            base_output = torch.empty_like(output)
            ops.moe_active(input,
                           act_mode,
                           is_gated,
                           output,
                           bias_real,
                           cusum_token_count.mlu() if has_bias or is_ep else None,
                           start_expert_id,
                           expert_size)
            deal_token_num = cusum_token_count[start_expert_id+expert_size] - cusum_token_count[start_expert_id]
            output = output[:deal_token_num, :]
            self.op_impl_base(input,
                              act_mode,
                              is_gated,
                              base_output,
                              bias_real,
                              cusum_token_count.mlu() if has_bias or is_ep else None,
                              start_expert_id,
                              expert_size)
            base_output = base_output[:deal_token_num, :]
            self.assertTensorsEqual(base_output.cpu().float(), output.cpu().float(),
                                    0.003, use_MSE=True, use_RAE=True)

    def test_inplace(self):
        num_expert, total_tokens, inner_size, dtype, is_gated, is_ep, has_bias, act_mode = \
                        32, 80, 1024, torch.half, True, False, False, 'gelu'
        input, bias, token_count, cusum_token_count, output, start_expert_id, expert_size = \
            gen_data(num_expert, total_tokens, inner_size, inner_size, dtype, is_gated, has_bias, is_ep)
        input_bak = input.clone()
        # test output inpalce with stride
        output = input.as_strided((total_tokens, inner_size), (inner_size * 2, 1))
        base_output = torch.empty_like(output)
        ops.moe_active(input,
                       act_mode,
                       is_gated,
                       output,
                       bias,
                       None,
                       start_expert_id,
                       expert_size)
        self.op_impl_base(input_bak, act_mode, is_gated, base_output, bias, None, start_expert_id, expert_size)
        self.assertTensorsEqual(base_output.cpu().float(), output.cpu().float(),
                                0.003, use_MSE=True, use_RAE=True)

    # 随机遍历测试
    def test_random(self):
        dtype_list = [torch.half, torch.float]
        if torch_mlu.mlu.is_bf16_supported():
            dtype_list.append(torch.bfloat16)

        for i in range(500):
            num_expert = random.randint(1, 64)
            total_tokens = random.randint(1, 32768)
            inner_size = random.randint(1, 8192)
            dtype = random.sample(dtype_list, 1)[0]
            is_gated = random.sample([True, False], 1)[0]
            is_ep = random.sample([True, False], 1)[0]
            # change True when test
            has_bias = random.sample([False, True], 1)[0]
            act_mode = random.sample(['gelu', 'silu'], 1)[0]
            print("===============================================================================")
            print(f"[{i}]: num_expert: {num_expert}, total_tokens: {total_tokens}")
            print(f"       inner_size: {inner_size}, dtype: {dtype}, is_gated: {is_gated}")
            print(f"       is_ep: {is_ep}, has_bias: {has_bias}, act_mode: {act_mode}")

            input, bias, token_count, cusum_token_count, output, start_expert_id, expert_size = \
                gen_data(num_expert, total_tokens, inner_size, inner_size, dtype, is_gated, has_bias, is_ep)
            bias_real = bias[start_expert_id:start_expert_id + expert_size] if has_bias else None
            base_output = torch.empty_like(output)
            ops.moe_active(input,
                           act_mode,
                           is_gated,
                           output,
                           bias_real,
                           cusum_token_count.mlu() if has_bias or is_ep else None,
                           start_expert_id,
                           expert_size)
            deal_token_num = cusum_token_count[start_expert_id+expert_size] - cusum_token_count[start_expert_id]
            output = output[:deal_token_num, :]
            self.op_impl_base(input,
                              act_mode,
                              is_gated,
                              base_output,
                              bias_real,
                              cusum_token_count.mlu() if has_bias or is_ep else None,
                              start_expert_id,
                              expert_size)
            base_output = base_output[:deal_token_num, :]
            self.assertTensorsEqual(base_output.cpu().float(), output.cpu().float(),
                                    0.003, use_MSE=True, use_RAE=True)

    def test_single(self):
        num_expert, total_tokens, inner_size, dtype, is_gated, is_ep, has_bias, act_mode = \
                        32, 64, 8192, torch.float, True, False, False, 'gelu'

        input, bias, token_count, cusum_token_count, output, start_expert_id, expert_size = \
            gen_data(num_expert, total_tokens, inner_size, inner_size, dtype, is_gated, has_bias, is_ep)
        bias_real = bias[start_expert_id:start_expert_id + expert_size] if has_bias else None
        base_output = torch.empty_like(output)
        ops.moe_active(input,
                       act_mode,
                       is_gated,
                       output,
                       bias_real,
                       cusum_token_count.mlu() if has_bias or is_ep else None,
                       start_expert_id,
                       expert_size)
        deal_token_num = cusum_token_count[start_expert_id+expert_size] - cusum_token_count[start_expert_id]
        output = output[:deal_token_num, :]

        self.op_impl_base(input, act_mode, is_gated, base_output, bias_real,
                       cusum_token_count.mlu() if has_bias or is_ep else None,
                       start_expert_id,
                       expert_size)
        base_output = base_output[:deal_token_num, :]
        self.assertTensorsEqual(base_output.cpu().float(), output.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 = ops.moe_active
        num_expert, total_tokens, inner_size, dtype, is_gated, is_ep, has_bias, act_mode = \
                        32, 64, 8192, torch.float, True, False, False, 'gelu'

        input, bias, token_count, cusum_token_count, output, start_expert_id, expert_size = \
            gen_data(num_expert, total_tokens, inner_size, inner_size, dtype, is_gated, has_bias, is_ep)
        bias_real = bias[start_expert_id:start_expert_id + expert_size] if has_bias else None
        act_mode = 'abc'
        self.assertException("act_mode must be 'silu', 'gelu', 'quick_gelu' or 'swish'.", func,
                             input, act_mode, is_gated, output, bias_real,
                             cusum_token_count.mlu() if has_bias or is_ep else None,
                             start_expert_id, expert_size)
        act_mode = 'gelu'
        self.assertException("input.dim() >= 2", func,
                             input.reshape(-1), act_mode, is_gated, output, bias_real,
                             cusum_token_count.mlu() if has_bias or is_ep else None,
                             start_expert_id, expert_size)
        is_gated = True
        self.assertException("in_channel % 2 == 0 if is_gated is true", func,
                             torch.randn(total_tokens, inner_size * 2 - 1, dtype=dtype, device='mlu'),
                             act_mode, is_gated, output, bias_real,
                             cusum_token_count.mlu() if has_bias or is_ep else None,
                             start_expert_id, expert_size)

    def test_inductor(self):
        num_expert, total_tokens, inner_size, dtype, is_gated, is_ep, has_bias, act_mode = \
                        32, 64, 8192, torch.float, True, False, True, 'gelu'

        input, bias, _, cusum_token_count, output, start_expert_id, expert_size = \
            gen_data(num_expert, total_tokens, inner_size, inner_size, dtype, is_gated, has_bias, is_ep)
        bias_real = bias[start_expert_id:start_expert_id + expert_size] if has_bias else None
        args = (input, output, bias_real,
                cusum_token_count.mlu() if has_bias or is_ep else None,
                act_mode, is_gated, start_expert_id, expert_size)
        self.base_opcheck(torch.ops.torch_mlu_ops.active, args)

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