# AUTOGENERATED FROM: babilim/training/losses.ipynb

# Cell: 0
"""doc
# babilim.training.losses

> A package containing all losses.
"""

# Cell: 1
from collections import defaultdict
from typing import Any
import json
import numpy as np
import babilim
from babilim.core.itensor import ITensor
from babilim.core.logging import info
from babilim.core.tensor import Tensor
from babilim.core.module import Module


# Cell: 2
class Loss(Module):
    def __init__(self, reduction: str = "mean"):
        """
        A loss is a statefull object which computes the difference between the prediction and the target.
        
        :param log_std: When true the loss will log its standard deviation. (default: False)
        :param log_min: When true the loss will log its minimum values. (default: False)
        :param log_max: When true the loss will log its maximal values. (default: False)
        :param reduction: Specifies the reduction to apply to the output: `'none'` | `'mean'` | `'sum'`. `'none'`: no reduction will be applied, `'mean'`: the sum of the output will be divided by the number of elements in the output, `'sum'`: the output will be summed. Default: 'mean'.
        """
        super().__init__()
        self._accumulators = defaultdict(list)
        self.reduction = reduction
        if reduction not in ["none", "mean", "sum"]:
            raise NotImplementedError()

    def call(self, y_pred: Any, y_true: Any) -> ITensor:
        """
        Implement a loss function between preds and true outputs.
        
        **DO NOT**:
        * Overwrite this function (overwrite `self.loss(...)` instead)
        * Call this function (call the module instead `self(y_pred, y_true)`)

        Arguments:
        :param y_pred: The predictions of the network. Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
        :param y_true: The desired outputs of the network (labels). Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
        """
        loss = self.loss(y_pred, y_true)
        if loss.is_nan().any():
            raise ValueError("Loss is nan. Loss value: {}".format(loss))
        if self.reduction == "mean":
            loss = loss.mean()
        elif self.reduction == "sum":
            loss = loss.sum()
        return loss

    def loss(self, y_pred: Any, y_true: Any) -> ITensor:
        """
        Implement a loss function between preds and true outputs.
        
        **`loss` must be overwritten by subclasses.**
        
        **DO NOT**:
        * Call this function (call the module instead `self(y_pred, y_true)`)

        Arguments:
        :param y_pred: The predictions of the network. Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
        :param y_true: The desired outputs of the network (labels). Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
        """
        raise NotImplementedError("Every loss must implement the call method.")

    def log(self, name: str, value: ITensor) -> None:
        """
        Log a tensor under a name.
        
        These logged values then can be used for example by tensorboard loggers.
        
        :param name: The name under which to log the tensor.
        :param value: The tensor that should be logged.
        """
        if isinstance(value, ITensor):
            val = value.numpy()
            if len(val.shape) > 0:
                self._accumulators[name].append(val)
            else:
                self._accumulators[name].append(np.array([val]))
        else:
            self._accumulators[name].append(np.array([value]))

    def reset_avg(self) -> None:
        """
        Reset the accumulation of tensors in the logging.
        
        Should only be called by a tensorboard logger.
        """
        self._accumulators = defaultdict(list)

    def summary(self, samples_seen, summary_writer=None, summary_txt=None, log_std=False, log_min=False, log_max=False) -> None:
        """
        Write a summary of the accumulated logs into tensorboard.
        
        :param samples_seen: The number of samples the training algorithm has seen so far (not iterations!).
            This is used for the x axis in the plot. If you use the samples seen it is independant of the batch size.
            If the network was trained for 4 batches with 32 batch size or for 32 batches with 4 batchsize does not matter.
        :param summary_writer: The summary writer to use for writing the summary. If none is provided it will use the tensorflow default.
        :param summary_txt: The file where to write the summary in csv format.
        """
        results = {}
        if summary_writer is not None:
            for k in self._accumulators:
                if not self._accumulators[k]:
                    continue
                combined = np.concatenate(self._accumulators[k], axis=0)
                summary_writer.add_scalar("{}".format(k), combined.mean(), global_step=samples_seen)
                results[f"{k}"] = combined.mean()
                if log_std:
                    results[f"{k}_std"] = combined.std()
                    summary_writer.add_scalar("{}_std".format(k), results[f"{k}_std"], global_step=samples_seen)
                if log_min:
                    results[f"{k}_min"] = combined.min()
                    summary_writer.add_scalar("{}_min".format(k), results[f"{k}_min"], global_step=samples_seen)
                if log_max:
                    results[f"{k}_max"] = combined.max()
                    summary_writer.add_scalar("{}_max".format(k), results[f"{k}_max"], global_step=samples_seen)
        else:
            import tensorflow as tf
            for k in self._accumulators:
                if not self._accumulators[k]:
                    continue
                combined = np.concatenate(self._accumulators[k], axis=0)
                tf.summary.scalar("{}".format(k), combined.mean(), step=samples_seen)
                results[f"{k}"] = combined.mean()
                if log_std:
                    results[f"{k}_std"] = combined.std()
                    tf.summary.scalar("{}_std".format(k), results[f"{k}_std"], step=samples_seen)
                if log_min:
                    results[f"{k}_min"] = combined.min()
                    tf.summary.scalar("{}_min".format(k), results[f"{k}_min"], step=samples_seen)
                if log_max:
                    results[f"{k}_max"] = combined.max()
                    tf.summary.scalar("{}_max".format(k), results[f"{k}_max"], step=samples_seen)

        if summary_txt is not None:
            results["samples_seen"] = samples_seen
            for k in results:
                results[k] = f"{results[k]:.5f}"
            with open(summary_txt, "a") as f:
                f.write(json.dumps(results)+"\n")


    @property
    def avg(self):
        """
        Get the average of the loged values.
        
        This is helpfull to print values that are more stable than values from a single iteration.
        """
        avgs = {}
        for k in self._accumulators:
            if not self._accumulators[k]:
                continue
            combined = np.concatenate(self._accumulators[k], axis=0)
            avgs[k] = combined.mean()
        return avgs


