# copyright (c) 2021 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.nn as nn

from paddleseg import utils
from paddleseg.cvlibs import manager


@manager.MODELS.add_component
class PortraitNet(nn.Layer):
    """
    The PortraitNet implementation based on PaddlePaddle.

    The original article refers to
    Song-Hai Zhanga, Xin Donga, Jia Lib, Ruilong Lia, Yong-Liang Yangc
    "PortraitNet: Real-time Portrait Segmentation Network for Mobile Device"
    (https://www.yongliangyang.net/docs/mobilePotrait_c&g19.pdf).

    Args:
        num_classes (int, optional): The unique number of target classes.  Default: 2.
        backbone (Paddle.nn.Layer): Backbone network, currently support MobileNetV2.
        add_edge (bool, optional): Whether output to edge. Default: False
        pretrained (str, optional): The path or url of pretrained model. Default: None
    """

    def __init__(self,
                 num_classes,
                 backbone,
                 min_channel=16,
                 channel_ratio=1.0,
                 add_edge=False,
                 pretrained=None):
        super(PortraitNet, self).__init__()
        self.backbone = backbone
        self.head = PortraitNetHead(num_classes, min_channel, channel_ratio,
                                    add_edge)
        self.pretrained = pretrained
        self.init_weight()

    def forward(self, x):
        feat_list = self.backbone(x)
        logits_list = self.head(feat_list)
        return [logits_list]

    def init_weight(self):
        if self.pretrained is not None:
            utils.load_entire_model(self, self.pretrained)


class PortraitNetHead(nn.Layer):
    def __init__(self,
                 num_classes,
                 min_channel=16,
                 channel_ratio=1.0,
                 add_edge=False):
        super().__init__()
        self.min_channel = min_channel
        self.channel_ratio = channel_ratio
        self.add_edge = add_edge
        self.deconv1 = nn.Conv2DTranspose(
            self.depth(96),
            self.depth(96),
            groups=1,
            kernel_size=4,
            stride=2,
            padding=1,
            bias_attr=False)
        self.deconv2 = nn.Conv2DTranspose(
            self.depth(32),
            self.depth(32),
            groups=1,
            kernel_size=4,
            stride=2,
            padding=1,
            bias_attr=False)
        self.deconv3 = nn.Conv2DTranspose(
            self.depth(24),
            self.depth(24),
            groups=1,
            kernel_size=4,
            stride=2,
            padding=1,
            bias_attr=False)
        self.deconv4 = nn.Conv2DTranspose(
            self.depth(16),
            self.depth(16),
            groups=1,
            kernel_size=4,
            stride=2,
            padding=1,
            bias_attr=False)
        self.deconv5 = nn.Conv2DTranspose(
            self.depth(8),
            self.depth(8),
            groups=1,
            kernel_size=4,
            stride=2,
            padding=1,
            bias_attr=False)

        self.transit1 = ResidualBlock(self.depth(320), self.depth(96))
        self.transit2 = ResidualBlock(self.depth(96), self.depth(32))
        self.transit3 = ResidualBlock(self.depth(32), self.depth(24))
        self.transit4 = ResidualBlock(self.depth(24), self.depth(16))
        self.transit5 = ResidualBlock(self.depth(16), self.depth(8))

        self.pred = nn.Conv2D(
            self.depth(8), num_classes, 3, 1, 1, bias_attr=False)
        if self.add_edge:
            self.edge = nn.Conv2D(
                self.depth(8), num_classes, 3, 1, 1, bias_attr=False)

    def depth(self, channels):
        min_channel = min(channels, self.min_channel)
        return max(min_channel, int(channels * self.channel_ratio))

    def forward(self, feat_list):
        feature_1_4, feature_1_8, feature_1_16, feature_1_32 = feat_list
        up_1_16 = self.deconv1(self.transit1(feature_1_32))
        up_1_8 = self.deconv2(self.transit2(feature_1_16 + up_1_16))
        up_1_4 = self.deconv3(self.transit3(feature_1_8 + up_1_8))
        up_1_2 = self.deconv4(self.transit4(feature_1_4 + up_1_4))
        up_1_1 = self.deconv5(self.transit5(up_1_2))

        pred = self.pred(up_1_1)
        if self.add_edge:
            edge = self.edge(up_1_1)
            return pred, edge
        else:
            return pred


class ConvDw(nn.Layer):
    def __init__(self, inp, oup, kernel, stride):
        super(ConvDw, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2D(
                inp,
                inp,
                kernel,
                stride, (kernel - 1) // 2,
                groups=inp,
                bias_attr=False),
            nn.BatchNorm2D(
                num_features=inp, epsilon=1e-05, momentum=0.1),
            nn.ReLU(),
            nn.Conv2D(
                inp, oup, 1, 1, 0, bias_attr=False),
            nn.BatchNorm2D(
                num_features=oup, epsilon=1e-05, momentum=0.1),
            nn.ReLU(), )

    def forward(self, x):
        return self.conv(x)


class ResidualBlock(nn.Layer):
    def __init__(self, inp, oup, stride=1):
        super(ResidualBlock, self).__init__()

        self.block = nn.Sequential(
            ConvDw(
                inp, oup, 3, stride=stride),
            nn.Conv2D(
                in_channels=oup,
                out_channels=oup,
                kernel_size=3,
                stride=1,
                padding=1,
                groups=oup,
                bias_attr=False),
            nn.BatchNorm2D(
                num_features=oup, epsilon=1e-05, momentum=0.1),
            nn.ReLU(),
            nn.Conv2D(
                in_channels=oup,
                out_channels=oup,
                kernel_size=1,
                stride=1,
                padding=0,
                bias_attr=False),
            nn.BatchNorm2D(
                num_features=oup, epsilon=1e-05, momentum=0.1), )
        if inp == oup:
            self.residual = None
        else:
            self.residual = nn.Sequential(
                nn.Conv2D(
                    in_channels=inp,
                    out_channels=oup,
                    kernel_size=1,
                    stride=1,
                    padding=0,
                    bias_attr=False),
                nn.BatchNorm2D(
                    num_features=oup, epsilon=1e-05, momentum=0.1), )
        self.relu = nn.ReLU()

    def forward(self, x):
        residual = x

        out = self.block(x)
        if self.residual is not None:
            residual = self.residual(x)

        out += residual
        out = self.relu(out)
        return out