# -*- 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.functional as F
import megengine.module as M

import official.vision.classification.resnet.model as resnet


class ASPP(M.Module):
    def __init__(self, in_channels, out_channels, dr=1):
        super().__init__()

        self.conv1 = M.Sequential(
            M.Conv2d(
                in_channels, out_channels, 1, 1, padding=0, dilation=dr, bias=False
            ),
            M.BatchNorm2d(out_channels),
            M.ReLU(),
        )
        self.conv2 = M.Sequential(
            M.Conv2d(
                in_channels,
                out_channels,
                3,
                1,
                padding=6 * dr,
                dilation=6 * dr,
                bias=False,
            ),
            M.BatchNorm2d(out_channels),
            M.ReLU(),
        )
        self.conv3 = M.Sequential(
            M.Conv2d(
                in_channels,
                out_channels,
                3,
                1,
                padding=12 * dr,
                dilation=12 * dr,
                bias=False,
            ),
            M.BatchNorm2d(out_channels),
            M.ReLU(),
        )
        self.conv4 = M.Sequential(
            M.Conv2d(
                in_channels,
                out_channels,
                3,
                1,
                padding=18 * dr,
                dilation=18 * dr,
                bias=False,
            ),
            M.BatchNorm2d(out_channels),
            M.ReLU(),
        )
        self.conv_gp = M.Sequential(
            M.Conv2d(in_channels, out_channels, 1, 1, 0, bias=False),
            M.BatchNorm2d(out_channels),
            M.ReLU(),
        )
        self.conv_out = M.Sequential(
            M.Conv2d(out_channels * 5, out_channels, 1, 1, padding=0, bias=False),
            M.BatchNorm2d(out_channels),
            M.ReLU(),
        )

    def forward(self, x):
        conv1 = self.conv1(x)
        conv31 = self.conv2(x)
        conv32 = self.conv3(x)
        conv33 = self.conv4(x)

        gp = 
F.mean(x, [2, 3], True)