# Cell: 3
class NativeLossWrapper(Loss):
    def __init__(self, loss, reduction: str = "mean"):
        """
        Wrap a native loss as a babilim loss.

        The wrapped object must have the following signature:

        ```python
        Callable(y_pred, y_true, log_val) -> Tensor
        ```

        where log_val will be a function which can be used for logging scalar tensors/values.

        :param loss: The loss that should be wrapped.
        :param reduction: Specifies the reduction to apply to the output: `'none'` | `'mean'` | `'sum'`. `'none'`: no reduction will be applied, `'mean'`: the sum of the output will be divided by the number of elements in the output, `'sum'`: the output will be summed. Default: 'mean'.
        """
        super().__init__(reduction=reduction)
        self.native_loss = loss
        self._auto_device()

    def _auto_device(self):
        if babilim.is_backend(babilim.PYTORCH_BACKEND):
            import torch
            self.native_loss = self.native_loss.to(torch.device(self.device))
            return self

    def loss(self, y_pred: Any, y_true: Any) -> ITensor:
        """
        Compute the loss using the native loss function provided in the constructor.
        
        :param y_pred: The predictions of the network. Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
        :param y_true: The desired outputs of the network (labels). Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
        """
        # Unwrap arguments
        tmp = y_true._asdict()
        y_true_tmp = {k: tmp[k].native for k in tmp}
        y_true = type(y_true)(**y_true_tmp)

        tmp = y_pred._asdict()
        y_pred_tmp = {k: tmp[k].native for k in tmp}
        y_pred = type(y_pred)(**y_pred_tmp)

        # call function
        result = self.native_loss(y_pred=y_pred, y_true=y_true,
                           log_val=lambda name, tensor: self.log(name, Tensor(data=tensor, trainable=True)))

        return Tensor(data=result, trainable=True)


