# -*- 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)