import time
import neat
import numpy as np
from bisect import bisect_right
from collections import defaultdict
from copy import deepcopy
from functools import partial
from itertools import chain
from operator import eq
from deap import algorithms
from deap.tools import HallOfFame, Logbook, Statistics, MultiStatistics

# From DEAP implementation https://deap.readthedocs.io/en/master/_modules/deap/tools/support.html#HallOfFame.update

class HallOfFameGeneral(HallOfFame):

    def __init__(self, maxsize, similar=eq, fit_attr='fitness'):
        HallOfFame.__init__(self, maxsize, similar)
        self.fit_attr = fit_attr

    def update(self, population):
        for ind in population:
            if len(self) == 0 and self.maxsize != 0:
                self.insert(population[0])
                continue
            if getattr(ind, self.fit_attr) > getattr(self[-1], self.fit_attr) or len(self) < self.maxsize:
                for hofer in self:
                    if self.similar(ind, hofer):
                        break
                else:
                    if len(self) >= self.maxsize:
                        self.remove(-1)
                    self.insert(ind)

    def insert(self, item):
        item = deepcopy(item)
        i = bisect_right(self.keys, getattr(item, self.fit_attr))
        self.items.insert(len(self) - i, item)
        self.keys.insert(i, getattr(item, self.fit_attr))

class LogbookManager:
    
    def __init__(self, general, stats) -> None:
        self.simple_statisctic = LogbookManager.__is_simple_statistic(stats)
        self.logbook = Logbook()
        
        # Initialize Header
        if self.simple_statisctic:
            self.logbook.header = *general.keys(), *(stats.fields if stats else []) 
        else:
            self.logbook.header = "general", *(stats.fields if stats else []) 
    
    @staticmethod
    def __is_simple_statistic(stats) -> bool:
        if(stats==None):
            return True
        elif (isinstance(stats, Statistics)):
            return True
        elif (isinstance(stats, MultiStatistics)):
            return False
        else:
            raise TypeError("'stats' parameter has incorrect Type")   
    
    def get_logbook(self) -> Logbook:
        return self.logbook
    
    def record(self, general, record):
        if self.simple_statisctic:
            self.logbook.record(**general, **record)  
        else:
            self.logbook.record(general=general, **record)

class GeneralStatistics:

    @staticmethod
    def register_general_stats(stats):
        stats.register("avg", lambda pop: round(np.nanmean(pop),4))
        stats.register("std", lambda pop: round(np.nanstd(pop),4))
        stats.register("min", lambda pop: round(np.nanmin(pop),4))
        stats.register("max", lambda pop: round(np.nanmax(pop),4))
    
    @staticmethod
    def create_novelty_stats() -> MultiStatistics:
        stats_fit = Statistics(key=lambda ind: ind.fitness.values)
        stats_obj = Statistics(key=lambda ind: ind.objective.values)
        stats = MultiStatistics(fitness=stats_fit, objective=stats_obj)
        GeneralStatistics.register_general_stats(stats)
        
        return stats

    @staticmethod
    def create_default_stats() -> Statistics:
        stats = Statistics(key=lambda ind: ind.fitness.values)
        GeneralStatistics.register_general_stats(stats)
        
        return stats
       
