#
# Copyright 2020-present by A. Mathis Group and contributors. All rights reserved.
#
# This project and all its files are licensed under GNU AGPLv3 or later version. 
# A copy is included in dlc2action/LICENSE.AGPL.
#
"""Implementations of `dlc2action.metric.base_metric.Metric`."""

import warnings
from abc import abstractmethod
from collections import defaultdict
from copy import copy, deepcopy
from typing import Any, Dict, List, Set, Tuple, Union

import editdistance
import numpy as np
import torch
from dlc2action.metric.base_metric import Metric
from sklearn import metrics


class _ClassificationMetric(Metric):
    """The base class for all metric that are calculated from true and false negative and positive rates."""

    needs_raw_data = True
    segmental = False
    """
    If `True`, the metric will be calculated over segments; otherwise over frames.
    """

    def __init__(
        self,
        num_classes: int,
        ignore_index: int = -100,
        average: str = "macro",
        ignored_classes: Set = None,
        exclusive: bool = True,
        iou_threshold: float = 0.5,
        tag_average: str = "micro",
        threshold_values: List = None,
        integration_interval: int = 0,
    ):
        """Initialize the metric.

        Parameters
        ----------
        ignore_index : int, default -100
            the class index that indicates ignored samples
        average: {'macro', 'micro', 'none'}
            method for averaging across classes
        num_classes : int, optional
            number of classes (not necessary if main_class is not None)
        ignored_classes : set, optional
            a set of class ids to ignore in calculation
        exclusive: bool, default True
            set to False for multi-label classification tasks
        iou_threshold : float, default 0.5
            if segmental is true, intervals with IoU larger than this threshold are considered correct
        tag_average: {'micro', 'macro', 'none'}
            method for averaging across meta tags (if given)
        threshold_values : List, optional
            a list of float values between 0 and 1 that will be used as decision thresholds
        integration_interval : int, default 0
            if not 0, the metric will be calculated over a window of this size (in frames) around the current frame

        """
        super().__init__()
        self.ignore_index = ignore_index
        self.average = average
        self.tag_average = tag_average
        self.tags = set()
        self.exclusive = exclusive
        self.threshold = iou_threshold
        self.integration_interval = integration_interval
        self.classes = list(range(int(num_classes)))
        self.optimize = False
        if threshold_values is None:
            threshold_values = [0.5]
        elif threshold_values == ["optimize"]:
            threshold_values = list(np.arange(0.25, 0.75, 0.05))
            self.optimize = True
        if self.exclusive and threshold_values != [0.5]:
            raise ValueError(
                "Cannot set threshold values for exclusive classification!"
            )
        self.threshold_values = threshold_values

        if ignored_classes is not None:
            for c in ignored_classes:
                if c in self.classes:
                    self.classes.remove(c)
        if len(self.classes) == 0:
            warnings.warn("No classes are followed!")

    def reset(self) -> None:
        """Reset the internal state (at the beginning of an epoch)."""
        self.tags = set()
        self.tp = defaultdict(lambda: defaultdict(lambda: 0))
        self.fp = defaultdict(lambda: defaultdict(lambda: 0))
        self.fn = defaultdict(lambda: defaultdict(lambda: 0))
        self.tn = defaultdict(lambda: defaultdict(lambda: 0))

    def update(
        self,
        predicted: torch.Tensor,
        target: torch.Tensor,
        tags: torch.Tensor,
    ) -> None:
        """Update the internal state (with a batch).

        Parameters
        ----------
        predicted : torch.Tensor
            the main prediction tensor generated by the model
        ssl_predicted : torch.Tensor
            the SSL prediction tensor generated by the model
        target : torch.Tensor
            the corresponding main target tensor
        ssl_target : torch.Tensor
            the corresponding SSL target tensor
        tags : torch.Tensor
            the tensor of meta tags (or `None`, if tags are not given)

        """
        if self.segmental:
            self._update_segmental(predicted, target, tags)
        else:
            self._update_normal(predicted, target, tags)

    def _key(self, tag: torch.Tensor, c: int) -> str:
        """Get a key for the intermediate value dictionaries from tag and class indices."""
        if tag is None or self.tag_average == "micro":
            return c
        else:
            return f"tag{int(tag)}_{c}"

    def _update_normal(
        self,
        predicted: torch.Tensor,
        target: torch.Tensor,
        tags: torch.Tensor,
    ) -> None:
        """Update the internal state (with a batch), calculating over frames."""
        if self.integration_interval != 0:
            pr = 0
            denom = torch.ones((1, 1, predicted.shape[-1])) * (
                2 * self.integration_interval + 1
            )
            for i in range(self.integration_interval):
                pr += torch.cat(
                    [
                        torch.zeros((predicted.shape[0], predicted.shape[1], i + 1)),
                        predicted[:, :, i + 1 :],
                    ],
                    dim=-1,
                )
                pr += torch.cat(
                    [
                        predicted[:, :, : -i - 1],
                        torch.zeros((predicted.shape[0], predicted.shape[1], i + 1)),
                    ],
                    dim=-1,
                )
                denom[:, :, i] = i + 1
                denom[:, :, -i] = i + 1
            predicted = pr / denom
        if self.exclusive:
            predicted = torch.max(predicted, 1)[1]
        if self.tag_average == "micro" or tags is None or tags[0] is None:
            tag_set = {None}
        else:
            tag_set = set(tags)
            self.tags.update(tag_set)
        for thr in self.threshold_values:
            for t in tag_set:
                if t is None:
                    predicted_t = predicted
                    target_t = target
                else:
                    predicted_t = predicted[tags == t]
                    target_t = target[tags == t]
                for c in self.classes:
                    if self.exclusive:
                        pos = (predicted_t == c) * (target_t != self.ignore_index)
                        neg = (predicted_t != c) * (target_t != self.ignore_index)
                        self.tp[self._key(t, c)][thr] += ((target_t == c) * pos).sum()
                        self.fp[self._key(t, c)][thr] += ((target_t != c) * pos).sum()
                        self.fn[self._key(t, c)][thr] += ((target_t == c) * neg).sum()
                        self.tn[self._key(t, c)][thr] += ((target_t != c) * neg).sum()
                    else:
                        if isinstance(thr, list):
                            key = ", ".join(map(str, thr))
                            for i, tt in enumerate(thr):
                                predicted_t[:, i, :] = (predicted_t[:, i, :] > tt).int()
                        else:
                            key = thr
                            predicted_t = (predicted_t > thr).int()
                        pos = predicted_t[:, c, :] * (
                            target_t[:, c, :] != self.ignore_index
                        )
                        neg = (predicted_t[:, c, :] != 1) * (
                            target_t[:, c, :] != self.ignore_index
                        )
                        self.tp[self._key(t, c)][key] += (target_t[:, c, :] * pos).sum()
                        self.fp[self._key(t, c)][key] += (
                            (target_t[:, c, :] != 1) * pos
                        ).sum()
                        self.fn[self._key(t, c)][key] += (target_t[:, c, :] * neg).sum()
                        self.tn[self._key(t, c)][key] += (
                            (target_t[:, c, :] != 1) * neg
                        ).sum()

    def _get_intervals(self, tensor: torch.Tensor) -> torch.Tensor:
        """Get a list of `True` group beginning and end indices from a boolean tensor."""
        output, indices = torch.unique_consecutive(tensor, return_inverse=True)
        true_indices = torch.where(output)[0]
        starts = torch.tensor(
            [(indices == i).nonzero(as_tuple=True)[0][0] for i in true_indices]
        )
        ends = torch.tensor(
            [(indices == i).nonzero(as_tuple=True)[0][-1] + 1 for i in true_indices]
        )
        return torch.stack([starts, ends]).T

    def _smooth(self, tensor: torch.Tensor, smooth_interval: int = 1) -> torch.Tensor:
        """Get rid of jittering in a non-exclusive classification tensor.

        First, remove intervals of 0 shorter than `smooth_interval`. Then, remove intervals of 1 shorter than
        `smooth_interval`.

        """
        for c in self.classes:
            intervals = self._get_intervals(tensor[:, c] == 0)
            interval_lengths = torch.tensor(
                [interval[1] - interval[0] for interval in intervals]
            )
            short_intervals = intervals[interval_lengths <= smooth_interval]
            for start, end in short_intervals:
                tensor[start:end, c] = 1
            intervals = self._get_intervals(tensor[:, c] == 1)
            interval_lengths = torch.tensor(
                [interval[1] - interval[0] for interval in intervals]
            )
            short_intervals = intervals[interval_lengths <= smooth_interval]
            for start, end in short_intervals:
                tensor[start:end, c] = 0
        return tensor

    def _sigmoid_threshold_function(
        self,
        low_threshold: float,
        high_threshold: float,
        low_length: int,
        high_length: int,
    ):
        """Generate a sigmoid threshold function.

        The resulting function outputs an intersection threshold given the length of the interval.

        """
        a = 2 / (high_length - low_length)
        b = 1 - a * high_length
        return lambda x: low_threshold + torch.sigmoid(4 * (a * x + b)) * (
            high_threshold - low_threshold
        )

    def _update_segmental(
        self,
        predicted: torch.tensor,
        target: torch.tensor,
        tags: torch.Tensor,
    ) -> None:
        """Update the internal state (with a batch), calculating over segments."""
        if self.exclusive:
            predicted = torch.max(predicted, 1)[1]
        predicted = torch.cat(
            [
                copy(predicted),
                -100 * torch.ones((*predicted.shape[:-1], 1)).to(predicted.device),
            ],
            dim=-1,
        )
        target = torch.cat(
            [
                copy(target),
                -100 * torch.ones((*target.shape[:-1], 1)).to(target.device),
            ],
            dim=-1,
        )
        if self.exclusive:
            predicted = predicted.flatten()
            target = target.flatten()
        else:
            num_classes = predicted.shape[1]
            predicted = predicted.transpose(1, 2).reshape(-1, num_classes)
            target = target.transpose(1, 2).reshape(-1, num_classes)
        if self.tag_average == "micro" or tags is None or tags[0] is None:
            tag_set = {None}
        else:
            tag_set = set(tags)
            self.tags.update(tag_set)
        for thr in self.threshold_values:
            key = thr
            for t in tag_set:
                if t is None:
                    predicted_t = predicted
                    target_t = target
                else:
                    predicted_t = predicted[tags == t]
                    target_t = target[tags == t]
                if not self.exclusive:
                    if isinstance(thr, list):
                        for i, tt in enumerate(thr):
                            predicted_t[i, :] = (predicted_t[i, :] > tt).int()
                        key = ", ".join(map(str, thr))
                    else:
                        predicted_t = (predicted_t > thr).int()
                for c in self.classes:
                    if self.exclusive:
                        predicted_intervals = self._get_intervals(predicted_t == c)
                        target_intervals = self._get_intervals(target_t == c)
                    else:
                        predicted_intervals = self._get_intervals(
                            predicted_t[:, c] == 1
                        )
                        target_intervals = self._get_intervals(target_t[:, c] == 1)
                    true_used = torch.zeros(target_intervals.shape[0])
                    for interval in predicted_intervals:
                        if len(target_intervals) > 0:
                            # Compute IoU against all others
                            intersection = torch.minimum(
                                interval[1], target_intervals[:, 1]
                            ) - torch.maximum(interval[0], target_intervals[:, 0])
                            union = torch.maximum(
                                interval[1], target_intervals[:, 1]
                            ) - torch.minimum(interval[0], target_intervals[:, 0])
                            IoU = intersection / union

                            # Get the best scoring segment
                            idx = IoU.argmax()

                            # If the IoU is high enough and the true segment isn't already used
                            # Then it is a true positive. Otherwise is it a false positive.
                            if IoU[idx] >= self.threshold and not true_used[idx]:
                                self.tp[self._key(t, c)][key] += 1
                                true_used[idx] = 1
                            else:
                                self.fp[self._key(t, c)][key] += 1
                        else:
                            self.fp[self._key(t, c)][key] += 1
                    self.fn[self._key(t, c)][key] += len(true_used) - torch.sum(
                        true_used
                    )

    def calculate(self) -> Union[float, Dict]:
        """Calculate the metric (at the end of an epoch).

        Returns
        -------
        result : float | dict
            either the single value of the metric or a dictionary where the keys are class indices and the values
            are class metric values

        """
        if self.tag_average == "micro" or self.tags == {None}:
            self.tags = {None}
        result = {}
        self.tags = sorted(list(self.tags))
        for tag in self.tags:
            if self.average == "macro":
                metric_list = []
                for c in self.classes:
                    metric_list.append(
                        self._calculate_metric(
                            self.tp[self._key(tag, c)],
                            self.fp[self._key(tag, c)],
                            self.fn[self._key(tag, c)],
                            self.tn[self._key(tag, c)],
                        )
                    )
                if tag is not None:
                    tag = int(tag)
                result[f"tag{tag}"] = sum(metric_list) / len(metric_list)
            elif self.average == "micro":
                # tp = sum([self.tp[self._key(tag, c)] for c in self.classes])
                # fn = sum([self.fn[self._key(tag, c)] for c in self.classes])
                # fp = sum([self.fp[self._key(tag, c)] for c in self.classes])
                # tn = sum([self.tn[self._key(tag, c)] for c in self.classes])
                # if tag is not None:
                #     tag = int(tag)
                # result[f"tag{tag}"] = self._calculate_metric(tp, fp, fn, tn)
                metric_list = []
                for c in self.classes:
                    metric_list.append(
                        self._calculate_metric(
                            self.tp[self._key(tag, c)],
                            self.fp[self._key(tag, c)],
                            self.fn[self._key(tag, c)],
                            self.tn[self._key(tag, c)],
                        )
                    )

                if tag is not None:
                    tag = int(tag)
                result[f"tag{tag}"] = metric_list
            elif self.average == "none":
                metric_dict = {}
                for c in self.classes:
                    metric_dict[self._key(tag, c)] = self._calculate_metric(
                        self.tp[self._key(tag, c)],
                        self.fp[self._key(tag, c)],
                        self.fn[self._key(tag, c)],
                        self.tn[self._key(tag, c)],
                    )
                result.update(metric_dict)
            else:
                raise ValueError(
                    f"The {self.average} averaging method is not available, please choose from "
                    f'["none", "micro", "macro"]'
                )
        if len(self.tags) == 1 and self.average != "none":
            tag = self.tags[0]
            if tag is not None:
                tag = int(tag)
            result = result[f"tag{tag}"]
        elif self.tag_average == "macro":
            if self.average == "none":
                r = {}
                for c in self.classes:
                    r[c] = torch.mean(
                        torch.tensor([result[self._key(tag, c)] for tag in self.tags])
                    )
            else:
                r = torch.mean(
                    torch.tensor([result[f"tag{int(tag)}"] for tag in self.tags])
                )
            result = r
        return result

    @abstractmethod
    def _calculate_metric(self, tp: Dict, fp: Dict, fn: Dict, tn: Dict) -> float:
        """Calculate the metric value from true and false positive and negative rates.

        Parameters
        ----------
        tp : float
            true positive rate
        fp: float
            false positive rate
        fn: float
            false negative rate
        tn: float
            true negative rate

        Returns
        -------
        metric : float
            metric value

        """