# Cell: 4
class SparseCrossEntropyLossFromLogits(Loss):
    def __init__(self, reduction: str = "mean"):
        """
        Compute a sparse cross entropy.
        
        This means that the preds are logits and the targets are not one hot encoded.
        
        :param reduction: Specifies the reduction to apply to the output: `'none'` | `'mean'` | `'sum'`. `'none'`: no reduction will be applied, `'mean'`: the sum of the output will be divided by the number of elements in the output, `'sum'`: the output will be summed. Default: 'mean'.
        """
        super().__init__(reduction=reduction)
        if babilim.is_backend(babilim.PYTORCH_BACKEND):
            from torch.nn import CrossEntropyLoss
            self.loss_fun = CrossEntropyLoss(reduction="none")
        else:
            from tensorflow.nn import sparse_softmax_cross_entropy_with_logits
            self.loss_fun = sparse_softmax_cross_entropy_with_logits

    def loss(self, y_pred: ITensor, y_true: ITensor) -> ITensor:
        """
        Compute the sparse cross entropy assuming y_pred to be logits.
        
        :param y_pred: The predictions of the network. Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
        :param y_true: The desired outputs of the network (labels). Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
        """
        y_true = y_true.cast("int64")
        if babilim.is_backend(babilim.PYTORCH_BACKEND):
            return Tensor(data=self.loss_fun(y_pred.native, y_true.native), trainable=True)
        else:
            return Tensor(data=self.loss_fun(labels=y_true.native, logits=y_pred.native), trainable=True)


# Cell: 5
class BinaryCrossEntropyLossFromLogits(Loss):
    def __init__(self, reduction: str = "mean"):
        """
        Compute a binary cross entropy.
        
        This means that the preds are logits and the targets are a binary (1 or 0) tensor of same shape as logits.

        :param reduction: Specifies the reduction to apply to the output: `'none'` | `'mean'` | `'sum'`. `'none'`: no reduction will be applied, `'mean'`: the sum of the output will be divided by the number of elements in the output, `'sum'`: the output will be summed. Default: 'mean'.
        """
        super().__init__(reduction=reduction)
        if babilim.is_backend(babilim.PYTORCH_BACKEND):
            from torch.nn import BCEWithLogitsLoss
            self.loss_fun = BCEWithLogitsLoss(reduction="none")
        else:
            from tensorflow.nn import sigmoid_cross_entropy_with_logits
            self.loss_fun = sigmoid_cross_entropy_with_logits

    def loss(self, y_pred: ITensor, y_true: ITensor) -> ITensor:
        """
        Compute the sparse cross entropy assuming y_pred to be logits.
        
        :param y_pred: The predictions of the network. Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
        :param y_true: The desired outputs of the network (labels). Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
        """
        if babilim.is_backend(babilim.PYTORCH_BACKEND):
            return Tensor(data=self.loss_fun(y_pred.native, y_true.native), trainable=True)
        else:
            return Tensor(data=self.loss_fun(labels=y_true.native, logits=y_pred.native), trainable=True)


# Cell: 6
class SmoothL1Loss(Loss):
    def __init__(self, reduction: str = "mean"):
        """
        Compute a binary cross entropy.
        
        This means that the preds are logits and the targets are a binary (1 or 0) tensor of same shape as logits.

        :param reduction: Specifies the reduction to apply to the output: `'none'` | `'mean'` | `'sum'`. `'none'`: no reduction will be applied, `'mean'`: the sum of the output will be divided by the number of elements in the output, `'sum'`: the output will be summed. Default: 'mean'.
        """
        super().__init__(reduction=reduction)
        if babilim.is_backend(babilim.PYTORCH_BACKEND):
            from torch.nn import SmoothL1Loss
            self.loss_fun = SmoothL1Loss(reduction="none")
        else:
            from tensorflow.keras.losses import huber
            self.loss_fun = huber
            self.delta = 1.0

    def loss(self, y_pred: ITensor, y_true: ITensor) -> ITensor:
        """
        Compute the sparse cross entropy assuming y_pred to be logits.
        
        :param y_pred: The predictions of the network. Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
        :param y_true: The desired outputs of the network (labels). Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
        """
        if babilim.is_backend(babilim.PYTORCH_BACKEND):
            return Tensor(data=self.loss_fun(y_pred.native, y_true.native), trainable=True)
        else:
            return Tensor(data=self.loss_fun(labels=y_true.native, logits=y_pred.native, delta=self.delta), trainable=True)


