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_mix_w4w8_param(bc, seq, k, n, experts_num, topk, data_type, has_bias, quant_wise):
    bs = bc * 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
    token_count = torch.bincount(sorted_expert_id, minlength=experts_num).to(torch.int32)
    quant_group = k // quant_wise
    quant_flag = random.choices([4,8], k=experts_num * quant_group)
    b_count = (sum(quant_flag) // 4) * (quant_wise // 2) * n
    b = torch.randint(-128, 127, (b_count,), dtype=torch.int32, device="mlu").to(torch.int8)
    b_scale = torch.normal(0, 0.01, (quant_group, experts_num, n), device="mlu", dtype=torch.float32)
    a = torch.randint(-128, 127, (bs, k), device="mlu", dtype=torch.int32).to(torch.int8)
    a = a[gather_idx]
    a_scale = torch.normal(0, 0.01, (bs,), device="mlu", dtype=torch.float32)
    a_scale = a_scale[gather_idx]
    c = torch.randn(token_topk, n, device="mlu", dtype=data_type)
    bias = torch.randn(experts_num, n, device="mlu", dtype=data_type) if has_bias else None
    return a, b, token_count, None, c, None, None, a_scale, b_scale, data_type, bs, bias, quant_flag

def gen_tensor(batch, seq, k, n, experts_num, topk, data_type, idx_mode, has_bias=False, quant_bit = 8, quant_group = 1):
    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)

    a, a_scale = QuantByRow(a, 8)
    if idx_mode:
        a_scale = a_scale[gather_idx]
    b_shape = b.shape
    b, b_scale = QuantByRow(b.view(-1, b.shape[-1]), quant_bit, quant_group)
    b = b.view(b_shape)
    if quant_bit == 4:
        b = PairlyPackInt8(b)
    b_scale = b_scale.view(experts_num, -1) if quant_group == 1 else b_scale.view(experts_num, -1, quant_group).permute(2, 0, 1).contiguous()
    alpha = None
    beta = None
    bias = torch.randn(experts_num, n, device="mlu", dtype=data_type) if has_bias else None
    gather_idx_ = gather_idx if idx_mode else None
    quant_flag = None
    return a, b, token_count, gather_idx_, c, alpha, beta, a_scale, b_scale, data_type, bs, bias, quant_flag

class TestSmoothQuantGroupGemmOp(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'])
            a_scale = create_tensor_from_dic(dic['a_scale'], 0, 0.01)
            b_scale = create_tensor_from_dic(dic['b_scale'], 0, 0.01)
            dtype = dic['dtype']['data']
            max_m = dic['max_m']['data']
            bias = None if dic['bias']['data'] is None else create_tensor_from_dic(dic['bias'])
            quant_flag = dic['quant_flag']['data']
            self.launch(a, b, m_list, expand_idx, c, alpha, beta, a_scale, b_scale, dtype, max_m, bias, quant_flag)

    def launch(self, *args):
        total_m = args[2].sum().item()
        torch_out = self.op_impl_base(*args)
        tmo_out = ops.smooth_quant_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, a_scale, b_scale, dtype, max_m, bias, quant_flag = 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))

        if a_scale is not None:
            a_scale_list = a_scale[:total_m].split(tuple(m_list))

        k = a.shape[1]
        n = b.size(1) if quant_flag is None else b_scale.shape[2]
        if b_scale is not None and b_scale.dim() == 3:  # for quant_grouped
            b_scale = b_scale.transpose(0, 1).contiguous()
            if quant_flag is not None:
                quant_group = b_scale.shape[1]
                group_wise = k // quant_group
                quant_flag = torch.tensor(quant_flag).view(-1, quant_group)
                b_offset_cu = torch.cumsum(quant_flag.sum(dim=1), dim=0) // 4 * (group_wise // 2) * n
                b_offset_cu = torch.nn.functional.pad(b_offset_cu, (1,0), "constant", 0)

        output_list = []
        experts = b.size(0) if quant_flag is None else b_scale.size(0)
        for i in range(experts):
            if (a_list[i].size(0) > 0):
                if a_scale is not None and b_scale is not None:
                    if quant_flag is None:
                        gemm_out = smooth_quant_matmul(a_list[i], a_scale_list[i], b[i], b_scale[i], dtype)
                    else:
                        gemm_out = smooth_quant_matmul_w4w8_mixed(a_list[i], a_scale_list[i],
                                                    b[b_offset_cu[i]:b_offset_cu[i+1]],
                                                    b_scale[i], dtype, quant_flag = quant_flag[i])
                else:
                    gemm_out = F.linear(a_list[i], b[i])
                if bias is not None:
                    gemm_out += bias[i]
                if alpha is not None:
                    gemm_out *= alpha[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], n, device=real_res.device).to(real_res.dtype)
        output[:total_m] = real_res
        return output

    def test_smooth_quant_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_tensor(batch, seq, k, n, experts_num, topk, data_type, idx_mode, has_bias)
            torch_out = self.op_impl_base(*param)
            tmo_out = ops.smooth_quant_group_gemm(*param)
            self.assertTensorsEqual(tmo_out.cpu().float(), torch_out.cpu().float(), 0.006, use_MSE=True)


    def test_sq_group_gemm_quant_group(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]
        quant_bit_list = [4, 8]
        quant_group_size_list = [128, 256]
        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, quant_bit_list, quant_group_size_list)
        for batch, seq, k, n, experts_num, topk, data_type, idx_mode, has_bias, quant_bit, quant_group_size 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}, quant_bit: {quant_bit}, quant_group_size: {quant_group_size} testing...", flush=True)
            quant_group = k // quant_group_size
            param = gen_tensor(batch, seq, k, n, experts_num, topk, data_type, idx_mode, has_bias, quant_bit, quant_group)
            torch_out = self.op_impl_base(*param)
            tmo_out = ops.smooth_quant_group_gemm(*param)
            self.assertTensorsEqual(tmo_out.cpu().float(), torch_out.cpu().float(), 0.006, use_MSE=True)

    def test_sq_group_gemm_w4w8_mixed(self):
        bs_l = [1, 3]
        seq_l = [5, 8]
        k_l = [1024, 2048, 3072]
        n_l = [512, 768, 2048]
        expert_l = [8, 32]
        topk_l = [2, 5]
        dtype_l = [torch.half, torch.bfloat16] if mlu.is_bf16_supported() else [torch.half]
        has_bias_l = [True, False]
        group_wise_l = [128, 256, 512]

        args = product(bs_l, seq_l, k_l, n_l, expert_l, topk_l, dtype_l, has_bias_l, group_wise_l)
        for bc, seq, k, n, experts, topk, data_type, has_bias, group_wise in args:
            print(f"bs: {bc}, seq_len: {seq}, k: {k}, n: {n}, experts: {experts}, \
topk: {topk}, dtype: {data_type}, has_bias: {has_bias}, group_wise: {group_wise}, testing...", flush=True)
            param = gen_mix_w4w8_param(bc, seq, k, n, experts, topk, data_type, has_bias, group_wise)
            torch_out = self.op_impl_base(*param)
            tmo_out = ops.smooth_quant_group_gemm(*param)
            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_tensor(batch, seq, k, n, experts_num, topk, data_type, idx_mode, has_bias)
            new_args = list(args)[:9]
            new_args.extend([args[-2], "half" if data_type == torch.half else "bfloat16", args[-1], None, args[-3]])
            self.base_opcheck(torch.ops.torch_mlu_ops.group_gemm, new_args)

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