class PR_AUC(_ClassificationMetric):
    """Area under precision-recall curve (not advised for training)."""

    def __init__(
        self,
        num_classes: int,
        ignore_index: int = -100,
        average: str = "macro",
        ignored_classes: Set = None,
        exclusive: bool = False,
        tag_average: str = "micro",
        threshold_step: float = 0.1,
    ):
        """Initialize the metric.

        Parameters
        ----------
        ignore_index : int, default -100
            the class index that indicates ignored samples
        average: {'macro', 'micro', 'none'}
            method for averaging across classes
        num_classes : int, optional
            number of classes (not necessary if main_class is not None)
        ignored_classes : set, optional
            a set of class ids to ignore in calculation
        exclusive: bool, default True
            set to False for multi-label classification tasks
        tag_average: {'micro', 'macro', 'none'}
            method for averaging across meta tags (if given)
        threshold_step : float, default 0.1
            the decision threshold step

        """
        if exclusive:
            raise ValueError(
                "The PR-AUC metric is not implemented for exclusive classification!"
            )
        super().__init__(
            num_classes,
            ignore_index,
            average,
            ignored_classes,
            exclusive,
            tag_average=tag_average,
            threshold_values=list(np.arange(0, 1, threshold_step)),
        )

    def _calculate_metric(self, tp: Dict, fp: Dict, fn: Dict, tn: Dict) -> float:
        """Calculate the metric value from true and false positive and negative rates."""
        precisions = []
        recalls = []
        for k in sorted(self.threshold_values):
            precisions.append(tp[k] / (tp[k] + fp[k] + 1e-7))
            recalls.append(tp[k] / (tp[k] + fn[k] + 1e-7))
        return metrics.auc(x=recalls, y=precisions)