# Cell: 7
class MeanSquaredError(Loss):
    def __init__(self, reduction: str = "mean"):
        """
        Compute the mean squared error.
        
        :param reduction: Specifies the reduction to apply to the output: `'none'` | `'mean'` | `'sum'`. `'none'`: no reduction will be applied, `'mean'`: the sum of the output will be divided by the number of elements in the output, `'sum'`: the output will be summed. Default: 'mean'.
        """
        super().__init__(reduction=reduction)
    
    def loss(self, y_pred: ITensor, y_true: ITensor, axis: int=-1) -> ITensor:
        """
        Compute the mean squared error.
        
        :param y_pred: The predictions of the network. Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
        :param y_true: The desired outputs of the network (labels). Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
        :param axis: (Optional) The axis along which to compute the mean squared error.
        """
        return ((y_pred - y_true) ** 2).mean(axis=axis)


# Cell: 8
class SparseCategoricalAccuracy(Loss):
    def __init__(self, reduction: str = "mean"):
        """
        Compute the sparse mean squared error.
        
        Sparse means that the targets are not one hot encoded.
        
        :param reduction: Specifies the reduction to apply to the output: `'none'` | `'mean'` | `'sum'`. `'none'`: no reduction will be applied, `'mean'`: the sum of the output will be divided by the number of elements in the output, `'sum'`: the output will be summed. Default: 'mean'.
        """
        super().__init__(reduction=reduction)

    def loss(self, y_pred: ITensor, y_true: ITensor, axis: int=-1) -> ITensor:
        """
        Compute the sparse categorical accuracy.
        
        :param y_pred: The predictions of the network. Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
        :param y_true: The desired outputs of the network (labels). Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
        :param axis: (Optional) The axis along which to compute the sparse categorical accuracy.
        """
        pred_class = y_pred.argmax(axis=axis)
        true_class = y_true.cast("int64")
        correct_predictions = pred_class == true_class
        return correct_predictions.cast("float32").mean(axis=axis)


# Cell: 9
class NaNMaskedLoss(Loss):
    def __init__(self, loss):
        """
        Compute a sparse cross entropy.
        
        This means that the preds are logits and the targets are not one hot encoded.
        
        :param loss: The loss that should be wrapped and only applied on non nan values.
        """
        super().__init__(reduction="none")
        self.wrapped_loss = loss
        self.zero = Tensor(data=np.array(0), trainable=False)

    def loss(self, y_pred: ITensor, y_true: ITensor) -> ITensor:
        """
        Compute the loss given in the constructor only on values where the GT is not NaN.
        
        :param y_pred: The predictions of the network. Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
        :param y_true: The desired outputs of the network (labels). Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
        """
        binary_mask = (~y_true.is_nan())
        mask = binary_mask.cast("float32")
        masked_y_true = (y_true * mask)[binary_mask]

        if y_pred.shape[-1] != binary_mask.shape[-1] and binary_mask.shape[-1] == 1:
            new_shape = binary_mask.shape[:-1]
            binary_mask = binary_mask.reshape(new_shape)
        masked_y_pred = (y_pred * mask)[binary_mask]
        
        if masked_y_pred.shape[0] > 0:
            loss = self.wrapped_loss(masked_y_pred, masked_y_true)
        else:
            loss = self.zero

        return loss
