import numpy as np
import matplotlib.pyplot as plt

np.random.seed(4)


# def fitness_function(x):
#     sum = 2.0
#     for item in x:
#         sum -= item**2.0
#     return sum


def fitness_function(x):
    sum = np.float64(-20.0)
    for item in x:
        sum += -(item**2) + 10 * np.cos(2 * np.pi * item)
    sum = 1 / (np.e ** (-sum / 2))
    return sum


def get_zero_population(seed, count_population, demention_population):
    zero_population = np.random.uniform(
        -5.2, 5.2, (count_population, demention_population)
    )
    return zero_population


def get_psi(g, NP, Lambda):
    psi = ((g) * NP + 1) ** (1 / Lambda)
    return psi


def select_reference_vertor(clear_generation, NP, g):
    reference_vector = None
    array_fitness_value = 0.0
    sum_arr = 0.0
    arr_ver = 0.0
    psi = get_psi(g, NP, 100)
    for item in clear_generation:
        num = fitness_function(item)

        array_fitness_value = np.append(array_fitness_value, num**psi)
        sum_arr += num**psi
    for i in array_fitness_value:
        arr_ver = np.append(arr_ver, i / sum_arr)

    id = np.random.choice(len(arr_ver), p=arr_ver)
    # TODO change np.random
    reference_vector = clear_generation[id - 2]

    return reference_vector


def calculate_A(x_min, x_max, x_r, e):
    return np.arctan((x_max - x_r) / e) - np.arctan((x_min - x_r) / e)


def calculate_e(g, NP, D):
    return 1 / ((g) * (NP) + 1) ** (1 / (2 * D))


def generate_potential_offspring(x_r, e, A):
    return x_r + e * np.tan((np.random.rand() - 0.5) * A)


def sofa(zero_population, data_cloud, fitness, mod, steps_number, epsilon):
    # TODO add data_cloud,fitness,mod,epsilon,true_answer
    start_population = np.copy(zero_population)
    mutant_populaion = np.copy(zero_population)
    arr_value_best_item = np.array([fitness_function(start_population[0])])
    value_best_item = arr_value_best_item[0]
    best_item = None

    for item in range(steps_number):
        reference_vector = select_reference_vertor(
            start_population, len(start_population), item
        )
        # print(f"reference_vector {reference_vector}")
        e = calculate_e(item, len(start_population), len(start_population[0]))
        for i in range(len(start_population)):
            const_a = calculate_A(-5.2, 5.2, i, e)
            mutant_populaion[i] = reference_vector + np.tan(
                (np.random.rand() - 0.5) * const_a
            )

        # print(f"mutant_populaion {mutant_populaion}")
        for i in range(len(start_population)):
            fit_mutant_popul = fitness_function(mutant_populaion[i])
            fit_start_popul = fitness_function(start_population[i])
            # print(f"test {fit_mutant_popul} start {fit_start_popul}")
            if fit_mutant_popul > fit_start_popul:
                start_population[i] = mutant_populaion[i]
                if fit_mutant_popul > value_best_item:
                    value_best_item = fit_mutant_popul
                    best_item = mutant_populaion[i]
                    # print(f"best in moment{best_item}")
                    # print(f"best_value in moment{value_best_item}")
        arr_value_best_item = np.append(arr_value_best_item, value_best_item)
    # TODO move for up 3 lines
    # for item in start_population:
    #     value_item = fitness_function(item)
    #     if value_item > value_best_item:
    #         value_best_item = value_item
    #         best_item = item

    print(f"final population: {start_population}")
    print(f"best vector: {best_item}")
    print(f"global maximum: {value_best_item}")
    print(f"fitness_changing: {arr_value_best_item}")
    index = list(np.arange(1.0, len(arr_value_best_item) + 1, 1))
    # print(f"index = {index}")
    fig, ax = plt.subplots()
    ax = plt.plot(index, arr_value_best_item)

    plt.show()
    return start_population, best_item, value_best_item


if __name__ == "__main__":
    # TODO optimization
    # TODO graphics
    steps = 10000
    steps = int(input("steps = "))
    a = get_zero_population(1, 50, 2)
    print(f"zero_population {a}")
    sofa(a, None, None, None, steps, 0.0001)