class Precision(_ClassificationMetric):
    """Precision."""

    def __init__(
        self,
        num_classes: int,
        ignore_index: int = -100,
        average: str = "macro",
        ignored_classes: Set = None,
        exclusive: bool = True,
        tag_average: str = "micro",
        threshold_value: Union[float, List] = None,
    ):
        """Initialize the class.

        Parameters
        ----------
        ignore_index : int, default -100
            the class index that indicates ignored samples
        average: {'macro', 'micro', 'none'}
            method for averaging across classes
        num_classes : int, optional
            number of classes (not necessary if main_class is not None)
        ignored_classes : set, optional
            a set of class ids to ignore in calculation
        exclusive: bool, default True
            set to False for multi-label classification tasks
        tag_average: {'micro', 'macro', 'none'}
            method for averaging across meta tags (if given)
        threshold_value : float | list, optional
            the decision threshold value (cannot be defined for exclusive classification, and 0.5 be default
            for non-exclusive); if `threshold_value` is a list, every value should correspond to the class
            under the same index

        """
        if threshold_value is None:
            threshold_value = 0.5
        super().__init__(
            num_classes,
            ignore_index,
            average,
            ignored_classes,
            exclusive,
            tag_average=tag_average,
            threshold_values=[threshold_value],
        )

    def _calculate_metric(self, tp: Dict, fp: Dict, fn: Dict, tn: Dict) -> float:
        """Calculate the metric value from true and false positive and negative rates."""
        k = self.threshold_values[0]
        if isinstance(k, list):
            k = ", ".join(map(str, k))
        return tp[k] / (tp[k] + fp[k] + 1e-7)


class SegmentalPrecision(_ClassificationMetric):
    """Segmental precision (not advised for training)."""

    segmental = True

    def __init__(
        self,
        num_classes: int,
        ignore_index: int = -100,
        average: str = "macro",
        ignored_classes: Set = None,
        exclusive: bool = True,
        iou_threshold: float = 0.5,
        tag_average: str = "micro",
        threshold_value: Union[float, List] = None,
    ):
        """Initialize the class.

        Parameters
        ----------
        ignore_index : int, default -100
            the class index that indicates ignored samples
        average: {'macro', 'micro', 'none'}
            method for averaging across classes
        num_classes : int, optional
            number of classes (not necessary if main_class is not None)
        ignored_classes : set, optional
            a set of class ids to ignore in calculation
        exclusive: bool, default True
            set to False for multi-label classification tasks
        iou_threshold : float, default 0.5
            if segmental is true, intervals with IoU larger than this threshold are considered correct
        tag_average: {'micro', 'macro', 'none'}
            method for averaging across meta tags (if given)
        threshold_value : float | list, optional
            the decision threshold value (cannot be defined for exclusive classification, and 0.5 be default
            for non-exclusive); if `threshold_value` is a list, every value should correspond to the class
            under the same index

        """
        if threshold_value is None:
            threshold_value = 0.5
        super().__init__(
            num_classes,
            ignore_index,
            average,
            ignored_classes,
            exclusive,
            iou_threshold,
            tag_average,
            [threshold_value],
        )

    def _calculate_metric(self, tp: Dict, fp: Dict, fn: Dict, tn: Dict) -> float:
        """Calculate the metric value from true and false positive and negative rates."""
        k = self.threshold_values[0]
        if isinstance(k, list):
            k = ", ".join(map(str, k))
        return tp[k] / (tp[k] + fp[k] + 1e-7)


class Recall(_ClassificationMetric):
    """Recall."""

    def __init__(
        self,
        num_classes: int,
        ignore_index: int = -100,
        average: str = "macro",
        ignored_classes: Set = None,
        exclusive: bool = True,
        tag_average: str = "micro",
        threshold_value: Union[float, List] = None,
    ):
        """Initialize the class.

        Parameters
        ----------
        ignore_index : int, default -100
            the class index that indicates ignored samples
        average: {'macro', 'micro', 'none'}
            method for averaging across classes
        num_classes : int, optional
            number of classes (not necessary if main_class is not None)
        ignored_classes : set, optional
            a set of class ids to ignore in calculation
        exclusive: bool, default True
            set to False for multi-label classification tasks
        tag_average: {'micro', 'macro', 'none'}
            method for averaging across meta tags (if given)
        threshold_value : float | list, optional
            the decision threshold value (cannot be defined for exclusive classification, and 0.5 be default
            for non-exclusive); if `threshold_value` is a list, every value should correspond to the class
            under the same index

        """
        if threshold_value is None:
            threshold_value = 0.5
        super().__init__(
            num_classes,
            ignore_index,
            average,
            ignored_classes,
            exclusive,
            tag_average=tag_average,
            threshold_values=[threshold_value],
        )

    def _calculate_metric(self, tp: Dict, fp: Dict, fn: Dict, tn: Dict) -> float:
        """Calculate the metric value from true and false positive and negative rates."""
        k = self.threshold_values[0]
        if isinstance(k, list):
            k = ", ".join(map(str, k))
        return tp[k] / (tp[k] + fn[k] + 1e-7)


class SegmentalRecall(_ClassificationMetric):
    """Segmental recall (not advised for training)."""

    segmental = True

    def __init__(
        self,
        num_classes: int,
        ignore_index: int = -100,
        average: str = "macro",
        ignored_classes: Set = None,
        exclusive: bool = True,
        iou_threshold: float = 0.5,
        tag_average: str = "micro",
        threshold_value: Union[float, List] = None,
    ):
        """Initialize the class.

        Parameters
        ----------
        ignore_index : int, default -100
            the class index that indicates ignored samples
        average: {'macro', 'micro', 'none'}
            method for averaging across classes
        num_classes : int, optional
            number of classes (not necessary if main_class is not None)
        ignored_classes : set, optional
            a set of class ids to ignore in calculation
        exclusive: bool, default True
            set to False for multi-label classification tasks
        iou_threshold : float, default 0.5
            if segmental is true, intervals with IoU larger than this threshold are considered correct
        tag_average: {'micro', 'macro', 'none'}
            method for averaging across meta tags (if given)
        threshold_value : float | list, optional
            the decision threshold value (cannot be defined for exclusive classification, and 0.5 be default
            for non-exclusive); if `threshold_value` is a list, every value should correspond to the class
            under the same index

        """
        if threshold_value is None:
            threshold_value = 0.5
        super().__init__(
            num_classes,
            ignore_index,
            average,
            ignored_classes,
            exclusive,
            iou_threshold,
            tag_average,
            [threshold_value],
        )

    def _calculate_metric(self, tp: Dict, fp: Dict, fn: Dict, tn: Dict) -> float:
        """Calculate the metric value from true and false positive and negative rates."""
        k = self.threshold_values[0]
        if isinstance(k, list):
            k = ", ".join(map(str, k))
        return tp[k] / (tp[k] + fn[k] + 1e-7)


