# Copyright: 2017, Markus Abel, Julien Gout, Markus Quade
# Licence: LGPL

"""Convenience classes and functions that allow you to quickly build gp apps."""

import os
import sys
import time
import inspect
import random
import logging
import argparse
import operator
from contextlib import contextmanager

import dill
import numpy as np
import toolz
import deap
import deap.tools

from . import gp
from . import utils
from .assessment import SingleProcessFactoy


class GPRunner(object):
    """Runner for gp problem sets.

    Takes care of propper initialization, execution, and accounting of a gp run
    (i.e. population creation, random state, generation count, hall of fame, and
    logbook). The method init() has to be called once before stepping through
    the evolution process with method step(); init() and step() invoke the
    assessment runner.
    """

    def __init__(self, IndividualClass, algorithm_factory, assessment_runner, callbacks=()):
        """Init GPRunner.

        :param IndividualClass: Class inherited from gp.AExpressionTree.
        :param algorithm_factory: callable() -> gp algorithm, as defined in
                                   gp.algorithms.
        :param assessment_runner: callable(population) -> None, updates fitness
                                   values of each invalid individual in population.
        """
        self.IndividualClass = IndividualClass
        self.algorithm_factory = algorithm_factory
        self.assessment_runner = assessment_runner
        self.halloffame = []
        self.logbook = ''
        self.callbacks = callbacks

    def init(self, pop_size):
        """Initialize the gp run."""
        self.halloffame = deap.tools.ParetoFront()
        self.logbook = deap.tools.Logbook()
        self.mstats = None
        self.step_count = 0
        with random_state(self):
            self.population = self.IndividualClass.create_population(pop_size)

        self.algorithm = self.algorithm_factory()
        self._update()

    def step(self):
        """Step through the evolution process."""
        with random_state(self):
            self.population = self.algorithm.evolve(self.population)
        self.step_count += 1
        self._update()
        for cb in self.callbacks:
            cb(self)

    def _update(self):
        """Evaluate invalid individuals and update statistics accordingly."""
        evals = self.assessment_runner(self.population)
        self.halloffame.update(self.population)
        if not self.mstats:
            self.mstats = create_stats(len(self.population[0].fitness.values))
        record = self.mstats.compile(self.population)
        self.logbook.record(gen=self.step_count, evals=evals, **record)



@contextmanager
def random_state(cls, rng=random):
    """Do work inside this contextmanager with a random state defined by cls.

    Looks for _prev_state to seed the rng.
    On exit, it will write the current state of the rng as _tmp_state
    to the cls.

    :params cls: Any object.
    :params rng: Instance of a random number generator.
    """
    cls._tmp_state = rng.getstate()
    rng.setstate(getattr(cls, '_prev_state', rng.getstate()))
    yield
    cls._prev_state = rng.getstate()
    rng.setstate(getattr(cls, '_tmp_state', rng.getstate()))


def default_gprunner(Individual, assessment_runner, callbacks=(), **kwargs):
    """Create a default GPRunner instance.

    For config options see `MateFactory`, `MutateFactory`, `AlgorithmFactory`.
    """
    default_config = dict(mating='cxonepoint', mating_max_height=20,
                          mutation='mutuniform', mutation_max_height=20,
                          algorithm='nsga2', crossover_prob=0.5,
                          mutation_prob=0.2, tournament_size=2,
                          select='nsga2', create_method='halfandhalf',
                          create_min_height=1, create_max_height=4,
                          mutate_tree_max=2, mutate_tree_min=0)

    default_config.update(kwargs)
    mate = MateFactory.create(default_config, Individual)
    mutate = MutateFactory.create(default_config, Individual)
    select = SelectFactory.create(default_config)
    create_method = CreateFactory.create(default_config, Individual)
    AlgorithmFactory.create(default_config, mate, mutate, select, create_method)  # A test run to check config params.
    algorithm_factory = toolz.partial(AlgorithmFactory.create, default_config, mate, mutate, select, create_method)
    return GPRunner(Individual, algorithm_factory, assessment_runner, callbacks=callbacks)


def make_checkpoint(app):
    if app.valid_checkpointing and (app.gp_runner.step_count % app.args.checkpoint_frequency == 0):
        app.checkpoint()


DEFAULT_CALLBACKS = make_checkpoint,


