# population_manip.py
# This file handles interfacing any population selection operations


# Internal Imports
from src.genetic_embedding.utils.population_manip import selection, genetic_changes

# External Imports


'''
----------selection----------
- This function performs a roulette wheel selection on the given population
-----Inputs-----
- population - The population from which to select values
- amount - The number of items to select (Defaults to 1)
- method - The selection method to use (Defaults to roulette-wheel)
-----Output-----
- chromosomes - The list of chromosomes (or single chromosome) selected
'''
def select(population, amount=1, method="", selection_pressure=2):
    chromosomes = []
    # Make a selection based on the method supplied
    if method == "tournament":
        chromosomes = [selection.tournament(population, selection_pressure) for i in range(amount)]
    elif method == "top_k":
        chromosomes = population[0:amount]
    else:
        # Run the default (roulette wheel selection)
        chromosomes = selection.roulette_wheel(population, amount)

    # Depending on the number of chromosomes in the list, return a value
    if len(chromosomes) == 1:
        return chromosomes[0]
    else:
        return chromosomes


'''
----------crossover----------
- This function performs a roulette wheel selection on the given population
-----Inputs-----
- population - The population to run crossover on
- crossover_span - The number of bits to cross over (Defaults to 50)
-----Output-----
- result - The resultant chromosomes
'''
def crossover(population, selection_method = "roulette_wheel", crossover_span = 50, crossover_method="default", mutation_rate = 0):
    # Select two chromosomes
    chromosomes = select(population, amount=2, method=selection_method)

    # Run the crossover and return the chromosomes
    if crossover_method == "uniform":
        return genetic_changes.crossover(chromosomes, crossover_method, mutation_rate)
    else:
        return genetic_changes.old_crossover(chromosomes, crossover_span)


'''
----------mutation----------
- This function performs a roulette wheel selection on the given population
-----Inputs-----
- population - The population to run mutation on
- mutation_amt - The number of bits to mutate (defaults to 1)
-----Output-----
- result - The resultant chromosome
'''
def mutation(population, selection_method = "roulette_wheel", mutation_amt = 1):
    # Select one chromosome
    chromosome = select(population, method=selection_method)

    # Run the mutation and return the result
    return genetic_changes.mutation(chromosome, mutation_amt)


'''
----------sort----------
- This function sorts the population from highest to lowest fitness
-----Inputs-----
- population - The population to sort
-----Output-----
- sorted - The sorted population
'''
def sort(population):
    # Return the sorted population
    return sorted(population, key = lambda cr: cr.fitness, reverse=True)