import numpy as np

from pymoo.algorithms.genetic_algorithm import GeneticAlgorithm
from pymoo.docs import parse_doc_string
from pymoo.model.individual import Individual
from pymoo.model.survival import Survival
from pymoo.operators.crossover.point_crossover import PointCrossover
from pymoo.operators.mutation.polynomial_mutation import PolynomialMutation
from pymoo.operators.sampling.random_sampling import RandomSampling
from pymoo.operators.selection.tournament_selection import compare, TournamentSelection
from pymoo.util.display import disp_multi_objective
from pymoo.util.dominator import Dominator
from pymoo.util.non_dominated_sorting import NonDominatedSorting
from pymoo.util.randomized_argsort import randomized_argsort
from misc import utils
from loguru import logger
from pymoo.rand import random
from projectcode.evolution.crossover import TracedPointCrossover
from projectcode.evolution.selection import binary_tournament
from projectcode.evolution.selection import RankAndCrowdingSurvival
import os
from compress_pickle import dump, load
from pymoo.model.individual import Individual
from projectcode.server.database import Generation, pickle_to_str


class NSGANet(GeneticAlgorithm):
    def __init__(self, save_to=None, continue_from=None, args = None, **kwargs):
        kwargs["individual"] = Individual(rank=np.inf, crowding=-1)
        super().__init__(**kwargs)

        logger.info("new version")
        self.tournament_type = "comp_by_dom_and_crowding"
        self.func_display_attrs = disp_multi_objective


        self.continue_from = continue_from
        self.args = args

        self.save_to = os.path.join(args.save, args.code)
        utils.create_exp_dir(self.save_to)

    def _solve(self, problem, termination):

        # generation counter
        self.n_gen = 1

        if self.continue_from is None:
            # initialize the first population and evaluate it
            self.pop = self._initialize()
            self._each_iteration(self, first=True)

            if self.save_to is not None:
                self.save_state(self.save_to)

        else:
            logger.info("Restoring from Save State")

            self.restore_state(self.continue_from)

            dump(self, os.path.join(self.save_to, 'emergency.pkl'))

        # while termination criterium not fulfilled
        while termination.do_continue(self):
            self.n_gen += 1

            # do the next iteration
            self.pop = self._next(self.pop)

            # execute the callback function in the end of each generation
            self._each_iteration(self)

            if self.save_to is not None:
                self.save_state(self.save_to)

            dump(self, os.path.join(self.save_to, 'emergency.pkl'))

        self._finalize()

        return self.pop

    def restore_state(self, continue_from):

        state = load(continue_from)
        self.n_gen = state["generation"]
        self.pop = state["population"]
        self.pop.individual = Individual()

    def save_state(self, save_to):

        logger.info("Saving Algorithm State")

        state = {"generation": self.n_gen, "population": self.pop}

        save_path = os.path.join(save_to, f"{self.args.code}_state_{self.n_gen:03d}.pkl")

        dump(state, save_path)

        p = pickle_to_str(obj=state)
        g = Generation()
        g.run_code = self.args.code
        g.client_stream = self.args.stream
        g.generation = self.n_gen
        g.population_pkl = p
        g.save()



    def _next(self, pop):
        logger.info("mating")

        # do the mating using the current population
        self.off = self._mating(pop)

        logger.info("evaluating")
        # evaluate the offspring
        self.evaluator.eval(self.problem, self.off, algorithm=self)

        logger.info("merging")
        # merge the offsprings with the current population
        pop = pop.merge(self.off)

        logger.info("culling")
        # the do survival selection
        pop = self.survival.do(self.problem, pop, self.pop_size, algorithm=self)

        return pop

    def _mating(self, pop):
        # adjusted mating to preserve heritage

        # the population object to be used
        off = pop.new()

        # mating counter - counts how often the mating needs to be done to fill up n_offsprings
        n_matings = 0

        # iterate until enough offsprings are created
        while len(off) < self.n_offsprings:

            # # how many parents need to be select for the mating - depending on number of offsprings remaining
            # n_select = math.ceil((self.n_offsprings - len(off)) / self.crossover.n_offsprings)

            # set to one to simplify sources of error. 1 means two parents
            n_select = 1

            # select the parents for the mating - just an index array
            parents = self.selection.do(
                pop, n_select, self.crossover.n_parents, algorithm=self
            )

            # do the crossover using the parents index and the population - additional data provided if necessary
            _off = self.crossover.do(self.problem, pop, parents, algorithm=self)

            # do the mutation on the offsprings created through crossover

            muta_off = self.mutation.do(self.problem, _off, algorithm=self)

            for i, ind in enumerate(_off):
                mutation_map = ind.X == muta_off[i].X
                ind.X = muta_off[i].X
                ind.mutation_map = mutation_map

                # repair the individuals if necessary
            if self.repair:
                _off = self.repair.do(self.problem, _off, algorithm=self)

            if self.eliminate_duplicates:
                is_duplicate = self.eliminate_duplicates(_off, pop, off, algorithm=self)
                _off = _off[np.logical_not(is_duplicate)]

            # if more offsprings than necessary - truncate them
            if len(_off) > self.n_offsprings - len(off):
                I = random.perm(self.n_offsprings - len(off))
                _off = _off[I]

            off = off.merge(_off)

            n_matings += 1

            # if no new offsprings can be generated within 100 trails -> return the current result
            if n_matings > 100:
                print(
                    "WARNING: Recombination could not produce new offsprings which are not already in the population!"
                )
                break

        for ind in off:
            assert hasattr(ind, "parents")
            assert hasattr(ind, "mutation_map")

        return off


def nsganet(
    pop_size=100,
    sampling=RandomSampling(var_type=np.int),
    selection=TournamentSelection(func_comp=binary_tournament),
    crossover=TracedPointCrossover(n_points=2),
    mutation=PolynomialMutation(eta=3, var_type=np.int),
    eliminate_duplicates=True,
    n_offsprings=None,
    save_to=None,
    continue_from=None,
    **kwargs,
):
    """

    Parameters
    ----------
    pop_size : {pop_size}
    sampling : {sampling}
    selection : {selection}
    crossover : {crossover}
    mutation : {mutation}
    eliminate_duplicates : {eliminate_duplicates}
    n_offsprings : {n_offsprings}

    Returns
    -------
    nsganet : :class:`~pymoo.model.algorithm.Algorithm`
        Returns an NSGANet algorithm object.


    """
    logger.info("new version")
    return NSGANet(
        pop_size=pop_size,
        sampling=sampling,
        selection=selection,
        crossover=crossover,
        mutation=mutation,
        survival=RankAndCrowdingSurvival(),
        eliminate_duplicates=eliminate_duplicates,
        n_offsprings=n_offsprings,
        save_to=save_to,
        continue_from=continue_from,
        **kwargs,
    )
