import numpy as np
import ionbench
import ionbench.utils.population_optimisers as pop_opt
import copy


# noinspection PyShadowingNames
def run(bm, x0=None, nGens=50, eta_cross=10, eta_mut=20, popSize=50, debug=False):
    """
    Runs the genetic algorithm from Bot et al. 2012.

    Parameters
    ----------
    bm : Benchmarker
        A benchmarker to evaluate the performance of the optimisation algorithm.
    x0 : list, optional
        Initial parameter guess. Population is generated by randomly perturbing this initial guess +-50%, clamped to bounds if necessary. If x0=None (the default), then the population will be sampled using bm.sample().
    nGens : int, optional
        The number of generations to run the optimisation algorithm for. The default is 50.
    eta_cross : float, optional
        Crossover parameter. The default is 10.
    eta_mut : float, optional
        Mutation parameter. The default is 20.
    popSize : int, optional
        The size of the population in each generation. The default is 50.
    debug : bool, optional
        If True, debug information will be printed, reporting that status of the optimisation each generation. The default is False.

    Returns
    -------
    xbest : list
        The best parameters identified.

    """

    cost_func = ionbench.utils.cache.get_cached_cost(bm)

    pop = pop_opt.get_pop(bm, x0, popSize, cost_func)
    gen = None
    for gen in range(nGens):
        elites = pop_opt.get_elites(pop, 1)
        if debug:
            print("------------")
            print(f'Gen {gen}, Best cost: {elites[0].cost}')

        # Tournament selection
        pop = pop_opt.tournament_selection(pop)

        # Crossover SBX
        pop = pop_opt.sbx_crossover(pop, bm, cost_func, eta_cross)

        # Mutation
        pop = pop_opt.polynomial_mutation(pop, bm, cost_func, eta_mut)

        if debug:
            print(f'Finishing gen {gen}')

        # Find costs
        pop = pop_opt.find_pop_costs(pop)

        # Elitism
        pop = pop_opt.set_elites(pop, elites)

        if bm.is_converged():
            break

    if gen >= nGens-1:
        bm.set_max_iter_flag()

    elites = pop_opt.get_elites(pop, 1)
    bm.evaluate()
    return elites[0].x


# noinspection PyShadowingNames
def get_modification(modNum=1):
    """
    modNum = 1 -> Bot2012
    modNum = 2 -> Groenendaal2015

    Returns
    -------
    mod : modification
        Modification corresponding to inputted modNum. Default is modNum = 1, so Bot2012.

    """

    if modNum == 1:
        mod = ionbench.modification.Bot2012()
    elif modNum == 2:
        mod = ionbench.modification.Groenendaal2015()
    else:
        mod = ionbench.modification.Empty(name='GA_Bot2012')
    return mod