class F1(_ClassificationMetric):
    """F1 score."""

    def __init__(
        self,
        num_classes: int,
        ignore_index: int = -100,
        average: str = "macro",
        ignored_classes: Set = None,
        exclusive: bool = True,
        tag_average: str = "micro",
        threshold_value: Union[float, List] = None,
        integration_interval: int = 0,
    ):
        """Initialize the class.

        Parameters
        ----------
        ignore_index : int, default -100
            the class index that indicates ignored samples
        average: {'macro', 'micro', 'none'}
            method for averaging across classes
        num_classes : int, optional
            number of classes (not necessary if main_class is not None)
        ignored_classes : set, optional
            a set of class ids to ignore in calculation
        exclusive: bool, default True
            set to False for multi-label classification tasks
        tag_average: {'micro', 'macro', 'none'}
            method for averaging across meta tags (if given)
        threshold_value : float | list, optional
            the decision threshold value (cannot be defined for exclusive classification, and 0.5 be default
            for non-exclusive); if `threshold_value` is a list, every value should correspond to the class
            under the same index
        integration_interval : int, default 0
            the number of frames to integrate over (0 means no integration)

        """
        if threshold_value is None:
            threshold_value = 0.5
        super().__init__(
            num_classes,
            ignore_index,
            average,
            ignored_classes,
            exclusive,
            tag_average=tag_average,
            threshold_values=[threshold_value],
            integration_interval=integration_interval,
        )

    def _calculate_metric(self, tp: Dict, fp: Dict, fn: Dict, tn: Dict) -> float:
        """Calculate the metric value from true and false positive and negative rates."""
        if self.optimize:
            scores = []
            for k in self.threshold_values:
                recall = tp[k] / (tp[k] + fn[k] + 1e-7)
                precision = tp[k] / (tp[k] + fp[k] + 1e-7)
                scores.append(2 * recall * precision / (recall + precision + 1e-7))
            f1 = max(scores)
        else:
            k = self.threshold_values[0]
            if isinstance(k, list):
                k = ", ".join(map(str, k))
            recall = tp[k] / (tp[k] + fn[k] + 1e-7)
            precision = tp[k] / (tp[k] + fp[k] + 1e-7)
            f1 = 2 * recall * precision / (recall + precision + 1e-7)
        return f1


class SegmentalF1(_ClassificationMetric):
    """Segmental F1 score (not advised for training)."""

    segmental = True

    def __init__(
        self,
        num_classes: int,
        ignore_index: int = -100,
        average: str = "macro",
        ignored_classes: Set = None,
        exclusive: bool = True,
        iou_threshold: float = 0.5,
        tag_average: str = "micro",
        threshold_value: Union[float, List] = None,
    ):
        """Initialize the class.

        Parameters
        ----------
        ignore_index : int, default -100
            the class index that indicates ignored samples
        average: {'macro', 'micro', 'none'}
            method for averaging across classes
        num_classes : int, optional
            number of classes (not necessary if main_class is not None)
        ignored_classes : set, optional
            a set of class ids to ignore in calculation
        exclusive: bool, default True
            set to False for multi-label classification tasks
        iou_threshold : float, default 0.5
            if segmental is true, intervals with IoU larger than this threshold are considered correct
        tag_average: {'micro', 'macro', 'none'}
            method for averaging across meta tags (if given)
        threshold_value : float | list, optional
            the decision threshold value (cannot be defined for exclusive classification, and 0.5 be default
            for non-exclusive); if `threshold_value` is a list, every value should correspond to the class
            under the same index

        """
        if threshold_value is None:
            threshold_value = 0.5
        super().__init__(
            num_classes,
            ignore_index,
            average,
            ignored_classes,
            exclusive,
            iou_threshold,
            tag_average,
            [threshold_value],
        )

    def _calculate_metric(self, tp: Dict, fp: Dict, fn: Dict, tn: Dict) -> float:
        """Calculate the metric value from true and false positive and negative rates."""
        if self.optimize:
            scores = []
            for k in self.threshold_values:
                recall = tp[k] / (tp[k] + fn[k] + 1e-7)
                precision = tp[k] / (tp[k] + fp[k] + 1e-7)
                scores.append(2 * recall * precision / (recall + precision + 1e-7))
            f1 = max(scores)
        else:
            k = self.threshold_values[0]
            if isinstance(k, list):
                k = ", ".join(map(str, k))
            recall = tp[k] / (tp[k] + fn[k] + 1e-7)
            precision = tp[k] / (tp[k] + fp[k] + 1e-7)
            f1 = 2 * recall * precision / (recall + precision + 1e-7)
        return f1


class Fbeta(_ClassificationMetric):
    """F-beta score."""

    def __init__(
        self,
        beta: float = 1,
        ignore_index: int = -100,
        average: str = "macro",
        num_classes: int = None,
        ignored_classes: Set = None,
        tag_average: str = "micro",
        exclusive: bool = True,
        threshold_value: float = 0.5,
    ):
        """Initialize the class.

        Parameters
        ----------
        beta : float, default 1
            the beta parameter
        ignore_index : int, default -100
            the class index that indicates ignored samples
        average: {'macro', 'micro', 'none'}
            method for averaging across classes
        num_classes : int, optional
            number of classes (not necessary if main_class is not None)
        ignored_classes : set, optional
            a set of class ids to ignore in calculation
        exclusive: bool, default True
            set to False for multi-label classification tasks
        tag_average: {'micro', 'macro', 'none'}
            method for averaging across meta tags (if given)
        threshold_value : float | list, optional
            the decision threshold value (cannot be defined for exclusive classification, and 0.5 be default
            for non-exclusive); if `threshold_value` is a list, every value should correspond to the class
            under the same index

        """
        if threshold_value is None:
            threshold_value = 0.5
        self.beta2 = beta**2
        super().__init__(
            num_classes,
            ignore_index,
            average,
            ignored_classes,
            exclusive,
            tag_average=tag_average,
            threshold_values=[threshold_value],
        )

    def _calculate_metric(self, tp: Dict, fp: Dict, fn: Dict, tn: Dict) -> float:
        """Calculate the metric value from true and false positive and negative rates."""
        if self.optimize:
            scores = []
            for k in self.threshold_values:
                recall = tp[k] / (tp[k] + fn[k] + 1e-7)
                precision = tp[k] / (tp[k] + fp[k] + 1e-7)
                scores.append(
                    (
                        (1 + self.beta2)
                        * precision
                        * recall
                        / (self.beta2 * precision + recall + 1e-7)
                    )
                )
            f1 = max(scores)
        else:
            k = self.threshold_values[0]
            if isinstance(k, list):
                k = ", ".join(map(str, k))
            recall = tp[k] / (tp[k] + fn[k] + 1e-7)
            precision = tp[k] / (tp[k] + fp[k] + 1e-7)
            f1 = (
                (1 + self.beta2)
                * precision
                * recall
                / (self.beta2 * precision + recall + 1e-7)
            )
        return f1


class SegmentalFbeta(_ClassificationMetric):
    """Segmental F-beta score (not advised for training)."""

    segmental = True

    def __init__(
        self,
        beta: float = 1,
        ignore_index: int = -100,
        average: str = "macro",
        num_classes: int = None,
        ignored_classes: Set = None,
        iou_threshold: float = 0.5,
        tag_average: str = "micro",
        exclusive: bool = True,
        threshold_value: float = 0.5,
    ):
        """Initialize the class.

        Parameters
        ----------
        beta : float, default 1
            the beta parameter
        ignore_index : int, default -100
            the class index that indicates ignored samples
        average: {'macro', 'micro', 'none'}
            method for averaging across classes
        num_classes : int, optional
            number of classes (not necessary if main_class is not None)
        ignored_classes : set, optional
            a set of class ids to ignore in calculation
        exclusive: bool, default True
            set to False for multi-label classification tasks
        iou_threshold : float, default 0.5
            if segmental is true, intervals with IoU larger than this threshold are considered correct
        tag_average: {'micro', 'macro', 'none'}
            method for averaging across meta tags (if given)
        threshold_value : float | list, optional
            the decision threshold value (cannot be defined for exclusive classification, and 0.5 be default
            for non-exclusive); if `threshold_value` is a list, every value should correspond to the class
            under the same index

        """
        if threshold_value is None:
            threshold_value = 0.5
        self.beta2 = beta**2
        super().__init__(
            num_classes,
            ignore_index,
            average,
            ignored_classes,
            exclusive,
            iou_threshold=iou_threshold,
            tag_average=tag_average,
            threshold_values=[threshold_value],
        )

    def _calculate_metric(self, tp: Dict, fp: Dict, fn: Dict, tn: Dict) -> float:
        """Calculate the metric value from true and false positive and negative rates."""
        if self.optimize:
            scores = []
            for k in self.threshold_values:
                recall = tp[k] / (tp[k] + fn[k] + 1e-7)
                precision = tp[k] / (tp[k] + fp[k] + 1e-7)
                scores.append(
                    (
                        (1 + self.beta2)
                        * precision
                        * recall
                        / (self.beta2 * precision + recall + 1e-7)
                    )
                )
            f1 = max(scores)
        else:
            k = self.threshold_values[0]
            if isinstance(k, list):
                k = ", ".join(map(str, k))
            recall = tp[k] / (tp[k] + fn[k] + 1e-7)
            precision = tp[k] / (tp[k] + fp[k] + 1e-7)
            f1 = (
                (1 + self.beta2)
                * precision
                * recall
                / (self.beta2 * precision + recall + 1e-7)
            )
        return f1


