# This code is referenced from 
# https://github.com/facebookresearch/astmt/
# 
# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
# 
# License: Attribution-NonCommercial 4.0 International

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.modules.module import Module
import numpy as np


class SoftMaxwithLoss(Module):
    """
    This function returns cross entropy loss for semantic segmentation
    """

    def __init__(self):
        super(SoftMaxwithLoss, self).__init__()
        self.softmax = nn.LogSoftmax(dim=1)
        self.criterion = nn.NLLLoss(ignore_index=255)

    def forward(self, out, label):
        assert not label.requires_grad
        # out shape  batch_size x channels x h x w
        # label shape batch_size x 1 x h x w
        label = label[:, 0, :, :].long()
        loss = self.criterion(self.softmax(out), label)

        return loss


class BalancedCrossEntropyLoss(Module):
    """
    Balanced Cross Entropy Loss with optional ignore regions
    """

    def __init__(self, size_average=True, batch_average=True, pos_weight=None):
        super(BalancedCrossEntropyLoss, self).__init__()
        self.size_average = size_average
        self.batch_average = batch_average
        self.pos_weight = pos_weight

    def forward(self, output, label, void_pixels=None):
        assert (output.size() == label.size())
        labels = torch.ge(label, 0.5).float()

        # Weighting of the loss, default is HED-style
        if self.pos_weight is None:
            num_labels_pos = torch.sum(labels)
            num_labels_neg = torch.sum(1.0 - labels)
            num_total = num_labels_pos + num_labels_neg
            w = num_labels_neg / num_total
        else:
            w = self.pos_weight

        output_gt_zero = torch.ge(output, 0).float()
        loss_val = torch.mul(output, (labels - output_gt_zero)) - torch.log(
            1 + torch.exp(output - 2 * torch.mul(output, output_gt_zero)))

        loss_pos_pix = -torch.mul(labels, loss_val)
        loss_neg_pix = -torch.mul(1.0 - labels, loss_val)

        if void_pixels is not None and not self.pos_weight:
            w_void = torch.le(void_pixels, 0.5).float()
            loss_pos_pix = torch.mul(w_void, loss_pos_pix)
            loss_neg_pix = torch.mul(w_void, loss_neg_pix)
            num_total = num_total - torch.ge(void_pixels, 0.5).float().sum()
            w = num_labels_neg / num_total

        loss_pos = torch.sum(loss_pos_pix)
        loss_neg = torch.sum(loss_neg_pix)

        final_loss = w * loss_pos + (1 - w) * loss_neg

        if self.size_average:
            final_loss /= float(np.prod(label.size()))
        elif self.batch_average:
            final_loss /= label.size()[0]

        return final_loss


class BinaryCrossEntropyLoss(Module):
    """
    Binary Cross Entropy with ignore regions, not balanced.
    """

    def __init__(self, size_average=True, batch_average=True):
        super(BinaryCrossEntropyLoss, self).__init__()
        self.size_average = size_average
        self.batch_average = batch_average

    def forward(self, output, label, void_pixels=None):
        assert (output.size() == label.size())

        labels = torch.ge(label, 0.5).float()

        output_gt_zero = torch.ge(output, 0).float()
        loss_val = torch.mul(output, (labels - output_gt_zero)) - torch.log(
            1 + torch.exp(output - 2 * torch.mul(output, output_gt_zero)))

        loss_pos_pix = -torch.mul(labels, loss_val)
        loss_neg_pix = -torch.mul(1.0 - labels, loss_val)

        if void_pixels is not None:
            w_void = torch.le(void_pixels, 0.5).float()
            loss_pos_pix = torch.mul(w_void, loss_pos_pix)
            loss_neg_pix = torch.mul(w_void, loss_neg_pix)

        loss_pos = torch.sum(loss_pos_pix)
        loss_neg = torch.sum(loss_neg_pix)
        final_loss = loss_pos + loss_neg

        if self.size_average:
            final_loss /= float(np.prod(label.size()))
        elif self.batch_average:
            final_loss /= label.size()[0]

        return final_loss


class DepthLoss(nn.Module):
    """
    Loss for depth prediction. By default L1 loss is used.  
    """
    def __init__(self, loss='l1'):
        super(DepthLoss, self).__init__()
        if loss == 'l1':
            self.loss = nn.L1Loss()

        else:
            raise NotImplementedError('Loss {} currently not supported in DepthLoss'.format(loss))

    def forward(self, out, label):
        mask = (label != 255)
        return self.loss(torch.masked_select(out, mask), torch.masked_select(label, mask))


class Normalize(nn.Module):
    def __init__(self):
        super(Normalize, self).__init__()

    def forward(self, bottom):
        qn = torch.norm(bottom, p=2, dim=1).unsqueeze(dim=1) + 1e-12
        top = bottom.div(qn)

        return top


