import numpy
import pygad
import pygad.nn
import pygad.gann
import ele

def fitness_func(solution, sol_idx):
    global GANN_instance, lastbest, lastbesti
    #Run give NN data and ask it what to do next
    envos = [ele.simulation(sol_idx)]
    data_inputs = numpy.empty(shape=[1,86], dtype=int)

    #structure of inputs: current floor, target floor of elevator person[]*max in elevator, target floor of waiting person[floor[max people per floor]]
    data_inputs[0] = envos[0].run([0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]) #Run the first iteration to get a data input
    envos[0].score = 0

    #Simulate 20 iterations
    #indices:                  0   1     2     3 4 5 6 7                                8 9 10 11 12 13 14 15 16 17   
    #strucutre of predictions: up, down, stop, unload passenger[]*max pass in elevator, load passenger[]*max passenger on floor
    for i in range(20):
        predictions = pygad.nn.predict(last_layer=GANN_instance.population_networks[sol_idx], #What will the network do in this situation?
                                    data_inputs=data_inputs, problem_type="regression")
        data_inputs[0] = envos[0].run(predictions[0]) #Simulate according to the instructions of the network and save that as the next input

    #save best current network index manually
    if lastbest <= envos[0].score:
        lastbest = envos[0].score
        lastbesti = sol_idx
    return envos[0].score

def callback_generation(ga_instance):
    global GANN_instance

    population_matrices = pygad.gann.population_as_matrices(population_networks=GANN_instance.population_networks, 
                                                            population_vectors=ga_instance.population)

    GANN_instance.update_population_trained_weights(population_trained_weights=population_matrices)

    print("Generation = {generation}".format(generation=ga_instance.generations_completed))
    print("Accuracy         = {fitness}".format(fitness=ga_instance.best_solution()[1]))

lastbest = 0
lastbesti = 0

GANN_instance = pygad.gann.GANN(num_solutions=50,
                                num_neurons_input=86,    #Number of Datapoints to respect
                                num_neurons_hidden_layers=[2],
                                num_neurons_output=18,   #Number of possible Actions
                                hidden_activations=["relu"],
                                output_activation="sigmoid")

population_vectors = pygad.gann.population_as_vectors(population_networks=GANN_instance.population_networks)

ga_instance = pygad.GA(num_generations=20, 
                       num_parents_mating=3, 
                       initial_population=population_vectors.copy(),
                       fitness_func=fitness_func,
                       mutation_percent_genes=35,
                       callback_generation=callback_generation)

ga_instance.run()
ga_instance.plot_fitness()

solution, solution_fitness, solution_idx = ga_instance.best_solution()
#print(solution)
print(solution_fitness) #result: about 16
print("Solution index acc. to Library:", solution_idx, "Solution acc. to comparison:", lastbesti)

print("Score of solution_idx:", fitness_func(0,solution_idx)) #result of both: all over the place, tho always very bad (<0)
print("Score of lastbesti:", fitness_func(0,lastbesti))

    num_generations = 3
    num_parents_mating = 2
    fitness_function = fitness_func
    sol_per_pop = 4
    num_genes = 4
    parent_selection_type = "tournament"
    keep_elitism = 1

[0.7014856849397931, 0.7132997342518398, 0.7557456777209327, 0.7557456777209327, 0.7429285673868089, 0.7429285673868089, 0.7429285673868089, 0.7429285673868089, 0.7429285673868089, 0.7429285673868089, 0.7429285673868089]