class _SemiSegmentalMetric(_ClassificationMetric):
    def __init__(
        self,
        num_classes: int,
        ignore_index: int = -100,
        ignored_classes: Set = None,
        exclusive: bool = True,
        average: str = "macro",
        tag_average: str = "micro",
        delta: int = 0,
        smooth_interval: int = 0,
        iou_threshold_long: float = 0.5,
        iou_threshold_short: float = 0.5,
        short_length: int = 30,
        long_length: int = 300,
        threshold_values: List = None,
    ) -> None:
        """Initialize the class.

        Parameters
        ----------
        num_classes : int
            the number of classes in the dataset
        ignore_index : int, default -100
            the ground truth label to ignore
        ignored_classes : set, optional
            the class indices to ignore in computation
        exclusive : bool, default True
            `False` for multi-label classification tasks
        average : {"macro", "micro", "none"}
            the method to average the results over classes
        tag_average : {"macro", "micro", "none"}
            the method to average the results over meta tags (if given)
        delta : int, default 0
            the number of frames to add to each ground truth interval before computing the intersection,
            see description of the class for details
        smooth_interval : int, default 0
            intervals shorter than this number of frames will be ignored (both in prediction and in ground truth,
            see description of the class for details
        iou_threshold_long : float, default 0.5
            the intersection threshold for segments longer than `long_length` frames (between 0 and 1),
            see description of the class for details
        iou_threshold_short : float, default 0.5
            the intersection threshold for segments shorter than `short_length` frames (between 0 and 1),
            see description of the class for details
        short_length : int, default 30
            the threshold number of frames for short intervals that will have an intersection threshold of
            `iou_threshold_short`, see description of the class for details
        long_length : int, default 300
            the threshold number of frames for long intervals that will have an intersection threshold of
            `iou_threshold_long`, see description of the class for details
        threshold_values : list, optional
            the decision threshold values (cannot be defined for exclusive classification, and 0.5 be default)

        """
        if smooth_interval > 0 and exclusive:
            warnings.warn(
                "Smoothing is not implemented for datasets with exclusive classification! Setting smooth_interval to 0..."
            )
        if threshold_values == [None]:
            threshold_values = [0.5]
        super().__init__(
            num_classes,
            ignore_index,
            average,
            ignored_classes,
            exclusive,
            tag_average=tag_average,
            threshold_values=threshold_values,
        )
        self.delta = delta
        self.random_sampling = False
        self.smooth_interval = smooth_interval
        self.threshold_function = self._sigmoid_threshold_function(
            low_threshold=iou_threshold_short,
            high_threshold=iou_threshold_long,
            low_length=short_length,
            high_length=long_length,
        )

    def update(
        self,
        predicted: torch.Tensor,
        target: torch.Tensor,
        tags: torch.Tensor,
    ) -> None:
        """Update the internal state (with a batch).

        Parameters
        ----------
        predicted : torch.Tensor
            the main prediction tensor generated by the model
        ssl_predicted : torch.Tensor
            the SSL prediction tensor generated by the model
        target : torch.Tensor
            the corresponding main target tensor
        ssl_target : torch.Tensor
            the corresponding SSL target tensor
        tags : torch.Tensor
            the tensor of meta tags (or `None`, if tags are not given)

        """
        if self.exclusive:
            predicted = torch.max(predicted, 1)[1]
        predicted = torch.cat(
            [
                copy(predicted),
                -100
                * torch.ones((*predicted.shape[:-1], self.delta + 1)).to(
                    predicted.device
                ),
            ],
            dim=-1,
        )
        target = torch.cat(
            [
                copy(target),
                -100
                * torch.ones((*target.shape[:-1], self.delta + 1)).to(target.device),
            ],
            dim=-1,
        )
        if self.exclusive:
            predicted = predicted.flatten()
            target = target.flatten()
        else:
            predicted = predicted.transpose(1, 2).reshape(-1, target.shape[1])
            target = target.transpose(1, 2).reshape(-1, predicted.shape[1])
        if self.tag_average == "micro" or tags is None or tags[0] is None:
            tag_set = {None}
        else:
            tag_set = set(tags)
            self.tags.update(tag_set)
        if self.smooth_interval > 0:
            target = self._smooth(target, self.smooth_interval)
        for thr in self.threshold_values:
            key = thr
            for t in tag_set:
                if t is None:
                    predicted_t = deepcopy(predicted)
                    target_t = target
                else:
                    predicted_t = deepcopy(predicted[tags == t])
                    target_t = target[tags == t]
                if not self.exclusive:
                    if isinstance(thr, list):
                        for i, tt in enumerate(thr):
                            predicted_t[i, :] = (predicted_t[i, :] > tt).int()
                        key = ", ".join(map(str, thr))
                    else:
                        predicted_t = (predicted_t > thr).int()
                if self.smooth_interval > 0:
                    predicted_t = self._smooth(predicted_t, self.smooth_interval)
                for c in self.classes:
                    if not self.random_sampling:
                        if self.exclusive:
                            target_intervals = self._get_intervals(target_t == c)
                        else:
                            target_intervals = self._get_intervals(target_t[:, c] == 1)
                        target_lengths = [
                            end - start for start, end in target_intervals
                        ]
                        target_intervals = [
                            [
                                max(start - self.delta, 0),
                                min(end + self.delta, len(target_t)),
                            ]
                            for start, end in target_intervals
                        ]
                    else:
                        if self.exclusive:
                            target_arr = target_t == c
                        else:
                            target_arr = target_t[:, c] == 1
                        target_points = torch.where(target_arr)[0][::20]
                        target_intervals = []
                        for p in target_points:
                            target_intervals.append(
                                [
                                    max(0, p - self.delta),
                                    min(len(target_arr), p + self.delta),
                                ]
                            )
                        target_lengths = [
                            end - start for start, end in target_intervals
                        ]
                    if self.exclusive:
                        predicted_arr = predicted_t == c
                    else:
                        predicted_arr = predicted_t[:, c] == 1
                    for interval, l in zip(target_intervals, target_lengths):
                        intersection = torch.sum(
                            predicted_arr[interval[0] : interval[1]]
                        )
                        IoU = intersection / l
                        if IoU >= self.threshold_function(l):
                            self.tp[self._key(t, c)][key] += 1
                        else:
                            self.fn[self._key(t, c)][key] += 1

                    if not self.random_sampling:
                        if self.exclusive:
                            predicted_intervals = self._get_intervals(predicted_t == c)
                        else:
                            predicted_intervals = self._get_intervals(
                                predicted_t[:, c] == 1
                            )
                        predicted_intervals_delta = [
                            [
                                max(start - self.delta, 0),
                                min(end + self.delta, len(target_t)),
                            ]
                            for start, end in predicted_intervals
                        ]
                    else:
                        if self.exclusive:
                            predicted_arr = predicted_t == c
                        else:
                            predicted_arr = predicted_t[:, c] == 1
                        predicted_points = torch.where(predicted_arr)[0][::20]
                        predicted_intervals = []
                        for p in predicted_points:
                            predicted_intervals.append(
                                [
                                    max(0, p - self.delta),
                                    min(len(predicted_arr), p + self.delta),
                                ]
                            )
                        predicted_intervals_delta = predicted_intervals
                    if self.exclusive:
                        target_arr = target_t == c
                        target_arr[target_t == -100] = -100
                    else:
                        target_arr = target_t[:, c]
                    for interval, interval_delta in zip(
                        predicted_intervals, predicted_intervals_delta
                    ):
                        if torch.sum(
                            target_arr[interval_delta[0] : interval_delta[1]] != -100
                        ) < 0.3 * (interval[1] - interval[0]):
                            continue
                        l = torch.sum(target_arr[interval[0] : interval[1]] != -100)
                        intersection = torch.sum(
                            target_arr[interval_delta[0] : interval_delta[1]] == 1
                        )
                        IoU = intersection / l
                        if IoU >= self.threshold_function(l):
                            self.tn[self._key(t, c)][key] += 1
                        else:
                            self.fp[self._key(t, c)][key] += 1


