# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import pytest
import torch
from deepspeed.ops import op_builder
from deepspeed.accelerator import get_accelerator

quantize_module = None


def int4x2to2xint4(int4X2tensor):
    high = int4X2tensor >> 4
    low = (int4X2tensor << 4) >> 4
    return torch.stack((high, low), dim=-1).flatten()


def run_quantize(data, num_groups, q_bits, is_symmetric_quant):
    global quantize_module
    if quantize_module is None:
        quantize_module = op_builder.QuantizerBuilder().load()

    return quantize_module.quantize(data, num_groups, q_bits,
                                    quantize_module.Symmetric if is_symmetric_quant else quantize_module.Asymmetric)


def run_dequantize(quantized_data, params, num_groups, q_bits, is_symmetric_quant):
    global quantize_module
    if quantize_module is None:
        quantize_module = op_builder.QuantizerBuilder().load()

    return quantize_module.dequantize(quantized_data, params, num_groups, q_bits,
                                      quantize_module.Symmetric if is_symmetric_quant else quantize_module.Asymmetric)


def run_ref_dequantize(quantized_data, params, num_groups, q_bits, is_symmetric_quant):

    if (q_bits == 4):
        quantized_data = int4x2to2xint4(quantized_data)

    quantized_data = quantized_data.reshape(num_groups, -1).to(torch.float32)

    if is_symmetric_quant:
        return (quantized_data * params).to(torch.float16)
    else:
        scales = params[:, 0].reshape(-1, 1)
        offsets = params[:, 1].reshape(-1, 1)
        return (quantized_data * scales + offsets).to(torch.float16)


@pytest.mark.inference_ops
@pytest.mark.parametrize("num_groups", [1, 13, 512])
@pytest.mark.parametrize("num_elems", [8, 16, 32, 64, 128, 256, 4096, 8192, 12288, 16384])
@pytest.mark.parametrize("is_symmetric_quant", [True, False])
@pytest.mark.parametrize("q_bits", [4, 8])
def test_dequantize(num_elems, num_groups, is_symmetric_quant, q_bits):

    activations = torch.randn((num_groups, num_elems), dtype=torch.float16, device=get_accelerator().device_name())
    quantized_data, params = run_quantize(activations, num_groups, q_bits, is_symmetric_quant)

    ds_dequant = run_dequantize(quantized_data, params, num_groups, q_bits, is_symmetric_quant)
    ref_dequant = run_ref_dequantize(quantized_data, params, num_groups, q_bits, is_symmetric_quant)

    assert (torch.allclose(ds_dequant.flatten(), ref_dequant.flatten(), rtol=3e-2, atol=2e-3))
