"""General interface definitions for distribution of Evolution Strategies."""
from abc import ABC, abstractmethod
from typing import List, Generator, Optional, Tuple, Dict, Union

import numpy as np

from maze.core.agent.policy import Policy
from maze.core.agent.torch_model import TorchModel
from maze.core.log_stats.log_stats import LogStats


class ESRolloutResult:
    """Result structure for distributed rollouts."""

    def __init__(self, is_eval: bool):
        """:param is_eval: If True, the result was generated by evaluation rollouts, in which case the noise indices
                           won't be set.
        """
        self.is_eval: bool = is_eval
        self.noise_indices: List[int] = list()
        self.episode_stats: List[LogStats] = list()


class ESDistributedRollouts(ABC):
    """Abstract base class of ES rollout distribution."""

    @abstractmethod
    def generate_rollouts(self,
                          policy: Union[Policy, TorchModel],
                          max_steps: Optional[int],
                          noise_stddev: float,
                          normalization_stats: Dict[str, Tuple[np.ndarray, np.ndarray]]
                          ) -> Generator[ESRolloutResult, None, None]:
        """Declare a new rollout task and start producing results that can be obtained from the returned generator.

        Note that different distribution strategies have different ways of balancing evaluation and training rollouts.

        :param policy: Multi-step policy encapsulating the policy networks
        :param max_steps: Optionally limit the rollout to a number of environment steps (horizon).
        :param noise_stddev: The standard deviation of the applied parameter noise.
        :param normalization_stats: Normalization statistics as calculated by the NormalizeObservationWrapper.
        """
