import torch
import torch_mlu
import unittest
import torch_mlu_ops as ops
from typing import Union, List, Tuple
from torch.nn.parameter import Parameter
from torch.nn import functional as F
import torch.multiprocessing as mp
import math
import sys
import os

work_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(work_dir))
from common_utils import *

def moe_split_inner_size(rank, *args):
    torch.manual_seed(0)
    batch, seq, hidden_size, inner_size, expert_num, topk, gated, renormalize, act_mode, dtype, block_n, world_size = args
    assert inner_size % world_size == 0, f"inner_size{inner_size} must be divisible by tp_num{world_size}"
    block_k = inner_size // world_size
    start_k = rank * block_k
    end_k = start_k + block_k

    setup(rank, world_size)
    pg = get_default_group()
    cncl_comm = pg._get_backend(torch.device("mlu")).get_cncl_comm(rank)

    scale_s = 0.01 # avoid the occurrence of inf
    hidden_states = torch.randn(batch, seq, hidden_size, device="mlu", dtype=dtype)
    router_logit = torch.randn(batch, seq, expert_num, device="mlu", dtype=torch.float32)
    residual = None # torch.randn(batch, seq, hidden_size, device="mlu", dtype=dtype)
    weight1 = torch.randn(expert_num, inner_size*(1+gated), hidden_size, device="mlu", dtype=dtype) * scale_s
    bias1 = None # torch.randn(expert_num, inner_size*(1+gated), device="mlu", dtype=dtype)
    weight2 = torch.randn(expert_num, hidden_size, inner_size, device="mlu", dtype=dtype) * scale_s
    bias2 = None # torch.randn(expert_num, hidden_size, device="mlu", dtype=dtype)

    weight1 = weight1.view(expert_num, -1, inner_size, hidden_size)
    weight1 = weight1[:, :, start_k:end_k]
    weight1 = weight1.reshape(expert_num, -1, hidden_size).contiguous()
    weight2 = weight2[..., start_k:end_k].contiguous()
    if bias1 is not None:
        bias1 = bias1.view(expert_num, -1, inner_size)
        bias1 = bias1[..., start_k:end_k]
        bias1 = bias1.reshape(expert_num, -1).contiguous()
    residual = residual if (residual is not None and rank == 0) else None

    param = [hidden_states, router_logit, weight1, weight2, bias1, bias2, residual,
             None, None, None, None, topk, renormalize, gated, act_mode]
    output = ops.fused_moe(*param)
    all_reduce(output, ReduceOp.SUM, group=pg)

    param = [hidden_states, router_logit, weight1, weight2, bias1, bias2, residual,
             None, None, None, None, topk, renormalize, gated, act_mode, 0, block_n, cncl_comm]
    output1 = ops.fused_moe(*param)

    new_inner_size = weight2.shape[-1]
    block_e = 4096 // new_inner_size
    if block_e * new_inner_size == 4096 and expert_num % block_e == 0:
        param = [hidden_states, router_logit, weight1, weight2.view(-1, block_e, hidden_size, new_inner_size).transpose(1,2).contiguous(),
                    bias1, bias2, residual,
                    None, None, None, None, topk, renormalize, gated, act_mode, 0, block_n, cncl_comm]
        output2 = ops.fused_moe(*param)
    if rank == 0:
        assertTensorsEqual(output.cpu().float(), output1.cpu().float(), 0.006, use_MSE=True, use_RAE=True)
        if block_e * new_inner_size == 4096 and expert_num % block_e == 0:
            assertTensorsEqual(output1.cpu().float(), output2.cpu().float(), 0.006, use_MSE=True, use_RAE=True)
    cleanup()

