"""Interface for environments to expose statistics logging capabilities."""
from abc import abstractmethod, ABC
from typing import Callable, Optional

from maze.core.log_stats.log_stats import LogStatsLevel, LogStatsValue, LogStatsAggregator


class LogStatsEnv(ABC):
    """Interface to access logging statistics generated by the environment. Most envs won't implement this directly,
    but use the LogStatsWrapper to implement this functionality.
    """

    @abstractmethod
    def get_stats(self, level: LogStatsLevel) -> LogStatsAggregator:
        """Get access to the statistics aggregator, which can be used to receive the latest statistics dictionary
        (via :meth:`~maze.core.log_stats.log_stats.LogStatsAggregator.last_stats`) or to register log consumers.

        :param level: The statistics aggregation level
                      (LogStatsLevel.STEP, LogStatsLevel.EPISODE or LogStatsLevel.EPOCH)

        :return: The statistics aggregator
        """

    @abstractmethod
    def write_epoch_stats(self) -> None:
        """Use this if you do not want to wait for the next :func:`maze.core.log_stats.log_stats.increment_log_step`
        call to update the epoch statistics.

        This can be useful in a training scenario to get the results of the evaluation rollouts immediately, instead
        of waiting for the next log step increment to occur.
        """

    @abstractmethod
    def clear_epoch_stats(self) -> None:
        """Use this to clear episode statistics collected so far in this epoch (i.e., inputs for epoch stats),
         and start fresh.

        This can be useful for clearing left-over statistics. E.g., in later epochs, there might be an environment
        left in the middle of an unfinished episode, which should not be counted into the current epoch.
        """

    @abstractmethod
    def get_stats_value(self,
                        event: Callable,
                        level: LogStatsLevel,
                        name: Optional[str] = None) -> LogStatsValue:
        """Obtain a single value from the statistics dict.

        :param event: The event interface method of the value in question
        :param level: The statistics aggregation level
                      (LogStatsLevel.STEP, LogStatsLevel.EPISODE or LogStatsLevel.EPOCH)
        :param name: The *output_name* of the statistics in case it has been specified in
                     :func:`maze.core.log_stats.event_decorators.define_epoch_stats`

        :return: The statistics dictionary
        """
