import io
from functools import partial
from itertools import chain
from typing import Callable

import numpy as np

import megengine as mge
import megengine.functional as F
import megengine.module as M
import megengine.module.qat as QM
import megengine.quantization as Q
from megengine import Tensor
from megengine.module.qat.module import QATModule
from megengine.traced_module import TracedModule, trace_module
from megengine.traced_module.utils import get_subattr


class MyConvBnRelu2d(M.ConvBnRelu2d):
    pass


class MyQATConvBnRelu2d(QM.ConvBnRelu2d):
    pass


class Myblcok(M.Module):
    def __init__(self,):
        super().__init__()
        self.conv0 = MyConvBnRelu2d(3, 3, 3, 1, 1)
        self.conv1 = M.ConvBn2d(3, 3, 1, 1, 0)
        self.conv2 = M.ConvBn2d(3, 3, 1, 1, 0)
        self.add = M.Elemwise("FUSE_ADD_RELU")

    def forward(self, x):
        x = self.conv0(x)
        x0 = self.conv1(x)
        x1 = self.conv2(x)
        o = self.add(x0, x1)
        return o


class MyModule(M.Module):
    def __init__(self):
        super().__init__()
        self.block0 = Myblcok()
        self.block1 = Myblcok()

    def forward(self, x):
        x = self.block0(x)
        x = self.block1(x)
        return x


class MyMinMaxObserver(Q.MinMaxObserver):
    pass


class MyTQT(Q.TQT):
    pass


def get_lsq_config(lsq_cls):
    return Q.QConfig(
        weight_observer=None,
        act_observer=None,
        weight_fake_quant=partial(lsq_cls, dtype="qint8_narrow"),
        act_fake_quant=partial(lsq_cls, dtype="qint8"),
    )


def get_observer_config(observer_cls):
    return Q.QConfig(
        weight_observer=partial(observer_cls, dtype="qint8_narrow"),
        act_observer=partial(observer_cls, dtype="qint8"),
        weight_fake_quant=None,
        act_fake_quant=None,
    )


def get_qparams(mod: QATModule):
    weight_qparams, act_qparams = None, None
    if mod.act_observer is not None:
        act_qparams = mod.act_observer.get_qparams()
    if mod.act_fake_quant:
        act_qparams = mod.act_fake_quant.get_qparams()

    if mod.weight_observer is not None:
        weight_qparams = mod.weight_observer.get_qparams()
    if mod.weight_fake_quant:
        weight_qparams = mod.weight_fake_quant.get_qparams()

    return weight_qparams, act_qparams


def check_qparams(qparmsa: Q.QParams, qparmsb: Q.QParams):
    assert qparmsa.dtype_meta == qparmsb.dtype_meta
    assert qparmsa.mode == qparmsb.mode
    np.testing.assert_equal(qparmsa.scale.numpy(), qparmsb.scale.numpy())
    if qparmsa.zero_point is not None:
        np.testing.assert_equal(qparmsa.zero_point.numpy(), qparmsb.zero_point.numpy())


def build_observered_net(net: M.Module, observer_cls):
    qat_net = Q.quantize_qat(
        net,
        qconfig=get_observer_config(observer_cls),
        mapping={MyConvBnRelu2d: MyQATConvBnRelu2d},
    )
    Q.enable_observer(qat_net)
    inp = Tensor(np.random.random(size=(5, 3, 32, 32)))
    qat_net.eval()
    qat_net(inp)
    Q.disable_observer(qat_net)
    return qat_net


def build_fakequanted_net(net: QATModule, fakequant_cls):
    qat_net = Q.reset_qconfig(net, get_lsq_config(fakequant_cls))
    qat_net.eval()
    return qat_net


def test_trace_qat():
    def _check_qat_module(qat_net: QATModule):
        inp = Tensor(np.random.random(size=(5, 3, 32, 32)))
        traced_net = trace_module(qat_net, inp)

        for name, qat_module in qat_net.named_modules():
            if not isinstance(qat_module, QATModule):
                continue
            traced_qat_module = get_subattr(traced_net, name)
            weight_qparams, act_qparams = get_qparams(qat_module)
            traced_weight_qparams, traced_act_qparams = get_qparams(traced_qat_module)
            if weight_qparams:
                check_qparams(weight_qparams, traced_weight_qparams)
            if act_qparams:
                check_qparams(act_qparams, traced_act_qparams)
        flatten_traced_net = traced_net.flatten()
        conv0_node = flatten_traced_net.graph.get_node_by_name(
            "MyModule_block0_conv0"
        ).as_unique()
        conv0_out_node = flatten_traced_net.graph.get_node_by_name(
            "MyModule_block0_conv0_out"
        ).as_unique()
        assert isinstance(conv0_node.owner, TracedModule)
        assert conv0_out_node.expr.inputs[0] is conv0_node

    _check_qat_module(build_observered_net(MyModule(), Q.MinMaxObserver))
    _check_qat_module(build_observered_net(MyModule(), MyMinMaxObserver))
    _check_qat_module(
        build_fakequanted_net(build_observered_net(MyModule(), Q.MinMaxObserver), Q.TQT)
    )
    _check_qat_module(
        build_fakequanted_net(build_observered_net(MyModule(), Q.MinMaxObserver), MyTQT)
    )


def test_load_param():
    def _check_param(moda: M.Module, modb: M.Module):
        for name, attr in chain(moda.named_parameters(), moda.named_buffers()):
            traced_attr = get_subattr(modb, name)
            np.testing.assert_equal(attr.numpy(), traced_attr.numpy())

    def _check_module(build_func: Callable):
        net = build_func()
        net.eval()
        buffer = io.BytesIO()
        mge.save(net.state_dict(), buffer)
        buffer.seek(0)

        inp = Tensor(np.random.random(size=(5, 3, 32, 32)))
        traced_net = trace_module(build_func(), inp)
        traced_net.load_state_dict(mge.load(buffer))

        _check_param(net, traced_net)

        buffer.seek(0)
        traced_net = trace_module(build_func(), inp).flatten()
        traced_net.load_state_dict(mge.load(buffer))

        _check_param(net, traced_net)

    _check_module(lambda: MyModule())
    _check_module(lambda: build_observered_net(MyModule(), Q.MinMaxObserver))


def test_qualname():
    def _check_qualname(net):
        inp = Tensor(np.random.random(size=(5, 3, 32, 32)))
        net.eval()
        traced_net = trace_module(net, inp)
        base_qualname = traced_net.graph.qualname
        for node in traced_net.graph.nodes():
            qualname = node.qualname
            qualname = qualname[len(base_qualname) + 1 :]
            if qualname.endswith("]"):
                qualname = qualname.rsplit(".", 1)[0]
            if qualname.startswith("["):
                qualname = ""
            traced_attr = get_subattr(traced_net, qualname)
            orig_attr = get_subattr(net, qualname)
            assert traced_attr is not None
            assert orig_attr is not None

    _check_qualname(MyModule())
    _check_qualname(build_observered_net(MyModule(), Q.MinMaxObserver))