class Application(object):
    """An application based on `GPRunner`.

    Controls execution of the runner and adds checkpointing and logging
    functionality; also defines a set of available command line options and
    their default values.

    To create a full console application one can use the factory function
    create_console_app().
    """

    def __init__(self, config, gp_runner, checkpoint_file, callbacks=()):
        """
        :param config: Container holding all configs
        :type config: dict or argparse.Namespace
        :param gp_runner: Instance of `GPRunner`
        :param checkpoint_file: Path to checkpoint_file
        :param callbacks:
        """
        self.args = to_argparse_namespace(config)
        self.gp_runner = gp_runner
        self.checkpoint_file = checkpoint_file
        self.pareto_fronts = []
        self.logger = logging.getLogger(__name__)
        self._initialized = False
        self.callbacks = callbacks + DEFAULT_CALLBACKS

    @property
    def assessment_runner(self):
        return self.gp_runner.assessment_runner

    def run(self, break_condition=None):
        """Run gp app.

        :param break_condition: is called after every evolutionary step.
        :type break_condition: callable(application)

        :return: number of iterations executed during run.
        """
        if break_condition is None:
            break_condition = lambda app: False
        iterations = 0
        if self.args.pop_size < 1:
            return iterations
        if not self._initialized:
            random.seed(self.args.seed)
            self.gp_runner.init(self.args.pop_size)
            self._update()
            iterations += 1
        while self.gp_runner.step_count < self.args.num_generations and not break_condition(self):
            self.gp_runner.step()
            self._update()
            iterations += 1
        return iterations

    def _update(self):
        for cb in self.callbacks:
            cb(self)

    @property
    def valid_checkpointing(self):
        return self.checkpoint_file is not None and isinstance(self.args.checkpoint_frequency, int)

    def checkpoint(self):
        """Checkpoint current state of evolution."""
        safe(self.checkpoint_file, args=self.args, runner=self.gp_runner,
             random_state=random.getstate(), pareto_fronts=self.pareto_fronts, callbacks=self.callbacks)
        self.logger.debug('Saved checkpoint to {}'.format(self.checkpoint_file))

    @property
    def workdir(self):
        return os.path.dirname(os.path.abspath(self.checkpoint_file))

    @classmethod
    def from_checkpoint(cls, file_name):
        """Create application from checkpoint file."""
        cp = load(file_name)
        app = cls(cp['args'], cp['runner'], file_name, callbacks=cp['callbacks'])
        app.pareto_fronts = cp['pareto_fronts']
        app._initialized = True
        random.setstate(cp['random_state'])
        return app

    @staticmethod
    def add_options(parser):
        """Add available parser options."""
        parser.add_argument('--pop_size', '-p', dest='pop_size', metavar='n',
                            type=utils.argparse.non_negative_int, default=10,
                            help='initial population size (default: 10)')
        parser.add_argument('--num_generations', '-n', dest='num_generations', metavar='n',
                            type=utils.argparse.non_negative_int, default=10,
                            help='number of generations to evolve (default: 10)')
        parser.add_argument('--seed', dest='seed', metavar='n', type=utils.argparse.non_negative_int,
                            default=random.randint(0, sys.maxsize),
                            help='a seed for the random genrator (default: random.randint(0, sys.maxsize))')
        parser.add_argument('--checkpoint_frequency', '-f', dest='checkpoint_frequency', metavar='n',
                            type=utils.argparse.positive_int, default=1,
                            help='do checkpointing every n generations (default: 1)')


def default_console_app(IndividualClass, AssessmentRunnerClass, parser=argparse.ArgumentParser(), callbacks=()):
    """Factory function for a console application."""
    Application.add_options(parser)
    cp_group = parser.add_mutually_exclusive_group(required=False)
    cp_group.add_argument('--resume', dest='resume_file', metavar='FILE', type=str,
                          help='continue previous run from a checkpoint file')
    cp_group.add_argument('-o', dest='checkpoint_file', metavar='FILE', type=str,
                          default=os.path.join('.', 'checkpoint.pickle'),
                          help='checkpoint to FILE (default: ./checkpoint.pickle)')
    parser.add_argument('--verbose', '-v', dest='verbosity', action='count', default=0,
                        help='set verbose output; raise verbosity level with -vv, -vvv, -vvvv')
    parser.add_argument('--logging_config', '-l', type=str, default='logging.yaml',
                        help='set config file for logging; overides --verbose (default: logging.yaml)')
    AlgorithmFactory.add_options(parser.add_argument_group('algorithm'))
    group_breeding = parser.add_argument_group('breeding')
    MateFactory.add_options(group_breeding)
    MutateFactory.add_options(group_breeding)
    SelectFactory.add_options(group_breeding)
    CreateFactory.add_options(group_breeding)
    ParallelizationFactory.add_options(parser.add_argument_group('parallel execution'))

    args = parser.parse_args()

    workdir = os.path.dirname(os.path.abspath(args.checkpoint_file))
    if not os.path.exists(workdir):
        raise RuntimeError('Path does not exist: "{}"'.format(workdir))
    log_level = utils.logging.log_level(args.verbosity)
    utils.logging.load_config(config_file=args.logging_config, default_level=log_level,
                              placeholders=dict(workdir=workdir))
    logger = logging.getLogger(__name__)

    if args.resume_file is not None:
        logger.debug('Loading checkpoint {}'.format(args.resume_file))
        app = Application.from_checkpoint(args.resume_file)
        return app, args
    else:
        mate = MateFactory.create(args, IndividualClass)
        mutate = MutateFactory.create(args, IndividualClass)
        select = SelectFactory.create(args)
        create_method = CreateFactory.create(args, IndividualClass)
        algorithm_factory = toolz.partial(AlgorithmFactory.create, args, mate, mutate, select, create_method)
        parallel_factory = toolz.partial(ParallelizationFactory.create, args)
        assessment_runner = AssessmentRunnerClass(parallel_factory)
        gp_runner = GPRunner(IndividualClass, algorithm_factory, assessment_runner)
        app = Application(args, gp_runner, args.checkpoint_file, callbacks=callbacks)
        return app, args


