# Adapted from:
# J. Blank and K. Deb, pymoo: Multi-Objective Optimization in Python, in IEEE Access, vol. 8, pp. 89497-89509, 2020, doi: 10.1109/ACCESS.2020.2990567; https://github.com/anyoptimization/pymoo

import numpy as np
from pymoo.core.algorithm import Algorithm
from pymoo.core.duplicate import DefaultDuplicateElimination, NoDuplicateElimination
from pymoo.core.initialization import Initialization
from pymoo.core.mating import Mating
from pymoo.core.population import Population
from pymoo.core.repair import NoRepair


class GeneticAlgorithm(Algorithm):
    def __init__(
        self,
        pop_size=None,
        sampling=None,
        selection=None,
        crossover=None,
        mutation=None,
        survival=None,
        n_offsprings=None,
        eliminate_duplicates=DefaultDuplicateElimination(),
        repair=None,
        mating=None,
        advance_after_initial_infill=False,
        **kwargs
    ):

        super().__init__(**kwargs)

        # the population size used
        self.pop_size = pop_size

        # whether the algorithm should be advanced after initialization of not
        self.advance_after_initial_infill = advance_after_initial_infill

        # the survival for the genetic algorithm
        self.survival = survival

        # number of offsprings to generate through recombination
        self.n_offsprings = n_offsprings

        # if the number of offspring is not set - equal to population size
        if self.n_offsprings is None:
            self.n_offsprings = pop_size

        # set the duplicate detection class - a boolean value chooses the default duplicate detection
        if isinstance(eliminate_duplicates, bool):
            if eliminate_duplicates:
                self.eliminate_duplicates = DefaultDuplicateElimination()
            else:
                self.eliminate_duplicates = NoDuplicateElimination()
        else:
            self.eliminate_duplicates = eliminate_duplicates

        # simply set the no repair object if it is None
        self.repair = repair if repair is not None else NoRepair()

        self.initialization = Initialization(
            sampling, repair=self.repair, eliminate_duplicates=self.eliminate_duplicates
        )

        if mating is None:
            mating = Mating(
                selection,
                crossover,
                mutation,
                repair=self.repair,
                eliminate_duplicates=self.eliminate_duplicates,
                n_max_iterations=100,
            )
        self.mating = mating

        # other run specific data updated whenever solve is called - to share them in all algorithms
        self.n_gen = None
        self.pop = None
        self.off = None

    def _initialize_infill(self):
        pop = self.initialization.do(self.problem, self.pop_size, algorithm=self)
        pop.set("n_gen", self.n_gen)

        if self.problem.variables.num_cat_var > 0:
            for individual_idx, individual in enumerate(pop):
                X_temp = (
                    individual.X
                    * (np.array(self.problem.variables.upper_bounds) - np.array(self.problem.variables.lower_bounds))
                ) + np.array(self.problem.variables.lower_bounds)
                X_temp = self.problem.variables.categorical_values_euc(X_temp.reshape(1, -1))
                pop[individual_idx].X = (X_temp.flatten() - np.array(self.problem.variables.lower_bounds)) / (
                    np.array(self.problem.variables.upper_bounds) - np.array(self.problem.variables.lower_bounds)
                )

        return pop

    def _initialize_advance(self, infills=None, **kwargs):
        if self.advance_after_initial_infill:
            self.pop = self.survival.do(self.problem, infills, n_survive=len(infills))

    def _infill(self):

        # do the mating using the current population
        off = self.mating.do(self.problem, self.pop, self.n_offsprings, algorithm=self)

        # if the mating could not generate any new offspring (duplicate elimination might make that happen)
        if len(off) == 0:
            self.termination.force_termination = True
            return

        # if not the desired number of offspring could be created
        elif len(off) < self.n_offsprings:
            if self.verbose:
                print("WARNING: Mating could not produce the required number of (unique) offsprings!")

        if self.problem.variables.num_cat_var > 0:
            for individual_idx, individual in enumerate(off):
                X_temp = (
                    individual.X
                    * (np.array(self.problem.variables.upper_bounds) - np.array(self.problem.variables.lower_bounds))
                ) + np.array(self.problem.variables.lower_bounds)
                X_temp = self.problem.variables.categorical_values_euc(X_temp.reshape(1, -1))
                off[individual_idx].X = (X_temp.flatten() - np.array(self.problem.variables.lower_bounds)) / (
                    np.array(self.problem.variables.upper_bounds) - np.array(self.problem.variables.lower_bounds)
                )

        return off

    def _advance(self, infills=None, **kwargs):

        # merge the offsprings with the current population
        if infills is not None:
            self.pop = Population.merge(self.pop, infills)

        # execute the survival to find the fittest solutions
        self.pop = self.survival.do(self.problem, self.pop, n_survive=self.pop_size, algorithm=self)
