import math
import time
from matplotlib import pyplot as plt
import numpy as np
from GenTunic.gen_math import project_bit
from GenTunic.gen_util import calc_population_fitness, create_population, crossover, mutation, turnament_selection


POPULATIUON_SIZE = 200
MUTATION_RATE = 0.05
CROSSOVER_RATE = 0.65

GEN_SIZE = 8 * 3
THRESHOLD = 0.5

def gen_tuning_main(AMOUNT_TRIES, AMOUNT_RUNS, REWARD_ON_WIN, REWARD_ON_LOSE):
    start_time = time.time()
    
    population = create_population(POPULATIUON_SIZE, GEN_SIZE)

    best_fintess_values = []
    best_fitness = 0
    counter = 0

    while True:
        print(f"Starting eveloution round {counter + 1}")
        #? Calc fitness
        population_propability, fintess_values = calc_population_fitness(population, AMOUNT_TRIES, AMOUNT_RUNS, REWARD_ON_WIN, REWARD_ON_LOSE)

        _, best_fitness = fintess_values
        best_fintess_values.append(best_fitness)
        print(best_fitness)

        if best_fitness > THRESHOLD:
            print("Breaking")
            break

        #? Selection
        amount_selections = math.floor((1 - CROSSOVER_RATE) * len(population_propability))
        amount_crossover = POPULATIUON_SIZE - amount_selections

        new_population = turnament_selection(population_propability, amount_selections)

        #? Crossover
        new_population = crossover(population_propability, new_population, amount_crossover, GEN_SIZE)

        #? Mutation
        population = mutation(new_population, MUTATION_RATE, GEN_SIZE)

        counter += 1


    population_propability, fintess_values = calc_population_fitness(population, AMOUNT_TRIES, AMOUNT_RUNS, REWARD_ON_WIN, REWARD_ON_LOSE)
    best_fintess_index, best_fitness = fintess_values

    print("\n=== BEST PARAMETERS ===")
    gen = population[best_fintess_index]["population"]
    parameter_names = ["Alpha: ", "Epsilon: ", "Gamma: "]
    parameters = [project_bit(x) for x in np.split(gen, 3)]
    for index, name in enumerate(parameter_names):
        print(f"{name}{parameters[index]}")

    time_amount = time.time() - start_time
    print(f"\nTook {time_amount}s")

    plt.plot(best_fintess_values)
    plt.show()