# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# 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 numpy as np

import paddle
import paddle.nn as nn
import paddle.nn.functional as F

from paddleseg.models import layers
from paddleseg.cvlibs import manager
from paddleseg.utils import utils


@manager.MODELS.add_component
class PFPNNet(nn.Layer):
    """
    The Panoptic Feature Pyramid Networks implementation based on PaddlePaddle.

    The original article refers to
    Alexander Kirillov, Ross Girshick, Kaiming He, Piotr Dollár, et al. "Panoptic Feature Pyramid Networks"
    (https://arxiv.org/abs/1901.02446)

    Args:
        num_classes (int): The unique number of target classes.
        backbone (Paddle.nn.Layer): Backbone network, currently support Resnet50/101.
        backbone_indices (tuple): Four values in the tuple indicate the indices of output of backbone.
        enable_auxiliary_loss (bool, optional): A bool value indicates whether adding auxiliary loss. Default: False.
        align_corners (bool, optional): An argument of F.interpolate. It should be set to False when the feature size is even,
            e.g. 1024x512, otherwise it is True, e.g. 769x769. Default: False.
        pretrained (str, optional): The path or url of pretrained model. Default: None.
    """

    def __init__(self,
                 num_classes,
                 backbone,
                 backbone_indices,
                 channels,
                 enable_auxiliary_loss=False,
                 align_corners=False,
                 dropout_ratio=0.1,
                 fpn_inplanes=[256, 512, 1024, 2048],
                 pretrained=None):
        super(PFPNNet, self).__init__()
        self.backbone = backbone
        self.backbone_indices = backbone_indices
        self.in_channels = [
            self.backbone.feat_channels[i] for i in backbone_indices
        ]
        self.align_corners = align_corners
        self.pretrained = pretrained
        self.enable_auxiliary_loss = enable_auxiliary_loss

        self.head = PFPNHead(
            num_class=num_classes,
            fpn_inplanes=fpn_inplanes,
            dropout_ratio=dropout_ratio,
            channels=channels,
            fpn_dim=channels,
            enable_auxiliary_loss=self.enable_auxiliary_loss)
        self.init_weight()

    def forward(self, x):
        feats = self.backbone(x)
        feats = [feats[i] for i in self.backbone_indices]
        logit_list = self.head(feats)
        return [
            F.interpolate(
                logit,
                paddle.shape(x)[2:],
                mode='bilinear',
                align_corners=self.align_corners) for logit in logit_list
        ]

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


class PFPNHead(nn.Layer):
    """
    The PFPNHead implementation.

    Args:
        inplane (int): Input channels of PPM module.
        num_class (int): The unique number of target classes.
        fpn_inplanes (list): The feature channels from backbone.
        fpn_dim (int, optional): The input channels of FPN module. Default: 512.
        enable_auxiliary_loss (bool, optional): A bool value indicates whether adding auxiliary loss. Default: False.
    """

    def __init__(self,
                 num_class,
                 fpn_inplanes,
                 channels,
                 dropout_ratio=0.1,
                 fpn_dim=256,
                 enable_auxiliary_loss=False,
                 align_corners=False):
        super(PFPNHead, self).__init__()
        self.enable_auxiliary_loss = enable_auxiliary_loss
        self.align_corners = align_corners
        self.lateral_convs = nn.LayerList()
        self.fpn_out = nn.LayerList()

        for fpn_inplane in fpn_inplanes:
            self.lateral_convs.append(
                nn.Sequential(
                    nn.Conv2D(fpn_inplane, fpn_dim, 1),
                    layers.SyncBatchNorm(fpn_dim), nn.ReLU()))
            self.fpn_out.append(
                nn.Sequential(
                    layers.ConvBNReLU(
                        fpn_dim, fpn_dim, 3, bias_attr=False)))

        self.scale_heads = nn.LayerList()
        for index in range(len(fpn_inplanes)):
            head_length = max(
                1,
                int(np.log2(fpn_inplanes[index]) - np.log2(fpn_inplanes[0])))
            scale_head = nn.LayerList()
            for head_index in range(head_length):
                scale_head.append(
                    layers.ConvBNReLU(
                        fpn_dim,
                        channels,
                        3,
                        padding=1, ))
                if fpn_inplanes[index] != fpn_inplanes[0]:
                    scale_head.append(
                        nn.Upsample(
                            scale_factor=2,
                            mode='bilinear',
                            align_corners=align_corners))
            self.scale_heads.append(nn.Sequential(*scale_head))

        if dropout_ratio:
            self.dropout = nn.Dropout2D(dropout_ratio)
            if self.enable_auxiliary_loss:
                self.dsn = nn.Sequential(
                    layers.ConvBNReLU(
                        fpn_inplanes[2], fpn_inplanes[2], 3, padding=1),
                    nn.Dropout2D(dropout_ratio),
                    nn.Conv2D(
                        fpn_inplanes[2], num_class, kernel_size=1))
        else:
            self.dropout = None
            if self.enable_auxiliary_loss:
                self.dsn = nn.Sequential(
                    layers.ConvBNReLU(
                        fpn_inplanes[2], fpn_inplanes[2], 3, padding=1),
                    nn.Conv2D(
                        fpn_inplanes[2], num_class, kernel_size=1))

        self.conv_last = nn.Sequential(
            layers.ConvBNReLU(
                len(fpn_inplanes) * fpn_dim, fpn_dim, 3, bias_attr=False),
            nn.Conv2D(
                fpn_dim, num_class, kernel_size=1))
        self.conv_seg = nn.Conv2D(channels, num_class, kernel_size=1)

    def cls_seg(self, feat):
        if self.dropout is not None:
            feat = self.dropout(feat)
        output = self.conv_seg(feat)
        return output

    def forward(self, conv_out):
        last_out = self.lateral_convs[-1](conv_out[-1])
        f = last_out
        fpn_feature_list = [last_out]
        for i in reversed(range(len(conv_out) - 1)):
            conv_x = conv_out[i]
            conv_x = self.lateral_convs[i](conv_x)
            prev_shape = paddle.shape(conv_x)[2:]
            f = conv_x + F.interpolate(
                f, prev_shape, mode='bilinear', align_corners=True)
            fpn_feature_list.append(self.fpn_out[i](f))

        output_size = paddle.shape(fpn_feature_list[-1])[2:]

        x = self.scale_heads[0](fpn_feature_list[-1])
        for index in range(len(self.scale_heads) - 2, 0, -1):
            x = x + F.interpolate(
                self.scale_heads[index](fpn_feature_list[index]),
                size=output_size,
                mode='bilinear',
                align_corners=self.align_corners)
        x = self.cls_seg(x)
        if self.enable_auxiliary_loss:
            dsn = self.dsn(conv_out[2])
            return [x, dsn]
        else:
            return [x]
