# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/050_losses.ipynb (unless otherwise specified).

__all__ = ['HuberLoss', 'LogCoshLoss', 'MaskedLossWrapper']

# Cell
from .imports import *

# Cell

## Available in Pytorch 1.9
class HuberLoss(nn.Module):
    """Huber loss

    Creates a criterion that uses a squared term if the absolute
    element-wise error falls below delta and a delta-scaled L1 term otherwise.
    This loss combines advantages of both :class:`L1Loss` and :class:`MSELoss`; the
    delta-scaled L1 region makes the loss less sensitive to outliers than :class:`MSELoss`,
    while the L2 region provides smoothness over :class:`L1Loss` near 0. See
    `Huber loss <https://en.wikipedia.org/wiki/Huber_loss>`_ for more information.
    This loss is equivalent to nn.SmoothL1Loss when delta == 1.
    """
    def __init__(self, reduction='mean', delta=1.0):
        assert reduction in ['mean', 'sum', 'none'], "You must set reduction to 'mean', 'sum' or 'none'"
        self.reduction, self.delta = reduction, delta
        super().__init__()

    def forward(self, input: Tensor, target: Tensor) -> Tensor:
        diff = input - target
        abs_diff = torch.abs(diff)
        mask = abs_diff < self.delta
        loss = torch.cat([(.5*diff[mask]**2), self.delta * (abs_diff[~mask] - (.5 * self.delta))])
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        else:
            return loss

# Cell

class LogCoshLoss(nn.Module):
    def __init__(self, reduction='mean', delta=1.0):
        assert reduction in ['mean', 'sum', 'none'], "You must set reduction to 'mean', 'sum' or 'none'"
        self.reduction, self.delta = reduction, delta
        super().__init__()

    def forward(self, input: Tensor, target: Tensor) -> Tensor:
        loss = torch.log(torch.cosh(input - target + 1e-12))
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        else:
            return loss

# Cell
class MaskedLossWrapper(Module):
    def __init__(self, crit):
        self.loss = crit

    def forward(self, inp, targ):
        inp = inp.flatten(1)
        targ = targ.flatten(1)
        mask = torch.isnan(targ)
        inp, targ = inp[~mask], targ[~mask]
        return self.loss(inp, targ)