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

__all__ = ['LossRouting', 'CombinedLoss']

# Cell
from fastai.basics import *

# Cell

from typing import Callable
from dataclasses import dataclass

@dataclass
class LossRouting:
    loss_func: Callable
    pred_idx: int
    target_idx: int
    weight: float = 1.0

class CombinedLoss:
    """Applies loss functions to multiple model outputs and sums them.
    If applicable, it can decode and compute activations for each model output."""

    def __init__(self, *loss_routings):
        self.loss_routings = loss_routings

    def __call__(self, outs, *targets, **kwargs):
        return sum([
            routing.weight*routing.loss_func(outs[routing.pred_idx], targets[routing.target_idx])
            for routing in self.loss_routings
        ])

    def activation(self, outs):
        return [getattr(routing.loss_func, 'activation', noop)(outs[routing.pred_idx]) for routing in self.loss_routings]

    def decodes(self, outs):
        return [getattr(routing.loss_func, 'decodes', noop)(outs[routing.pred_idx]) for routing in self.loss_routings]

    @classmethod
    def from_one_to_one_routing(cls, *loss_funcs):
        return cls(*[LossRouting(loss_func, pred_idx=i, target_idx=i, weight=1.0) for i, loss_func in enumerate(loss_funcs)])
