import time


class Job:
    def __init__(self, model, persistence):
        self.model = model
        self.persistence = persistence

    def run(self):
        result = self.model()
        self.persistence.store(self.model, result)


def _execute(job):
    """
    Use this function instead of a lambda expression
    since standard python multiprocessing cannot
    pickle lambda expresison.
    """
    job.run()


class ModelWithPrior:
    """
    A wrapper around a model and its prior for sampling.

    Parameters
    ----------

    prior: Distribution
        The model's prior distribution.

    model: Node
        The model.
    """
    def __init__(self, prior, model, name="NoNameDefined"):
        self.prior = prior
        self.model = model
        self.name = name

    def __call__(self):
        parameter = self.prior.rvs()
        analysis_result = self.model(parameter)
        return {"parameter": dict(parameter), "result": dict(analysis_result)}

    def __repr__(self):
        return "<ModelWithPrior model={model}, prior={prior}>".format(model=self.model, prior=self.prior)


class TimedModelWithPriorSample(ModelWithPrior):
    """
    A subclass of :class:`ModelWithPrior <connectome.builder.sampler.ModelWithPrior>`

    This wrapper stores in addition "time_elapsed", i.e. the time for model generation, in the results.
    """
    def __call__(self):
        start_time = time.time()
        result = super().__call__()
        end_time = time.time()
        result["time_elapsed"] = end_time - start_time
        return result


class Sampler:
    """
    Sample models.

    Parameters
    ----------

    map_callable: map-like
        A callable which behaves like the built-in map function.
        Applies function to each of the entries in argument_list.
        It should behave like the built in map but computation can be executed
        locally or distributed on a cluster.

    persistence: InMemoryStorer or DistributedFileStorer or alike
        Has to implement the methods:
          * `store(model, result)`: Store the result of a model
          * fetch_all(): Return a list of dictionaries of the form :code:`[{model_name: [results]}, ...]`.
    """

    def __init__(self, map_callable, persistence):
        self.map_callable = map_callable
        self.persistence = persistence
        self.models = []

    def add_model(self, model):
        """
        Parameters
        ----------
        model: ModelWithPrior
            Callable and has to have a property :code:`model.name`.
        """
        self.models.append(model)

    def sample(self, nr_samples_per_model=0):
        if not self.models:
            raise Exception("No models defined. Cannot sample")
        jobs = []
        for model in self.models:
            jobs.append(Job(model, self.persistence))
        # make sure things get executed
        for _ in self.map_callable(_execute, jobs*nr_samples_per_model):
            pass

    def results(self):
        return self.persistence.fetch_all()
