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

from paddlers.models.ppseg.cvlibs import manager


@manager.LOSSES.add_component
class DetailAggregateLoss(nn.Layer):
    """
    DetailAggregateLoss's implementation based on PaddlePaddle.

    The original article refers to Meituan
    Fan, Mingyuan, et al. "Rethinking BiSeNet For Real-time Semantic Segmentation."
    (https://arxiv.org/abs/2104.13188)

    Args:
        ignore_index (int64, optional): Specifies a target value that is ignored
            and does not contribute to the input gradient. Default ``255``.

    """

    def __init__(self, ignore_index=255):
        super(DetailAggregateLoss, self).__init__()
        self.ignore_index = ignore_index
        self.laplacian_kernel = paddle.to_tensor([-1, -1, -1, -1, 8, -1, -1, -1, -1], dtype='float32').reshape(
            (1, 1, 3, 3))
        self.fuse_kernel = paddle.create_parameter([1, 3, 1, 1], dtype='float32')

    def forward(self, logits, label):
        """
        Args:
            logits (Tensor): Logit tensor, the data type is float32, float64. Shape is
                (N, C), where C is number of classes, and if shape is more than 2D, this
                is (N, C, D1, D2,..., Dk), k >= 1.
            label (Tensor): Label tensor, the data type is int64. Shape is (N), where each
                value is 0 <= label[i] <= C-1, and if shape is more than 2D, this is
                (N, D1, D2,..., Dk), k >= 1.
        Returns: loss
        """
        boundary_targets = F.conv2d(paddle.unsqueeze(label, axis=1).astype('float32'), self.laplacian_kernel,
                                    padding=1)
        boundary_targets = paddle.clip(boundary_targets, min=0)
        boundary_targets = boundary_targets > 0.1
        boundary_targets = boundary_targets.astype('float32')

        boundary_targets_x2 = F.conv2d(paddle.unsqueeze(label, axis=1).astype('float32'), self.laplacian_kernel,
                                       stride=2, padding=1)
        boundary_targets_x2 = paddle.clip(boundary_targets_x2, min=0)
        boundary_targets_x4 = F.conv2d(paddle.unsqueeze(label, axis=1).astype('float32'), self.laplacian_kernel,
                                       stride=4, padding=1)
        boundary_targets_x4 = paddle.clip(boundary_targets_x4, min=0)

        boundary_targets_x8 = F.conv2d(paddle.unsqueeze(label, axis=1).astype('float32'), self.laplacian_kernel,
                                       stride=8, padding=1)
        boundary_targets_x8 = paddle.clip(boundary_targets_x8, min=0)

        boundary_targets_x8_up = F.interpolate(boundary_targets_x8, boundary_targets.shape[2:], mode='nearest')
        boundary_targets_x4_up = F.interpolate(boundary_targets_x4, boundary_targets.shape[2:], mode='nearest')
        boundary_targets_x2_up = F.interpolate(boundary_targets_x2, boundary_targets.shape[2:], mode='nearest')

        boundary_targets_x2_up = boundary_targets_x2_up > 0.1
        boundary_targets_x2_up = boundary_targets_x2_up.astype('float32')

        boundary_targets_x4_up = boundary_targets_x4_up > 0.1
        boundary_targets_x4_up = boundary_targets_x4_up.astype('float32')

        boundary_targets_x8_up = boundary_targets_x8_up > 0.1
        boundary_targets_x8_up = boundary_targets_x8_up.astype('float32')

        boudary_targets_pyramids = paddle.stack((boundary_targets, boundary_targets_x2_up, boundary_targets_x4_up),
                                                axis=1)

        boudary_targets_pyramids = paddle.squeeze(boudary_targets_pyramids, axis=2)
        boudary_targets_pyramid = F.conv2d(boudary_targets_pyramids, self.fuse_kernel)

        boudary_targets_pyramid = boudary_targets_pyramid > 0.1
        boudary_targets_pyramid = boudary_targets_pyramid.astype('float32')

        if logits.shape[-1] != boundary_targets.shape[-1]:
            logits = F.interpolate(
                logits, boundary_targets.shape[2:], mode='bilinear', align_corners=True)

        bce_loss = F.binary_cross_entropy_with_logits(logits, boudary_targets_pyramid)
        dice_loss = self.fixed_dice_loss_func(F.sigmoid(logits), boudary_targets_pyramid)
        detail_loss = bce_loss + dice_loss

        label.stop_gradient = True
        return detail_loss

    def fixed_dice_loss_func(self, input, target):
        """
            simplified diceloss for DetailAggregateLoss.
        """
        smooth = 1.
        n = input.shape[0]
        iflat = paddle.reshape(input, [n, -1])
        tflat = paddle.reshape(target, [n, -1])
        intersection = paddle.sum((iflat * tflat), axis=1)
        loss = 1 - ((2. * intersection + smooth) /
                    (paddle.sum(iflat, axis=1) + paddle.sum(tflat, axis=1) + smooth))
        return paddle.mean(loss)