"""
Main ABC algorithm
------------------
"""

import datetime
import logging
import sys
import time
from typing import List, Callable, Iterable, Union

from util.parameters import Parameter
from util.random_variables import RV, ModelPerturbationKernel, Distribution, Kernel
from .distance_functions import DistanceFunction
from .epsilon import Epsilon
from .storage import History


class ABCSMC:
    """
    Approximate Bayesian Computation - Sequential Monte Carlo (ABCSMC).

    This is an implementation of an ABCSMC algorithm similar to [#toni-stumpf]_


    Parameters
    ----------

    models: List[Callable[[Parameter], dict]]
       Calling ``models[m](par)`` returns the calculated summary statistics
       of model ``m`` with the corresponding parameters ``par``.

       Each callable represents thus one single model.

    model_prior_distribution: RV
        A random variable giving the prior weights of the model classes.
        If the prior is uniform over the model classes
        this is something like ``RV("randint", 0, len(models))``.

    model_perturbation_kernel: ModelPerturbationKernel
        Kernel which governs with which probability to switch the model
        for a given sample.

    parameter_given_model_prior_distribution: List[Distribution]
        A list of prior distributions for the models' parameters.
        Each list entry is the prior distribution for the corresponding model.

    adaptive_parameter_perturbation_kernels: List[Callable[[int, dict], Kernel]]
        A list of functions mapping ``(t, stat) -> Kernel``, where

            * ``t`` is the population nr
            * ``stat`` a dictionary of summary statistics.
               E.g. ``stat['std']['parameter_1']`` is the standard deviation of ``parameter_1``.

                .. warning:: If a model has only one particle left the standard deviation is zero.

        This callable is called at the beginning of a new population with the statistics dictionary
        from the last population to determine the new parameter perturbation kernel for the next population.

    distance_function: DistanceFunction
        Measures the distance of the tentatively sampled particle to the measured data.

    eps: Epsilon
        Returns the current acceptance epsilon.
        This epsilon changes from population to population.
        The eps instance provides the strategy according to which to change it.

    mapper: map like
        A callable which behaves like the built-in map function.
        I.e. mapper(f, args) takes a callable ``f`` and applies it to the arguments in the list ``args``.
        This mapper is used for particle sampling.
        It can be a distributed mapper such as the :class:`parallel.sge.SGE` class.

    debug: bool
        Whether to output additional debug information.

    max_nr_allowed_sample_attempts_per_particle: int
        The maximum number of sample attempts allowed for each particle.
        If this number is reached, the sampling for a particle is stopped.
        Hence, a population may return with less particles than started.
        This is an approximation to the ABCSMC algorithm which ensures, that
        the algorithm terminates.

    min_nr_particles_per_population: int
        Minimum number of samples which have to be accepted for a population.
        If this number is not reached, the algorithm stops.
        This option, together with the ``max_nr_allowed_sample_attempts_per_particle``
        ensures that the algorithm terminates.
        This parameter determines to which extent the
        ABCSMC algorithm is approximated.


    .. [#toni-stumpf] Toni, Tina, and Michael P. H. Stumpf.
                  “Simulation-Based Model Selection for Dynamical
                  Systems in Systems and Population Biology.”
                  Bioinformatics 26, no. 1 (2010):
                  104–10. doi:10.1093/bioinformatics/btp619.
    """
    def __init__(self,
                 models: List[Callable[[Parameter], dict]],
                 model_prior_distribution: RV,
                 model_perturbation_kernel: ModelPerturbationKernel,
                 parameter_given_model_prior_distribution: List[Distribution],
                 adaptive_parameter_perturbation_kernels: List[Callable[[int, dict], Kernel]],
                 distance_function: DistanceFunction,
                 eps: Epsilon,
                 nr_particles: int,
                 mapper=map,
                 debug: bool =False,
                 max_nr_allowed_sample_attempts_per_particle: int =500,
                 min_nr_particles_per_population: int =1):

        # sanity checks
        self.models = list(models)
        if not (len(self.models)
                == len(parameter_given_model_prior_distribution)
                == len(adaptive_parameter_perturbation_kernels)):
            raise Exception("Nr of models has to be equal to the number of parameter prior distributions has to be equal"
                            " to the number of parameter perturbation kernels")
        self.model_prior_distribution = model_prior_distribution
        self.model_perturbation_kernel = model_perturbation_kernel
        self.parameter_given_model_prior_distribution = parameter_given_model_prior_distribution  # this cannot be serialized by dill
        self.adaptive_parameter_perturbation_kernels = adaptive_parameter_perturbation_kernels
        self.distance_function = distance_function
        self.eps = eps
        self.nr_particles = nr_particles
        self.mapper = mapper
        self.debug = debug
        self.stop_if_only_single_model_alive = True
        self.x_0 = None
        self.history = None #History
        self._points_sampled_from_prior = None
        self.max_nr_allowed_sample_attempts_per_particle = max_nr_allowed_sample_attempts_per_particle
        self.min_nr_particles_per_population = min_nr_particles_per_population

    def do_not_stop_when_only_single_model_alive(self):
        """
        Calling this method causes the ABCSMC to still continue if only
        a single model is still alive. This is useful if the interest lies in
        estimating the model parameter as compared to performing model selection.

        The default behavior is to stop when only a single model is alive.
        """
        self.stop_if_only_single_model_alive = False

    def set_data(self, observed_summary_statistics: dict,
                 ground_truth_model_nr_or_name: Union[int, str],
                 ground_truth_parameter: dict,
                 abc_options: dict,
                 model_names: Iterable[str]):
        """
        Set the data to be fitted.

        Parameters
        ----------

        observed_summary_statistics : dict
               **This is the really important parameter here**. It is of the form
               ``{'statistic_1' : val_1, 'statistic_2': val_2, ... }``.

               The dictionary provided here represents the measured data.
               Particles during ABCSMC sampling are compared with the summary statistics
               provided here.

        ground_truth_model_nr_or_name: Union[int, str]
            This is only meta data stored to the database, but not actually used for the ABCSMC algorithm.
            To evaluate the ABCSMC procedure against synthetic samples, this parameter can be used to indicate the
            ground truth model number or name. This helps with further analysis. If actually measured data is used,
            it is recommended to set this parameter to ``-1``.

        ground_truth_parameter: dict
            Similar to ``ground_truth_model_nr_or_name``, this is only for recording purposes, but not used in the
            ABCSMC algorithm. This stores the parameters of the ground truth model if it was synthetically
            obtained.

        abc_options: dict
            Has to contain the key "db_path" which has to be a valid SQLAlchemy database identifier.
            Can contain an arbitrary number of additional keys, only for recording purposes.
            Store arbitrary meta information in this dictionary.

        model_names: List[str]
            Only for recording purposes. Record names of the models.
        """
        # initialize
        self.x_0 = observed_summary_statistics
        self.history = History(abc_options['db_path'], len(self.models), list(model_names),
                               self.min_nr_particles_per_population)

        # initialize distance function and epsilon
        sample_from_prior = self.sample_from_prior()
        self.distance_function.initialize(sample_from_prior)

        def distance_to_ground_truth_function(x):
            return self.distance_function(x, self.x_0)

        self.eps.initialize(sample_from_prior, distance_to_ground_truth_function)
        self.history.store_initial_data(ground_truth_model_nr_or_name, abc_options,
                                        observed_summary_statistics, ground_truth_parameter,
                                        self.distance_function.to_json(),
                                        self.eps.to_json())

    def sample_from_prior(self) -> List[dict]:
        """
        Only sample from prior and return results without changing
        the history. This can be used to get initial samples
        for the distance function or the epsilon to calibrate them.

        .. warning::

            The sample is cached.
        """
        if self._points_sampled_from_prior is None:
            # not saved as attribute b/c Mapper of type "ipython_cluster" is not pickable
            def sample(_):
                m = self.model_prior_distribution.rvs()
                par = self.parameter_given_model_prior_distribution[m].rvs()
                x = self.models[m](par)
                return x
            if self.debug:
                print('sample from prior')
            self._points_sampled_from_prior = list(self.mapper(sample, list(range(self.nr_particles))))
            if self.debug:
                print('sample from prior done')
        if self.debug:
            print('return sample from prior')
        return self._points_sampled_from_prior

    def _sample_single_particle(self, parameter_perturbation_kernels,
                               nr_samples_per_particle: int,
                               t: int,
                               t0: int,
                               current_eps: float):
        """
        This is where the actual stuff happens.

        Parameters
        ----------
        parameter_perturbation_kernels
        nr_samples_per_particle
        t: int
            population number
        t0: int
            initial population
        current_eps

        Returns
        -------

        """
        simulation_counter = 0
        theta_s, theta_ss = "initial", "initial"
        while True:  # find valid theta_ss and (corresponding b) according to data x_0
            while True:  # find m_s and theta_ss, valid according to prior
                if t == 0:  # sample from prior
                    m_ss = self.model_prior_distribution.rvs()
                    theta_ss = self.parameter_given_model_prior_distribution[m_ss].rvs()
                else:
                    m_s = self.history.sample_from_models(t-1)
                    m_ss = self.model_perturbation_kernel.rvs(m_s)
                    theta_s = self.history.sample_from_population(t-1, m_ss)
                    if theta_s is not None:   # theta_s is None if the population m_s has died out
                        theta_ss = parameter_perturbation_kernels[m_ss].rvs(theta_s)
                    else:
                        theta_ss = None
                if theta_ss is not None and (self.model_prior_distribution.pmf(m_ss)
                                             * self.parameter_given_model_prior_distribution[m_ss].pdf(theta_ss) > 0):
                    break

            # from here, theta_ss is valid according to the prior
            distance_list = []
            summary_statistics_list = []
            for __ in range(nr_samples_per_particle[t-t0]):
                ##### MODEL SIMULATION - THIS IS THE EXPENSIVE PART ######
                simulation_counter += 1
                # stop builder if it takes too long
                if simulation_counter > self.max_nr_allowed_sample_attempts_per_particle:
                    print("Max number of samples (={n_max}) for particle theta_s={theta_s} reached."
                          .format(theta_s=theta_s, n_max=self.max_nr_allowed_sample_attempts_per_particle), file=sys.stderr)
                    return None
                start_time = time.time()
                x_s = self.models[m_ss](theta_ss)  # the actual simulation
                end_time = time.time()
                duration = end_time - start_time
                if self.debug:
                    print("Sampled model={}-{}, delta_time={}s, end_time={}, theta_s={}, theta_ss={}"
                          .format(m_ss, self.history.model_names[m_ss], duration, end_time,
                                  theta_s, theta_ss))
                distance = self.distance_function(x_s, self.x_0)
                if distance <= current_eps:
                    distance_list.append(distance)
                    summary_statistics_list.append(x_s)
            if len(distance_list) > 0:
                break

        if t == 0:
            weight = len(distance_list) / nr_samples_per_particle[t-t0] # b_t in Toni et al.
        else:
            normalization = (sum(self.history.get_model_probabilities(t-1)[j] * self.model_perturbation_kernel.pmf(m_ss, j)
                                 for j in range(len(self.models)))
                             * sum(particle['weight']  # this is already conditioned on m,
                                                       # so do not divide by P_{t-1}(m_{t-1} = m_t^{(i)}) = P_{t-1}(m_t^{(i)})
                                   * parameter_perturbation_kernels[m_ss].pdf(theta_ss, particle['parameter'])
                                   for particle in self.history.store[t-1][m_ss])
                            )
            if normalization == 0:
                print('normalization is zero!')
            weight = (self.model_prior_distribution.pmf(m_ss)                             # P(m_t^i)
                      * self.parameter_given_model_prior_distribution[m_ss].pdf(theta_ss) # × P(theta_t^i | m_t^i)
                                                                                          # = P(m_t^i, theta_t^i) in Toni et al.
                      * len(distance_list) / nr_samples_per_particle[t-t0] # b_t in Toni et al.
                      / normalization) # S in Toni et al.
        if self.debug:
            print('.', end='')
        return m_ss, theta_ss, weight, distance_list, simulation_counter, summary_statistics_list

    def run(self, nr_samples_per_particle: List[int], minimum_epsilon: float) -> History:
        """
        Run the ABCSMC model selection. This method can be called many times. It makes another
        step continuing where it has stopped before.

        It is stopped when the maximum number of populations is reached
        or the ``minimum_epsilon`` value is reached.

        Parameters
        ----------

        nr_samples_per_particle: List[int]
            The length of the list determines the maximal number of populations.

            The entries of the list the number of iterated simulations
            in the notation from [#toni-stumpf-nine]_ these are the :math:`B_t`.
            Usually, the entries are all ones:
            ``nr_samples_per_particle = [1] * nr_populations``.

        minimum_epsilon: float
            Stop if epsilon is smaller than minimum epsilon specified here.


        .. [#toni-stumpf-nine] Toni, Tina, David Welch, Natalja Strelkowa, Andreas Ipsen, and Michael P. H. Stumpf. “Approximate Bayesian Computation Scheme for Parameter Inference and Model Selection in Dynamical Systems.” Journal of The Royal Society Interface 6, no. 31 (2009): 187–202. doi:10.1098/rsif.2008.0172.

        """
        t0 = self.history.t
        self.history.start_time = datetime.datetime.now()

        t = t0
        while True:
            current_eps = self.eps(t, self.history)  # this is calculated here to avoid double initialization of medians
            if self.debug:
                print('t:', t, 'eps:', current_eps)
            statistics = self.history.get_statistics(t-1)
            parameter_perturbation_kernels = self._make_parameter_perturbation_kernels(statistics, t)
            if self.debug:
                print('now submitting population', t)
            new_particle_population = list(
                    self.mapper(lambda _: self._sample_single_particle(parameter_perturbation_kernels,
                                                                    nr_samples_per_particle,
                                                                    t,
                                                                    t0,
                                                                    current_eps),
                                [None] * self.nr_particles))
            new_particle_population = [particle for particle in new_particle_population
                                       if not isinstance(particle, Exception)]
            if self.debug:
                print('population', t, 'done')
            new_particle_population_non_empty = self.history.append_population(t, current_eps, new_particle_population)
            if self.debug:
                print('\ntotal number simulations up to t =', t, 'is', self.history.total_nr_simulations)
                sys.stdout.flush()

            t += 1
            terminate = False

            if not t < t0 + len(nr_samples_per_particle):
                logging.info('Stopping because maximal number of population is exceeded')
                terminate = True

            if not new_particle_population_non_empty:
                logging.info('Stopping because new population of particles is empty')
                terminate = True

            if self.stop_if_only_single_model_alive and self.history.nr_of_models_alive() <= 1:
                logging.info('Stopping because at most one model remains')
                terminate = True

            if minimum_epsilon is not None and current_eps <= minimum_epsilon:
                logging.info('Stopping because epislon is too low')
                terminate = True

            if terminate:
                break


        self.history.done()
        return self.history

    def _make_parameter_perturbation_kernels(self, statistics, t):
        parameter_perturbation_kernels = [apk(t, stat) if stat is not None else None
                                          for apk, stat in
                                          zip(self.adaptive_parameter_perturbation_kernels, statistics)]
        return parameter_perturbation_kernels