class SemiSegmentalRecall(_SemiSegmentalMetric):
    """Semi-segmental recall (not advised for training).

    A metric in-between segmental and frame-wise recall.

    This metric follows the following algorithm:
    1) smooth over too-short intervals, both in ground truth and in prediction (first remove
        groups of zeros shorter than `smooth_interval` and then do the same with groups of ones),
    2) add `delta` frames to each ground truth interval at both ends and count the number of predicted
        positive frames at the resulting intervals (intersection),
    3) calculate the threshold for each interval as
        `t = sigmoid(4 * (a * x + b)) * (iou_threshold_long - iou_threshold_short))`, where
        `a = 2 / (long_length - short_length)`, `b = 1 - a * long_length`, `x` is the length of the interval
        before `delta` was added,
    4) for each interval, if intersection is higher than `t * x`, the interval is labeled as true positive (`TP`),
        and otherwise as false negative (`FN`),
    5) the final metric value is computed as `TP / (TP + FN)`.
    """

    def __init__(
        self,
        num_classes: int,
        ignore_index: int = -100,
        ignored_classes: Set = None,
        exclusive: bool = True,
        average: str = "macro",
        tag_average: str = "micro",
        delta: int = 0,
        smooth_interval: int = 0,
        iou_threshold_long: float = 0.5,
        iou_threshold_short: float = 0.5,
        short_length: int = 30,
        long_length: int = 300,
        threshold_value: Union[float, List] = None,
    ) -> None:
        """Initialize the class.

        Parameters
        ----------
        num_classes : int
            the number of classes in the dataset
        ignore_index : int, default -100
            the ground truth label to ignore
        ignored_classes : set, optional
            the class indices to ignore in computation
        exclusive : bool, default True
            `False` for multi-label classification tasks
        average : {"macro", "micro", "none"}
            the method to average the results over classes
        tag_average : {"macro", "micro", "none"}
            the method to average the results over meta tags (if given)
        delta : int, default 0
            the number of frames to add to each ground truth interval before computing the intersection,
            see description of the class for details
        smooth_interval : int, default 0
            intervals shorter than this number of frames will be ignored (both in prediction and in ground truth,
            see description of the class for details
        iou_threshold_long : float, default 0.5
            the intersection threshold for segments longer than `long_length` frames (between 0 and 1),
            see description of the class for details
        iou_threshold_short : float, default 0.5
            the intersection threshold for segments shorter than `short_length` frames (between 0 and 1),
            see description of the class for details
        short_length : int, default 30
            the threshold number of frames for short intervals that will have an intersection threshold of
            `iou_threshold_short`, see description of the class for details
        long_length : int, default 300
            the threshold number of frames for long intervals that will have an intersection threshold of
            `iou_threshold_long`, see description of the class for details
        threshold_value : float | list, optional
            the decision threshold value (cannot be defined for exclusive classification, and 0.5 be default)

        """
        super().__init__(
            num_classes,
            ignore_index,
            ignored_classes,
            exclusive,
            average,
            tag_average,
            delta,
            smooth_interval,
            iou_threshold_long,
            iou_threshold_short,
            short_length,
            long_length,
            [threshold_value],
        )

    def _calculate_metric(self, tp: Dict, fp: Dict, fn: Dict, tn: Dict) -> float:
        """Calculate the metric value from true and false positive and negative rates."""
        k = self.threshold_values[0]
        if isinstance(k, list):
            k = ", ".join(map(str, k))
        return tp[k] / (tp[k] + fn[k] + 1e-7)


class SemiSegmentalPrecision(_SemiSegmentalMetric):
    """Semi-segmental precision (not advised for training).

    A metric in-between segmental and frame-wise precision.

    This metric follows the following algorithm:
    1) smooth over too-short intervals, both in ground truth and in prediction (first remove
        groups of zeros shorter than `smooth_interval` and then do the same with groups of ones),
    2) add `delta` frames to each predicted interval at both ends and count the number of ground truth
        positive frames at the resulting intervals (intersection),
    3) calculate the threshold for each interval as
        `t = sigmoid(4 * (a * x + b)) * (iou_threshold_long - iou_threshold_short))`, where
        `a = 2 / (long_length - short_length)`, `b = 1 - a * long_length`, `x` is the length of the interval
        before `delta` was added,
    4) for each interval, if intersection is higher than `t * x`, the interval is labeled as true positive (`TP`),
        and otherwise as false positive (`FP`),
    5) the final metric value is computed as `TP / (TP + FP)`.

    """
    def __init__(
        self,
        num_classes: int,
        ignore_index: int = -100,
        ignored_classes: Set = None,
        exclusive: bool = True,
        average: str = "macro",
        tag_average: str = "micro",
        delta: int = 0,
        smooth_interval: int = 0,
        iou_threshold_long: float = 0.5,
        iou_threshold_short: float = 0.5,
        short_length: int = 30,
        long_length: int = 300,
        threshold_value: Union[float, List] = None,
    ) -> None:
        """Initialize the class.

        Parameters
        ----------
        num_classes : int
            the number of classes in the dataset
        ignore_index : int, default -100
            the ground truth label to ignore
        ignored_classes : set, optional
            the class indices to ignore in computation
        exclusive : bool, default True
            `False` for multi-label classification tasks
        average : {"macro", "micro", "none"}
            the method to average the results over classes
        tag_average : {"macro", "micro", "none"}
            the method to average the results over meta tags (if given)
        delta : int, default 0
            the number of frames to add to each ground truth interval before computing the intersection,
            see description of the class for details
        smooth_interval : int, default 0
            intervals shorter than this number of frames will be ignored (both in prediction and in ground truth,
            see description of the class for details
        iou_threshold_long : float, default 0.5
            the intersection threshold for segments longer than `long_length` frames (between 0 and 1),
            see description of the class for details
        iou_threshold_short : float, default 0.5
            the intersection threshold for segments shorter than `short_length` frames (between 0 and 1),
            see description of the class for details
        short_length : int, default 30
            the threshold number of frames for short intervals that will have an intersection threshold of
            `iou_threshold_short`, see description of the class for details
        long_length : int, default 300
            the threshold number of frames for long intervals that will have an intersection threshold of
            `iou_threshold_long`, see description of the class for details
        threshold_value : float | list, optional
            the decision threshold value (cannot be defined for exclusive classification, and 0.5 be default)

        """
        super().__init__(
            num_classes,
            ignore_index,
            ignored_classes,
            exclusive,
            average,
            tag_average,
            delta,
            smooth_interval,
            iou_threshold_long,
            iou_threshold_short,
            short_length,
            long_length,
            [threshold_value],
        )

    def _calculate_metric(self, tp: Dict, fp: Dict, fn: Dict, tn: Dict) -> float:
        """Calculate the metric value from true and false positive and negative rates."""
        k = self.threshold_values[0]
        if isinstance(k, list):
            k = ", ".join(map(str, k))
        return tn[k] / (tn[k] + fp[k] + 1e-7)


