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

def gen_args(index: int = 0):
    token_num = random.randint(1, 32768)
    expert_num = random.randint(1, 256)
    topk = random.randint(1, expert_num)
    expert_id = torch.randint(low=0, high=expert_num, size=(token_num, topk)).to(torch.int32).to('mlu')
    print("===============================================================================")
    print(f"[{index}]: token_num: {token_num}, expert_num: {expert_num}, topk: {topk}")
    return expert_id, expert_num

class TestGenIdx(BtTestCase):
    def op_impl_base(self, *args):
        expert_id, expert_num = args
        token_num, topk = expert_id.size(0), expert_id.size(1)
        sorted_expert_id, indices = expert_id.int().flatten().sort()
        expand_idx_out = indices.int() // topk
        combine_idx_out = torch.zeros((token_num * topk,), dtype=torch.int, device="mlu")
        combine_idx_out.scatter_(0, indices, torch.arange(token_num * topk, dtype=torch.int, device="mlu"))
        token_count_out = torch.bincount(sorted_expert_id, minlength=expert_num)
        cusum_token_count_out = torch.cat((torch.tensor([0]).to('mlu'), torch.cumsum(token_count_out, dim=0))).to(torch.int32)
        return tuple([expand_idx_out, combine_idx_out, token_count_out, cusum_token_count_out])

    def test_kernel_random(self):
        for i in range(1500):
            expert_id, expert_num = gen_args(i)
            base_gather_expand_idx, base_gather_combine_idx, base_token_count, base_cusum_token_count = self.op_impl_base(expert_id, expert_num)
            gather_expand_idx_out, gather_combine_idx_out, token_count_out, cusum_token_count_out = \
              torch_mlu_ops.moe_gen_idx(expert_id, expert_num)
            self.assertTensorsEqual(base_gather_expand_idx.cpu(), gather_expand_idx_out.cpu(), 0.00, use_MSE=True, use_RAE=True)
            self.assertTensorsEqual(base_gather_combine_idx.cpu(), gather_combine_idx_out.cpu(), 0.00, use_MSE=True, use_RAE=True)
            self.assertTensorsEqual(base_token_count.cpu(), token_count_out.cpu(), 0.00, use_MSE=True, use_RAE=True)
            self.assertTensorsEqual(base_cusum_token_count.cpu(), cusum_token_count_out.cpu(), 0.00, use_MSE=True, use_RAE=True)

    def test_inductor(self):
        args = gen_args()
        self.base_opcheck(torch.ops.torch_mlu_ops.moe_gen_idx, args)

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_prevent due to ASan issues")
    def test_prevent(self):
        func = torch_mlu_ops.moe_gen_idx
        token_num, expert_num, topk = 128, 32, 16
        expert_id = torch.randint(low=0, high=expert_num, size=(token_num, topk, 1)).to(torch.int32).to('mlu')
        self.assertException("expert_id dim must be equal to 2.", func, expert_id, expert_num)
        expert_id = expert_id.reshape(token_num, topk)
        self.assertException("data type of expert_id must be int32.", func, expert_id.to(torch.int64), expert_num)

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