"""
Distance functions
------------------
"""

import json
import numpy as np
import logging
from scipy import linalg as la
from abc import ABC, abstractmethod
from typing import List


class DistanceFunction(ABC):
    """
    Abstract case class for distance functions.

    Any other distance function should inherit from this class.
    """
    @abstractmethod
    def __call__(self, x: dict, x_0: dict) -> float:
        """
        Abstract method. This method has to be overwritten by all concrete implementations.

        Evaluate the distance of the tentatively sampled particles relative to the measured data.

        Parameters
        ----------

        x: dict
            Summary statistics of the tentatively sampled parameter.

        x_0: dict
            Summary statistics of the measured data.

        Returns
        -------

        distance: float
            Distance of the tentatively sampled particles to the measured data.
        """

    def initialize(self, sample_from_prior: List[dict]):
        """
        This method is called by the ABCSMC framework before the first usage of the distance function
        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 containig the summary statistics.
        """

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

        Returns
        -------

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

    def to_json(self) -> str:
        """
        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())


class DistanceFunctionWithMeasureList(DistanceFunction):
    """
    Base class for distance functions with measure list.

    Parameters
    ----------

    measures_to_use: Union[str, List[str]].
        * If set to "all", all measures are used. This is the default.
        * If a list is provided, the measures in the list are used.
    """

    def __init__(self, measures_to_use='all'):
        self._measures_to_use_passed_to_init = measures_to_use
        self.measures_to_use = None  #: The measures (summary statistics) to use for distance calculation.

    def initialize(self, sample_from_prior):
        super().initialize(sample_from_prior)
        if self._measures_to_use_passed_to_init == 'all':
            self.measures_to_use == sample_from_prior[0].keys()
            raise Exception("distance function from all measures not implemented.")
        else:
            self.measures_to_use = self._measures_to_use_passed_to_init

    def sanitize_sample_from_prior(self, sample):
        """
        Remove samples in which any of the measures is NaN.
        Added by Alessandro Motta <alessandro.motta@brain.mpg.de>
        """

        assert self.measures_to_use is not None
        is_ok = lambda s: not any(np.isnan(s[m]) for m in self.measures_to_use)

        before_count = len(sample)
        sample = [s for s in sample if is_ok(s)]
        nok_count = before_count - len(sample)

        if nok_count > 0:
            logging.warning("{} out of {} prior samples have NaN measures. " \
                           "Ignoring them...".format(nok_count, before_count))

        return sample

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


class ZScoreDistanceFunction(DistanceFunctionWithMeasureList):
    """
    Calculate distance as sum of ZScores over the selected measures.
    The measured data is the reference for the ZScore.

    Hence

    .. math::

        d(x, y) = \\sum_{i \\in \\text{measures}} \\left| \\frac{x_i-y_i}{y_i} \\right|.
    """
    def __call__(self, x, y):
        return sum(abs((x[key]-y[key])/y[key]) if y[key] != 0 else
                   (0 if x[key] == 0 else np.inf)
                   for key in self.measures_to_use) / len(self.measures_to_use)


class PCADistanceFunction(DistanceFunctionWithMeasureList):
    """
    Calculate distance in whitened coordinates.

    A whitening transformation :math:`W` is calculated from an initial sample.
    The distance is measured as Euclidean distance in the transformed space. I.e

    .. math::

        d(x,y) = \\| Wx - Wy \\|.
    """
    def __init__(self, measures_to_use='all'):
        super().__init__(measures_to_use)
        self._whitening_transformation_matrix = None

    def _dict_to_to_vect(self, x):
        return np.asarray([x[key] for key in self.measures_to_use])

    def _calculate_whitening_transformation_matrix(self, sample_from_prior):
        samples_vec = np.asarray([self._dict_to_to_vect(x) for x in sample_from_prior])
        # samples_vec is an array of shape nr_samples x nr_features
        means = samples_vec.mean(axis=0)
        centered = samples_vec - means
        covariance = centered.T.dot(centered)
        w, v = la.eigh(covariance)
        self._whitening_transformation_matrix = v.dot(np.diag(1. / np.sqrt(w))).dot(v.T)

    def initialize(self, sample_from_prior):
        super().initialize(sample_from_prior)
        sample_from_prior = self.sanitize_sample_from_prior(sample_from_prior)
        self._calculate_whitening_transformation_matrix(sample_from_prior)

    def __call__(self, x, y):
        x_vec, y_vec = self._dict_to_to_vect(x), self._dict_to_to_vect(y)
        distance = la.norm(self._whitening_transformation_matrix.dot(x_vec - y_vec), 2)
        return distance


class RangeEstimatorDistanceFunction(DistanceFunctionWithMeasureList):
    """
    Abstract base class for distance functions whose estimate is based on a range.

    It defines the two template methods ``lower`` and ``upper``.

    Hence

    .. math::

        d(x, y) = \\sum_{i \\in \\text{measures}} \\left | \\frac{x_i - y_i}{u_i - l_i}  \\right |,

    where :math:`l_i` and :math:`u_i` are the lower and upper margins for measure :math:`i`.
    """
    @staticmethod
    def lower(parameter_list: List[float]):
        """
        Calculate the lower margin form a list of parameter values.

        Parameters
        ----------
        parameter_list: List[float]
            List of values of a parameter.

        Returns
        -------

        lower_margin: float
            The lower margin of the range calculated from these parameters.
        """

    @staticmethod
    def upper(parameter_list: List[float]):
        """
        Calculate the upper margin form a list of parameter values.

        Parameters
        ----------
        parameter_list: List[float]
            List of values of a parameter.

        Returns
        -------

        upper_margin: float
            The upper margin of the range calculated from these parameters.
        """

    def __init__(self, measures_to_use='all'):
        super().__init__(measures_to_use)
        self.normalization = None

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

    def _calculate_normalization(self, sample_from_prior):
        measures = {name: [] for name in self.measures_to_use}
        for sample in sample_from_prior:
            for measure in self.measures_to_use:
                measures[measure].append(sample[measure])
        self.normalization = {measure: self.upper(measures[measure]) - self.lower(measures[measure])
                              for measure in self.measures_to_use}

        # HACKHACKHACK(amotta): The measures may have very little (or no)
        # variance under extreme conditions. For example, the relative
        # reciprocity may be zero across all samples from the prior when
        # simulating the locally dense reconstruction of a minor part of the
        # barrel volume.
        #  In these cases, the upper and lower margins may be equal, resulting
        # in a zero normalization constant. Upon division by zero, NumPy will
        # display a warning and return an infinite value.
        #  To avoid that all samples are rejected due to infinite distance,
        # let's replace zero normalization constants by tiny positive values.
        epsilon = np.finfo(np.float).eps
        for measure in self.measures_to_use:
            if not self.normalization[measure]:
                logging.warning("Normalization for {} is zero. Replacing with tiny epsilon...".format(measure))
                self.normalization[measure] = epsilon

    def initialize(self, sample_from_prior):
        super().initialize(sample_from_prior)
        sample_from_prior = self.sanitize_sample_from_prior(sample_from_prior)
        self._calculate_normalization(sample_from_prior)

    def __call__(self, x, y):
        distance = sum(abs((x[key]-y[key])/self.normalization[key]) for key in self.measures_to_use)
        return distance


class MinMaxDistanceFunction(RangeEstimatorDistanceFunction):
    """
    Calculate upper and lower margins as max and min of the parameters.
    """
    @staticmethod
    def upper(parameter_list):
        return max(parameter_list)

    @staticmethod
    def lower(parameter_list):
        return min(parameter_list)


class PercentileDistanceFunction(RangeEstimatorDistanceFunction):
    """
    Calculate normalization 20% and 80% from percentiles as lower and upper margins.
    """

    PERCENTILE = 20  #: The percentiles

    @staticmethod
    def upper(measures):
        measures = np.asarray(measures)
        percentile = np.percentile(measures, 100 - PercentileDistanceFunction.PERCENTILE)
        assert not np.isnan(percentile), "Upper limit is NaN"
        return percentile

    @staticmethod
    def lower(measures):
        measures = np.asarray(measures)
        percentile = np.percentile(measures, PercentileDistanceFunction.PERCENTILE)
        assert not np.isnan(percentile), "Lower limit is NaN"
        return percentile

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


class AcceptAllDistance(DistanceFunction):
    """
    Just a mock distance function which always returns -1.
    So any sample should be accepted for any sane epsilon object.

    Can be used for testing.
    """
    def __call__(self, x, y):
        """
        Parameters
        ----------
        x: dictionary
            Sampled point,
        y: dictionary
            Measured point,
        """
        return -1
