import random


def fitness_function(x, y):
    return (x ** 2) + (y ** 2)+1


def evaluate_generation(population):
    scores = []
    total = 0
    for individual in population:
        if len(individual) == 2:
            r = -fitness_function(individual[0], individual[1])
            scores.append(r)
            total += r
        else:
            print("error: Wrong number of arguments received")
    avg = total / len(scores)
    return scores, avg


# Create child from parent
def mutate(individual):
    new = []
    for attribute in individual:
        new.append(attribute + random.normalvariate(0, attribute + .1))  # Random factor of normal distribution
    return new


def find_best(population):
    best = None
    val = None
    for individual in population:
        if len(individual) == 2:
            r = fitness_function(individual[0], individual[1])
            try:
                if r < val:
                    best = individual
                    val = r
            except:  # On the first run, set the result as best
                best = individual
                val = r
    else:
        print("error: Wrong number of arguments received")

    return best, val


# Create a population of p lists of [0, 0, ..., 0] of length n
def initialize(n, p):
    pop = [[0] * n]
    for i in range(p):
        pop.append(mutate(pop[0]))
    return pop


# Handle the output of the genetic algorithm
def termination(best, val, total_iterations, population_size, num_attributes):
    best = [round(x, 3) for x in best]  # Round for printing
    print("Ran", total_iterations, "iterations on a population of", population_size)
    print("The optimal input is", best, "with a value of", round(val, 3))


if __name__ == "__main__":
    num_attributes = 2
    population_size = 100
    total_iterations = 100
    population = initialize(num_attributes, population_size)
    for iteration in range(total_iterations):
        scores, avg = evaluate_generation(population)
        deleted = 0
        new_population = []
        for i in range(len(population)):
            if scores[i] < avg:
                deleted += 1
            else:
                new_population.append(population[i])
        for i in range(deleted):
            new_population.append(
                mutate(new_population[i % len(new_population)]))  # iterate over population with overflow protection
        population = new_population
    best, val = find_best(population)
    termination(best, val, total_iterations, population_size, num_attributes)