# 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 paddle
import paddle.nn as nn
import paddle.nn.functional as F

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


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

    The original article refers to
    Yu, Changqian, et al. "BiSeNet V2: Bilateral Network with Guided Aggregation for Real-time Semantic Segmentation"
    (https://paperswithcode.com/paper/bisenet-bilateral-segmentation-network-for)

    Args:
        num_classes (int): The unique number of target classes.
        backbone (paddle.nn.Layer): Backbone network, currently support Resnet18_vd/Resnet34_vd/Resnet50_vd/Resnet101_vd.
        pretrained (str, optional): The path or url of pretrained model. Default: None.
    """

    def __init__(self, num_classes, backbone, conv_channel=128,
                 pretrained=None):
        super().__init__()
        self.backbone = backbone
        self.spatial_path = SpatialPath(3, 128)
        self.global_context = nn.Sequential(
            nn.AdaptiveAvgPool2D(1),
            layers.ConvBNReLU(
                512, conv_channel, 1, bias_attr=False), )

        self.arms = nn.LayerList([
            AttentionRefinement(512, conv_channel),
            AttentionRefinement(256, conv_channel),
        ])
        self.refines = nn.LayerList([
            layers.ConvBNReLU(
                conv_channel,
                conv_channel,
                3,
                stride=1,
                padding=1,
                bias_attr=False),
            layers.ConvBNReLU(
                conv_channel,
                conv_channel,
                3,
                stride=1,
                padding=1,
                bias_attr=False),
        ])

        self.heads = nn.LayerList([
            BiSeNetHead(conv_channel, num_classes, 8, True),
            BiSeNetHead(conv_channel, num_classes, 8, True),
            BiSeNetHead(conv_channel * 2, num_classes, 8, False),
        ])

        self.ffm = FeatureFusion(conv_channel * 2, conv_channel * 2, 1)

        self.pretrained = pretrained

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

    def forward(self, x):
        spatial_out = self.spatial_path(x)
        context_blocks = self.backbone(x)
        context_blocks.reverse()

        global_context = self.global_context(context_blocks[0])
        global_context = F.interpolate(
            global_context,
            size=paddle.shape(context_blocks[0])[2:],
            mode='bilinear',
            align_corners=True)
        last_fm = global_context
        pred_out = []

        for i, (
                fm, arm, refine
        ) in enumerate(zip(context_blocks[:2], self.arms, self.refines)):
            fm = arm(fm)
            fm += last_fm
            last_fm = F.interpolate(
                fm,
                size=paddle.shape(context_blocks[i + 1])[2:],
                mode='bilinear',
                align_corners=True)
            last_fm = refine(last_fm)
            pred_out.append(last_fm)
        context_out = last_fm

        concate_fm = self.ffm(spatial_out, context_out)
        pred_out.append(concate_fm)

        output = []
        if self.training:
            for i, head in enumerate(self.heads):
                out = head(pred_out[i])
                output.append(out)
        else:
            out = self.heads[-1](pred_out[-1])
            output.append(out)
        return output


class SpatialPath(nn.Layer):
    """
    SpatialPath module of BiseNetV1 model

    Args:
        in_channels (int): The number of input channels in spatial path module.
        out_channels (int): The number of output channels in spatial path module.
    """

    def __init__(self, in_channels, out_channels, inner_channel=64):
        super().__init__()
        self.conv_7x7 = layers.ConvBNReLU(
            in_channels, inner_channel, 7, stride=2, padding=3, bias_attr=False)
        self.conv_3x3_1 = layers.ConvBNReLU(
            inner_channel,
            inner_channel,
            3,
            stride=2,
            padding=1,
            bias_attr=False)
        self.conv_3x3_2 = layers.ConvBNReLU(
            inner_channel,
            inner_channel,
            3,
            stride=2,
            padding=1,
            bias_attr=False)
        self.conv_1x1 = layers.ConvBNReLU(
            inner_channel, out_channels, 1, bias_attr=False)

    def forward(self, x):
        x = self.conv_7x7(x)
        x = self.conv_3x3_1(x)
        x = self.conv_3x3_2(x)
        x = self.conv_1x1(x)
        return x


class BiSeNetHead(nn.Layer):
    """
    BiSeNet head of BiseNetV1 model

    Args:
        in_channels (int): The number of input channels in spatial path module.
        out_channels (int): The number of output channels in spatial path module.
        scale (int, float): The scale factor of interpolation.
    """

    def __init__(self, in_channels, out_channels, scale, is_aux=False):
        super().__init__()
        inner_channel = 128 if is_aux else 64
        self.conv_3x3 = layers.ConvBNReLU(
            in_channels, inner_channel, 3, stride=1, padding=1, bias_attr=False)
        self.conv_1x1 = nn.Conv2D(inner_channel, out_channels, 1)
        self.scale = scale

    def forward(self, x):
        x = self.conv_3x3(x)
        x = self.conv_1x1(x)
        if self.scale > 1:
            x = F.interpolate(
                x, scale_factor=self.scale, mode='bilinear', align_corners=True)
        return x


class AttentionRefinement(nn.Layer):
    """
    AttentionRefinement module of BiseNetV1 model

    Args:
        in_channels (int): The number of input channels in spatial path module.
        out_channels (int): The number of output channels in spatial path module.
    """

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv_3x3 = layers.ConvBNReLU(
            in_channels, out_channels, 3, stride=1, padding=1, bias_attr=False)
        self.channel_attention = nn.Sequential(
            nn.AdaptiveAvgPool2D(1),
            layers.ConvBNReLU(
                out_channels, out_channels, 1, bias_attr=False),
            nn.Sigmoid(), )

    def forward(self, x):
        x = self.conv_3x3(x)
        se = self.channel_attention(x)
        x = x * se
        return x


class FeatureFusion(nn.Layer):
    """
    AttentionRefinement module of BiseNetV1 model

    Args:
        in_channels (int): The number of input channels in spatial path module.
        out_channels (int): The number of output channels in spatial path module.
        reduction (int): A factor shrinks convolutional channels. Default: 1.
    """

    def __init__(self, in_channels, out_channels, reduction=1):
        super().__init__()
        self.conv_1x1 = layers.ConvBNReLU(
            in_channels, out_channels, 1, bias_attr=False)
        self.channel_attention = nn.Sequential(
            nn.AdaptiveAvgPool2D(1),
            layers.ConvBNReLU(
                out_channels, out_channels // reduction, 1, bias_attr=False),
            layers.ConvBNReLU(
                out_channels // reduction, out_channels, 1, bias_attr=False),
            nn.Sigmoid(), )

    def forward(self, x1, x2):
        fm = paddle.concat([x1, x2], axis=1)
        fm = self.conv_1x1(fm)
        fm_se = self.channel_attention(fm)
        output = fm + fm * fm_se
        return output