import torch
import torch_mlu
import unittest
import torch_mlu_ops as tmo
from common_utils import *
import random


class TestPerTokenSmoothQuantizeOp(BtTestCase):
    def op_impl_base(self, *args):
        x, smooth, zero, m_list = args
        x_shape = x.size()
        scale_shape = x.size()[0:-1]
        if m_list is None:
            smoothed = x * smooth
        else:
            input_list = x.split(tuple(m_list))
            experts = len(input_list)
            result = []
            for i in range(experts):
                result.append(input_list[i] * smooth[i])
            smoothed = torch.concat(result, dim=0)
        output, scale = QuantByRow(smoothed, 8)
        return output.reshape(x_shape), scale.reshape(scale_shape)

    def test_random_case(self):
        torch.manual_seed(0)
        case_list = set()
        while(len(case_list) < 200):
            dtype_list = [torch.half, torch.float]
            if torch_mlu.mlu.is_bf16_supported():
                dtype_list.append(torch.bfloat16)
            dtype = random.choice(dtype_list)
            has_group = random.choice([False, True])
            if has_group:
                experts = random.randint(1, 40)
                m_list = torch.randint(1, 100, (experts,), device="mlu", dtype=torch.int32)
                ci = m_list.sum().item()
            else:
                experts = None
                ci = random.randint(1, 4096)
            co = random.randint(1, 4096)
            case = (experts, ci, co, dtype)
            if case in case_list:
                continue
            else:
                case_list.add(case)
                x = torch.randn(ci, co, device="mlu", dtype=dtype)
                if has_group:
                    scale = torch.randn(experts, co, device="mlu", dtype=torch.float32)
                else:
                    scale = torch.randn(co, device="mlu", dtype=torch.float32)
                print("experts={}, ci={}, co={}, dtype={}, testing...".format(experts, ci, co, dtype), flush=True)
                param = (x, scale, None, m_list if has_group else None)
                tmo_output, tmo_scale = tmo.per_token_smooth_quantize(*param)
                torch_output, torch_scale = self.op_impl_base(*param)
                self.assertTensorsEqual(torch_output.cpu().float(), tmo_output.cpu().float(), 0.01, use_MSE=True, use_RAE=True)
                self.assertTensorsEqual(torch_scale.cpu().float(), tmo_scale.cpu().float(), 0.01, use_MSE=True, use_RAE=True)

    def test_inductor(self):
        m_list = torch.randint(1, 100, (8,), device="mlu", dtype=torch.int32)
        total_m = m_list.sum().item()
        x = torch.randn(total_m, 1024, device="mlu", dtype=torch.half)
        scale = torch.randn(8, 1024, device="mlu", dtype=torch.float32)
        output = torch.empty(x.size(), dtype=torch.int8, device="mlu")
        output_scale = torch.empty(x.size()[:-1], dtype=torch.float32, device="mlu")
        args = (x, scale, output, output_scale, None, m_list, None, None, 'per_token', True)
        self.base_opcheck(torch.ops.torch_mlu_ops.smooth_quant, args)


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