import numpy as np
import pygame
import tiles

#Parameters for genetic algorithm
n_trashes = tiles.howManyMesses
n_population = 100
mutation_rate = 0.5
listDraft = []

#Some settings for the whole project window
pygame.display.set_caption('Inteligentny Kelner - projekt SI 2022')
Icon = pygame.image.load('waiter.png')
pygame.display.set_icon(Icon)

#Take data from tiles.py
for i in range(len(tiles.messesArray)):
    # Generating a list of coordinates representing each trash
    coordinates_list = [[x*32,y*32] for x,y in tiles.messesArray]
    listDraft.append(tiles.coordsToNum(coordinates_list[i]))

#Add some data
names_list = np.array(listDraft)
trashes_dict = {x:y for x, y in zip(names_list, coordinates_list)}

#Testing:
#print(trashes_dict)

#Function computing the distance between two points by its coordinates
def compute_trash_distance_coordinates(a, b):
    return ((a[0]-b[0])**2+(a[1]-b[1])**2)**0.5

def compute_trash_distance_numbers(trash_a, trash_b, trash_dict):
    return compute_trash_distance_coordinates(trash_dict[trash_a], trash_dict[trash_b])

#First step is to create the first population set
def genesis(trash_list, n_population):

    population_set = []
    for i in range(n_population):
        #Randomly generated new solution
        sol_i = trash_list[np.random.choice(list(range(n_trashes)), n_trashes, replace=False)]
        population_set.append(sol_i)
    return np.array(population_set)

population_set = genesis(names_list, n_population)
#Testing
#print(population_set)


def fitness_eval(trash_list, trash_dict):
    total = 0
    for i in range(n_trashes - 1):
        a = trash_list[i]
        b = trash_list[i + 1]
        total += compute_trash_distance_numbers(a, b, trash_dict)
    return total


def get_all_fitnes(population_set, trash_dict):
    fitnes_list = np.zeros(n_population)

    #Looping over all solutions computing the fitness for each solution
    for i in  range(n_population):
        fitnes_list[i] = fitness_eval(population_set[i], trash_dict)

    return fitnes_list

fitnes_list = get_all_fitnes(population_set, trashes_dict)
#Testing
#print(fitnes_list)


def ancestors_choose(population_set, fitnes_list):
    total_fit = fitnes_list.sum()
    prob_list = fitnes_list / total_fit
    pop_len = len(population_set)
    ancestor_list_a = np.random.choice(list(range(pop_len)), pop_len, p=prob_list, replace=True)
    ancestor_list_b = np.random.choice(list(range(pop_len)), pop_len, p=prob_list, replace=True)

    ancestor_list_a = population_set[ancestor_list_a]
    ancestor_list_b = population_set[ancestor_list_b]

    return np.array([ancestor_list_a, ancestor_list_b])


ancestor_list = ancestors_choose(population_set, fitnes_list)
#print(ancestor_list[0][2])


def mate_ancestors(prog_a, prog_b):
    offspring = prog_a[0:5]

    for trash in prog_b:
        if not trash in offspring:
            offspring = np.concatenate((offspring, [trash]))

    return offspring


def mate_population(progenitor_list):
    new_population_set = []
    for i in range(progenitor_list.shape[1]):
        prog_a, prog_b = progenitor_list[0][i], progenitor_list[1][i]
        offspring = mate_ancestors(prog_a, prog_b)
        new_population_set.append(offspring)

    return new_population_set

new_population_set = mate_population(ancestor_list)
#print(new_population_set[0])


def mutate_offspring(offspring):
    for q in range(int(n_trashes * mutation_rate)):
        a = np.random.randint(0, n_trashes)
        b = np.random.randint(0, n_trashes)

        offspring[a], offspring[b] = offspring[b], offspring[a]

    return offspring


def mutate_population(new_population_set):
    mutated_pop = []
    for offspring in new_population_set:
        mutated_pop.append(mutate_offspring(offspring))
    return mutated_pop


mutated_pop = mutate_population(new_population_set)
#print(mutated_pop[0])

best_solution = [-1, np.inf, np.array([])]
for i in range(1001):
    #What I put here: iteration number, minimal fitness value, average fitness value
    if i % 100 == 0: print(i, fitnes_list.min(), fitnes_list.mean())
    fitnes_list = get_all_fitnes(mutated_pop, trashes_dict)

    #Saving the best solution
    if fitnes_list.min() < best_solution[1]:
        best_solution[0] = i
        best_solution[1] = fitnes_list.min()
        best_solution[2] = np.array(mutated_pop)[fitnes_list.min() == fitnes_list]

    ancestor_list = ancestors_choose(population_set, fitnes_list)
    new_population_set = mate_population(ancestor_list)

    mutated_pop = mutate_population(new_population_set)
    trashesList = []

    for i in range(len(best_solution[2][0])):
        trashesList.append(0)

    for i in range(len(best_solution[2][0])):
        trashesList[i] = best_solution[2][0][i]

#Prints some specs to recreate if needed
print(trashesList)
print(tiles.messesArray)

#Pass solution to draw
tiles.main(trashesList)