import torch
import unittest
import torch_mlu_ops as ops
import random
from common_utils import *
from itertools import product
import copy
from typing import Optional

def generate_token_count(num_expert,
                         total_token_count):
    token_count = torch.randint(low=1, high=1024, size=(num_expert, ), \
                                dtype=torch.int32).to(dtype=torch.float32)
    sum = torch.sum(token_count, dim=-1) * 1.0
    token_count *= total_token_count / sum.item()
    token_count = token_count.to(dtype=torch.int32)
    cusum_token_count = torch.zeros(num_expert + 1, dtype=torch.int32)
    end_expert_token_count = torch.ones(num_expert + 1, dtype=torch.int32) * total_token_count
    cusum_token_count[1:] = torch.cumsum(token_count, dim=-1)
    cusum_token_count = torch.minimum(cusum_token_count, end_expert_token_count)
    cusum_token_count[-1] = total_token_count
    return cusum_token_count


def gen_case(num_tokens,
             topk,
             hidden_size,
             num_expert,
             expert_size,
             has_bias,
             has_residual,
             dtype,
             device):
    input = torch.randn((num_tokens * topk, hidden_size), dtype=dtype, device=device)
    reduce_weight = torch.randn((num_tokens, topk), dtype=torch.float32, device=device)
    gather_ids = torch.randperm(num_tokens * topk, dtype=torch.int32, device=device)
    bias = None
    residual = None

    cusum_token_count = None

    if has_bias:
        bias = torch.randn((num_expert, hidden_size), dtype=dtype, device=device)
    if has_residual:
        residual = torch.randn((num_tokens, hidden_size), dtype=dtype, device=device)

    if has_bias or expert_size < num_expert:
        cusum_token_count = generate_token_count(num_expert, num_tokens * topk)
        cusum_token_count = cusum_token_count.to(device=device)
    return input, reduce_weight, gather_ids, residual, bias, cusum_token_count

