import numpy as np
import pytest

import megengine as mge
import megengine.functional as F
import megengine.module as Float
import megengine.module.qat as QAT
import megengine.module.quantized as Q
from megengine.core.tensor import dtype
from megengine.quantization import min_max_fakequant_qconfig
from megengine.quantization.quantize import (
    disable_fake_quant,
    disable_observer,
    propagate_qconfig,
)

"""
Calculate testing scales based on ``min_max_fakequant_qconfig``
"""

inp_scale = np.float32(np.random.rand() + 1)

min_val = np.random.randint(-127, 0, size=(2,)).astype("float32")
max_val = np.random.randint(1, 127, size=(2,)).astype("float32")
weight_scale = np.float32(np.max([-min_val[0], max_val[0]]) / 254 * 2)
act_scale = np.float32(np.max([-min_val[1], max_val[1]]) / 255 * 2)


def quant(x, scale):
    inp_dtype = dtype.qint8(scale)
    return x.astype(inp_dtype)


def fake_quant(x, scale):
    x = x / scale
    x = F.round(x)
    x = F.clip(x, -128, 127)
    x = x * scale
    return x


def init_qat_net(net):
    if net.with_weight:
        net.weight_observer.min_val.set_value(min_val[0])
        net.weight_observer.max_val.set_value(max_val[0])
    if net.with_act:
        net.act_observer.min_val.set_value(min_val[1])
        net.act_observer.max_val.set_value(max_val[1])


def test_quant_stub():
    normal_net = Float.QuantStub()
    normal_net.eval()

    qat_from_float = QAT.QuantStub.from_float_module(normal_net)
    qat_from_float.eval()
    
disable_observer(qat_from_float)