# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/multiloss.ipynb.

# %% ../nbs/multiloss.ipynb 2
from __future__ import annotations
from types import FunctionType

from torch.distributions.beta import Beta

from fastai.callback.core import Callback
from fastai.callback.mixup import reduce_loss
from fastai.learner import Recorder
from fastai.layers import NoneReduce

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

# %% auto 0
__all__ = ['MultiLoss', 'MultiTargetLoss', 'MixHandlerX', 'MultiLossCallback']

# %% ../nbs/multiloss.ipynb 5
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)

# %% ../nbs/multiloss.ipynb 6
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:Listy[Numeric|Tensor]|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), f"{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), f"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), f"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), f"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,{}

        for loss in self.loss_funcs:
            if getattr(loss, 'y_int', False):
                self.y_int = True

    def forward(self, pred, targ):
        for i, loss_func in enumerate(self.loss_funcs):
            l = self.weights[i]*loss_func(pred, targ)
            if i == 0: loss = 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 forward_mixup(self, pred, targ1, targ2, lam):
        for i, loss_func in enumerate(self.loss_funcs):
            l = self.weights[i]*torch.lerp(loss_func(pred, targ1), loss_func(pred, targ2), lam)
            if i == 0:
                loss = torch.zeros_like(targ1).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)

# %% ../nbs/multiloss.ipynb 11
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:Listy[Numeric|Tensor]|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 forward_mixup(self, **kwargs):
        raise NotImplementedError("Mixup doesn't support Multi-Target training")

    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)]

# %% ../nbs/multiloss.ipynb 16
class MultiAvgLoss(AvgLossX):
    "Average the MultiLoss losses taking into account potential different batch sizes"
    def __init__(self,
        i:int, # `Multiloss` loss function location
        name:str, # 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

# %% ../nbs/multiloss.ipynb 17
class MultiAvgSmoothLoss(AvgSmoothLossX):
    "Smooth average of the MultiLoss losses (exponentially weighted with `beta`)"
    def __init__(self,
        i:int, # `Multiloss` loss function location
        name:str, # 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)

# %% ../nbs/multiloss.ipynb 18
class MultiAvgSmoothLossMixup(AvgSmoothLossX):
    "Smooth average of the MultiLoss losses (exponentially weighted with `beta`)"
    def __init__(self,
        i:int, # `Multiloss` loss function location
        name:str, # 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_mixup.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)

# %% ../nbs/multiloss.ipynb 20
# MixHandlerX is based on fastai's MixHandler:
# fastai - Apache License 2.0 - Copyright (c) 2023 fast.ai

class MixHandlerX(Callback):
    "A handler class for implementing `MixUp` style scheduling. Like fastai's `MixHandler` but supports `MultiLoss`."
    run_valid = False
    def __init__(self,
        alpha:float=0.5, # Alpha & beta parametrization for `Beta` distribution
        interp_label:bool|None=None # Blend or stack labels. Defaults to `loss_func.y_int` if None
    ):
        store_attr()
        self.distrib = Beta(tensor(alpha), tensor(alpha))

    def before_fit(self):
        "Determine whether to stack or interpolate labels"
        self.multiloss = isinstance(self.learn.loss_func, MultiLoss)
        if self.interp_label is None:
            self.stack_y = getattr(self.learn.loss_func, 'y_int', False)
        else:
            self.stack_y = not self.interp_label

    def before_train(self):
        "Determine whether to stack y"
        if self.stack_y:
            if self.multiloss:
                self.learn.loss_func_mixup = self.learn.loss_func
                self.learn.loss_func = self.multi_lf
            else:
                self.old_lf = self.learn.loss_func
                self.learn.loss_func = self.solo_lf

    def after_train(self):
        "Set the loss function back to the original loss"
        if self.stack_y:
            if self.multiloss:
                self.learn.loss_func = self.learn.loss_func_mixup
            else:
                self.learn.loss_func = self.old_lf

    def after_cancel_train(self):
        "If training is canceled, still set the loss function back"
        self.after_train()

    def after_cancel_fit(self):
        "If fit is canceled, still set the loss function back"
        self.after_train()

    def solo_lf(self, pred, *yb):
        "Interpolates losses on stacked labels by `self.lam` during training"
        if not self.training:
            return self.old_lf(pred, *yb)

        with NoneReduce(self.old_lf) as lf:
            loss = torch.lerp(lf(pred,*self.yb1), lf(pred,*yb), self.lam)

        return reduce_loss(loss, getattr(self.old_lf, 'reduction', 'mean'))

    def multi_lf(self, pred, *yb):
        if not self.training:
            return self.learn.loss_func_mixup(pred, *yb)
        else:
            return self.learn.loss_func_mixup.forward_mixup(pred, *self.yb1, *yb, self.lam)

# %% ../nbs/multiloss.ipynb 23
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")

        mixup = len(self.learn._grab_cbs(MixHandlerX)) > 0 and getattr(self.learn.loss_func, 'y_int', False)

        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)):
            if mixup:
                mets += MultiAvgSmoothLossMixup(i, self.loss_func.loss_names[i], self.beta, reduction)
            else:
                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
