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


from __future__ import annotations


__all__ = ['MultiLoss', 'MultiTargetLoss', 'MultiLossCallback']

# Cell
#nbdev_comment from __future__ import annotations
from types import FunctionType

from fastai.callback.core import Callback
from fastai.learner import Recorder

from .basics import is_listish
from .metrics import AvgLossX, AvgSmoothLossX
from .imports import *

# Internal Cell
def init_loss(l, **kwargs):
    "Initiatize loss class or partial loss function"
    return partialler(l, reduction='none') if isinstance(l, FunctionType) else l(reduction='none', **kwargs)

# Cell
class MultiLoss(Module):
    """
    Combine multiple `loss_funcs` on one prediction & target via `reduction`, with optional weighting.

    Log `loss_funcs` as metrics via `MultiLossCallback`, optionally using `loss_names`.
    """
    def __init__(self,
        loss_funcs:listy[Callable[...,nn.Module]|FunctionType], # Uninitialized loss functions or classes. Must support PyTorch `reduction` string.
        weights:listified[Number]|None=None, # Weight per loss. Defaults to uniform weighting.
        loss_kwargs:listy[dict[str,Any]]|None=None, # kwargs to pass to each loss function. Defaults to None.
        loss_names:listy[str]|None=None, # Loss names to log using `MultiLossCallback`. Defaults to loss `__name__`.
        reduction:str|None='mean' # PyTorch loss reduction
    ):
        store_attr(but='loss_names')
        assert is_listy(loss_funcs), "`loss_funcs` must be list-like"

        if weights is None or len(weights)==0:
            self.weights = [1]*len(loss_funcs)
        else:
            assert len(loss_funcs) == len(weights), "Must provide same number of `weights` as `loss_funcs`"
            self.weights = weights

        if loss_kwargs is None or len(loss_kwargs)==0: loss_kwargs = [{}]*len(loss_funcs)
        else: assert len(loss_funcs) == len(loss_kwargs), "Must provide same number of `loss_kwargs` as `loss_funcs`"

        if loss_names is None or len(loss_names)==0: loss_names = [l.__name__ for l in loss_funcs]
        else: assert len(loss_funcs) == len(loss_names), "Must provide same number of `loss_names` as `loss_funcs`"

        self.loss_funcs = [init_loss(l, **k) for l, k in zip(loss_funcs, loss_kwargs)]
        self.loss_names = L(loss_names)
        self._reduction,self._loss = reduction,{}

    def forward(self, pred, targ):
        for i, loss_func in enumerate(self.loss_funcs):
            l = TensorBase(self.weights[i]*loss_func(pred, targ))
            if i == 0: loss = TensorBase(torch.zeros_like(targ)).float()
            loss += l
            self._loss[i] = l

        return loss.mean() if self._reduction=='mean' else loss.sum() if self._reduction=='sum' else loss

    @property
    def losses(self): return self._loss

    @property
    def reduction(self): return self._reduction

    @reduction.setter
    def reduction(self, r): self._reduction = r

    @delegates(Module.to)
    def to(self, *args, **kwargs):
        device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
        if is_listish(self.weights) or not isinstance(self.weights, torch.Tensor): self.weights = torch.Tensor(self.weights)
        if self.weights.device != device: self.weights = self.weights.to(device=device)
        super().to(*args, **kwargs)

    def activation(self, pred):
        "Returns first `loss_funcs` `activation`"
        return getattr(self.loss_funcs[0], 'activation', noop)(pred)

    def decodes(self, pred):
        "Returns first `loss_funcs` `decodes`"
        return getattr(self.loss_funcs[0], 'decodes', noop)(pred)

# Cell
class MultiTargetLoss(MultiLoss):
    """
    Combine `loss_funcs` from multiple predictions & targets via `reduction`, with optional weighting.

    Log `loss_funcs` as metrics via `MultiLossCallback`, optionally using `loss_names`.
    """
    def __init__(self,
        loss_funcs:listy[Callable[...,nn.Module]|FunctionType], # Uninitialized loss functions or classes. One per prediction and target. Must support PyTorch `reduction` string.
        weights:listified[Number]|None=None, # Weight per loss. Defaults to uniform weighting.
        loss_kwargs:listy[dict[str,Any]]|None=None, # kwargs to pass to each loss function. Defaults to None.
        loss_names:listy[str]|None=None, # Loss names to log using `MultiLossCallback`. Defaults to loss `__name__`.
        reduction:str|None='mean' # PyTorch loss reduction
    ):
        super().__init__(loss_funcs, weights, loss_kwargs, loss_names, reduction)

    def forward(self, preds, targs):
        for i, (loss_func, pred, targ) in enumerate(zip(self.loss_funcs, preds, targs)):
            l = TensorBase(self.weights[i]*loss_func(pred, targ))
            if i == 0: loss = TensorBase(torch.zeros_like(targ)).float()
            loss += l
            self._loss[i] = l

        return loss.mean() if self._reduction=='mean' else loss.sum() if self._reduction=='sum' else loss

    def activation(self, preds):
        "Returns list of `activation`"
        return [getattr(self.loss_funcs[i], 'activation', noop)(pred) for i, pred in enumerate(preds)]

    def decodes(self, preds):
        "Returns list of `decodes`"
        return [getattr(self.loss_funcs[i], 'decodes', noop)(pred) for i, pred in enumerate(preds)]

# Internal Cell
class MultiAvgLoss(AvgLossX):
    "Average the MultiLoss losses taking into account potential different batch sizes"
    def __init__(self,
        i, # `Multiloss` loss function location
        name, # Loss function name
        reduction:str|None='mean' # Override loss reduction for logging
    ):
        store_attr(but='name')
        self._name = name

    def accumulate(self, learn):
        bs = find_bs(learn.yb)
        loss = learn.loss_func.losses[self.i]
        loss = loss.mean() if self.reduction=='mean' else loss.sum() if self.reduction=='sum' else loss
        self.total += learn.to_detach(loss)*bs
        self.count += bs

# Internal Cell
class MultiAvgSmoothLoss(AvgSmoothLossX):
    "Smooth average of the MultiLoss losses (exponentially weighted with `beta`)"
    def __init__(self,
        i, # `Multiloss` loss function location
        name, # Loss function name
        beta:float=0.98, # Smoothing beta
        reduction:str|None='mean' # Override loss reduction for logging
    ):
        super().__init__()
        store_attr(but='name')
        self._name = name

    def accumulate(self, learn):
        self.count += 1
        loss = learn.loss_func.losses[self.i]
        loss = loss.mean() if self.reduction=='mean' else loss.sum() if self.reduction=='sum' else loss
        self.val = torch.lerp(to_detach(loss, gather=False), self.val, self.beta)

# Cell
class MultiLossCallback(Callback):
    "Callback to automatically log and name `MultiLoss` losses as fastxtend metrics"
    run_valid,order = False,Recorder.order-1

    def __init__(self,
        beta:float=0.98, # Smoothing beta
        reduction:str|None='mean' # Override loss reduction for logging
    ):
        store_attr()

    def before_fit(self):
        if not isinstance(self.loss_func, MultiLoss):
            raise ValueError("`MultiLossCallback` requires loss to be `MultiLoss` class")

        mets= L()
        reduction = self.loss_func.reduction if self.reduction is None else self.reduction
        for i in range(len(self.loss_func.loss_funcs)):
            mets += MultiAvgSmoothLoss(i, self.loss_func.loss_names[i], self.beta, reduction)
            mets += MultiAvgLoss(i, self.loss_func.loss_names[i], reduction)

        self.learn.metrics = mets + self.learn.metrics