# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2020 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 megengine as mge
import megengine.functional as F
import megengine.hub as hub
import megengine.module as M
import math
import official.vision.classification.resnet.model as resnet

import numpy as np


class ResnetBody(M.Module):
    def __init__(
        self,
        block,
        init_channel,
        layers,
        channels,
        zero_init_residual=False,
        norm=M.BatchNorm2d,
    ):
        super(ResnetBody, self).__init__()
        self.in_channels = init_channel
        self.layer1 = self._make_layer(
            block, channels[0], layers[0], stride=1, norm=norm
        )

        self.layer2 = self._make_layer(
            block, channels[1], layers[1], stride=2, norm=norm
        )

        self.layer3 = self._make_layer(
            block, channels[2], layers[2], stride=2, norm=norm,
        )

        self.layer4 = self._make_layer(
            block, channels[3], layers[3], stride=2, norm=norm,
        )

        for m in self.modules():
            if isinstance(m, M.Conv2d):
                M.init.msra_normal_(m.weight, mode="fan_out", nonlinearity="relu")
                if m.bias is not None:
                    fan_in, _ = M.init.calculate_fan_in_and_fan_out(m.weight)
                    bound = 1 / math.sqrt(fan_in)
                    M.init.uniform_(m.bias, -bound, bound)
            elif isinstance(m, M.BatchNorm2d):
                M.init.ones_(m.weight)
                M.init.zeros_(m.bias)
            elif isinstance(m, M.Linear):
                M.init.msra_uniform_(m.weight, a=math.sqrt(5))
                if m.bias is not None:
                    fan_in, _ = M.init.calculate_fan_in_and_fan_out(m.weight)
                    bound = 1 / math.sqrt(fan_in)
                    M.init.uniform_(m.bias, -bound, bound)

    def _make_layer(self, block, channels, blocks, stride=1, norm=M.BatchNorm2d):
        layers = []
        layers.append(block(self.in_channels, channels, stride, norm=norm))
        self.in_channels = channels * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.in_channels, channels, norm=norm))

        return M.Sequential(*layers)

    def forward(self, x):
        outputs = []

        x = self.layer1(x)
        outputs.append(x)
        x = self.layer2(x)
        outputs.append(x)
        x = self.layer3(x)
        outputs.append(x)
        x = self.layer4(x)
        outputs.append(x)

        return outputs


class SingleStage(M.Module):
    def __init__(
        self, block, init_channel, layers, channels, mid_channel, norm=M.BatchNorm2d
    ):
        super(SingleStage, self).__init__()
        self.down = ResnetBody(block, init_channel, layers, channels, norm)
        channel = block.expansion * channels[-1]
        self.up1 = M.Sequential(
            M.Conv2d(channel, mid_channel, 1, 1, 0), norm(mid_channel)
        )
        self.deconv1 = M.Sequential(
            M.ConvTranspose2d(mid_channel, mid_channel, 4, 2, 1), norm(mid_channel)
        )

        channel = block.expansion * channels[-2]
        self.up2 = M.Sequential(
            M.Conv2d(channel, mid_channel, 1, 1, 0), norm(mid_channel)
        )
        self.deconv2 = M.Sequential(
            
M.ConvTranspose2d(mid_channel, mid_channel, 4, 2, 1)