# Native
from typing import Optional

# pymoo imports
from pymoo.core.algorithm import Algorithm
from pymoo.core.duplicate import DefaultDuplicateElimination, NoDuplicateElimination
from pymoo.core.initialization import Initialization
from pymoo.core.infill import InfillCriterion
from pymoo.core.population import Population
from pymoo.core.repair import NoRepair, Repair
from pymoo.core.survival import Survival
from pymoo.core.sampling import Sampling
from pymoo.core.mating import Mating


# =========================================================================================================
# Implementation
# =========================================================================================================

class EvolutionaryAlgorithm(Algorithm):

    def __init__(
        self,
        pop_size: Optional[int] = None,
        sampling: Optional[Sampling] = None,
        mating: Optional[Mating] = None,
        survival: Optional[Survival] = None,
        n_offsprings: Optional[int] = None,
        eliminate_duplicates: bool = True,
        repair: Optional[Repair] = None,
        advance_after_initial_infill: bool = False,
        **kwargs
    ):
        """Base class for Evolutionary Algorithms

        Parameters
        ----------
        pop_size : int, optional
            Population size, by default None

        sampling : Sampling, optional
            pymoo Sampling instance, by default None

        mating : InfillCriterion, optional
            pymoo mating operator, by default None

        survival : Survival, optional
            pymoo survival operator, by default None

        n_offsprings : int, optional
            Number of offspring individuals created at each generation, by default None

        eliminate_duplicates : DuplicateElimination | bool | None, optional
            Eliminate duplicates in mating, by default True

        repair : Repair, optional
            pymoo repair operator. In the algorithm level it should be called when sampling.
            It is recommended that mating operators also have some repair associated with. By default None

        advance_after_initial_infill : bool, optional
            Either or not apply survival after initialization, by default False
        """

        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 = InfillCriterion(
                repair=self.repair,
                eliminate_duplicates=self.eliminate_duplicates,
        )
        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)
        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), algorithm=self, **kwargs)

    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!")

        return off

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

        # the current population
        pop = self.pop

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

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