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

__all__ = ['after_loss', 'cross_entropy', 'cross_entropy_mixup', 'binary_cross_entropy', 'binary_cross_entropy_mixup',
           'binary_cross_entropy_scaled_mixup', 'focal_loss', 'get_loss']

# Cell
from fastai.vision.all import *

# Cell
def after_loss(loss, y=None):
    return loss

def cross_entropy(x, y, before=lambda *o:o, after=after_loss):
    f = nn.CrossEntropyLoss()
    loss = after(f(*before(x,y)), y=y)
    return loss

def cross_entropy_mixup(x, y, before=lambda *o:o, after=after_loss):
    if isinstance(y, dict):
        f = nn.CrossEntropyLoss(reduction='none')
        x0, y0 = before(x, y['yb'])
        x1, y1 = before(x, y['yb1'])
        a = y['a'].unsqueeze(1)
        loss = a*f(x0, y0) + (1-a)*f(x1, y1)
        loss = after(loss, y=y)
        return loss.mean()
    else: return cross_entropy(x, y, before=before, after=after)

def binary_cross_entropy(x, y, before=lambda *o:o, after=after_loss):
    f = nn.BCEWithLogitsLoss()
    x, y = before(x, y)
    if x.shape != y.shape:
        y = F.one_hot(y, x.shape[1]).float()
    loss = after(f(x,y), y=y)
    return loss

def binary_cross_entropy_mixup(x, y, before=lambda *o:o, after=after_loss):
    if isinstance(y, dict):
        f = nn.BCEWithLogitsLoss(reduction='none')
        x0, y0 = before(x, y['yb'])
        x1, y1 = before(x, y['yb1'])
        a = y['a'].unsqueeze(1)
        y0 = F.one_hot(y0, x0.shape[1]).float()
        y1 = F.one_hot(y1, x1.shape[1]).float()
        loss = a*f(x0, y0) + (1-a)*f(x1, y1)
        loss = after(loss, y=y)
        return loss.mean()
    else: return binary_cross_entropy(x, y, before=before, after=after)

def binary_cross_entropy_scaled_mixup(x, y, before=lambda *o:o, after=after_loss, weights=[0.1,1]):
    if isinstance(y, dict):
        f = nn.BCEWithLogitsLoss(reduction='none')
        x0, y0 = before(x, y['yb'])
        x1, y1 = before(x, y['yb1'])
        a = y['a'].unsqueeze(1)
        y0 = F.one_hot(y0, x0.shape[1]).float()
        y1 = F.one_hot(y1, x1.shape[1]).float()
        y = y0*a + (1-a)*y1
        loss = a*f(x0, y0) + (1-a)*f(x1, y1)
        loss = torch.cat([loss[y<0.5]*weights[0],loss[y>=0.5]*weights[1]])
        loss = after(loss, y=y)
        return loss.mean()
    else: return binary_cross_entropy(x, y, before=before, after=after)

def focal_loss(x, y, alpha=1, gamma=2, before=lambda *o:o, after=after_loss):
    "Adapted from https://www.kaggle.com/c/rfcx-species-audio-detection/discussion/213075"
    f = nn.BCEWithLogitsLoss(reduction='none')
    x, y = before(x, y)
    if x.shape != y.shape:
        y = F.one_hot(y, x.shape[1]).float()
    bce_loss = f(x, y)
    p = bce_loss.sigmoid()
    loss = torch.where(y >= 0.5, alpha * (1. - p)**gamma * bce_loss, p**gamma * bce_loss)
    loss = after(loss, y=y)
    return loss.mean()

# Cell
def get_loss(loss_func, before=lambda *o:o, after=after_loss):
    if loss_func == 'cross_entropy':
        return partial(cross_entropy, before=before, after=after)
    if loss_func == 'cross_entropy_mixup':
        return partial(cross_entropy_mixup, before=before, after=after)
    if loss_func == 'binary_cross_entropy':
        return partial(binary_cross_entropy, before=before, after=after)
    if loss_func == 'binary_cross_entropy_mixup':
        return partial(binary_cross_entropy_mixup, before=before, after=after)
    if loss_func == 'focal_loss':
        return partial(focal_loss, before=before, after=after)
    if loss_func == 'binary_cross_entropy_scaled_mixup':
        return partial(binary_cross_entropy_scaled_mixup, before=before, after=after)
    raise Exception(f'{loss_func} not defined.')