class AFactory(object):
    _mapping = {}

    @classmethod
    def create(cls, config, *args, **kwargs):
        config = to_argparse_namespace(config)
        return cls._create(config, *args, **kwargs)

    @staticmethod
    def add_options(parser):
        """Add available parser options."""
        raise NotImplementedError

    @classmethod
    def get_from_mapping(cls, key):
        try:
            func = cls._mapping[key]
        except KeyError:
            raise RuntimeError('Option {} not supported'.format(key))
        return func


def get_mapping(group):
    return {obj.__name__.lower(): obj for obj in group}


class AlgorithmFactory(AFactory):
    """Factory class for gp algorithms."""

    _mapping = get_mapping(gp.all_algorithms)

    @classmethod
    def _create(cls, args, mate_func, mutate_func, select, create_func):
        """Setup gp algorithm."""
        kwargs = locals().copy()
        args.algorithm = args.algorithm.lower()
        algorithm_class = cls.get_from_mapping(args.algorithm)
        signature = inspect.signature(algorithm_class)
        algorithm = algorithm_class(*tuple(kwargs.get(x) for x in signature.parameters))
        algorithm.crossover_prob = args.crossover_prob
        algorithm.mutation_prob = args.mutation_prob
        return algorithm

    @staticmethod
    def add_options(parser):
        """Add available parser options."""
        parser.add_argument('--algorithm', dest='algorithm', type=str, default='nsga2',
                            choices=list(AlgorithmFactory._mapping.keys()),
                            help='the gp algorithm (default: nsga2)')
        parser.add_argument('--cxpb', dest='crossover_prob', metavar='p', type=utils.argparse.unit_interval,
                            default=0.5,
                            help='crossover probability for mating (default: 0.5)')
        parser.add_argument('--mutpb', dest='mutation_prob', metavar='p', type=utils.argparse.unit_interval,
                            default=0.2,
                            help='mutation probability (default: 0.2)')
        parser.add_argument('--tournament-size', dest='tournament_size', metavar='n', type=utils.argparse.unit_interval,
                            default=2,
                            help='tournament size for tournament selection (default: 2)')


class MateFactory(AFactory):
    """Factory class for gp mating functions."""

    _mapping = get_mapping(gp.all_crossover)

    @staticmethod
    def _create(args, IndividualClass):
        """Setup mating function."""
        args.mating = args.mating.lower()
        mate = MateFactory.get_from_mapping(args.mating)(**vars(args))
        static_limit_decorator = deap.gp.staticLimit(key=operator.attrgetter("height"),
                                                     max_value=args.mating_max_height)
        mate = static_limit_decorator(mate)
        return mate

    @staticmethod
    def add_options(parser):
        """Add available parser options."""
        parser.add_argument('--mating', dest='mating', type=str, default='cxonepoint',
                            choices=list(MateFactory._mapping.keys()),
                            help='the mating method (default: cxonepoint)')
        parser.add_argument('--mating-max-height', dest='mating_max_height', metavar='n',
                            type=utils.argparse.positive_int, default=20,
                            help='limit for the expression tree height as a result of mating (default: 20)')


