# 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, 'torch.nn.LogSoftmax', 'nn.LogSoftmax', ([], {'dim': '(1)'}), True, 'import torch.nn as nn\n'), (24, 'torch.nn.NLLLoss', 'nn.NLLLoss', ([], {'ignore_index': '(255)'}), True, 'import torch.nn as nn\n'), (74, 'torch.sum', 'torch.sum', (['loss_pos_pix'], {}), False, 'import torch\n'), (75, 'torch.sum', 'torch.sum', (['loss_neg_pix'], {}), False, 'import torch\n'), (114, 'torch.sum', 'torch.sum', (['loss_pos_pix'], {}), False, 'import torch\n'), (115, 'torch.sum', 'torch.sum', (['loss_neg_pix'], {}), False, 'import torch\n'), (53, 'torch.sum', 'torch.sum', (['labels'], {}), False, 'import torch\n'), (54, 'torch.sum', 'torch.sum', (['(1.0 - labels)'], {}), False, 'import torch\n'), (61, 'torch.mul', 'torch.mul', (['output', '(labels - output_gt_zero)'], {}), False, 'import torch\n'), (64, 'torch.mul', 'torch.mul', (['labels', 'loss_val'], {}), False, 'import torch\n'), (65, 'torch.mul', 'torch.mul', (['(1.0 - labels)', 'loss_val'], {}), False, 'import torch\n'), (69, 'torch.mul', 'torch.mul', (['w_void', 'loss_pos_pix'], {}), False, 'import torch\n'), (70, 'torch.mul', 'torch.mul', (['w_void', 'loss_neg_pix'], {}), False, 'import torch\n'), (103, 'torch.mul', 'torch.mul', (['output', '(labels - output_gt_zero)'], {}), False, 'import torch\n'), (106, 'torch.mul', 'torch.mul', (['labels', 'loss_val'], {}), False, 'import torch\n'), (107, 'torch.mul', 'torch.mul', (['(1.0 - labels)', 'loss_val'], {}), False, 'import torch\n'), (111, 'torch.mul', 'torch.mul', (['w_void', 'loss_pos_pix'], {}), False, 'import torch\n'), (112, 'torch.mul', 'torch.mul', (['w_void', 'loss_neg_pix'], {}), False, 'import torch\n'), (133, 'torch.nn.L1Loss', 'nn.L1Loss', ([], {}), True, 'import torch.nn as nn\n'), (140, 'torch.masked_select', 'torch.masked_select', (['out', 'mask'], {}), False, 'import torch\n'), (140, 'torch.masked_select', 'torch.masked_select', (['label', 'mask'], {}), False, 'import torch\n'), (49, 'torch.ge', 'torch.ge', (['label', '(0.5)'], {}), False, 'import torch\n'), (60, 'torch.ge', 'torch.ge', (['output', '(0)'], {}), False, 'import torch\n'), (100, 'torch.ge', 'torch.ge', (['label', '(0.5)'], {}), False, 'import torch\n'), (102, 'torch.ge', 'torch.ge', (['output', '(0)'], {}), False, 'import torch\n'), (181, 'torch.sum', 'torch.sum', (['mask'], {}), False, 'import torch\n'), (185, 'torch.masked_select', 'torch.masked_select', (['out_norm', 'mask'], {}), False, 'import torch\n'), (185, 'torch.masked_select', 'torch.masked_select', (['label', 'mask'], {}), False, 'import torch\n'), (187, 'torch.masked_select', 'torch.masked_select', (['out', 'mask'], {}), False, 'import torch\n'), (187, 'torch.masked_select', 'torch.masked_select', (['label', 'mask'], {}), False, 'import torch\n'), (68, 'torch.le', 'torch.le', (['void_pixels', '(0.5)'], {}), False, 'import torch\n'), (110, 'torch.le', 'torch.le', (['void_pixels', '(0.5)'], {}), False, 'import torch\n'), (148, 'torch.norm', 'torch.norm', (['bottom'], {'p': '(2)', 'dim': '(1)'}), False, 'import torch\n'), (62, 'torch.mul', 'torch.mul', (['output', 'output_gt_zero'], {}), False, 'import torch\n'), (71, 'torch.ge', 'torch.ge', (['void_pixels', '(0.5)'], {}), False, 'import torch\n'), (104, 'torch.mul', 'torch.mul', (['output', 'output_gt_zero'], {}), False, 'import torch\n')]