class SemiSegmentalF1(_SemiSegmentalMetric):
    """The F1 score for semi-segmental recall and precision (not advised for training)."""

    def __init__(
        self,
        num_classes: int,
        ignore_index: int = -100,
        ignored_classes: Set = None,
        exclusive: bool = True,
        average: str = "macro",
        tag_average: str = "micro",
        delta: int = 0,
        smooth_interval: int = 0,
        iou_threshold_long: float = 0.5,
        iou_threshold_short: float = 0.5,
        short_length: int = 30,
        long_length: int = 300,
        threshold_value: Union[float, List] = None,
    ) -> None:
        """Initialize the class.

        Parameters
        ----------
        num_classes : int
            the number of classes in the dataset
        ignore_index : int, default -100
            the ground truth label to ignore
        ignored_classes : set, optional
            the class indices to ignore in computation
        exclusive : bool, default True
            `False` for multi-label classification tasks
        average : {"macro", "micro", "none"}
            the method to average the results over classes
        tag_average : {"macro", "micro", "none"}
            the method to average the results over meta tags (if given)
        delta : int, default 0
            the number of frames to add to each ground truth interval before computing the intersection,
            see description of the class for details
        smooth_interval : int, default 0
            intervals shorter than this number of frames will be ignored (both in prediction and in ground truth,
            see description of the class for details
        iou_threshold_long : float, default 0.5
            the intersection threshold for segments longer than `long_length` frames (between 0 and 1),
            see description of the class for details
        iou_threshold_short : float, default 0.5
            the intersection threshold for segments shorter than `short_length` frames (between 0 and 1),
            see description of the class for details
        short_length : int, default 30
            the threshold number of frames for short intervals that will have an intersection threshold of
            `iou_threshold_short`, see description of the class for details
        long_length : int, default 300
            the threshold number of frames for long intervals that will have an intersection threshold of
            `iou_threshold_long`, see description of the class for details
        threshold_value : float | list, optional
            the decision threshold value (cannot be defined for exclusive classification, and 0.5 be default)

        """
        super().__init__(
            num_classes,
            ignore_index,
            ignored_classes,
            exclusive,
            average,
            tag_average,
            delta,
            smooth_interval,
            iou_threshold_long,
            iou_threshold_short,
            short_length,
            long_length,
            [threshold_value],
        )

    def _calculate_metric(self, tp: Dict, fp: Dict, fn: Dict, tn: Dict) -> float:
        """Calculate the metric value from true and false positive and negative rates."""
        if self.optimize:
            scores = []
            for k in self.threshold_values:
                recall = tp[k] / (tp[k] + fn[k] + 1e-7)
                precision = tn[k] / (tn[k] + fp[k] + 1e-7)
                scores.append(2 * recall * precision / (recall + precision + 1e-7))
            f1 = max(scores)
        else:
            k = self.threshold_values[0]
            if isinstance(k, list):
                k = ", ".join(map(str, k))
            recall = tp[k] / (tp[k] + fn[k] + 1e-7)
            precision = tn[k] / (tn[k] + fp[k] + 1e-7)
            f1 = 2 * recall * precision / (recall + precision + 1e-7)
        return f1


class SemiSegmentalPR_AUC(_SemiSegmentalMetric):
    """The area under the precision-recall curve for semi-segmental metrics (not advised for training)."""

    def __init__(
        self,
        num_classes: int,
        ignore_index: int = -100,
        ignored_classes: Set = None,
        exclusive: bool = True,
        average: str = "macro",
        tag_average: str = "micro",
        delta: int = 0,
        smooth_interval: int = 0,
        iou_threshold_long: float = 0.5,
        iou_threshold_short: float = 0.5,
        short_length: int = 30,
        long_length: int = 300,
        threshold_step: float = 0.1,
    ) -> None:
        super().__init__(
            num_classes,
            ignore_index,
            ignored_classes,
            exclusive,
            average,
            tag_average,
            delta,
            smooth_interval,
            iou_threshold_long,
            iou_threshold_short,
            short_length,
            long_length,
            list(np.arange(0, 1, threshold_step)),
        )

    def _calculate_metric(self, tp: Dict, fp: Dict, fn: Dict, tn: Dict) -> float:
        """Calculate the metric value from true and false positive and negative rates."""
        precisions = []
        recalls = []
        for k in sorted(self.threshold_values):
            precisions.append(tn[k] / (tn[k] + fp[k] + 1e-7))
            recalls.append(tp[k] / (tp[k] + fn[k] + 1e-7))
        return metrics.auc(x=recalls, y=precisions)


class Accuracy(Metric):
    """Accuracy."""

    def __init__(self, ignore_index=-100):
        """Initialize the class.

        Parameters
        ----------
        ignore_index: int
            the class index that indicates ignored sample

        """
        super().__init__()
        self.ignore_index = ignore_index

    def reset(self) -> None:
        """Reset the internal state (at the beginning of an epoch)."""
        self.total = 0
        self.correct = 0

    def calculate(self) -> float:
        """Calculate the metric value.

        Returns
        -------
        metric : float
            metric value

        """
        return self.correct / (self.total + 1e-7)

    def update(
        self,
        predicted: torch.Tensor,
        target: torch.Tensor,
        tags: torch.Tensor = None,
    ) -> None:
        """Update the internal state (with a batch).

        Parameters
        ----------
        predicted : torch.Tensor
            the main prediction tensor generated by the model
        ssl_predicted : torch.Tensor
            the SSL prediction tensor generated by the model
        target : torch.Tensor
            the corresponding main target tensor
        ssl_target : torch.Tensor
            the corresponding SSL target tensor
        tags : torch.Tensor
            the tensor of meta tags (or `None`, if tags are not given)

        """
        mask = target != self.ignore_index
        self.total += torch.sum(mask)
        self.correct += torch.sum((target == predicted)[mask])


class Count(Metric):
    """Fraction of samples labeled by the model as a class."""

    def __init__(self, classes: Set, exclusive: bool = True):
        """Initialize the class.

        Parameters
        ----------
        classes : set
            the set of classes to count
        exclusive: bool, default True
            set to False for multi-label classification tasks

        """
        super().__init__()
        self.classes = classes
        self.exclusive = exclusive

    def reset(self) -> None:
        """Reset the internal state (at the beginning of an epoch)."""
        self.count = defaultdict(lambda: 0)
        self.total = 0

    def update(
        self,
        predicted: torch.Tensor,
        target: torch.Tensor,
        tags: torch.Tensor,
    ) -> None:
        """Update the internal state (with a batch).

        Parameters
        ----------
        predicted : torch.Tensor
            the main prediction tensor generated by the model
        ssl_predicted : torch.Tensor
            the SSL prediction tensor generated by the model
        target : torch.Tensor
            the corresponding main target tensor
        ssl_target : torch.Tensor
            the corresponding SSL target tensor
        tags : torch.Tensor
            the tensor of meta tags (or `None`, if tags are not given)

        """
        if self.exclusive:
            for c in self.classes:
                self.count[c] += torch.sum(predicted == c)
            self.total += torch.numel(predicted)
        else:
            for c in self.classes:
                self.count[c] += torch.sum(predicted[:, c, :] == 1)
            self.total += torch.numel(predicted[:, 0, :])

    def calculate(self) -> Dict:
        """Calculate the metric (at the end of an epoch).

        Returns
        -------
        result : dict
            a dictionary where the keys are class indices and the values are class metric values

        """
        for c in self.classes:
            self.count[c] = self.count[c] / (self.total + 1e-7)
        return dict(self.count)


