import numpy as np

from typing import List

from catalyst.core import Callback, CallbackOrder, IRunner

from ._metrics import lwlrap


class LWLRAPCallback(Callback):
    def __init__(self,
                 input_key: str = "targets",
                 output_key: str = "logits",
                 model_output_key: str = "clipwise_output",
                 prefix: str = "lwlrap"):
        super().__init__(CallbackOrder.Metric)

        self.input_key = input_key
        self.output_key = output_key
        self.model_output_key = model_output_key
        self.prefix = prefix

    def on_loader_start(self, state: IRunner):
        self.prediction: List[np.ndarray] = []
        self.target: List[np.ndarray] = []

    def on_batch_end(self, state: IRunner):
        targ = state.input[self.input_key]["weak"].detach().cpu().numpy()
        out = state.output[self.output_key]

        clipwise_output = out[self.model_output_key].detach().cpu().numpy()

        self.prediction.append(clipwise_output)
        self.target.append(targ)

        score_class, weight = lwlrap(targ, clipwise_output)
        score = (score_class * weight).sum()
        state.batch_metrics[self.prefix] = score

    def on_loader_end(self, state: IRunner):
        y_pred = np.concatenate(self.prediction, axis=0)
        y_true = np.concatenate(self.target, axis=0)
        score_class, weight = lwlrap(y_true, y_pred)
        score = (score_class * weight).sum()
        if state.is_valid_loader:
            state.epoch_metrics[state.valid_loader + "_epoch_" + self.prefix] = score
        else:
            state.epoch_metrics["train_epoch_" + self.prefix] = score

import numpy as np

from typing import List

from catalyst.core import Callback, CallbackOrder, IRunner

from ._metrics import lwlrap


class LWLRAPCallback(Callback):
    def __init__(self,
                 input_key: str = "targets",
                 output_key: str = "logits",
                 model_output_key: str = "clipwise_output",
                 prefix: str = "lwlrap"):
        super().__init__(CallbackOrder.Metric)

        self.input_key = input_key
        self.output_key = output_key
        self.model_output_key = model_output_key
        self.prefix = prefix

    def on_loader_start(self, state: IRunner):
        self.prediction: List[np.ndarray] = []
        self.target: List[np.ndarray] = []

    def on_batch_end(self, state: IRunner):
        targ = state.input[self.input_key]["weak"].detach().cpu().numpy()
        out = state.output[self.output_key]

        clipwise_output = out[self.model_output_key].detach().cpu().numpy()

        self.prediction.append(clipwise_output)
        self.target.append(targ)

        score_class, weight = lwlrap(targ, clipwise_output)
        score = (score_class * weight).sum()
        state.batch_metrics[self.prefix] = score

    def on_loader_end(self, state: IRunner):
        y_pred = np.concatenate(self.prediction, axis=0)
        y_true = np.concatenate(self.target, axis=0)
        score_class, weight = lwlrap(y_true, y_pred)
        score = (score_class * weight).sum()
        if state.is_valid_loader:
            state.epoch_metrics[state.valid_loader + "_epoch_" + self.prefix] = score
        else:
            state.epoch_metrics["train_epoch_" + self.prefix] = score