class NormalsLoss(Module):
    """
    L1 loss with ignore labels
    normalize: normalization for surface normals
    """
    def __init__(self, size_average=True, normalize=False, norm=1):
        super(NormalsLoss, self).__init__()

        self.size_average = size_average

        if normalize:
            self.normalize = Normalize()
        else:
            self.normalize = None

        if norm == 1:
            print('Using L1 loss for surface normals')
            self.loss_func = F.l1_loss
        elif norm == 2:
            print('Using L2 loss for surface normals')
            self.loss_func = F.mse_loss
        else:
            raise NotImplementedError

    def forward(self, out, label, ignore_label=255):
        assert not label.requires_grad
        mask = (label != ignore_label)
        n_valid = torch.sum(mask).item()

        if self.normalize is not None:
            out_norm = self.normalize(out)
            loss = self.loss_func(torch.masked_select(out_norm, mask), torch.masked_select(label, mask), reduction='sum')
        else:
            loss = self.loss_func(torch.masked_select(out, mask), torch.masked_select(label, mask), reduction='sum')

        if self.size_average:
            if ignore_label:
                ret_loss = torch.div(loss, max(n_valid, 1e-6))
                return ret_loss
            else:
                ret_loss = torch.div(loss, float(np.prod(label.size())))
                return ret_loss

        return loss

