"""Contains the DensityArchive."""
import numpy as np
from scipy.spatial.distance import cdist
from sklearn.neighbors import KernelDensity

from ribs._utils import check_batch_shape, check_finite, readonly
from ribs.archives._archive_base import ArchiveBase
from ribs.archives._utils import parse_dtype


def gkern(x):
    """Gaussian kernel."""
    gauss = np.exp(-0.5 * np.square(x))
    return gauss / np.sqrt(2 * np.pi)


def gaussian_kde_measures(measures, buffer, h):
    """Evaluates kernel density estimation with a Gaussian kernel.

    The density is defined as zero if the buffer is empty.

    Args:
        measures (numpy.ndarray): (measures_batch_size, measure_dim) array
            of measures at which to estimate density.
        buffer (np.ndarray): (buffer_batch_size, measure_dim) batch of
            measures that parameterize the KDE.
        h (float): Kernel bandwidth.
    Returns:
        Evaluation of KDE(m).
    """
    if buffer.shape[0] == 0:
        return np.zeros(measures.shape[0], dtype=measures.dtype)

    # (measures_batch_size, buffer_batch_size)
    norms = cdist(measures, buffer) / h

    # (measures_batch_size,)
    t = np.sum(gkern(norms), axis=1)

    return t / (buffer.shape[0] * h)


# Developer Notes:
# - The documentation for this class is hacked -- to add new methods, manually
#   modify the template in docs/_templates/autosummary/class.rst