class EditDistance(Metric):
    """Edit distance (not advised for training).

    Normalized by the length of the sequences.
    """

    def __init__(self, ignore_index: int = -100, exclusive: bool = True) -> None:
        """Initialize the class.

        Parameters
        ----------
        ignore_index : int, default -100
            the class index that indicates samples that should be ignored
        exclusive : bool, default True
            set to False for multi-label classification tasks

        """
        super().__init__()
        self.ignore_index = ignore_index
        self.exclusive = exclusive

    def reset(self) -> None:
        """Reset the internal state (at the beginning of an epoch)."""
        self.edit_distance = 0
        self.total = 0

    def update(
        self,
        predicted: torch.Tensor,
        target: torch.Tensor,
        tags: torch.Tensor,
    ) -> None:
        """Update the internal state (with a batch).

        Parameters
        ----------
        predicted : torch.Tensor
            the main prediction tensor generated by the model
        ssl_predicted : torch.Tensor
            the SSL prediction tensor generated by the model
        target : torch.Tensor
            the corresponding main target tensor
        ssl_target : torch.Tensor
            the corresponding SSL target tensor
        tags : torch.Tensor
            the tensor of meta tags (or `None`, if tags are not given)

        """
        mask = target != self.ignore_index
        self.total += torch.sum(mask)
        if self.exclusive:
            predicted = predicted[mask].flatten()
            target = target[mask].flatten()
            self.edit_distance += editdistance.eval(
                predicted.detach().cpu().numpy(), target.detach().cpu().numpy()
            )
        else:
            for c in range(target.shape[1]):
                predicted_class = predicted[:, c, :][mask[:, c, :]].flatten()
                target_class = target[:, c, :][mask[:, c, :]].flatten()
                self.edit_distance += editdistance.eval(
                    predicted_class.detach().cpu().tolist(),
                    target_class.detach().cpu().tolist(),
                )

    def _is_equal(self, a, b):
        """Compare while ignoring samples marked with ignore_index."""
        if self.ignore_index in [a, b] or a == b:
            return True
        else:
            return False

    def calculate(self) -> float:
        """Calculate the metric (at the end of an epoch).

        Returns
        -------
        result : float
            the metric value

        """
        return self.edit_distance / (self.total + 1e-7)


class mAP(Metric):
    """Mean average precision (segmental) (not advised for training)."""

    needs_raw_data = True

    def __init__(
        self,
        exclusive,
        num_classes,
        average="macro",
        iou_threshold=0.5,
        threshold_value=0.5,
        ignored_classes=None,
    ):
        """Initialize the class.

        Parameters
        ----------
        exclusive : bool
            set to False for multi-label classification tasks
        num_classes : int
            the number of classes
        average : {"macro", "micro", "none"}
            the type of averaging to perform
        iou_threshold : float, default 0.5
            the IoU threshold for matching
        threshold_value : float, default 0.5
            the threshold value for binarization
        ignored_classes : list, default None
            the list of classes to ignore

        """
        if ignored_classes is None:
            ignored_classes = []
        self.average = average
        self.iou_threshold = iou_threshold
        self.threshold = threshold_value
        self.exclusive = exclusive
        self.classes = [x for x in list(range(num_classes)) if x not in ignored_classes]
        super().__init__()

    def match(self, lst, ratio, ground):
        """Given a list of proposals, match them to the ground truth boxes."""
        lst = sorted(lst, key=lambda x: x[2])

        def overlap(prop, ground):
            s_p, e_p, _ = prop
            s_g, e_g, _ = ground
            return (min(e_p, e_g) - max(s_p, s_g)) / (max(e_p, e_g) - min(s_p, s_g))

        cos_map = [-1 for x in range(len(lst))]
        count_map = [0 for x in range(len(ground))]

        for x in range(len(lst)):
            for y in range(len(ground)):
                if overlap(lst[x], ground[y]) < ratio:
                    continue
                if overlap(lst[x], ground[y]) < overlap(lst[x], ground[cos_map[x]]):
                    continue
                cos_map[x] = y
            if cos_map[x] != -1:
                count_map[cos_map[x]] += 1
        positive = sum([(x > 0) for x in count_map])
        return cos_map, count_map, positive, [x[2] for x in lst]

    def reset(self) -> None:
        """Reset the internal state (at the beginning of an epoch)."""
        self.count_map = defaultdict(lambda: [])
        self.positive = defaultdict(lambda: 0)
        self.cos_map = defaultdict(lambda: [])
        self.confidence = defaultdict(lambda: [])

    def calc_pr(self, positive, proposal, ground):
        """Get precision."""
        if proposal == 0:
            return 0, 0
        if ground == 0:
            return 0, 0
        return (1.0 * positive) / proposal, (1.0 * positive) / ground

    def calculate(self) -> Union[float, Dict]:
        """Calculate the metric (at the end of an epoch)."""
        if self.average == "micro":
            confidence = []
            count_map = []
            cos_map = []
            positive = sum(self.positive.values())
            for key in self.count_map.keys():
                confidence += self.confidence[key]
                cos_map += list(np.array(self.cos_map[key]) + len(count_map))
                count_map += self.count_map[key]
            return self.ap(cos_map, count_map, positive, confidence)
        results = {
            key: self.ap(
                self.cos_map[key],
                self.count_map[key],
                self.positive[key],
                self.confidence[key],
            )
            for key in self.count_map.keys()
        }
        if self.average == "none":
            return results
        else:
            return float(np.mean(list(results.values())))

    def ap(self, cos_map, count_map, positive, confidence):
        """Compute average precision."""
        indices = np.argsort(confidence)
        cos_map = list(np.array(cos_map)[indices])
        score = 0
        number_proposal = len(cos_map)
        number_ground = len(count_map)
        old_precision, old_recall = self.calc_pr(
            positive, number_proposal, number_ground
        )

        for x in range(len(cos_map)):
            number_proposal -= 1
            if cos_map[x] == -1:
                continue
            count_map[cos_map[x]] -= 1
            if count_map[cos_map[x]] == 0:
                positive -= 1

            precision, recall = self.calc_pr(positive, number_proposal, number_ground)
            if precision > old_precision:
                old_precision = precision
            score += old_precision * (old_recall - recall)
            old_recall = recall
        return score

    def _get_intervals(
        self, tensor: torch.Tensor, probability: torch.Tensor = None
    ) -> Union[Tuple, torch.Tensor]:
        """Get `True` group beginning and end indices from a boolean tensor and average probability over these intervals."""
        output, indices = torch.unique_consecutive(tensor, return_inverse=True)
        true_indices = torch.where(output)[0]
        starts = torch.tensor(
            [(indices == i).nonzero(as_tuple=True)[0][0] for i in true_indices]
        )
        ends = torch.tensor(
            [(indices == i).nonzero(as_tuple=True)[0][-1] + 1 for i in true_indices]
        )
        confidence = torch.tensor(
            [probability[indices == i].mean() for i in true_indices]
        )
        return torch.stack([starts, ends, confidence]).T

    def update(
        self,
        predicted: torch.Tensor,
        target: torch.Tensor,
        tags: torch.Tensor,
    ) -> None:
        """Update the state (at the end of each batch)."""
        predicted = torch.cat(
            [
                copy(predicted),
                -100 * torch.ones((*predicted.shape[:-1], 1)).to(predicted.device),
            ],
            dim=-1,
        )
        target = torch.cat(
            [
                copy(target),
                -100 * torch.ones((*target.shape[:-1], 1)).to(target.device),
            ],
            dim=-1,
        )
        num_classes = predicted.shape[1]
        predicted = predicted.transpose(1, 2).reshape(-1, num_classes)
        if self.exclusive:
            target = target.flatten()
        else:
            target = target.transpose(1, 2).reshape(-1, num_classes)
        probability = copy(predicted)
        if not self.exclusive:
            predicted = (predicted > self.threshold).int()
        else:
            predicted = torch.max(predicted, 1)[1]
        for c in self.classes:
            if self.exclusive:
                predicted_intervals = self._get_intervals(
                    predicted == c, probability=probability[:, c]
                )
                target_intervals = self._get_intervals(
                    target == c, probability=probability[:, c]
                )
            else:
                predicted_intervals = self._get_intervals(
                    predicted[:, c] == 1, probability=probability[:, c]
                )
                target_intervals = self._get_intervals(
                    target[:, c] == 1, probability=probability[:, c]
                )
            cos_map, count_map, positive, confidence = self.match(
                predicted_intervals, self.iou_threshold, target_intervals
            )
            cos_map = np.array(cos_map)
            cos_map[cos_map != -1] += len(self.count_map[c])
            self.cos_map[c] += list(cos_map)
            self.count_map[c] += count_map
            self.confidence[c] += confidence
            self.positive[c] += positive
