import torch
from torch.autograd import Variable
import torch.nn.functional as F

import ipdb

"""
Lovasz-Softmax and Jaccard hinge loss in PyTorch
<NAME> 2018 ESAT-PSI KU Leuven (MIT License)
"""

def lovasz_grad(gt_sorted):
    """
    Computes gradient of the Lovasz extension w.r.t sorted errors
    See Alg. 1 in paper
    """
    p = len(gt_sorted)
    gts = gt_sorted.sum()
    intersection = gts - gt_sorted.float().cumsum(0)
    union = gts + (1 - gt_sorted).float().cumsum(0)
    jaccard = 1. - intersection / union
    if p > 1: # cover 1-pixel case
        jaccard[1:p] = jaccard[1:p] - jaccard[0:-1]
    return jaccard

def lovasz_softmax(probas, labels, classes='present', per_image=False, ignore=None):
    """
    Multi-class Lovasz-Softmax loss
      probas: [B, C, H, W] Variable, class probabilities at each prediction (between 0 and 1).
              Interpreted as binary (sigmoid) output with outputs of size [B, H, W].
      labels: [B, H, W] Tensor, ground truth labels (between 0 and C - 1)
      classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average.
      per_image: compute the loss per image instead of per batch
      ignore: void class labels
    """
    if per_image:
        loss = mean(lovasz_softmax_flat(*flatten_probas(prob.unsqueeze(0), lab.unsqueeze(0), ignore), classes=classes)
                          for prob, lab in zip(probas, labels))
    else:
        loss = lovasz_softmax_flat(*flatten_probas(probas, labels, ignore), classes=classes)
    return loss


def lovasz_softmax_flat(probas, labels, classes='present'):
    """
    Multi-class Lovasz-Softmax loss
      probas: [P, C] Variable, class probabilities at each prediction (between 0 and 1)
      labels: [P] Tensor, ground truth labels (between 0 and C - 1)
      classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average.
    """
    if probas.numel() == 0:
        # only void pixels, the gradients should be 0
        return probas * 0.
    num_classes = probas.size(1)
    losses = []
    class_to_sum = list(range(num_classes)) if classes in ['all', 'present'] else classes
    for _class in class_to_sum:
        fg = (labels == _class).float() # foreground for class c
        if (classes is 'present' and fg.sum() == 0):
            continue
        if num_classes == 1:
            if len(classes) > 1:
                raise ValueError('Sigmoid output possible only with 1 class')
            class_pred = probas[:, 0]
        else:
            class_pred = probas[:, _class]
        errors = (Variable(fg) - class_pred).abs()
        errors_sorted, perm = torch.sort(errors, 0, descending=True)
        perm = perm.data
        fg_sorted = fg[perm]
        losses.append(torch.dot(errors_sorted, Variable(lovasz_grad(fg_sorted))))

    losses = torch.stack(losses)
    return losses.mean()


def flatten_probas(probas, labels, ignore=None, sparse=True):
    """
    Flattens predictions in the batch
    """
    if probas.dim() != 2:
        if probas.dim() == 3:
            # assumes output of a sigmoid layer
            B, H, W = probas.size()
            probas = probas.view(B, 1, H, W)
        elif probas.dim() == 5:
            #3D segmentation
            B, C, L, H, W = probas.size()
            probas = probas.contiguous().view(B, C, L, H*W)
        B, C, H, W = probas.size()
        probas = probas.permute(0, 2, 3, 1).contiguous().view(-1, C)  # B * H * W, C = P, C
    labels = labels.view(-1)
    if ignore is None:
        return probas, labels
    valid = (labels != ignore)
    valid_probas = probas[valid]
    valid_labels = labels[valid]
    return valid_probas, valid_labels
import torch
from torch.autograd import Variable
import torch.nn.functional as F

import ipdb

"""
Lovasz-Softmax and Jaccard hinge loss in PyTorch
<NAME> 2018 ESAT-PSI KU Leuven (MIT License)
"""

