#
# 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.
#
"""Abstract parent class for all metrics."""

from abc import ABC, abstractmethod
from typing import Any, Dict, Union

import torch

torch.manual_seed(0)


class Metric(ABC):
    """Base class for all metrics.

    Metrics are reset at the beginning of each epoch, updated with batch data and then calculated at the end of the epoch.
    If needs_raw_data is True for a metric class, it should expect to receive raw model output as the predicted vector;
    otherwise it should be the final class prediction
    """

    needs_raw_data = False
    """
    If `True`, `dlc2action.task.universal_task.Task` will pass raw data to the metric (only primary predict
    function applied).
    Otherwise it will pass a prediction for the classes.
    """

    def __init__(self) -> None:
        """Initialize the class."""
        self.reset()

    @abstractmethod
    def update(
        self,
        predicted: torch.Tensor,
        target: torch.Tensor,
        tags: torch.Tensor,
    ) -> None:
        """Update the intrinsic parameters (with a batch).

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

        """

    @abstractmethod
    def reset(self) -> None:
        """Reset the intrinsic parameters (at the beginning of an epoch)."""

    @abstractmethod
    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

        """