def sq_moe_split_inner_size(rank, *args):
    torch.manual_seed(0)
    batch, seq, hidden_size, inner_size, expert_num, topk, gated, renormalize, act_mode, dtype, block_n, world_size = args
    assert inner_size % world_size == 0, f"inner_size{inner_size} must be divisible by tp_num{world_size}"
    block_k = inner_size // world_size
    start_k = rank * block_k
    end_k = start_k + block_k

    setup(rank, world_size)
    pg = get_default_group()
    cncl_comm = pg._get_backend(torch.device("mlu")).get_cncl_comm(rank)

    scale_s = 0.1 # avoid the occurrence of inf
    eps = 0.1       # Avoid the occurrence of nan
    residual = None # torch.randn(batch, seq, hidden_size, device="mlu", dtype=dtype)
    bias1, bias2 = None, None
    hidden_states = torch.normal(0, 0.1, size=(batch, seq, hidden_size), device="mlu", dtype=dtype)
    weight1 = torch.normal(0, 0.1, size=(expert_num, inner_size*(1+gated), hidden_size), device="mlu", dtype=dtype)
    weight2 = torch.normal(0, 0.01, size=(expert_num, hidden_size, inner_size), device="mlu", dtype=dtype)
    router_logit = torch.normal(0, 0.1, size=(batch, seq, expert_num), device="mlu", dtype=torch.float32)
    input_smooth = torch.normal(0, 0.1, size=(expert_num, hidden_size), device="mlu", dtype=torch.float32).abs() + eps
    act_smooth = torch.normal(0, 0.1, size=(expert_num, inner_size), device="mlu", dtype=torch.float32).abs() + eps
    weight1_shape, weight2_shape = weight1.shape, weight2.shape
    weight1 = weight1 / input_smooth.unsqueeze(1)
    weight2 = weight2 / act_smooth.unsqueeze(1)
    quant_w1, w1_scale = QuantByRow(weight1.view(-1, weight1.shape[-1]), 8)
    quant_w2, w2_scale = QuantByRow(weight2.view(-1, weight2.shape[-1]), 8)
    quant_w1, quant_w2 = quant_w1.view(weight1_shape), quant_w2.view(weight2_shape)
    w1_scale, w2_scale = w1_scale.view(expert_num, -1), w2_scale.view(expert_num, -1)
    quant_w1 = quant_w1.view(expert_num, -1, inner_size, hidden_size)
    quant_w1 = quant_w1[:, :, start_k:end_k]
    quant_w1 = quant_w1.reshape(expert_num, -1, hidden_size).contiguous()
    quant_w2 = quant_w2[..., start_k:end_k].contiguous()

    if bias1 is not None:
        bias1 = bias1.view(expert_num, -1, inner_size)
        bias1 = bias1[..., start_k:end_k]
        bias1 = bias1.reshape(expert_num, -1).contiguous()
    if w1_scale is not None:
        w1_scale = w1_scale.view(expert_num, -1, inner_size)
        w1_scale = w1_scale[..., start_k:end_k]
        w1_scale = w1_scale.reshape(expert_num, -1).contiguous()
    if act_smooth is not None:
        act_smooth = act_smooth[..., start_k:end_k].contiguous()
    residual = residual if (residual is not None and rank == 0) else None

    param = [hidden_states, router_logit, quant_w1, quant_w2, bias1, bias2, residual,
             input_smooth, act_smooth, w1_scale, w2_scale, topk, renormalize, gated, act_mode]
    output = ops.fused_moe(*param)
    all_reduce(output, ReduceOp.SUM, group=pg)

    param = [hidden_states, router_logit, quant_w1, quant_w2, bias1, bias2, residual,
             input_smooth, act_smooth, w1_scale, w2_scale, topk, renormalize, gated, act_mode, 0,
             block_n, cncl_comm]
    output1 = ops.fused_moe(*param)

    new_inner_size = quant_w2.shape[-1]
    block_e = 4096 // new_inner_size
    if block_e * new_inner_size == 4096 and expert_num % block_e == 0:
        param = [hidden_states, router_logit, quant_w1,
                  quant_w2.view(-1, block_e, hidden_size, new_inner_size).transpose(1,2).contiguous(),
                  bias1, bias2, residual, input_smooth, act_smooth, w1_scale, w2_scale, topk,
                  renormalize, gated, act_mode, 0, block_n, cncl_comm]
        output2 = ops.fused_moe(*param)
    if rank == 0:
        assertTensorsEqual(output.cpu().float(), output1.cpu().float(), 0.006, use_MSE=True, use_RAE=True)
        if block_e * new_inner_size == 4096 and expert_num % block_e == 0:
            assertTensorsEqual(output1.cpu().float(), output2.cpu().float(), 0.006, use_MSE=True, use_RAE=True)
    cleanup()


