# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

from functools import partial

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 import Parameter, Tensor
from megengine.core.tensor import dtype
from megengine.quantization import (
    FakeQuantize,
    MinMaxObserver,
    QConfig,
    QuantMode,
    create_qparams,
)
from megengine.quantization.quantize import (
    disable_fake_quant,
    disable_observer,
    propagate_qconfig,
)

min_max_fakequant_qconfig = QConfig(
    weight_observer=partial(MinMaxObserver, dtype="qint8_narrow"),
    act_observer=partial(MinMaxObserver, dtype="qint8"),
    weight_fake_quant=partial(FakeQuantize, dtype="qint8_narrow"),
    act_fake_quant=partial(FakeQuantize, dtype="qint8"),
)


def gen_inp_scale():
    return 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, qmin, qmax):
    x = x / scale
    x = F.round(x)
    x = F.clip(x, qmin, qmax)
    x = x * scale
    return x


fake_quant_act = partial(fake_quant, qmin=-128, qmax=127)
fake_quant_weight = partial(fake_quant, qmin=-127, qmax=127)
fake_quant_bias = partial(fake_quant, qmin=-(2 ** 31), qmax=2 ** 31 - 1)


def init_qat_net(net):
    if net.with_weight:
        net.weight_observer.min_val[...] = Tensor(min_val[0])
        net.weight_observer.max_val[...] = Tensor(max_val[0])
    if net.with_act:
        net.act_observer.min_val[...] = Tensor(min_val[1])
        net.act_observer.max_val[...] = Tensor(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)
    disable_fake_quant(qat_from_float)

    qat_net = QAT.QuantStub()
    qat_net.eval()
    disable_observer(qat_net)

    propagate_qconfig(qat_net, min_max_fakequant_qconfig)
    init_qat_net(qat_net)

    q_net = Q.QuantStub.from_qat_module(qat_net)
    q_net.eval()

    x = mge.tensor(np.random.normal(size=(3, 3)).astype("float32"))

    normal = normal_net(x)
    qat_without_fakequant = qat_from_float(x)
    fake_quant_normal = fake_quant_act(normal_net(x), act_scale)
    qat = qat_net(x)
    q = q_net(x).numpy() * act_scale
    np.testing.assert_allclose(qat_without_fakequant, normal)
    np.testing.assert_allclose(qat, fake_quant_normal)
    np.testing.assert_allclose(q, fake_quant_normal.numpy())


def test_dequant_stub():
    normal_net = Float.DequantStub()
    normal_net.eval()

    qat_from_float = QAT.DequantStub.from_float_module(normal_net)
    qat_from_float.eval()
    disable_fake_quant(qat_from_float)
    disable_observer(qat_from_float)

    qat_net = QAT.DequantStub()
    qat_net.eval()
    
disable_observer(qat_net)