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

__all__ = ['ClassBalancedCrossEntropyLoss', 'ClassBalancedBCEWithLogitsLoss']

# Cell
import torch.nn.modules.loss as TL

from .imports import *

# Internal Cell
class ClassBalanced(TL._WeightedLoss):
    "Class Balanced weight calculation, from https://arxiv.org/abs/1901.05555."
    def __init__(self,
        samples_per_class:Tensor,
        beta:float=0.99,
        size_average=None,
        reduce=None,
        reduction:str='mean',
    ):
        num_classes = len(samples_per_class)
        if not isinstance(samples_per_class, Tensor):
            samples_per_class = Tensor(samples_per_class)

        effective_num = 1.0 - torch.pow(beta, samples_per_class)
        weight = (1.0 - beta) / effective_num
        weight = weight / weight.sum() * num_classes

        super().__init__(weight, size_average, reduce, reduction)

    @delegates(nn.Module.to)
    def to(self, *args, **kwargs):
        device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
        if self.weight.device != device: self.weight = self.weight.to(device=device)
        super().to(*args, **kwargs)

# Refactored from https://github.com/vandit15/Class-balanced-loss-pytorch by Vandit Jain

# Original code, Copyright 2019 by Vandit Jain
# Licensed under the MIT License

# Cell
class ClassBalancedCrossEntropyLoss(ClassBalanced):
    "Class Balanced Cross Entropy Loss, from https://arxiv.org/abs/1901.05555."
    __constants__ = ['ignore_index', 'reduction', 'label_smoothing']
    ignore_index: int
    label_smoothing: float

    def __init__(self,
        samples_per_class:Tensor,
        beta:float=0.99,
        size_average=None,
        ignore_index:int=-100,
        reduce=None,
        reduction:str='mean',
        label_smoothing:float=0.0,
        axis:int=-1
    ):
        super().__init__(samples_per_class, beta, size_average, reduce, reduction)
        self.ignore_index = ignore_index
        self.label_smoothing = label_smoothing
        self.axis=axis

    def forward(self, input:Tensor, target:Tensor) -> Tensor:
        return F.cross_entropy(input, target, weight=self.weight,
                               ignore_index=self.ignore_index, reduction=self.reduction,
                               label_smoothing=self.label_smoothing)

    def decodes(self, x:Tensor):    return x.argmax(dim=self.axis)
    def activation(self, x:Tensor): return F.softmax(x, dim=self.axis)

# Cell
class ClassBalancedBCEWithLogitsLoss(ClassBalanced):
    "Class Balanced BCE With Logits Loss, from https://arxiv.org/abs/1901.05555."

    def forward(self, input:Tensor, target:Tensor) -> Tensor:
        return F.binary_cross_entropy_with_logits(input, target, weight=self.weight, reduction=self.reduction)