class TestCombineResult(BtTestCase):
    def op_impl_base(self, *args):
        input, reduce_weight, gather_ids, residual, cusum_token_count, start_expert_id, \
            expert_size, bias = args
        input = input.to(dtype=torch.float32).cpu()
        reduce_weight = reduce_weight.cpu()
        gather_ids = gather_ids.cpu()
        bias = None

        dtype = input.dtype
        num_expert = expert_size
        hidden_size = input.shape[1]

        if bias is not None:
            bias = bias.to(dtype=torch.float32).cpu()
            num_expert = bias.shape[0]
        if cusum_token_count is not None:
            num_expert = cusum_token_count.shape[0] - 1
            cusum_token_count = cusum_token_count.cpu()

        if cusum_token_count is not None and expert_size < num_expert:
            gathered_input = input[gather_ids - cusum_token_count[start_expert_id].item()]
        else:
            gathered_input = input[gather_ids]

        if bias is not None and cusum_token_count is not None:
            for i in range(start_expert_id, start_expert_id + expert_size):
                gathered_input[cusum_token_count[i] : cusum_token_count[i+1]] += bias[i]
        gathered_input = gathered_input.reshape(*reduce_weight.shape, hidden_size)

        if cusum_token_count is not None:
            filtered_ids = (gather_ids >= cusum_token_count[start_expert_id]) * \
                (gather_ids < cusum_token_count[start_expert_id + expert_size])
            filtered_ids = filtered_ids.to(dtype=torch.float32)
            reduce_weight = reduce_weight * filtered_ids.reshape(reduce_weight.shape)

        gathered_input *= reduce_weight.reshape(*reduce_weight.shape, -1)
        output = torch.sum(gathered_input, dim=1, keepdim=False)
        if residual is not None:
            residual = residual.to(dtype=torch.float32).cpu()
            output += residual
        return output.to(dtype=dtype)

    def test_random_case(self):
        torch.manual_seed(444)
        test_cases = 200
        num_tokens_list = torch.randint(low=1, high=2048, size=(test_cases, ), dtype=torch.int32)
        topk_list = torch.randint(low=1, high=33, size=(test_cases, ), dtype=torch.int32)
        hidden_size_list = torch.randint(low=256, high=8193, size=(test_cases, ), dtype=torch.int32)
        num_expert_list = torch.randint(low=1, high=129, size=(test_cases, ), dtype=torch.int32)
        num_expert_list = torch.maximum(topk_list, num_expert_list)
        expert_size_list = torch.randint(low=1, high=129, size=(test_cases, ), dtype=torch.int32)
        expert_size_list = torch.minimum(expert_size_list, num_expert_list)
        start_expert_id_list = torch.randint(low=0, high=129, size=(test_cases, ), dtype=torch.int32)
        start_expert_id_list = torch.minimum(start_expert_id_list, num_expert_list - expert_size_list)
        start_expert_id_list = torch.maximum(start_expert_id_list, torch.zeros(test_cases, dtype=torch.int32))
        has_bias_list = torch.randint(low=0, high=2, size=(test_cases, ), dtype=torch.int32)
        has_residual_list = torch.randint(low=0, high=2, size=(test_cases, ), dtype=torch.int32)
        dtype_list = torch.randint(low=0, high=10, size=(test_cases, ), dtype=torch.int32)
        dtype_map = [torch.half, torch.bfloat16, torch.float]

        device = 'mlu'
        mlu_name = torch.mlu.get_device_name()
        is_mlu370 = "MLU3" in mlu_name
        max_num_tokens = 128 * 1024
        for i in range(test_cases):
            num_tokens = num_tokens_list[i].item()
            topk = topk_list[i].item()
            hidden_size = hidden_size_list[i].item()
            num_expert = num_expert_list[i].item()
            expert_size = expert_size_list[i].item()
            start_expert_id = start_expert_id_list[i].item()
            has_bias = False  # has_bias_list[i].item()
            has_residual = has_residual_list[i].item()

            topk = min(topk, (int)(max_num_tokens / num_tokens))

            if dtype_list[i].item() < 4:
                dtype = dtype_map[0]
            elif  dtype_list[i].item() < 9:
                dtype = dtype_map[1]
            else:
                dtype = dtype_map[2]

            if is_mlu370 and dtype is torch.bfloat16:
                continue

            inputs = gen_case(num_tokens,
                              topk,
                              hidden_size,
                              num_expert,
                              expert_size,
                              has_bias,
                              has_residual,
                              dtype,
                              device)
            input = inputs[0]
            reduce_weight = inputs[1]
            gather_ids = inputs[2]
            residual = inputs[3]
            bias = inputs[4]
            cusum_token_count = inputs[5]

            print("num_tokens={}, topk={}, hidden_size={}, num_expert={}, expert_size={}, "
                  "start_expert_id={}, has_bias={}, has_residual={}, dtype={}, testing...".format(
                      num_tokens, topk, hidden_size, num_expert, expert_size, start_expert_id, \
                      has_bias, has_residual, dtype))

            golden_output = self.op_impl_base(input, reduce_weight, gather_ids, residual,
                                                   cusum_token_count, start_expert_id, expert_size, None)

            output = ops.moe_combine_result(input, reduce_weight, gather_ids, residual,
                                            cusum_token_count, start_expert_id, expert_size)

            self.assertTensorsEqual(output.cpu().float(), golden_output.cpu().float(), 0.003,
                                    "golden_output must equal output", True, True, True, True)


    def test_perf_case(self):
        num_tokens_list = [1, 72, 512]
        hidden_size_list = [2048, 4096, 5120, 8192]
        # [num_expert, topk, start_expert_id, expert_size]
        expert_options_list = [[1, 1, 0, 1], [8, 2, 0, 8], [32, 5, 0, 32], [32, 5, 24, 8]]
        has_residual_list = [True, False]
        dtype_list = [torch.half, torch.bfloat16]

        device = 'mlu'
        mlu_name = torch.mlu.get_device_name()
        is_mlu370 = "MLU3" in mlu_name
        args = product(num_tokens_list, hidden_size_list, expert_options_list,\
                       has_residual_list, dtype_list)
        for num_tokens, hidden_size, expert_options, has_residual, dtype in args:
            num_expert = expert_options[0]
            topk = expert_options[1]
            start_expert_id = expert_options[2]
            expert_size = expert_options[3]
            has_bias = False

            if is_mlu370 and dtype is torch.bfloat16:
                continue

            torch.manual_seed(444)
            inputs = gen_case(num_tokens,
                              topk,
                              hidden_size,
                              num_expert,
                              expert_size,
                              has_bias,
                              has_residual,
                              dtype,
                              device)
            input = inputs[0]
            reduce_weight = inputs[1]
            gather_ids = inputs[2]
            residual = inputs[3]
            bias = inputs[4]
            cusum_token_count = inputs[5]

            golden_output = self.op_impl_base(input, reduce_weight, gather_ids, residual,
                                              cusum_token_count, start_expert_id, expert_size, None)

            notify_start = torch.mlu.Event(enable_timing=True)
            notify_end = torch.mlu.Event(enable_timing=True)
            notify_start.record()
            loop = 10
            for _ in range(loop):
                output = ops.moe_combine_result(input, reduce_weight, gather_ids, residual,
                                            cusum_token_count, start_expert_id, expert_size)
            notify_end.record()
            notify_end.synchronize()
            time = notify_start.hardware_time(notify_end) / loop

            print("num_tokens={}, topk={}, hidden_size={}, num_expert={}, expert_size={}, "
                  "start_expert_id={}, has_bias={}, has_residual={}, dtype={}, time={:.1f}".format(
                      num_tokens, topk, hidden_size, num_expert, expert_size, start_expert_id, \
                      has_bias, has_residual, dtype, time))

            self.assertTensorsEqual(output.cpu().float(), golden_output.cpu().float(), 0.003,
                                    "golden_output must equal output", True, True, True, True)

    def test_inductor(self):
        num_tokens, hidden_size, has_bias, has_residual, dtype = 1, 2048, False, True, torch.float16
        num_expert, topk, start_expert_id, expert_size = 8, 2, 0, 8
        input, reduce_weight, gather_ids, residual, \
            bias, cusum_token_count = gen_case(num_tokens, topk, hidden_size, num_expert, expert_size,
                                               has_bias, has_residual, dtype, 'mlu')
        args = (input, reduce_weight, gather_ids, residual, cusum_token_count, start_expert_id, expert_size, bias)
        self.base_opcheck(torch.ops.torch_mlu_ops.moe_combine_result, args)

if __name__ == '__main__':
    exit(run_unittest(TestCombineResult))
