import numpy as np
import matplotlib.pyplot as plt


DO_VIS = True


# in-place
def clip_params(population, clipper_function):
    for i, individual in enumerate(population):
        population[i] = clipper_function(individual)


def evolution_strategies_optimizer(objective_function, clipper_function,
                                   init_mean, init_scale,
                                   number_of_generations,
                                   population_size):
    # Initialize parameters
    mutation_scale = 0.1
    selection_ratio = 0.5
    selected_size = int(population_size * selection_ratio)

    # Initialize population (randomly)
    population = np.random.normal(loc=init_mean, scale=init_scale, size=(population_size, len(init_mean)))
    clip_params(population, clipper_function)

    for generation in range(number_of_generations):
        # Evaluate fitness
        fitness = np.array([objective_function(individual) for individual in population])

        # Select the best individuals
        selected_indices = np.argsort(fitness)[:selected_size]
        selected = population[selected_indices]

        # Reproduce (mutate)
        offspring = selected + np.random.normal(loc=0, scale=init_scale * mutation_scale, size=(selected_size, len(init_mean)))
        clip_params(offspring, clipper_function) # in-place

        # Logging
        best_fitness = fitness[selected_indices[0]]
        best_solution = selected[0].copy()
        print(f"Generation {generation + 1}: Best Fitness = {best_fitness}", f"Best solution so far: {best_solution}")

        # Replacement: Here we simply generate new candidates around the selected ones
        population[:selected_size] = selected
        population[selected_size:] = offspring

    print(f"Best solution found: {best_solution} with loss {best_fitness}")
    return best_solution


def toy_objective_function(x):
    return (x[..., 0] - 3)**2 + (x[..., 1] + 2)**2


def toy_clipper_function(x):
    return x


def main():
    init_mean = np.array([0.0, 0.0])
    init_scale = np.array([10.0, 10.0])
    best_solution = evolution_strategies_optimizer(toy_objective_function, toy_clipper_function, init_mean, init_scale)


if __name__ == '__main__':
    main()