import numpy as np
import itertools
import random
import matplotlib.pyplot as plt
import networkx as nx
from networkx.drawing.layout import fruchterman_reingold_layout
from sklearn.cluster import KMeans

np.random.seed(4)

VERTEX = 1.0
EGDE = 0.006
WEGDE = 0.0001
POPUL_LENGTH = 50
GRAPH_LENGTH = 4
N_CLUSTERS = 4  # должно равняться GRAPH_LENGTH


def data(n):

    prufer_sequence = np.random.randint(0, n, n - 2)

    G1 = nx.from_prufer_sequence(prufer_sequence)

    pos = fruchterman_reingold_layout(G1, dim=2)

    edge_points = [pos[node] for node in G1.nodes]

    vectors = []
    for point in edge_points:
        for _ in range(10):
            vectors.append(point + np.random.normal(size=2) * 0.05)
    return vectors


def cluster_data(vectors, n_clusters):
    # Преобразование списка векторов в массив NumPy для кластеризации
    vectors = np.array(vectors)

    # Создание модели K-средних с заданным количеством кластеров
    kmeans = KMeans(n_clusters=n_clusters)

    # Выполнение кластеризации
    kmeans.fit(vectors)

    # Получение центроидов кластеров
    centroids = kmeans.cluster_centers_

    return centroids


def visualize_graph(vectors, selected_vectors, G):
    for point in vectors:
        plt.scatter(point[0], point[1], color="b")
    for node in selected_vectors:
        plt.scatter(node[0], node[1], color="r")
    pos = {
        node: (selected_vectors[node][0], selected_vectors[node][1])
        for node in G.nodes()
    }
    nx.draw(G, pos, with_labels=False, node_size=0, edge_color="k")
    plt.show()


def get_zero_population_graph(data, length_population=50, graph_length=3):
    zero_population_graph = []
    data_list = list(data)
    for i in range(length_population):
        selected_elements = np.array(random.sample(data_list, graph_length))
        prufer_sequence = np.random.randint(0, graph_length, graph_length - 2)
        G1 = nx.from_prufer_sequence(prufer_sequence)
        if i == 0:
            zero_population = np.copy(selected_elements)
        else:
            zero_population = np.concatenate(
                [zero_population, selected_elements], axis=0
            )
        zero_population_graph.append(G1)

    zero_population = np.vsplit(zero_population, length_population)
    return zero_population, zero_population_graph


def calculate_vertex_energy(vertex, vectors, k):
    distances = [np.linalg.norm(vertex - vector) for vector in vectors]

    if 0.0 in distances:
        distances[distances.index(0.0)] = np.inf

    sorted_indices = np.argsort(distances)

    nearest_vector = sorted_indices[0]

    energy = 0
    energy += k * np.linalg.norm(vertex - vectors[nearest_vector]) ** 2
    return energy


def calculate_edge_energy(edge, vectors, k):
    v1, v2 = vectors[edge[0] - 1], vectors[edge[1] - 1]
    return k * np.linalg.norm(v1 - v2) ** 2


def calculate_wedge_energy(wedge, vectors, k):
    v1, v2, v3 = vectors[wedge[0] - 1], vectors[wedge[1] - 1], vectors[wedge[2] - 1]
    return k * np.linalg.norm(v1 + v3 - 2 * v2) ** 2


def fitness_function(data, points, graph):
    fitness = 0
    for node in graph.nodes:
        vertex_energy = calculate_vertex_energy(points[int(node) - 1], data, VERTEX)
        fitness += vertex_energy
    for edge in graph.edges:
        edge_energy = calculate_edge_energy(edge, points, EGDE)
        fitness += edge_energy
    sum_wedge_energy = 0.0
    for nodes in itertools.combinations(graph.nodes, 3):
        if graph.has_edge(nodes[0], nodes[1]) and graph.has_edge(nodes[0], nodes[2]):
            wedge_energy = calculate_wedge_energy(
                [nodes[1], nodes[0], nodes[2]], points, WEGDE
            )
            sum_wedge_energy += wedge_energy
        elif graph.has_edge(nodes[0], nodes[1]) and graph.has_edge(nodes[1], nodes[2]):
            wedge_energy = calculate_wedge_energy(nodes, points, WEGDE)
            sum_wedge_energy += wedge_energy
        elif graph.has_edge(nodes[0], nodes[2]) and graph.has_edge(nodes[1], nodes[2]):
            wedge_energy = calculate_wedge_energy(
                [nodes[1], nodes[2], nodes[0]], points, WEGDE
            )
            sum_wedge_energy += wedge_energy
    fitness += sum_wedge_energy
    return 1000 - fitness


def get_zero_population(seed, count_population, demention_population):
    zero_population = np.random.uniform(
        -2.0, 2.0, (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(data, clear_generation, graph, NP, g):
    reference_vector = None
    array_fitness_value = 0.0
    sum_arr = 0.0
    arr_ver = 0.0
    psi = get_psi(g, NP, POPUL_LENGTH)
    for item in range(len(clear_generation)):
        num = fitness_function(data, clear_generation[item], graph[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)
    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, graph, 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(data_cloud, zero_population[0], graph[0])]
    )
    value_best_item = np.copy(arr_value_best_item[0])
    print(f"first_value_best: {value_best_item}")

    best_item = np.copy(mutant_populaion[0])
    print(f"best_vector1: {best_item}")
    for item in range(steps_number):
        reference_vector = select_reference_vertor(
            data_cloud, start_population, graph, len(start_population), item
        )
        e = calculate_e(item, len(start_population), len(start_population[0]))
        for i in range(len(start_population)):
            for j in range(len(start_population[0])):
                const_a = calculate_A(-1.0, 1.0, reference_vector[j], e)
                mutant_populaion[i][j] = reference_vector[j] + np.tan(
                    (np.random.random(2) - 0.5) * const_a
                )

        for i in range(len(start_population)):
            fit_mutant_popul = fitness_function(data, mutant_populaion[i], graph[i])
            fit_start_popul = fitness_function(data, start_population[i], graph[i])
            if fit_mutant_popul > fit_start_popul:
                start_population[i] = np.copy(mutant_populaion[i])
                if fit_mutant_popul > value_best_item:
                    value_best_item = fit_mutant_popul
                    print(f"value_best_item {value_best_item}")
                    best_item = np.copy(mutant_populaion[i])
        print(f"step: {item}")

        arr_value_best_item = np.append(arr_value_best_item, value_best_item)

    print(f"best vector2: {best_item}")
    print(f"global maximum: {value_best_item}")
    index = list(np.arange(1.0, len(arr_value_best_item) + 1, 1))
    return start_population, best_item, value_best_item


if __name__ == "__main__":
    data = data(GRAPH_LENGTH)
    steps = int(input("steps = "))
    # Кластеризация данных
    centroids = cluster_data(data, N_CLUSTERS)
    a, graph = get_zero_population_graph(centroids, POPUL_LENGTH, GRAPH_LENGTH)
    a, item, b = sofa(a, graph, data, None, None, steps, 0.0001)
    visualize_graph(data, item, graph[0])