#!/usr/bin/env python
# -*- coding: utf-8 -*-

##############
# GaudiMM: Genetic Algorithms with Unrestricted
# Descriptors for Intuitive Molecular Modeling
#
# https://github.com/insilichem/gaudi
#
# Copyright 2017 Jaime Rodriguez-Guerra, Jean-Didier Marechal
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##############

"""
This module implements evolutionary algorithms as seen in DEAP,
and extends their functionality to make use of GAUDI goodies.

.. todo::

    * Genealogy

"""

from __future__ import print_function, division
import sys
import os
import logging
from time import time, strftime
from datetime import timedelta
from deap import tools
from deap.algorithms import varOr
import yaml
import gaudi
import random

logger = logging.getLogger(__name__)

if sys.version_info.major == 3:
    xrange = range
    raw_input = input


def ea_mu_plus_lambda(population, toolbox, mu, lambda_, cxpb, mutpb, ngen, cfg,
                      stats=None, halloffame=None, verbose=True,
                      prompt_on_exception=True):
    """This is the :math:`(\mu + \lambda)` evolutionary algorithm.

    :param population: A list of individuals.
    :param toolbox: A :class:`~deap.base.Toolbox` that contains the evolution
                    operators.
    :param mu: The number of individuals to select for the next generation.
    :param lambda\_: The number of children to produce at each generation.
    :param cxpb: The probability that an offspring is produced by crossover.
    :param mutpb: The probability that an offspring is produced by mutation.
    :param ngen: The number of generation.
    :param stats: A :class:`~deap.tools.Statistics` object that is updated
                  inplace, optional.
    :param halloffame: A :class:`~deap.tools.HallOfFame` object that will
                       contain the best individuals, optional.
    :param verbose: Whether or not to log the statistics.
    :returns: The final population.

    First, the individuals having an invalid fitness are evaluated. Then, the
    evolutionary loop begins by producing *lambda_* offspring from the
    population, the offspring are generated by a crossover, a mutation or a
    reproduction proportionally to the probabilities *cxpb*, *mutpb* and 1 -
    (cxpb + mutpb). The offspring are then evaluated and the next generation
    population is selected from both the offspring **and** the population.
    Briefly, the operators are applied as following ::

        evaluate(population)
        for i in range(ngen):
            offspring = varOr(population, toolbox, lambda_, cxpb, mutpb)
            evaluate(offspring)
            population = select(population + offspring, mu)

    This function expects :meth:`toolbox.mate`, :meth:`toolbox.mutate`,
    :meth:`toolbox.select` and :meth:`toolbox.evaluate` aliases to be
    registered in the toolbox. This algorithm uses the :func:`varOr`
    variation.
    """
    t0 = time()
    population_ = population[:]
    logbook = tools.Logbook()
    logbook.header = ['gen', 'progress', 'nevals', 'speed', 'eta'] + (stats.fields if stats else [])

    # Evaluate the individuals with an invalid fitness
    invalid_ind = [ind for ind in population if not ind.fitness.valid]
    fitnesses = toolbox.map(toolbox.evaluate, invalid_ind)
    for ind, fit in zip(invalid_ind, fitnesses):
        ind.fitness.values = fit

    # This is just to assign the crowding distance to the individuals
    # no actual selection is done
    population = toolbox.select(population, len(population))
    
    if halloffame is not None:
        halloffame.update(population)

    record = stats.compile(population) if stats is not None else {}
    nevals = len(invalid_ind)
    t1 = time()
    speed = nevals / (t1-t0)
    performed_evals = nevals
    estimated_evals = (ngen + 1) * lambda_ * (cxpb + mutpb)
    remaining_evals = estimated_evals - performed_evals
    remaining = timedelta(seconds=int(remaining_evals / speed))
    progress = '{:.2f}%'.format(100/(ngen+1))
    logbook.record(gen=0, progress=progress, nevals=nevals,
                   speed='{:.2f} ev/s'.format(speed), eta=remaining, **record)
    if verbose:
        logger.log(100, logbook.stream)
    # Begin the generational process
    for gen in xrange(1, ngen + 1):
        try:
            # Vary the population
            t0 = time()
            offspring = varOr(population, toolbox, lambda_, cxpb, mutpb)

            # Evaluate the individuals with an invalid fitness
            t1 = time()
            invalid_ind = [ind for ind in offspring if not ind.fitness.valid]
            fitnesses = toolbox.map(toolbox.evaluate, invalid_ind)
            
            for ind, fit in zip(invalid_ind, fitnesses):
                ind.fitness.values = fit

            #Discard similar inds
            new_population = []
            for i in population + offspring:
                to_delete = [] #Delete those similar inds with worse fitness
                for n_i in new_population:
                    # Loop through the new_population to check for any
                    # similar individual
                    if toolbox.similarity(i, n_i):
                        if i.fitness <= n_i.fitness:
                            break
                        else:
                            to_delete.append(new_population.index(n_i))
                else:
                    if to_delete:
                        new_population = [i for j, i in enumerate(new_population) if j not in to_delete]
                    # The child is different enough
                    new_population.append(i)
            # Complete new_population with random ind if there are not enough
            while len(new_population) < mu:
                new_population.append(random.choice(population + offspring))

            # Update the hall of fame with the generated individuals
            if halloffame is not None:
                halloffame.clear()
                halloffame.update(new_population)

            # Select the next generation population
            population[:] = toolbox.select(new_population, mu)
            # Update the statistics with the new population
            nevals = len(invalid_ind)
            t2 = time()
            ev_speed = nevals / (t2-t1)
            speed = nevals / (t2-t0)
            performed_evals += nevals
            remaining_evals = estimated_evals - performed_evals
            remaining = timedelta(seconds=int(remaining_evals / speed))
            record = stats.compile(population) if stats is not None else {}
            progress = '{:.2f}%'.format(100*(gen+1)/(ngen+1))
            logbook.record(gen=gen, progress=progress, nevals=nevals,
                           speed='{:.2f} ev/s'.format(ev_speed), eta=remaining, **record)
            if verbose:
                logger.log(100, logbook.stream)
        except (Exception, KeyboardInterrupt) as e:
            logger.error(e)
            if prompt_on_exception:
                answer = raw_input('\nInterruption detected. Write results so far? (y/N): ')
                if answer.lower() not in ('y', 'yes'):
                    sys.exit('Ok, bye!')
            for individual in population:
                try:
                    individual.unexpress()
                except Exception:  # individual was already unexpressed
                    pass
            break
        else:
            # Save a copy of an fully evaluated population, in case the
            # simulation is stopped in next generation.
            population_ = population[:]
            if halloffame and cfg.output.check_every and not gen % cfg.output.check_every:
                try:
                    dump_population(halloffame, cfg, subdir='check{}'.format(gen))
                except Exception:
                    logger.warn('Could not write checkpoing for gen #%s', gen)
    return population_, logbook


def dump_population(population, cfg, subdir=None):
    logger.log(100, 'Writing %s results to disk', len(population))
    results = {'GAUDI.objectives': ['{} ({})'.format(obj.name, obj.module) for obj in cfg.objectives]}
    results['GAUDI.results'] = {}
    basepath = cfg.output.path
    if subdir is not None:
        try:
            basepath = os.path.join(cfg.output.path, subdir)
            os.mkdir(basepath)
        except (IOError, OSError):
            pass
    for i, ind in enumerate(population):
        filename = ind.write(i, path=basepath)
        results['GAUDI.results'][os.path.basename(filename)] = map(float, ind.fitness.values)
    gaudi_output = os.path.join(basepath, cfg.output.name + '.gaudi-output')
    with open(gaudi_output, 'w+') as out:
        out.write('# Generated by GAUDI v{} on {}\n\n'.format(gaudi.__version__,
                                                              strftime("%Y-%m-%d %H:%M:%S")))
        out.write(yaml.safe_dump(results, default_flow_style=False))