def lovasz_grad(gt_sorted):
    """
    Computes gradient of the Lovasz extension w.r.t sorted errors
    See Alg. 1 in paper
    """
    p = len(gt_sorted)
    gts = gt_sorted.sum()
    intersection = gts - gt_sorted.float().cumsum(0)
    union = gts + (1 - gt_sorted).float().cumsum(0)
    jaccard = 1. - intersection / union
    if p > 1: # cover 1-pixel case
        jaccard[1:p] = jaccard[1:p] - jaccard[0:-1]
    return jaccard

def lovasz_softmax(probas, labels, classes='present', per_image=False, ignore=None):
    """
    Multi-class Lovasz-Softmax loss
      probas: [B, C, H, W] Variable, class probabilities at each prediction (between 0 and 1).
              Interpreted as binary (sigmoid) output with outputs of size [B, H, W].
      labels: [B, H, W] Tensor, ground truth labels (between 0 and C - 1)
      classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average.
      per_image: compute the loss per image instead of per batch
      ignore: void class labels
    """
    if per_image:
        loss = mean(lovasz_softmax_flat(*flatten_probas(prob.unsqueeze(0), lab.unsqueeze(0), ignore), classes=classes)
                          for prob, lab in zip(probas, labels))
    else:
        loss = lovasz_softmax_flat(*flatten_probas(probas, labels, ignore), classes=classes)
    return loss


def lovasz_softmax_flat(probas, labels, classes='present'):
    """
    Multi-class Lovasz-Softmax loss
      probas: [P, C] Variable, class probabilities at each prediction (between 0 and 1)
      labels: [P] Tensor, ground truth labels (between 0 and C - 1)
      classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average.
    """
    if probas.numel() == 0:
        # only void pixels, the gradients should be 0
        return probas * 0.
    num_classes = probas.size(1)
    losses = []
    class_to_sum = list(range(num_classes)) if classes in ['all', 'present'] else classes
    for _class in class_to_sum:
        fg = (labels == _class).float() # foreground for class c
        if (classes is 'present' and fg.sum() == 0):
            continue
        if num_classes == 1:
            if len(classes) > 1:
                raise ValueError('Sigmoid output possible only with 1 class')
            class_pred = probas[:, 0]
        else:
            class_pred = probas[:, _class]
        errors = (Variable(fg) - class_pred).abs()
        errors_sorted, perm = torch.sort(errors, 0, descending=True)
        perm = perm.data
        fg_sorted = fg[perm]
        losses.append(torch.dot(errors_sorted, Variable(lovasz_grad(fg_sorted))))

    losses = torch.stack(losses)
    return losses.mean()


def flatten_probas(probas, labels, ignore=None, sparse=True):
    """
    Flattens predictions in the batch
    """
    if probas.dim() != 2:
        if probas.dim() == 3:
            # assumes output of a sigmoid layer
            B, H, W = probas.size()
            probas = probas.view(B, 1, H, W)
        elif probas.dim() == 5:
            #3D segmentation
            B, C, L, H, W = probas.size()
            probas = probas.contiguous().view(B, C, L, H*W)
        B, C, H, W = probas.size()
        probas = probas.permute(0, 2, 3, 1).contiguous().view(-1, C)  # B * H * W, C = P, C
    labels = labels.view(-1)
    if ignore is None:
        return probas, labels
    valid = (labels != ignore)
    valid_probas = probas[valid]
    valid_labels = labels[valid]
    return valid_probas, valid_labels
Lovasz-Softmax and Jaccard hinge loss in PyTorch <NAME> 2018 ESAT-PSI KU Leuven (MIT License) Computes gradient of the Lovasz extension w.r.t sorted errors     See Alg. 1 in paper # cover 1-pixel case Multi-class Lovasz-Softmax loss       probas: [B, C, H, W] Variable, class probabilities at each prediction (between 0 and 1).               Interpreted as binary (sigmoid) output with outputs of size [B, H, W].       labels: [B, H, W] Tensor, ground truth labels (between 0 and C - 1)       classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average.       per_image: compute the loss per image instead of per batch       ignore: void class labels Multi-class Lovasz-Softmax loss       probas: [P, C] Variable, class probabilities at each prediction (between 0 and 1)       labels: [P] Tensor, ground truth labels (between 0 and C - 1)       classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average. # only void pixels, the gradients should be 0 # foreground for class c Flattens predictions in the batch # assumes output of a sigmoid layer #3D segmentation # B * H * W, C = P, C