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

def linear_combination(x, y, epsilon):
    return epsilon * x + (1 - epsilon) * y

def reduce_loss(loss, reduction="mean"):
    if reduction == "mean":
        return loss.mean()
    
    elif reduction == "sum":
        return loss.sum()
    
    else:
        return loss
    
class LabelSmoothingCrossEntropy(nn.Module):
    def __init__(self, epsilon: float = 0.1, reduction="mean"):
        super(LabelSmoothingCrossEntropy, self).__init__()
        self.epsilon = epsilon
        self.reduction = reduction
    def forward(self, preds, target):
        n_classes = preds.size()[-1]
        
        log_preds = F.log_softmax(preds, dim=-1)
        loss = reduce_loss(-log_preds.sum(dim=-1), self.reduction)
        nll = F.nll_loss(log_preds, target, reduction=self.reduction)
        
        return linear_combination(loss / n_classes, nll, self.epsilon)
