import numpy as np
import logging
import json
from abc import ABC, abstractmethod
from .storage import History
from typing import List, Callable, Union
logging.basicConfig(level=logging.DEBUG)
eps_logger = logging.getLogger("Epsilon")


class Epsilon(ABC):
    """
    Abstract epsilon base class.

    This class encapsulates a strategy for setting a new epsilon for each new population.
    """
    def initialize(self, sample_from_prior: List[dict], distance_to_ground_truth_function: Callable[[dict], float]):
        """
        This method is called by the ABCSMC framework before the first usage of the epsilon
        and can be used to calibrate it to the statistics of the samples.

        Per default, no calibration is made.

        Parameters
        ----------

        sample_from_prior: List[dict]
            List of dictionaries containing the summary statistics.

        distance_to_ground_truth_function: Callable[[dict], float]
            One of the distance functions pre-evaluated at its second argument
            (the one representing the measured data).
            E.g. similar to ``lambda x: distance_function(x, x_measured)``.

        """

    def get_config(self):
        """
        Return configuration of the distance function.

        Returns
        -------

        config: dict
            Dictionary describing the distance function.
        """
        return {"name": self.__class__.__name__}

    def to_json(self):
        """
        Return JSON encoded configuration of the distance function.

        Returns
        -------

        json_str: str
            JSON encoded string describing the distance function.
            The default implementation is to try to convert the dictionary
            returned by ``get_config``.
        """
        return json.dumps(self.get_config())

    @abstractmethod
    def __call__(self, t: int, history: History):
        """

        Parameters
        ----------
        t: int
            The population number.

        history: History
            ABC history object. Can be used to query summary statistics to set the epsilon.

        Returns
        -------

        eps: float
            The new epsilon for population ``t``.
        """


class ConstantEpsilon(Epsilon):
    """
    Keep epsilon constant over all populations.

    Parameters
    ----------

    constant_epsilon_value: float
        The epsilon value for all populations.
    """
    def __init__(self, constant_epsilon_value: float):
        super().__init__()
        self.constant_epsilon_value = constant_epsilon_value

    def get_config(self):
        config = super().get_config()
        config["constant_epsilon_value"] = self.constant_epsilon_value
        return config

    def __call__(self, t, history):
        return self.constant_epsilon_value


class ListEpsilon(Epsilon):
    """
    Return epsilon values from a predefined list.

    Parameters
    ----------

    values: List[float]
        List of epsilon values.
        ``values[k]`` is the value for population k.
    """
    def __init__(self, values: List[float]):
        super().__init__()
        self.epsilon_values = list(values)

    def get_config(self):
        config = super().get_config()
        config["epsilon_values"] = self.epsilon_values
        return config

    def __call__(self, t, history):
        return self.epsilon_values[t]


class MedianEpsilon(Epsilon):
    """
    Calculate epsilon as median from the last population.

    Parameters
    ----------

    initial_epsilon: Union[str, int]
        * If 'from_sample', then the initial median is calculated from samples as its median.
        * If a number is given, this number is used.

    median_multiplier: float
        Multiplies the median by that number. Also applies it
        to the initial median if it is calculated from samples.
        However, it does **not** apply to the initial median if
        it is given as a number.
    """

    def __init__(self, initial_epsilon: Union[str, int]='from_sample', median_multiplier: float =1):
        eps_logger.debug("init medianepsilon initial_epsilon={}, median_multiplier={}".format(initial_epsilon, median_multiplier))
        super().__init__()
        self._initial_epsilon = initial_epsilon
        self.median_multiplier = median_multiplier

    def get_config(self):
        config = super().get_config()
        config.update({"initial_epsilon": self._initial_epsilon,
                       "median_multiplier": self.median_multiplier})
        return config

    def initialize(self, sample_from_prior, distance_to_ground_truth_function):
        super().initialize(sample_from_prior, distance_to_ground_truth_function)
        eps_logger.debug("calc initial epsilon")
        if self._initial_epsilon == 'from_sample':  # calculate initial epsilon if not given
            distances = np.asarray([distance_to_ground_truth_function(x) for x in sample_from_prior])
            nan_mask = np.isnan(distances)

            if nan_mask.any():
                # TODO(amotta): Find the source of these NaNs and try to solve the root cause.
                eps_logger.warning("{} out of {} distances are NaN. Ignoring them...".format(nan_mask.sum(), nan_mask.size))
                distances = distances[~nan_mask]
                assert distances.size > 0

            eps_t0 = np.median(distances) * self.median_multiplier
            self._look_up = {0: eps_t0}
        else:
            self._look_up = {0: self._initial_epsilon}

        assert not np.isnan(self._look_up[0]), "Initial epsilon is NaN"
        assert self._look_up[0] > 0, "Initial epsilon is non-positive"
        eps_logger.debug("initial epsilon is {}".format(self._look_up[0]))

    def __call__(self, t, history):
        try:
            return self._look_up[t]
        except KeyError:
            median = history.get_complete_population_median(t-1)
            self._look_up[t] = median * self.median_multiplier
            eps_logger.debug("new eps, t={}, eps={}".format(t, self._look_up[t]))
            return self._look_up[t]
