from catalyst.dl.core import Callback, RunnerState, CallbackOrder
from catalyst.dl.callbacks import CriterionCallback
from catalyst.contrib.criterion import FocalLossBinary
from catalyst.dl.utils.criterion import accuracy
from catalyst.utils import get_activation_fn
import torch
import torch.nn as nn
import numpy as np
from typing import List

import torch
from catalyst.utils import get_activation_fn


class MultiTaskCriterionCallback(Callback):
    def __init__(
        self,
        input_seg_key: str = "targets",
        input_cls_key: str = "labels",
        output_seg_key: str = "logits",
        output_cls_key: str = "cls_logits",
        prefix: str = "loss",
        criterion_key: str = None,
        loss_key: str = None,
        multiplier: float = 1.0,
        loss_weights: List[float] = None,
    ):
        super(MultiTaskCriterionCallback, self).__init__(CallbackOrder.Criterion)
        self.input_seg_key = input_seg_key
        self.input_cls_key = input_cls_key
        self.output_seg_key = output_seg_key
        self.output_cls_key = output_cls_key
        self.prefix = prefix
        self.criterion_key = criterion_key
        self.loss_key = loss_key
        self.multiplier = multiplier
        self.loss_weights = loss_weights

        self.criterion_cls = nn.BCEWithLogitsLoss()

    def _add_loss_to_state(self, state: RunnerState, loss):
        if self.loss_key is None:
            if state.loss is not None:
                if isinstance(state.loss, list):
                    state.loss.append(loss)
                else:
                    state.loss = [state.loss, loss]
            else:
                state.loss = loss
        else:
            if state.loss is not None:
                assert isinstance(state.loss, dict)
                state.loss[self.loss_key] = loss
            else:
                state.loss = {self.loss_key: loss}

    def _compute_loss(self, state: RunnerState, criterion):
        output_seg = state.output[self.output_seg_key]
        output_cls = state.output[self.output_cls_key]
        input_seg = state.input[self.input_seg_key]
        input_cls = state.input[self.input_cls_key]

        # assert len(self.loss_weights) == len(outputs)
        loss = 0

        # Segmentation loss
        loss += criterion(output_seg, input_seg) * self.loss_weights[0]
        # Classification loss
        loss += self.criterion_cls(output_cls, input_cls) * self.loss_weights[1]

        return loss

    def on_stage_start(self, state: RunnerState):
        assert state.criterion is not None

    def on_batch_end(self, state: RunnerState):
        criterion = state.get_key(
            key="criterion", inner_key=self.criterion_key
        )

        loss = self._compute_loss(state, criterion) * self.multiplier

        state.metrics.add_batch_value(
            metrics_dict={
                self.prefix: loss.item(),
            }
        )

        self._add_loss_to_state(state, loss)