from nnga import nnGA, GaussianInitializationStrategy, \
    GaussianMutationStrategy, BasicCrossoverStrategy, \
    PopulationParameters

def make_network(parameters=None):
    ''' Function that creates a network given a set of parameters '''
    neural_network = ...
    return neural_network


def fitness(idx, parameters):
    ''' Fitness function to evaluate a set of parameters '''
    # Evaluate parameters
    network = make_network(parameters)
    return evaluate_network(network)


if __name__ == '__main__':
    # Initialize GA parameters
    network = make_initial_network()
    network_structure = [list(layer.shape) for layer in network]  # List of tuples, containing the shape of each layer
    
    # Population parameters
    population = PopulationParameters(population_size=200)
    
    # Mutation strategy
    mutation = GaussianMutationStrategy(network_structure, 1e-1)
    
    # Crossover strategy
    crossover = BasicCrossoverStrategy(network_structure)
    
    # Initialization strategy
    init = GaussianInitializationStrategy(
        mean=0., std=1., network_structure=network_structure)

    ga = nnGA(
        epochs=50,  # Number of epochs
        fitness_function=fitness,
        population_parameters=population,
        mutation_strategy=mutation,
        initialization_strategy=init,
        crossover_strategy=crossover,
        num_processors=8)  # Number of cores

    # Run GA
    network_parameters, best_result, results = ga.run()