# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import paddle
import paddle.nn as nn

from arch.base_module import SNConv, SNConvTranspose, ResBlock


class Encoder(nn.Layer):
    def __init__(self, name, in_channels, encode_dim, use_bias, norm_layer,
                 act, act_attr, conv_block_dropout, conv_block_num,
                 conv_block_dilation):
        super(Encoder, self).__init__()
        self._pad2d = paddle.nn.Pad2D([3, 3, 3, 3], mode="replicate")
        self._in_conv = SNConv(
            name=name + "_in_conv",
            in_channels=in_channels,
            out_channels=encode_dim,
            kernel_size=7,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        self._down1 = SNConv(
            name=name + "_down1",
            in_channels=encode_dim,
            out_channels=encode_dim * 2,
            kernel_size=3,
            stride=2,
            padding=1,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        self._down2 = SNConv(
            name=name + "_down2",
            in_channels=encode_dim * 2,
            out_channels=encode_dim * 4,
            kernel_size=3,
            stride=2,
            padding=1,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        self._down3 = SNConv(
            name=name + "_down3",
            in_channels=encode_dim * 4,
            out_channels=encode_dim * 4,
            kernel_size=3,
            stride=2,
            padding=1,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        conv_blocks = []
        for i in range(conv_block_num):
            conv_blocks.append(
                ResBlock(
                    name="{}_conv_block_{}".format(name, i),
                    channels=encode_dim * 4,
                    norm_layer=norm_layer,
                    use_dropout=conv_block_dropout,
                    use_dilation=conv_block_dilation,
                    use_bias=use_bias))
        self._conv_blocks = nn.Sequential(*conv_blocks)

    def forward(self, x):
        out_dict = dict()
        x = self._pad2d(x)
        out_dict["in_conv"] = self._in_conv.forward(x)
        out_dict["down1"] = self._down1.forward(out_dict["in_conv"])
        out_dict["down2"] = self._down2.forward(out_dict["down1"])
        out_dict["down3"] = self._down3.forward(out_dict["down2"])
        out_dict["res_blocks"] = self._conv_blocks.forward(out_dict["down3"])
        return out_dict


class EncoderUnet(nn.Layer):
    def __init__(self, name, in_channels, encode_dim, use_bias, norm_layer,
                 act, act_attr):
        super(EncoderUnet, self).__init__()
        self._pad2d = paddle.nn.Pad2D([3, 3, 3, 3], mode="replicate")
        self._in_conv = SNConv(
            name=name + "_in_conv",
            in_channels=in_channels,
            out_channels=encode_dim,
            kernel_size=7,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        self._down1 = SNConv(
            name=name + "_down1",
            in_channels=encode_dim,
            out_channels=encode_dim * 2,
            kernel_size=3,
            stride=2,
            padding=1,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        self._down2 = SNConv(
            name=name + "_down2",
            in_channels=encode_dim * 2,
            out_channels=encode_dim * 2,
            kernel_size=3,
            stride=2,
            padding=1,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        self._down3 = SNConv(
            name=name + "_down3",
            in_channels=encode_dim * 2,
            out_channels=encode_dim * 2,
            kernel_size=3,
            stride=2,
            padding=1,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        self._down4 = SNConv(
            name=name + "_down4",
            in_channels=encode_dim * 2,
            out_channels=encode_dim * 2,
            kernel_size=3,
            stride=2,
            padding=1,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        self._up1 = SNConvTranspose(
            name=name + "_up1",
            in_channels=encode_dim * 2,
            out_channels=encode_dim * 2,
            kernel_size=3,
            stride=2,
            padding=1,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        self._up2 = SNConvTranspose(
            name=name + "_up2",
            in_channels=encode_dim * 4,
            out_channels=encode_dim * 4,
            kernel_size=3,
            stride=2,
            padding=1,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)

    def forward(self, x):
        output_dict = dict()
        x = self._pad2d(x)
        output_dict['in_conv'] = self._in_conv.forward(x)
        output_dict['down1'] = self._down1.forward(output_dict['in_conv'])
        output_dict['down2'] = self._down2.forward(output_dict['down1'])
        output_dict['down3'] = self._down3.forward(output_dict['down2'])
        output_dict['down4'] = self._down4.forward(output_dict['down3'])
        output_dict['up1'] = self._up1.forward(output_dict['down4'])
        output_dict['up2'] = self._up2.forward(
            paddle.concat(
                (output_dict['down3'], output_dict['up1']), axis=1))
        output_dict['concat'] = paddle.concat(
            (output_dict['down2'], output_dict['up2']), axis=1)
        return output_dict