import torch
import torch_mlu_ops
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
from typing import Tuple, Optional
import os

class TestExpandInput(BtTestCase):
    def run_gen_case(self, dic):
        dump_data = dic.pop('dump_data')
        if dump_data:
            self.launch(*dic.values())
        else:
            input = create_tensor_from_dic(dic['input'])
            gather_idx = dic['gather_idx']['data']
            cusum_token_count = dic['cusum_token_count']['data']
            start_expert_id = dic['start_expert_id']['data']
            expert_size = dic['expert_size']['data']
            self.launch(input, gather_idx, cusum_token_count, start_expert_id, expert_size)

    def launch(self, *args):
        cusum_token_count = args[2]
        start_expert_id = args[3]
        expert_size = args[4]
        real_token_count = None if cusum_token_count is None else cusum_token_count[start_expert_id+expert_size-1] - cusum_token_count[start_expert_id]
        base_out = self.op_impl_base(*args)
        tmo_out = torch_mlu_ops.moe_expand_input(*args)
        self.assertTensorsEqual(base_out[:real_token_count].cpu().float(), tmo_out[:real_token_count].cpu().float(),
                                    0.00, use_MSE=True, use_RAE=True)

    def op_impl_base(self, *args):
        input, gather_idx, cusum_token_count, start_expert_id, expert_size = args
        if cusum_token_count is None:
            return input[gather_idx]
        else:
            idx = gather_idx[cusum_token_count[start_expert_id]:cusum_token_count[start_expert_id + expert_size]]
            return input[idx]

    def get_tensor(self, token_num, hidden_size, expert_num, topk, start_expert_id, expert_size, 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, _ = generate_token_count(expert_num, token_num * topk)
        cusum_token_count = cusum_token_count.to('mlu')
        use_all_experts = expert_num == expert_size
        if use_all_experts:
            cusum_token_count = None
            real_token_count = token_num * topk
        else:
            real_token_count = cusum_token_count[start_expert_id+expert_size-1] - cusum_token_count[start_expert_id]
        return input, gather_idx, cusum_token_count, real_token_count

    def test_kernel_random(self):
        for i in range(100):
            token_num = random.randint(1, 2048)
            hidden_size = random.randint(1, 4096)
            expert_num = random.randint(1, 32)
            topk = random.randint(1, expert_num)
            start_expert_id = random.randint(0, expert_num-1)
            expert_size = random.randint(1, expert_num-start_expert_id)
            dtype_list = [torch.half, torch.float, torch.int8]
            if torch_mlu.mlu.is_bf16_supported():
              dtype_list.append(torch.bfloat16)
            dtype = random.sample(dtype_list, 1)[0]
            print("===============================================================================")
            print(f"[{i}]: token_num: {token_num}, hidden_size: {hidden_size}, expert_num: {expert_num}")
            print(f"       topk: {topk}, start_expert_id: {start_expert_id}, expert_size: {expert_size}, dtype: {dtype}")
            input, gather_idx, cusum_token_count, real_token_count = self.get_tensor(token_num,
              hidden_size, expert_num, topk, start_expert_id, expert_size, dtype)
            base_expand_hidden_states = self.op_impl_base(input, gather_idx, None, 0, 0)
            expand_hidden_states = torch_mlu_ops.moe_expand_input(input, gather_idx)
            self.assertTensorsEqual(base_expand_hidden_states.cpu().float(), expand_hidden_states.cpu().float(),
                                    0.00, use_MSE=True, use_RAE=True)
            del base_expand_hidden_states, expand_hidden_states
            base_expand_hidden_states = self.op_impl_base(input, gather_idx, cusum_token_count, start_expert_id, expert_size)[:real_token_count]
            expand_hidden_states = torch_mlu_ops.moe_expand_input(input, gather_idx, cusum_token_count,
                                                                  start_expert_id, expert_size)[:real_token_count]
            self.assertTensorsEqual(base_expand_hidden_states.cpu().float(), expand_hidden_states.cpu().float(),
                                    0.00, use_MSE=True, use_RAE=True)
            del base_expand_hidden_states, expand_hidden_states

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_prevent due to ASan issues")
    def test_prevent(self):
        func = torch_mlu_ops.moe_expand_input
        token_num, hidden_size, topk, expert_num, start_expert_id, expert_size, dtype = 5, 5, 6, 32, 0, 16, torch.half
        input, gather_idx, cusum_token_count, _ = self.get_tensor(token_num, hidden_size, expert_num, \
          topk, start_expert_id, expert_size, dtype)
        input = torch.randn(token_num, 1, hidden_size).to(dtype).to('mlu')
        self.assertException("input dim must be equal to 2.", func, input, gather_idx)
        input = torch.randn(token_num, hidden_size).to(dtype).to('mlu')
        gather_idx = torch.randint(0, token_num, (token_num * topk, 1)).to(torch.int32).to('mlu')
        self.assertException("gather_idx dim must be equal to 1.", func, input, gather_idx)
        gather_idx = gather_idx.reshape(token_num * topk)
        self.assertException("input tensor must on mlu.", func, input.cpu(), gather_idx)
        self.assertException("gather_idx must on mlu.", func, input, gather_idx.cpu())
        self.assertException("data type of gather_idx must be int32.", func, input, gather_idx.to(torch.int64))
        input = torch.randn(token_num, hidden_size).to(dtype).to('mlu')
        gather_idx = torch.randint(0, token_num, (8,)).to(torch.int32).to('mlu')
        self.assertException("expand_token_num % token_num == 0.", func, input, gather_idx)
        gather_idx = torch.randint(1, token_num, (token_num * topk,)).to(torch.int32).to('mlu')
        self.assertException("cusum_token_count must on mlu.", func, input, gather_idx, cusum_token_count.cpu(),
          start_expert_id, expert_size)
        self.assertException("data type of cusum_token_count must be int32.", func, input, gather_idx, cusum_token_count.to(torch.int64),
          start_expert_id, expert_size)
        start_expert_id = -1
        self.assertException("start_expert_id >=0 && start_expert_id < expert_num.",
                            func, input, gather_idx, cusum_token_count, start_expert_id, expert_size)
        start_expert_id = expert_num
        self.assertException("start_expert_id >=0 && start_expert_id < expert_num.",
                            func, input, gather_idx, cusum_token_count, start_expert_id, expert_size)
        start_expert_id = 16
        expert_size = 17
        self.assertException("start_expert_id + expert_size <= expert_num.",
                            func, input, gather_idx, cusum_token_count, start_expert_id, expert_size)

    def test_inductor(self):
      dtype = torch.float16
      token_num, hidden_size, expert_num, topk, start_expert_id, expert_size = 64, 128, 16, 8, 3, 10
      input, gather_idx, cusum_token_count, _ = self.get_tensor(token_num, hidden_size, expert_num, \
          topk, start_expert_id, expert_size, dtype)
      args = (input, gather_idx, cusum_token_count, start_expert_id, expert_size)
      self.base_opcheck(torch.ops.torch_mlu_ops.moe_expand_input, args)

if __name__ == "__main__":
    random.seed(0)
    torch.manual_seed(0)
    exit(run_unittest(TestExpandInput))