class MutateFactory(AFactory):
    """Factory class for gp mutation functions."""

    _mapping = get_mapping(gp.all_mutations)

    @staticmethod
    def _create(args, IndividualClass):
        """Setup mutation function."""
        args.mutation = args.mutation.lower()
        mutate = MutateFactory.get_from_mapping(args.mutation)(IndividualClass.pset, **vars(args))
        static_limit_decorator = deap.gp.staticLimit(key=operator.attrgetter("height"),
                                                     max_value=args.mutation_max_height)
        mutate = static_limit_decorator(mutate)
        return mutate

    @staticmethod
    def add_options(parser):
        """Add available parser options."""
        parser.add_argument('--mutation', dest='mutation', type=str, default='mutuniform',
                            choices=list(MutateFactory._mapping.keys()),
                            help='the mutation method (default: mutuniform)')
        parser.add_argument('--mutation-max-height', dest='mutation_max_height', metavar='n',
                            type=utils.argparse.positive_int, default=20,
                            help='limit for the expression tree height as a result of mutation (default: 20)')
        parser.add_argument('--mutate_tree_min', dest='mutate_tree_min', default=0, metavar='min_',
                            type=utils.argparse.positive_int,
                            help="minimum value for tree based mutation methods (default: 0)")
        parser.add_argument('--mutate_tree_max', dest='mutate_tree_max', default=2, metavar='max_',
                            type=utils.argparse.positive_int,
                            help="maximum value for tree based mutation methods (default: 2)")


class SelectFactory(AFactory):
    """Factory class for selection"""

    _mapping = {"nsga2": deap.tools.selNSGA2,
                "spea2": deap.tools.selSPEA2,
                }

    @staticmethod
    def _create(args):
        args.select = args.select.lower()
        return SelectFactory.get_from_mapping(args.select)

    @staticmethod
    def add_options(parser):
        """Add available parser options."""
        parser.add_argument('--select', dest='select', type=str, default='nsga2',
                            choices=list(SelectFactory._mapping.keys()),
                            help='the selection method (default: nsga2)')


class CreateFactory(AFactory):
    """Factory class for creation"""
    _mapping = {"halfandhalf": deap.gp.genHalfAndHalf}

    @staticmethod
    def _create(args, IndividualClass):
        args.create_method = args.create_method.lower()
        m = CreateFactory.get_from_mapping(args.create_method)
        create_ = toolz.partial(IndividualClass.create_population, gen_method=m,
                                min=args.create_min_height, max=args.create_max_height)
        return create_

    def add_options(parser):
        parser.add_argument('--create_method', dest='create_method', type=str, default='halfandhalf',
                            choices=list(CreateFactory._mapping.keys()),
                            help='the create method (default: halfandhalf)')
        parser.add_argument('--create_max_height', dest='create_max_height', default=4,
                            type=utils.argparse.positive_int,
                            help="maximum value for tree based create methods (default: 4)")
        parser.add_argument('--create_min_height', dest='create_min_height', default=1,
                            type=utils.argparse.positive_int,
                            help="maximum value for tree based create methods (default: 1)")


class ParallelizationFactory(AFactory):
    """Factory class for parallel execution schemes."""

    @staticmethod
    def _create(args):
        return SingleProcessFactoy

    @staticmethod
    def add_options(parser):
        """Add available parser options."""
        # todo


def safe(file_name, **kwargs):
    """Dump kwargs to file."""
    with open(file_name, 'wb') as file:
        dill.dump(kwargs, file)


def load(file_name):
    """Load data saved with safe()."""
    with open(file_name, 'rb') as file:
        cp = dill.load(file)
    return cp


def create_tmp_dir(prefix='run-'):
    """Create directory with current time as signature."""
    start_date_str = time.strftime('%Y-%m-%d-%H%M%S', time.localtime())
    workdir = os.path.relpath(prefix + start_date_str)
    os.mkdir(workdir)
    return workdir


def _create_logger(verbosity, config_file, workdir):
    log_level = utils.logging.log_level(verbosity)
    utils.logging.load_config(config_file=config_file, default_level=log_level, placeholders=dict(workdir=workdir))
    return logging.getLogger(__name__)


def create_stats(n):
    """Create deap.tools.MultiStatistics object for n fitness values."""

    def val(i, ind):
        return ind.fitness.values[i]

    stats = dict()
    for i in range(n):
        stats['fit{}'.format(i)] = deap.tools.Statistics(toolz.partial(val, i))
    mstats = deap.tools.MultiStatistics(**stats)
    mstats.register("min", np.nanmin)
    mstats.register("max", np.nanmax)
    return mstats


def to_argparse_namespace(d):
    """Return argparse.Namespace object created from dictionary d."""
    if isinstance(d, argparse.Namespace):
        return d
    elif isinstance(d, dict):
        return argparse.Namespace(**d)
    else:
        raise RuntimeError('Cannot convert {} to argparse.Namespace.'.format(type(d)))