[((23, 23, 23, 43), 'torch.nn.LogSoftmax', 'nn.LogSoftmax', (), '', True, 'import torch.nn as nn\n'), ((24, 25, 24, 53), 'torch.nn.NLLLoss', 'nn.NLLLoss', (), '', True, 'import torch.nn as nn\n'), ((74, 19, 74, 42), 'torch.sum', 'torch.sum', ({(74, 29, 74, 41): 'loss_pos_pix'}, {}), '(loss_pos_pix)', False, 'import torch\n'), ((75, 19, 75, 42), 'torch.sum', 'torch.sum', ({(75, 29, 75, 41): 'loss_neg_pix'}, {}), '(loss_neg_pix)', False, 'import torch\n'), ((114, 19, 114, 42), 'torch.sum', 'torch.sum', ({(114, 29, 114, 41): 'loss_pos_pix'}, {}), '(loss_pos_pix)', False, 'import torch\n'), ((115, 19, 115, 42), 'torch.sum', 'torch.sum', ({(115, 29, 115, 41): 'loss_neg_pix'}, {}), '(loss_neg_pix)', False, 'import torch\n'), ((53, 29, 53, 46), 'torch.sum', 'torch.sum', ({(53, 39, 53, 45): 'labels'}, {}), '(labels)', False, 'import torch\n'), ((54, 29, 54, 52), 'torch.sum', 'torch.sum', ({(54, 39, 54, 51): '1.0 - labels'}, {}), '(1.0 - labels)', False, 'import torch\n'), ((61, 19, 61, 63), 'torch.mul', 'torch.mul', ({(61, 29, 61, 35): 'output', (61, 38, 61, 61): '(labels - output_gt_zero)'}, {}), '(output, labels - output_gt_zero)', False, 'import torch\n'), ((64, 24, 64, 51), 'torch.mul', 'torch.mul', ({(64, 34, 64, 40): 'labels', (64, 42, 64, 50): 'loss_val'}, {}), '(labels, loss_val)', False, 'import torch\n'), ((65, 24, 65, 57), 'torch.mul', 'torch.mul', ({(65, 34, 65, 46): '(1.0 - labels)', (65, 48, 65, 56): 'loss_val'}, {}), '(1.0 - labels, loss_val)', False, 'import torch\n'), ((69, 27, 69, 58), 'torch.mul', 'torch.mul', ({(69, 37, 69, 43): 'w_void', (69, 45, 69, 57): 'loss_pos_pix'}, {}), '(w_void, loss_pos_pix)', False, 'import torch\n'), ((70, 27, 70, 58), 'torch.mul', 'torch.mul', ({(70, 37, 70, 43): 'w_void', (70, 45, 70, 57): 'loss_neg_pix'}, {}), '(w_void, loss_neg_pix)', False, 'import torch\n'), ((103, 19, 103, 63), 'torch.mul', 'torch.mul', ({(103, 29, 103, 35): 'output', (103, 38, 103, 61): '(labels - output_gt_zero)'}, {}), '(output, labels - output_gt_zero)', False, 'import torch\n'), ((106, 24, 106, 51), 'torch.mul', 'torch.mul', ({(106, 34, 106, 40): 'labels', (106, 42, 106, 50): 'loss_val'}, {}), '(labels, loss_val)', False, 'import torch\n'), ((107, 24, 107, 57), 'torch.mul', 'torch.mul', ({(107, 34, 107, 46): '(1.0 - labels)', (107, 48, 107, 56): 'loss_val'}, {}), '(1.0 - labels, loss_val)', False, 'import torch\n'), ((111, 27, 111, 58), 'torch.mul', 'torch.mul', ({(111, 37, 111, 43): 'w_void', (111, 45, 111, 57): 'loss_pos_pix'}, {}), '(w_void, loss_pos_pix)', False, 'import torch\n'), ((112, 27, 112, 58), 'torch.mul', 'torch.mul', ({(112, 37, 112, 43): 'w_void', (112, 45, 112, 57): 'loss_neg_pix'}, {}), '(w_void, loss_neg_pix)', False, 'import torch\n'), ((133, 24, 133, 35), 'torch.nn.L1Loss', 'nn.L1Loss', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((140, 25, 140, 55), 'torch.masked_select', 'torch.masked_select', ({(140, 45, 140, 48): 'out', (140, 50, 140, 54): 'mask'}, {}), '(out, mask)', False, 'import torch\n'), ((140, 57, 140, 89), 'torch.masked_select', 'torch.masked_select', ({(140, 77, 140, 82): 'label', (140, 84, 140, 88): 'mask'}, {}), '(label, mask)', False, 'import torch\n'), ((49, 17, 49, 37), 'torch.ge', 'torch.ge', ({(49, 26, 49, 31): 'label', (49, 33, 49, 36): '0.5'}, {}), '(label, 0.5)', False, 'import torch\n'), ((60, 25, 60, 44), 'torch.ge', 'torch.ge', ({(60, 34, 60, 40): 'output', (60, 42, 60, 43): '0'}, {}), '(output, 0)', False, 'import torch\n'), ((100, 17, 100, 37), 'torch.ge', 'torch.ge', ({(100, 26, 100, 31): 'label', (100, 33, 100, 36): '0.5'}, {}), '(label, 0.5)', False, 'import torch\n'), ((102, 25, 102, 44), 'torch.ge', 'torch.ge', ({(102, 34, 102, 40): 'output', (102, 42, 102, 43): '0'}, {}), '(output, 0)', False, 'import torch\n'), ((181, 18, 181, 33), 'torch.sum', 'torch.sum', ({(181, 28, 181, 32): 'mask'}, {}), '(mask)', False, 'import torch\n'), ((185, 34, 185, 69), 'torch.masked_select', 'torch.masked_select', ({(185, 54, 185, 62): 'out_norm', (185, 64, 185, 68): 'mask'}, {}), '(out_norm, mask)', False, 'import torch\n'), ((185, 71, 185, 103), 'torch.masked_select', 'torch.masked_select', ({(185, 91, 185, 96): 'label', (185, 98, 185, 102): 'mask'}, {}), '(label, mask)', False, 'import torch\n'), ((187, 34, 187, 64), 'torch.masked_select', 'torch.masked_select', ({(187, 54, 187, 57): 'out', (187, 59, 187, 63): 'mask'}, {}), '(out, mask)', False, 'import torch\n'), ((187, 66, 187, 98), 'torch.masked_select', 'torch.masked_select', ({(187, 86, 187, 91): 'label', (187, 93, 187, 97): 'mask'}, {}), '(label, mask)', False, 'import torch\n'), ((68, 21, 68, 47), 'torch.le', 'torch.le', ({(68, 30, 68, 41): 'void_pixels', (68, 43, 68, 46): '0.5'}, {}), '(void_pixels, 0.5)', False, 'import torch\n'), ((110, 21, 110, 47), 'torch.le', 'torch.le', ({(110, 30, 110, 41): 'void_pixels', (110, 43, 110, 46): '0.5'}, {}), '(void_pixels, 0.5)', False, 'import torch\n'), ((148, 13, 148, 43), 'torch.norm', 'torch.norm', (), '', False, 'import torch\n'), ((62, 39, 62, 72), 'torch.mul', 'torch.mul', ({(62, 49, 62, 55): 'output', (62, 57, 62, 71): 'output_gt_zero'}, {}), '(output, output_gt_zero)', False, 'import torch\n'), ((71, 36, 71, 62), 'torch.ge', 'torch.ge', ({(71, 45, 71, 56): 'void_pixels', (71, 58, 71, 61): '(0.5)'}, {}), '(void_pixels, 0.5)', False, 'import torch\n'), ((104, 39, 104, 72), 'torch.mul', 'torch.mul', ({(104, 49, 104, 55): 'output', (104, 57, 104, 71): 'output_gt_zero'}, {}), '(output, output_gt_zero)', False, 'import torch\n')]