import torch
from torch_mlu import mlu
import unittest
import torch_mlu_ops as ops
from common_utils import *
from itertools import product
from torch.nn import functional as F

def gen_params(batch, seq, k, n, experts_num, topk, data_type, idx_mode, has_bias=False):
    bs = batch * seq
    token_topk = bs * topk
    expert_id = torch.randint(experts_num, (token_topk,), device="mlu")
    sorted_expert_id, indices = expert_id.sort()
    gather_idx = indices // topk
    gather_idx = gather_idx.to(torch.int32)
    token_count = torch.bincount(sorted_expert_id, minlength=experts_num).to(torch.int32)

    a = torch.randn(bs, k, device="mlu", dtype=data_type)
    if not idx_mode:
        a = a[gather_idx]
    b = torch.randn(experts_num, n, k, device="mlu", dtype=data_type)
    c = torch.randn(token_topk, n, device="mlu", dtype=data_type)
    alpha = torch.randn(experts_num, device="mlu", dtype=torch.float32)
    beta = torch.randn(experts_num, device="mlu", dtype=torch.float32)
    a_scale = None
    b_scale = None
    bias = None
    if has_bias:
        bias = torch.randn(experts_num, n, device="mlu", dtype=data_type)

    gather_idx_ = gather_idx if idx_mode else None
    return a, b, token_count, gather_idx_, c, alpha, beta, a_scale, b_scale, data_type, bias


class TestGroupGemmOp(BtTestCase):
    def run_gen_case(self, dic):
        dump_data = dic.pop('dump_data')
        if dump_data:
            self.launch(*dic.values())
        else:
            a = create_tensor_from_dic(dic['a'])
            b = create_tensor_from_dic(dic['b'])
            m_list = dic['m_list']['data']
            expand_idx = dic['expand_idx']['data']
            c = None if dic['c']['data'] is None else create_tensor_from_dic(dic['c'])
            alpha = None if dic['alpha']['data'] is None else create_tensor_from_dic(dic['alpha'])
            beta = None if dic['beta']['data'] is None else create_tensor_from_dic(dic['beta'])
            max_m = dic['max_m']['data']
            bias = None if dic['bias']['data'] is None else create_tensor_from_dic(dic['bias'])
            self.launch(a, b, m_list, expand_idx, c, alpha, beta, max_m, bias)

    def launch(self, *args):
        total_m = args[2].sum().item()
        torch_out = self.op_impl_base(*args)
        tmo_out = ops.group_gemm(*args)
        self.assertTensorsEqual(tmo_out.cpu().float()[0:total_m], torch_out.cpu().float()[0:total_m], 0.006, use_MSE=True)

    def op_impl_base(self, *args):
        a, b, m_list, expand_idx, c, alpha, beta, max_m, bias = args
        a = a.reshape(-1, a.size(-1))
        if expand_idx is not None:
            a = a[expand_idx]
        total_m = m_list.sum().item()
        a_list = a[:total_m].split(tuple(m_list))

        c_list = []
        if c is not None:
            c = c.reshape(-1, c.size(-1))
            c_list = c[:total_m].split(tuple(m_list))

        output_list = []
        for i in range(b.size(0)): # alpha*(a*b) + bias + beta*c
            if (a_list[i].size(0) > 0):
                gemm_out = torch.matmul(a_list[i], b[i].permute(1, 0))
                if alpha is not None:
                    gemm_out *= alpha[i]
                if bias is not None:
                    gemm_out += bias[i]
                if beta is not None and c_list != []:
                    gemm_out += c_list[i] * beta[i]
                output_list.append(gemm_out)
        real_res = torch.cat(output_list, dim=0)
        output = torch.empty(a.shape[0], b.shape[1], device=real_res.device).to(real_res.dtype)
        output[:total_m] = real_res
        return output

    def test_group_gemm(self):
        bs_list = [1, 3]
        seq_list = [5, 8]
        k_list = [512, 1024]
        n_list = [512, 768, 2048]
        expert_list = [8, 32]
        topk_list = [2, 5]
        dtype_list = [torch.half, torch.bfloat16] if mlu.is_bf16_supported() else [torch.half]
        idx_list =  [False] if mlu.get_device_name() == 'MLU370' else [False, True]
        has_bias_list = [True, False]

        args = product(bs_list, seq_list, k_list, n_list, expert_list, topk_list, dtype_list, idx_list, has_bias_list)
        for batch, seq, k, n, experts_num, topk, data_type, idx_mode, has_bias in args:
            print(f"bs: {batch}, seq_len: {seq}, k: {k}, n: {n}, experts_num: {experts_num}, topk: {topk}, \
dtype: {data_type}, idx_mode: {idx_mode}, has_bias: {has_bias} testing...", flush=True)
            param = gen_params(batch, seq, k, n, experts_num, topk, data_type, idx_mode, has_bias)
            torch_out = self.op_impl_base(*param[:7], batch * seq, param[10])
            tmo_out = ops.group_gemm(*param[:7],  batch * seq, param[10])
            self.assertTensorsEqual(tmo_out.cpu().float(), torch_out.cpu().float(), 0.006, use_MSE=True)

    def test_inductor(self):
        batch, seq, k, n, experts_num, topk = 1, 8, 1024, 2048, 8, 5
        dtype_list = [torch.half, torch.bfloat16] if mlu.is_bf16_supported() else [torch.half]
        idx_list =  [False] if mlu.get_device_name() == 'MLU370' else [False, True]
        has_bias_list = [True, False]
        args = product( dtype_list, idx_list, has_bias_list)
        for data_type, idx_mode, has_bias in args:
            args = gen_params(batch, seq, k, n, experts_num, topk, data_type, idx_mode, has_bias)
            args = list(args)
            args[-2] = args[-1] # bias
            args[-1] = None #dtype
            args.extend([None, None, batch * seq]) #b_offset, max_m
            self.base_opcheck(torch.ops.torch_mlu_ops.group_gemm, args)

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