import os
import random

import numpy

from abcsmc import Kernel, RV, ModelPerturbationKernel, PercentileDistanceFunction, MedianEpsilon, ABCSMC, History
from connectome.builder import PipelineBuilder
from typing import Iterable
from concurrent.futures import ProcessPoolExecutor, as_completed
import datetime
import time
from util.representation import ReprMixin
from connectome.model import Network
from parallel import sge_available
from typing import Union


class GroundTruthData(ReprMixin):
    """
    Ground truth data.

    This is an object encapsulating the data for a network ABC task.

    Parameters
    ----------

    network: Network
        The network sample.

    model_nr_or_name: Union[int, str]
        A  integer or string representing the model.

    parameter: dict
        Metadata for recording purposes. This helps to identify a run in the database.
        The data provided here is not used for the ABC inference.
    """
    def __init__(self, network: Network, model_nr_or_name: Union[int, str], parameter: dict):
        self.network = network
        self.model_nr_or_name = model_nr_or_name
        self.parameter = parameter


# not a lambda to be able to pickle it
def kernel_maker(t, stat):
    return Kernel(stat['cov'] * 2)


class ABCConnectomeTask:
    """
    Facade class for an ABC Task involving a network.
    This class runs a default ABC task with a network as input.


    Intended usage of this class is via the execute_task_list function.

    E.g.

    .. code::

        abc_tasks = [ABCConnectomeTask(model_definitions, gt_data_1, db, map),
                    ABCConnectomeTask(model_definitions, gt_data_1, db, map),
                    ...]

        execute_task_list(abc_tasks)


    However, a single task can be executed using its ``run_abc`` method.


    Parameters
    ----------

    model_definitions: dict
        The dictionary of model definitions. It contains the definitions for models, noise and the analysis.

    ground_truth: GroundTruthData
        This represents the measured data.

    db: str
        Path to an sqlite database. This is **not** an SQLAlchemy identifier.

    mapper: map like object
        This is used to execute the ABC tasks.
        However, for parallel task execution the concurrent.futures module is used.

    probability_to_stay_at_model: float, optional
        A number in [0,1]. Probability to not to perturb the model.

    measures_to_use: Iterable[str], optional
        An iterable containing the name of the measures (summary statistics) to use.

    nr_particles: int, optional
        Number of particles per population.

    max_nr_allowed_sample_attempts_per_particle: int, optional
        Maximum number of allowed sample attempts per particle.

    min_nr_particles_per_population: int, optional
        Minimum number of particles to be contained in a population.
        The max_nr_allowed_sample_attempts_per_particle restricts the number of sample attempts.
        So the number of particles in  a population can be less than nr_particles, but not less
        than min_nr_particles_per_population.

    max_nr_populations: int, optional
        Maximum number of populations. Stop ABC if this value is reached.

    minimum_epsilon: float, optional
        Minimum epsilon. Stop ABC if this value is reached.
    """
    def __init__(self, model_definitions: dict, ground_truth_data: GroundTruthData, db: str, mapper,
                 *,  # to avoid typos only allow keyword arguments
                 probability_to_stay_at_model=0.85,
                 measures_to_use: Iterable[str] =("relative_reciprocity_ee",
                   "relative_reciprocity_ei",
                   "relative_reciprocity_ie",
                   "relative_reciprocity_ii",
                   "relative_cycles_5",
                   "in_out_degree_correlation_exc"),
                 nr_particles=2000,
                 max_nr_allowed_sample_attempts_per_particle=2000,
                 min_nr_particles_per_population=1000,
                 max_nr_populations=8,
                 minimum_epsilon=.175):
        self.model_definitions = model_definitions
        self.ground_truth = ground_truth_data
        self.db = "sqlite:///" + db
        self.mapper = mapper
        self.probability_to_stay_at_model = probability_to_stay_at_model
        self.measures_to_use = list(measures_to_use)
        self.abc_nr_particles = nr_particles
        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
        self.max_nr_populations = max_nr_populations
        self.minimum_epsilon = minimum_epsilon

    def __call__(self):
        print("PID={}, Start ABCConnectomeTask.".format(os.getpid()), flush=True)
        random.seed()
        numpy.random.seed()
        result = self.run_abc()
        return result

    def get_params(self) -> dict:
        return {"db_path": self.db,
                "model_definitions": self.model_definitions,
                "probability_to_stay_at_model": self.probability_to_stay_at_model,
                "measures_to_use": self.measures_to_use,
                "nr_particles": self.abc_nr_particles,
                "max_nr_allowed_sample_attempts_per_particle": self.max_nr_allowed_sample_attempts_per_particle,
                "min_nr_particles_per_population": self.min_nr_particles_per_population,
                "max_nr_populations": self.max_nr_populations,
                "minimum_epsilon": self.minimum_epsilon}

    def run_abc(self) -> History:
        """
        Run the defined ABC task.

        Returns
        -------
        abc_history: History
            The abc History object.
        """
        builder = PipelineBuilder(self.model_definitions)
        analysis = builder.build_analysis()
        summary_statistics = analysis(network=self.ground_truth.network)

        models = [builder.build_pipeline(model_name) for model_name in builder.defined_models()]
        prior = RV("randint", 0, len(models))
        model_pert_kernel = ModelPerturbationKernel(len(models), self.probability_to_stay_at_model)

        parameter_prior = [builder.build_prior(model_name) for model_name in builder.defined_models()]

        adaptive_par_pert_kernels = [kernel_maker for _ in models]
        distance_function = PercentileDistanceFunction(self.measures_to_use)
        epsilon = MedianEpsilon()

        abc_smc = ABCSMC(models, prior, model_pert_kernel, parameter_prior,
                         adaptive_par_pert_kernels, distance_function, epsilon, self.abc_nr_particles,
                         self.mapper,
                         max_nr_allowed_sample_attempts_per_particle=self.max_nr_allowed_sample_attempts_per_particle,
                         min_nr_particles_per_population=self.min_nr_particles_per_population,
                         debug=False)

        abc_smc.set_data(summary_statistics,
                         self.ground_truth.model_nr_or_name,
                         self.ground_truth.parameter,
                         self.get_params(),
                         builder.defined_models())
        abc_history = abc_smc.run([1] * self.max_nr_populations, self.minimum_epsilon)
        print("ABC done")
        return abc_history


def execute_task_list(abc_tasks: Iterable[ABCConnectomeTask], nr_parallel_jobs: int):
    """
    Parameters
    ----------
    abc_tasks: Iterable[ABCConnectomeTask]
        The ABCTask instances to run.

    nr_parallel_jobs: int
        How many ABCTask instances to run in parallel.
    """
    if nr_parallel_jobs > 1:
        with ProcessPoolExecutor(nr_parallel_jobs) as pool:
            submitted_jobs = []
            for abc_task in abc_tasks:
                time.sleep(1)  # to prevent a sqlite race condition
                submitted_jobs.append(pool.submit(abc_task))
            print("Submitted all {} jobs.".format(len(submitted_jobs)), flush=True)
            for nr, future in enumerate(as_completed(submitted_jobs)):
                print("{}: Finished job {}/{}: {}".format(datetime.datetime.now(),
                                                          nr + 1, len(submitted_jobs), future.result()), flush=True)
    else:
        for abc_task in abc_tasks:
            abc_task()
