"""Provides EmitterBase."""
import itertools
from abc import ABC, abstractmethod

import numpy as np


class EmitterBase(ABC):
    """Base class for emitters.

    Every emitter has an :meth:`ask` method that generates a batch of solutions,
    and a :meth:`tell` method that inserts solutions into the emitter's archive.
    Child classes are only required to override :meth:`ask`.

    Args:
        archive (ribs.archives.ArchiveBase): An archive to use when creating and
            inserting solutions. For instance, this can be
            :class:`ribs.archives.GridArchive`.
        solution_dim (int): The dimension of solutions produced by this emitter.
        bounds (None or array-like): Bounds of the solution space. Pass None to
            indicate there are no bounds. Alternatively, pass an array-like to
            specify the bounds for each dim. Each element in this array-like can
            be None to indicate no bound, or a tuple of
            ``(lower_bound, upper_bound)``, where ``lower_bound`` or
            ``upper_bound`` may be None to indicate no bound.

            Unbounded upper bounds are set to +inf, and unbounded lower bounds
            are set to -inf.
    """

    def __init__(self, archive, solution_dim, bounds):
        self._archive = archive
        self._solution_dim = solution_dim
        (self._lower_bounds,
         self._upper_bounds) = self._process_bounds(bounds, self._solution_dim,
                                                    archive.dtype)

    @staticmethod
    def _process_bounds(bounds, solution_dim, dtype):
        """Processes the input bounds.

        Returns:
            tuple: Two arrays containing all the lower bounds and all the upper
                bounds.
        Raises:
            ValueError: There is an error in the bounds configuration.
        """
        lower_bounds = np.full(solution_dim, -np.inf, dtype=dtype)
        upper_bounds = np.full(solution_dim, np.inf, dtype=dtype)

        if bounds is None:
            return lower_bounds, upper_bounds

        # Handle array-like bounds.
        if len(bounds) != solution_dim:
            raise ValueError("If it is an array-like, bounds must have the "
                             "same length as x0")
        for idx, bnd in enumerate(bounds):
            if bnd is None:
                continue  # Bounds already default to -inf and inf.
            if len(bnd) != 2:
                raise ValueError("All entries of bounds must be length 2")
            lower_bounds[idx] = -np.inf if bnd[0] is None else bnd[0]
            upper_bounds[idx] = np.inf if bnd[1] is None else bnd[1]
        return lower_bounds, upper_bounds

    @property
    def archive(self):
        """ribs.archives.ArchiveBase: The archive which stores solutions
        generated by this emitter."""
        return self._archive

    @property
    def solution_dim(self):
        """int: The dimension of solutions produced by this emitter."""
        return self._solution_dim

    @property
    def lower_bounds(self):
        """numpy.ndarray: ``(solution_dim,)`` array with lower bounds of
        solution space.

        For instance, ``[-1, -1, -1]`` indicates that every dimension of the
        solution space has a lower bound of -1.
        """
        return self._lower_bounds

    @property
    def upper_bounds(self):
        """numpy.ndarray: ``(solution_dim,)`` array with upper bounds of
        solution space.

        For instance, ``[1, 1, 1]`` indicates that every dimension of the
        solution space has an upper bound of 1.
        """
        return self._upper_bounds

    @abstractmethod
    def ask(self, grad_estimate=False):
        """Generates an ``(n, solution_dim)`` array of solutions."""

    def tell(self, solutions, objective_values, behavior_values, jacobian=None, metadata=None):
        """Inserts entries into the archive.

        This base class implementation (in :class:`~ribs.emitters.EmitterBase`)
        simply inserts entries into the archive by calling
        :meth:`~ribs.archives.ArchiveBase.add`. It is enough for simple emitters
        like :class:`~ribs.emitters.GaussianEmitter`, but more complex emitters
        will almost certainly need to override it.

        Args:
            solutions (numpy.ndarray): Array of solutions generated by this
                emitter's :meth:`ask()` method.
            objective_values (numpy.ndarray): 1D array containing the objective
                function value of each solution.
            behavior_values (numpy.ndarray): ``(n, <behavior space dimension>)``
                array with the behavior space coordinates of each solution.
            jacobian (numpy.ndarray): Jacobian matrix for differentiable QD algorithms.           
            metadata (numpy.ndarray): 1D object array containing a metadata
                object for each solution.
        """
        metadata = itertools.repeat(None) if metadata is None else metadata
        for sol, obj, beh, meta in zip(solutions, objective_values,
                                       behavior_values, metadata):
            self.archive.add(sol, obj, beh, meta)