class DensityArchive(ArchiveBase):
    """An archive that models the density of solutions in measure space.

    This archive originates in Density Descent Search in `Lee 2024
    <https://dl.acm.org/doi/10.1145/3638529.3654001>`_. It maintains a buffer of
    measures, and using that buffer, it builds a density estimator such as a
    KDE. The density estimator indicates which areas of measure space have, for
    instance, a high density of solutions -- to improve exploration, an
    algorithm would need to target areas with a low density of solutions.

    Incoming solutions are added to the buffer with `reservoir sampling
    <https://en.wikipedia.org/wiki/Reservoir_sampling>`_, specifically as
    described in `Li 1994 <https://dl.acm.org/doi/abs/10.1145/198429.198435>`_.
    Reservoir sampling enables sampling uniformly from the incoming stream of
    solutions generated by the emitters.

    Unlike other archives, this archive does not store any elites, and as such,
    most methods from :class:`ArchiveBase` are not implemented. Rather, it is
    assumed that a separate ``result_archive`` (see
    :class:`~ribs.schedulers.Scheduler`) will store solutions when using this
    archive.

    Args:
        measure_dim (int): Dimension of the measure space.
        buffer_size (int): Size of the buffer of measures.
        density_method (str): Method for computing density. Currently supports
            ``"kde"`` (KDE -- kernel density estimator), ``"kde_sklearn"`` (KDE
            using :class:`sklearn.neighbors.KernelDensity`). Note that when
            ``"kde_sklearn"`` is used, this archive computes *log density*; see
            :meth:`sklearn.neighbors.KernelDensity.score_samples` for more info.
        bandwidth (float): Bandwidth when using ``kde`` or ``kde_sklearn`` as
            the ``density_method``.
        sklearn_kwargs (dict): kwargs for
            :class:`sklearn.neighbors.KernelDensity` when using
            ``"kde_sklearn"`` as the ``density_method``. Note that bandwidth is
            already passed in via the ``bandwidth`` parameter above.
        seed (int): Value to seed the random number generator. Set to None to
            avoid a fixed seed.
        dtype (str or data-type or dict): Data type of the measures.
            This can be ``"f"`` / ``np.float32``, ``"d"`` / ``np.float64``. For
            consistency with other archives, this can also be a dict specifying
            separate dtypes, of the form ``{"solution": <dtype>, "objective":
            <dtype>, "measures": <dtype>}``.
    Raises:
        ValueError: Unknown ``density_method`` provided.
    """

    def __init__(
        self,
        *,
        measure_dim,
        buffer_size=10000,
        density_method="kde",
        bandwidth=None,
        sklearn_kwargs=None,
        seed=None,
        dtype=np.float64,
    ):
        self._rng = np.random.default_rng(seed)
        dtypes = parse_dtype(dtype)
        self._measure_dtype = dtypes["measures"]

        ArchiveBase.__init__(
            self,
            solution_dim=0,
            objective_dim=(),
            measure_dim=measure_dim,
        )

        # Buffer for storing the measures.
        self._buffer = np.empty((buffer_size, measure_dim),
                                dtype=self._measure_dtype)
        # Number of occupied entries in the buffer.
        self._n_occupied = 0
        # Acceptance threshold for the buffer.
        self._w = np.exp(np.log(self._rng.uniform()) / buffer_size)
        # Number of solutions to skip.
        self._n_skip = int(np.log(self._rng.uniform()) / np.log(1 - self._w))

        # Set up density estimator.
        self._density_method = density_method
        if self._density_method == "kde":
            # Kernel density estimation
            self._bandwidth = bandwidth
        elif self._density_method == "kde_sklearn":
            self._bandwidth = bandwidth
            self._sklearn_kwargs = ({} if sklearn_kwargs is None else
                                    sklearn_kwargs.copy())
        else:
            raise ValueError(f"Unknown density_method '{self._density_method}'")

    ## Properties inherited from ArchiveBase ##

    # Necessary to implement this since `Scheduler` calls it.
    @property
    def empty(self):
        """bool: Whether the archive is empty. Since the archive does not store
        elites, we always mark it as not empty."""
        return False

    ## Properties that are not in ArchiveBase ##

    @property
    def buffer(self):
        """numpy.ndarray: Buffer of measures considered in the density
        estimator. Shape (n, :attr:`measure_dim`)."""
        return readonly(self._buffer[:self._n_occupied])

    ## Utilities ##

    def compute_density(self, measures):
        """Computes density at the given points in measure space.

        Args:
            measures (array-like): (batch_size, :attr:`measure_dim`) array with
                measure space coordinates of all the solutions.
        Returns:
            numpy.ndarray: Array of density values of the input solutions.
        """
        measures = np.asarray(measures, dtype=self._measure_dtype)

        if self._density_method == "kde":
            # Use self.buffer instead of self._buffer since self.buffer only
            # contains the valid entries of the buffer.
            return gaussian_kde_measures(
                measures,
                self.buffer,
                self._bandwidth,
            ).astype(self._measure_dtype)
        elif self._density_method == "kde_sklearn":
            if self.buffer.shape[0] == 0:
                return np.zeros(measures.shape[0], dtype=measures.dtype)
            # Note that this is log density with some normalization too.
            kde = KernelDensity(
                bandwidth=self._bandwidth,
                **self._sklearn_kwargs,
            ).fit(self.buffer)
            return kde.score_samples(measures).astype(self._measure_dtype)
        else:
            raise ValueError(f"Unknown density_method '{self._density_method}'")

    ## Methods for writing to the archive ##

    def add(
        self,
        solution,
        objective,
        measures,
        **fields,
    ):
        """Adds measures to the buffer and updates the density estimator if
        necessary.

        The measures are added to the buffer with reservoir sampling to enable
        sampling uniformly from the incoming solutions.

        Args:
            solution (None or array-like): Included for API consistency. Any
                value is ignored.
            objective (None or array-like): Included for API consistency. Any
                value is ignored.
            measures (array-like): (batch_size, :attr:`measure_dim`) array with
                measure space coordinates of all the solutions.
            fields (keyword arguments): Included for API consistency. Any value
                is ignored.

        Returns:
            dict: Information describing the result of the add operation. The
            dict contains the following keys:

            - ``"status"`` (:class:`numpy.ndarray` of :class:`np.int32`): An
              array of integers that represent the "status" obtained when
              attempting to insert each solution in the batch. Since this
              archive does not store any elites, all statuses are set to 2
              (which normally indicates the solution discovered a new cell in
              the archive -- see :class:`AddStatus`).

            - ``"density"`` (:class:`numpy.ndarray` of the dtype passed in at
              init): The density values of the measure passed in, before the
              buffer or density estimator was updated. Note that when
              ``"kde_sklearn"`` is used as the ``density_method``, *log density*
              is computed; see
              :meth:`sklearn.neighbors.KernelDensity.score_samples` for more
              info.

        Raises:
            ValueError: The array arguments do not match their specified shapes.
            ValueError: ``measures`` has non-finite values (inf or NaN).
        """
        measures = np.asarray(measures, dtype=self._measure_dtype)
        check_batch_shape(measures, "measures", self.measure_dim, "measure_dim",
                          "")
        check_finite(measures, "measures")
        batch_size = len(measures)
        buffer_size = len(self._buffer)

        add_info = {
            # Make all statuses be 2 as a placeholder value.
            "status": np.full(batch_size, 2, dtype=np.int32),
            # Note that density should be computed _before_ updating the buffer
            # or density estimator.
            "density": self.compute_density(measures),
        }

        # Add to the buffer using reservoir sampling as in Li 1994
        # (https://dl.acm.org/doi/pdf/10.1145/198429.198435).

        # First, fill the buffer if there are any slots available.
        n_fill = 0
        if buffer_size > self._n_occupied:
            n_fill = min(buffer_size - self._n_occupied, batch_size)
            self._buffer[self._n_occupied:self._n_occupied + n_fill] = \
                measures[:n_fill]
            remaining_measures = measures[n_fill:]
            self._n_occupied += n_fill
        else:
            remaining_measures = measures

        # Replace measures in the buffer using reservoir sampling.
        n_remaining = remaining_measures.shape[0]
        while n_remaining > 0:
            # Done with skipping, replace measures.
            if self._n_skip < n_remaining:
                replace = self._rng.integers(buffer_size)
                self._buffer[replace] = remaining_measures[self._n_skip]
                self._w *= np.exp(np.log(self._rng.uniform()) / buffer_size)
                self._n_skip = int(
                    np.log(self._rng.uniform()) / np.log(1 - self._w))
            skip = min(self._n_skip, n_remaining)
            n_remaining -= skip
            self._n_skip -= skip

        return add_info