class Algorithms:

    @staticmethod
    def ea_simple(population, toolbox, cxpb, mutpb, ngen, stats=None, halloffame:HallOfFame=None, verbose=False):
        # Initialize
        general = {'gen':0, 'neval':0, 'elapsed_time':0}
        logbook = LogbookManager(general, stats)
        
        s_time = time.time()
        # Main loop
        for gen in range(0, ngen + 1):
            
            # Execute
            invalid_idx = [i for i,ind in enumerate(population) if not ind.fitness.valid]
            toolbox.execute(population,invalid_idx)
            invalid_ind =[population[i] for i in invalid_idx]
            time.sleep(0.1)
            
            # Evaluate
            fitnesses = toolbox.evaluate(invalid_ind)
            for ind, fit in zip(invalid_ind, fitnesses):
                ind.fitness.values = fit
            if halloffame is not None:
                halloffame.update(population)
            
            # Record
            record = stats.compile(population) if stats else {}
            general.update({'gen':gen, 'neval':len(invalid_ind), 'elapsed_time':round(time.time() - s_time, 2)})
            logbook.record(general, record)
            if verbose:
                print(logbook.get_logbook().stream)
            
            # Evolve
            if gen < ngen:
                offspring = toolbox.select(population, len(population))
                offspring = algorithms.varAnd(offspring, toolbox, cxpb, mutpb)
                population[:] = offspring
        
        return population, logbook.get_logbook()

    @staticmethod
    def ea_novelty(population, toolbox, cxpb, mutpb, ngen, stats=None, halloffame:HallOfFame=None, verbose=False):
        # Initialize
        general = {'gen':0, 'neval':0, 'elapsed_time':0}
        logbook = LogbookManager(general, stats)
        
        s_time = time.time()
        # Main evolution loop
        for gen in range(0, ngen + 1):
            
            # Execute
            invalid_idx = [i for i,ind in enumerate(population) if not ind.fitness.valid]
            invalid_ind =[population[i] for i in invalid_idx]
            toolbox.execute(population, invalid_idx)
            
            # Safety wait for server
            time.sleep(0.1)
            
            # Evaluate
            fitnesses = toolbox.evaluate(population)
            for ind, fit in zip(population, fitnesses):
                ind.fitness.values = fit
            objectives = toolbox.evaluate_obj(invalid_ind)
            for ind, obj in zip(invalid_ind, objectives):
                ind.objective.values = obj
            toolbox.update_archive(population)
            if halloffame is not None:
                halloffame.update(population)
            
            # Record
            record = stats.compile(population) if stats else {}
            general.update({'gen':gen, 'neval':len(invalid_ind), 'elapsed_time':round(time.time() - s_time, 2)})
            logbook.record(general, record)
            if verbose:
                print(logbook.get_logbook().stream)
            
            # Evolve
            if gen < ngen:
                offspring = toolbox.select(population, len(population))
                offspring = algorithms.varAnd(offspring, toolbox, cxpb, mutpb)
                population[:] = offspring
        
        return population, logbook.get_logbook()

    @staticmethod
    def ea_general(population, toolbox, evaluate, cxpb, mutpb, ngen, stats=None, halloffame:HallOfFame=None, verbose=False):
        # Initialize
        general = {'gen':0, 'neval':0, 'elapsed_time':0}
        logbook = LogbookManager(general, stats)
        
        s_time = time.time()
        # Main evolution loop
        for gen in range(0, ngen + 1):
            
            # Execute
            results = toolbox.simulate(population)
            
            # Evaluate
            for func, attr in evaluate.items():
                values = toolbox.map(getattr(toolbox, func), population, *results)
                if attr!=None:
                    for ind, value in zip(population, values):
                        getattr(ind, attr).values = value
                else:
                    values = list(values)
            
            if halloffame is not None:
                halloffame.update(population)
            
            # Record
            record = stats.compile(population) if stats else {}
            general.update({'gen':gen, 'neval':len(population), 'elapsed_time':round(time.time() - s_time, 2)})
            logbook.record(general, record)
            if verbose:
                print(logbook.get_logbook().stream)
            
            # Evolve
            if gen < ngen:
                offspring = toolbox.select(population, len(population))
                offspring = algorithms.varAnd(offspring, toolbox, cxpb, mutpb)
                population[:] = offspring
        
        return population, logbook.get_logbook()

    evaluate_simple= {
        "fitness":"fitness",
    }

    evaluate_novelty= {
        "objective":"objective",
        "fitness":"fitness",
    }
    
    evaluate_novelty_update= {
        "objective":"objective",
        "update_archive": None,
        "fitness":"fitness",
    }
