import torch


def ModEntropy(softprob, labels):
    """Compute the modified entropy, described https://www.usenix.org/system/files/sec21fall-song.pdf, of a vector of
    soft probabilities.
    softprob: pytorch tensor. Vector of soft probabilities with shape [batch,num_classes]
    labels: pytorch tensor of shape [batch] containing the integer labels of the samples
    outputs -> pytorch tensor containing the modified entropy of each sample in the batch
    """
    epsilon = 1e-8
    confidence = torch.stack([softprob[i, labels[i]] for i in range(softprob.shape[0])])
    firstTerm = (confidence - 1) * torch.log(confidence + epsilon)
    secondTerm = - torch.sum(softprob * torch.log(1 - softprob + epsilon), 1)
    excessTerm = confidence * torch.log(1 - confidence + epsilon)
    return firstTerm + secondTerm + excessTerm