# 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.
import pickle
from collections import defaultdict
from itertools import chain

import numpy as np

import megengine.functional as F
import megengine.module as M
import megengine.module.qat as qat
from megengine.module.identity import Identity
from megengine.traced_module import trace_module
from megengine.traced_module.expr import CallFunction, CallMethod, Expr, GetAttr, Input
from megengine.traced_module.node import ModuleNode, Node, TensorNode


class IdentityMod(M.Module):
    def forward(self, x):
        return x


class MyBlock(M.Module):
    def __init__(self, in_channels=3, channels=3):
        super(MyBlock, self).__init__()
        self.conv1 = M.Conv2d(in_channels, channels, 3, 1, padding=1, bias=False)
        self.bn1 = M.BatchNorm2d(channels)
        self.nothing = IdentityMod()

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = F.relu(x) + 1
        x = self.nothing(x)
        return x


class MyModule(M.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.block0 = MyBlock()
        self.block1 = MyBlock()
        self.nothing = IdentityMod()

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


class MyBlock1(M.Module):
    def forward(self, a):
        y = F.concat([a, a])
        return a, y


class MyModule1(M.Module):
    def __init__(self):
        super().__init__()
        self.block0 = MyBlock1()
        self.block1 = MyBlock1()

    def forward(self, a):
        a, y1 = self.block0(a)
        a = a + 1
        a, y2 = self.block1(a)
        return a, y1 + y2


class NewModule(M.Module):
    def __init__(self, traced_module):
        super(NewModule, self).__init__()
        self.module = traced_module

    def forward(self, x):
        x = x - 1
        x = self.module(x)
        x = x + 1
        return x


def _check_expr_users(traced_module):
    node_user = defaultdict(list)
    for expr in traced_module.graph._exprs:
        for node in expr.inputs:
            node_user[node].append(expr)
    for node in traced_module.graph.nodes():
        node.users.sort(key=lambda m: m._id)
        node_user[node].sort(key=lambda m: m._id)
        assert node.users == node_user[node]


def _init_cls(cls):
    module = cls()
    x = 
F.ones((1, 3, 3, 3))