class TestFusedMOEAllReduceOp(BtTestCase):
    def op_impl_base(self, *args):
        return super().op_impl_base(*args)

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_single_fused_moe_allreduce due to ASan issues")
    def test_single_fused_moe_allreduce(self):
        print("test_single_fused_moe_allreduce")
        device_n = min(torch_mlu.mlu.device_count(), 8)
        if device_n < 2:
            print(f"device count is {device_n}, can not test communication between devices.", flush=True)
        block_n = 2

        batch, seq, hidden_size, inner_size = 1, 1024, 8192, 8192
        expert_num, topk, gated, renormalize, act_mode, dtype = 8, 2, True, True, 'silu', torch.float16
        print(f"bs: {batch}, seq_len: {seq}, hidden_size: {hidden_size}, inner_size: {inner_size}, experts_num: {expert_num}, \
topk: {topk}, gated: {gated}, renormalize: {renormalize}, dtype: {dtype}, act_mode: {act_mode}, \
tp_num: {device_n}, block_n: {block_n}, testing...", flush=True)
        param = [batch, seq, hidden_size, inner_size, expert_num, topk, gated, renormalize, act_mode, dtype, block_n, device_n]
        mp.spawn(moe_split_inner_size, param, nprocs=device_n, join=True)

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_random_fused_moe_allreduce due to ASan issues")
    def test_random_fused_moe_allreduce(self):
        print("test_random_fused_moe_allreduce")
        import random
        random.seed(0)
        device_n = min(torch_mlu.mlu.device_count(), 8)
        if device_n < 2:
            print(f"device count is {device_n}, can not test communication between devices.", flush=True)

        act_mode = 'gelu'
        case_list = set()
        dtype_list = [torch.half, torch.bfloat16] if torch_mlu.mlu.is_bf16_supported() else [torch.half]
        while(len(case_list) < 10):
            block_n = random.randint(-5, 5)
            tp_num = random.randint(1, device_n)
            batch = random.randint(1, 10)
            seq = random.randint(1, 10)
            hidden_size = random.randrange(512, 1024, 2)
            if block_n != 0:
                if hidden_size // abs(block_n) < 256:
                    continue
                hidden_size = (hidden_size // block_n) * block_n
            else:
                hidden_size = 1024 * random.randint(1, 10)
            k_start = (512 // tp_num) * tp_num
            inner_size = random.randrange(k_start, 1024, tp_num * 2)
            expert_num = random.randint(1, 32)
            topk = random.randint(1,expert_num)
            gated = random.choice([True, False])
            renormalize = random.choice([True, False])
            dtype = random.choice(dtype_list)
            case = (batch, seq, hidden_size, inner_size, expert_num, topk, gated, renormalize, act_mode, dtype, tp_num, block_n)
            if case in case_list:
                continue
            case_list.add(case)
            print(f"random bs: {batch}, seq_len: {seq}, hidden_size: {hidden_size}, inner_size: {inner_size}, experts_num: {expert_num}, \
topk: {topk}, gated: {gated}, renormalize: {renormalize}, dtype: {dtype}, act_mode: {act_mode}, \
tp_num: {tp_num}, block_n: {block_n}, testing...", flush=True)
            param = [batch, seq, hidden_size, inner_size, expert_num, topk, gated, renormalize, act_mode, dtype, block_n, tp_num]
            mp.spawn(moe_split_inner_size, param, nprocs=tp_num, join=True)

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_sq_fused_moe_allreduce due to ASan issues")
    def test_sq_fused_moe_allreduce(self):
        print("test_sq_fused_moe_allreduce")
        device_n = min(torch_mlu.mlu.device_count(), 8)
        if device_n < 2:
            print(f"device count is {device_n}, can not test communication between devices.", flush=True)
        block_n = 2
        batch, seq, hidden_size, inner_size = 5, 9, 8192, 8192
        expert_num, topk, gated, renormalize, act_mode, dtype = 20, 16, False, True, 'gelu', torch.float16
        print(f"sq bs: {batch}, seq_len: {seq}, hidden_size: {hidden_size}, inner_size: {inner_size}, experts_num: {expert_num}, \
topk: {topk}, gated: {gated}, renormalize: {renormalize}, dtype: {dtype}, act_mode: {act_mode}, \
tp_num: {device_n}, block_n: {block_n}, testing...", flush=True)
        param = [batch, seq, hidden_size, inner_size, expert_num, topk, gated, renormalize, act_mode, dtype, block_n, device_n]
        mp.spawn(sq_moe_split_inner_size, param, nprocs=device_n, join=True)

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_random_sq_fused_moe_allreduce due to ASan issues")
    def test_random_sq_fused_moe_allreduce(self):
        print("test_random_sq_fused_moe_allreduce")
        import random
        random.seed(0)
        device_n = min(torch_mlu.mlu.device_count(), 8)
        if device_n < 2:
            print(f"device count is {device_n}, can not test communication between devices.", flush=True)
        act_mode = 'gelu'
        case_list = set()
        while(len(case_list) < 10):
            block_n = random.randint(-5, 5)
            tp_num = random.randint(1, device_n)
            batch = random.randint(1, 10)
            seq = random.randint(1, 10)
            hidden_size = random.randrange(256, 512, 2)
            if block_n != 0:
                if hidden_size // abs(block_n) < 256:
                    continue
                hidden_size = (hidden_size // block_n) * block_n
            else:
                hidden_size = 1024 * random.randint(1, 10)
            k_start = (512 // tp_num) * tp_num
            inner_size = random.randrange(k_start, 1024, tp_num * 2)
            expert_num = random.randint(1, 32)
            topk = random.randint(1, expert_num)
            gated = random.choice([True, False])
            renormalize = random.choice([True, False])
            dtype = random.choice([ torch.float16])
            if torch_mlu.mlu.get_device_name() == 'MLU370':
                dtype = torch.float16
            case = (batch, seq, hidden_size, inner_size, expert_num, topk, gated, renormalize, act_mode, dtype, tp_num, block_n)
            if case in case_list:
                continue
            case_list.add(case)
            print(f"random sq bs: {batch}, seq_len: {seq}, hidden_size: {hidden_size}, inner_size: {inner_size}, experts_num: {expert_num}, \
topk: {topk}, gated: {gated}, renormalize: {renormalize}, dtype: {dtype}, act_mode: {act_mode}, \
tp_num: {tp_num}, block_n: {block_n}, testing...", flush=True)
            param = [batch, seq, hidden_size, inner_size, expert_num, topk, gated, renormalize, act_mode, dtype, block_n, tp_num]
            mp.spawn(sq_moe_split_inner_size, param, nprocs=tp_num, join=True)

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_random_sq_fused_moe_with_4D_w2_allreduce due to ASan issues")
    def test_random_sq_fused_moe_with_4D_w2_allreduce(self):
        print("test_random_sq_fused_moe_with_4D_w2_allreduce")
        import random
        random.seed(0)
        device_n = min(torch_mlu.mlu.device_count(), 8)
        if device_n < 2:
            print(f"device count is {device_n}, can not test communication between devices.", flush=True)
        act_mode = 'gelu'
        case_list = set()
        while (len(case_list) < 10):
            block_n = random.randint(-5, 5)
            tp_num = random.randint(1, device_n)
            batch = random.randint(1, 10)
            seq = random.randint(1, 10)
            hidden_size = random.randrange(512, 4096, 2)
            if block_n != 0:
                if hidden_size // abs(block_n) < 256:
                    continue
                hidden_size = (hidden_size // block_n) * block_n
            else:
                hidden_size = 1024 * random.randint(1, 10)
            inner_size_per_tp = random.choice([256, 512])
            inner_size = inner_size_per_tp * tp_num
            expert_num_base = random.randint(1, 4)
            expert_num_factor = 4096 // inner_size_per_tp
            expert_num = expert_num_base * expert_num_factor
            topk = random.randint(1, expert_num)
            gated = random.choice([True, False])
            renormalize = random.choice([True, False])
            dtype = random.choice([ torch.float16])
            if torch_mlu.mlu.get_device_name() == 'MLU370':
                dtype = torch.float16
            case = (batch, seq, hidden_size, inner_size, expert_num, topk, gated, renormalize, act_mode, dtype, tp_num, block_n)
            if case in case_list:
                continue
            case_list.add(case)
            print(f"random sq bs: {batch}, seq_len: {seq}, hidden_size: {hidden_size}, inner_size: {inner_size}, experts_num: {expert_num}, \
topk: {topk}, gated: {gated}, renormalize: {renormalize}, dtype: {dtype}, act_mode: {act_mode}, \
tp_num: {tp_num}, block_n: {block_n}, testing...", flush=True)
            param = [batch, seq, hidden_size, inner_size, expert_num, topk, gated, renormalize, act_mode, dtype, block_n, tp_num]
            mp.spawn(sq_moe_split_inner_size, param, nprocs=tp_num, join=True)

    def test_inductor(self):
        return super().test_inductor()

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