import numpy as np
import matplotlib.pyplot as plt
import networkx as nx
from scipy.spatial import distance_matrix
import itertools
from networkx.drawing.layout import fruchterman_reingold_layout


def add_unique_vector(vectors, selected_vectors, remaining_indices):

    # Выбираем случайный индекс из оставшихся
    new_index = np.random.choice(remaining_indices)

    # Добавляем новый выбранный вектор
    selected_vectors.append(vectors[new_index])
    remaining_indices = list(remaining_indices - new_index)

    return selected_vectors, remaining_indices


# Расчет энергии вершины
def calculate_vertex_energy(vertex, vectors, k, l):
    # Вычисляем расстояния от вершины до каждого вектора
    distances = [np.linalg.norm(vertex - vector) for vector in vectors]

    # Получаем индексы векторов, отсортированных по расстоянию
    sorted_indices = np.argsort(distances)

    # Выбираем num_nearest ближайших векторов
    nearest_vector = sorted_indices[0]

    # Вычисляем энергию вершины
    energy = 0

    energy += k * (1 / l) * np.linalg.norm(vertex - nearest_vector) ** 2
    return energy


# Расчет энергии ребра
def calculate_edge_energy(edge, vectors, k, s):
    v1, v2 = vectors[edge[0]], vectors[edge[1]]
    return k * s**2 * np.linalg.norm(v1 - v2) ** 2


# Расчет энергии клина
def calculate_wedge_energy(wedge, vectors, k, r):
    v1, v2, v3 = vectors[wedge[0]], vectors[wedge[1]], vectors[wedge[2]]
    return k * r**4 * np.linalg.norm(v1 + v3 - 2 * v2) ** 2


# Расчет энергии графа
def calculate_energy(
    graph,
    vectors,
    selected_vectors,
    max_vertex_energy,
    max_edge_energy,
    max_wedge_energy,
    max_energy_vertex,
    max_energy_edge,
    max_energy_wedge,
    edge_combination,
):
    # vertex_k = 0.5 # Замените на ваше значение для коэффициента аппроксимации вершин
    # edge_k = 0.5 # Замените на ваше значение для коэффициента растяжения ребер
    # wedge_k = 0.5 # Замените на ваше значение для коэффициента изгиба клиньев
    max_wedge_energy = (
        -np.inf
    )  # Инициализируем максимальную энергию клина как минус бесконечность
    max_energy_wedge = None  # Инициализируем клин с максимальной энергией как None
    edge_combination = 0
    max_vertex_energy = -np.inf
    max_edge_energy = -np.inf
    max_energy_vertex = 0
    max_energy_edge = None

    energy = 0
    for node in graph.nodes:
        vertex_energy = calculate_vertex_energy(
            selected_vectors[int(node)], vectors, vertex_k, len(selected_vectors)
        )
        if vertex_energy > max_vertex_energy:
            max_vertex_energy = vertex_energy
            max_energy_vertex = node
        energy += vertex_energy
    for edge in graph.edges:
        edge_energy = calculate_edge_energy(edge, selected_vectors, edge_k)
        if edge_energy > max_edge_energy:
            max_edge_energy = edge_energy
            max_energy_edge = edge
        energy += edge_energy

    # Добавляем энергию клина для всех подграфов из трех вершин
    for nodes in itertools.combinations(graph.nodes, 3):
        # print(nodes)
        if graph.has_edge(nodes[0], nodes[1]) and graph.has_edge(nodes[0], nodes[2]):
            wedge_energy = calculate_wedge_energy(nodes, selected_vectors, wedge_k)
            # print("Энергия клина:")
            # print(p1)
            if wedge_energy > max_wedge_energy:
                max_wedge_energy = wedge_energy
                max_energy_wedge = nodes
                edge_combination = 1
            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, selected_vectors, wedge_k)
            if wedge_energy > max_wedge_energy:
                max_wedge_energy = wedge_energy
                max_energy_wedge = nodes
                edge_combination = 2
                # print("Энергия клина:")
            # print(p1)
            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, selected_vectors, wedge_k)
            if wedge_energy > max_wedge_energy:
                max_wedge_energy = wedge_energy
                max_energy_wedge = nodes
                edge_combination = 3
            # print("Энергия клина:")
            # print(p1)
            energy += wedge_energy
    return (
        energy,
        max_vertex_energy,
        max_edge_energy,
        max_wedge_energy,
        max_energy_vertex,
        max_energy_edge,
        max_energy_wedge,
        edge_combination,
    )


# Расчет энергии графа
def calculate_energy1(graph, vectors, selected_vectors):
    vertex_k = 2  # Замените на ваше значение для коэффициента аппроксимации вершин
    edge_k = 1  # Замените на ваше значение для коэффициента растяжения ребер
    wedge_k = 0.25
    energy = 0
    for node in graph.nodes:
        vertex_energy = calculate_vertex_energy(
            selected_vectors[int(node)], vectors, vertex_k, len(graph.nodes)
        )
        energy += vertex_energy
    for edge in graph.edges:
        edge_energy = calculate_edge_energy(
            edge, selected_vectors, edge_k, len(graph.edges)
        )
        energy += edge_energy
    r = 0
    # Добавляем энергию клина для всех подграфов из трех вершин
    for nodes in itertools.combinations(graph.nodes, 3):
        # print(nodes)
        if graph.has_edge(nodes[0], nodes[1]) and graph.has_edge(nodes[0], nodes[2]):
            r = r + 1
        elif graph.has_edge(nodes[0], nodes[1]) and graph.has_edge(nodes[1], nodes[2]):
            r = r + 1
        elif graph.has_edge(nodes[0], nodes[2]) and graph.has_edge(nodes[1], nodes[2]):
            r = r + 1

    # Добавляем энергию клина для всех подграфов из трех вершин
    for nodes in itertools.combinations(graph.nodes, 3):
        # print(nodes)
        if graph.has_edge(nodes[0], nodes[1]) and graph.has_edge(nodes[0], nodes[2]):
            wedge_energy = calculate_wedge_energy(nodes, selected_vectors, wedge_k, r)
            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, selected_vectors, wedge_k, r)
            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, selected_vectors, wedge_k, r)
            energy += wedge_energy
    return energy


def visualize_graph(vectors, selected_vectors, G):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")

    # Отображаем все векторы на графике
    for vector in vectors:
        ax.scatter(vector[0], vector[1], vector[2], color="b")

    # Отображаем выбранные векторы на графике
    for vector in selected_vectors:
        ax.scatter(vector[0], vector[1], vector[2], color="r")

    # Добавляем ребра между вершинами
    for edge in G.edges:
        ax.plot(
            [selected_vectors[edge[0]][0], selected_vectors[edge[1]][0]],
            [selected_vectors[edge[0]][1], selected_vectors[edge[1]][1]],
            [selected_vectors[edge[0]][2], selected_vectors[edge[1]][2]],
            "r-",
        )

    # Добавляем оси
    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    ax.set_zlabel("Z")
    plt.show()


def modify_graph(
    G,
    selected_vectors,
    max_vertex_energy,
    max_edge_energy,
    max_wedge_energy,
    max_energy_wedge,
    edge_combination,
    vectors,
    remaining_indices,
):
    if max_vertex_energy >= max_edge_energy and max_vertex_energy >= max_wedge_energy:
        print("hello1")
        # Добавляем новую вершину к вершине с наибольшей энергией
        # max_energy_vertex = vertex_energies.index(max_vertex_energy)
        G.add_node(len(selected_vectors))
        G.add_edge(max_energy_vertex, len(selected_vectors))
        selected_vectors, remaining_indices = add_unique_vector(
            vectors, selected_vectors, remaining_indices
        )
        # selected_vectors.append(np.random.randint(-6, 7, 3))  # Добавляем новый вектор
    elif max_edge_energy >= max_vertex_energy and max_edge_energy >= max_wedge_energy:
        print("hello2")
        # Разбиваем ребро с наибольшей энергией новой вершиной
        # edges = list(G.edges)  # Получаем список всех ребер
        # max_energy_edge = edges[edge_energies.index(max_edge_energy)]  # Находим ребро с максимальной энергией
        G.add_node(len(selected_vectors))
        G.add_edge(max_energy_edge[0], len(selected_vectors))
        G.add_edge(max_energy_edge[1], len(selected_vectors))
        G.remove_edge(max_energy_edge[0], max_energy_edge[1])
        selected_vectors, remaining_indices = add_unique_vector(
            vectors, selected_vectors, remaining_indices
        )  # Добавляем новый вектор
    else:
        print("hello3")
        if edge_combination == 1:
            old_edge_length = np.linalg.norm(
                selected_vectors[max_energy_wedge[0]]
                - selected_vectors[max_energy_wedge[1]]
            )
            new_edge_length = np.linalg.norm(
                selected_vectors[max_energy_wedge[0]]
                - selected_vectors[max_energy_wedge[2]]
            )
            if new_edge_length < old_edge_length:
                # Удаляем старое ребро и добавляем новое
                G.remove_edge(max_energy_wedge[0], max_energy_wedge[1])
                G.add_edge(max_energy_wedge[2], max_energy_wedge[1])
            else:
                G.remove_edge(max_energy_wedge[0], max_energy_wedge[2])
                G.add_edge(max_energy_wedge[2], max_energy_wedge[1])
        elif edge_combination == 2:
            old_edge_length = np.linalg.norm(
                selected_vectors[max_energy_wedge[0]]
                - selected_vectors[max_energy_wedge[1]]
            )
            new_edge_length = np.linalg.norm(
                selected_vectors[max_energy_wedge[1]]
                - selected_vectors[max_energy_wedge[2]]
            )
            print(f"длина старого ребра: {old_edge_length}")
            print(f"длина нового ребра: {new_edge_length }")
            if new_edge_length < old_edge_length:
                # Удаляем старое ребро и добавляем новое
                G.remove_edge(max_energy_wedge[0], max_energy_wedge[1])
                G.add_edge(max_energy_wedge[0], max_energy_wedge[2])
            else:
                G.remove_edge(max_energy_wedge[1], max_energy_wedge[2])
                G.add_edge(max_energy_wedge[0], max_energy_wedge[2])
        elif edge_combination == 3:
            old_edge_length = np.linalg.norm(
                selected_vectors[max_energy_wedge[0]]
                - selected_vectors[max_energy_wedge[2]]
            )
            new_edge_length = np.linalg.norm(
                selected_vectors[max_energy_wedge[1]]
                - selected_vectors[max_energy_wedge[2]]
            )
            if new_edge_length < old_edge_length:
                # Удаляем старое ребро и добавляем новое
                G.remove_edge(max_energy_wedge[0], max_energy_wedge[2])
                G.add_edge(max_energy_wedge[0], max_energy_wedge[1])
            else:
                G.remove_edge(max_energy_wedge[2], max_energy_wedge[1])
                G.add_edge(max_energy_wedge[0], max_energy_wedge[1])
        # print(max_energy_wedge[0])
        # print(max_energy_wedge[1])
        # print(max_energy_wedge[2])
    return G, selected_vectors, remaining_indices
    # for edges in G.edges:
    # print(edges)


def calculate_psi(g, NP, lambda_):
    return ((g - 1) * (NP + 1) + 1) ** (1 / lambda_)


def calculate_probabilities(
    vectors, graphs, selected_vectors_list, vertex_k, edge_k, wedge_k, g, NP, lambda_
):
    # Вычисляем psi
    psi_g = calculate_psi(g, NP, lambda_)

    # Вычисляем энергии для всех векторов
    energies = [
        calculate_energy1(graphs[i], vectors, selected_vectors_list[i]) ** psi_g
        for i in range(len(selected_vectors_list))
    ]

    # Вычисляем сумму энергий
    sum_energies = sum(energies)

    # Вычисляем вероятности
    probabilities = [energy / sum_energies for energy in energies]
    # print(probabilities)
    return probabilities


def select_support_vector(
    vectors, graphs, selected_vectors_list, vertex_k, edge_k, wedge_k, g, NP, lambda_
):
    # Вычисляем вероятности
    probabilities = calculate_probabilities(
        vectors,
        graphs,
        selected_vectors_list,
        vertex_k,
        edge_k,
        wedge_k,
        g,
        NP,
        lambda_,
    )

    # Выбираем опорный вектор на основе вероятностей
    id = np.random.choice(len(selected_vectors_list), p=probabilities)
    support_vector = selected_vectors_list[id]

    return support_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)) ** -1


def calculate_e(i, g, NP, D):
    return 1 / ((g - 1) * (NP + 1) + i + 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 crossover(selected_vectors_list, potential_offspring, q):
    # Создаем mutant_vectors, равный по размеру selected_vectors
    mutant_vectors = potential_offspring

    # Заменяем случайные q векторов mutant_vectors на соответствующие векторы из selected_vectors

    # Для оставшихся векторов
    for l in range(q, len(selected_vectors_list)):
        for j in range(len(selected_vectors_list[l])):
            # Выбираем случайное число
            p = np.random.rand()
            # Вычисляем CR_{l,g}
            CR_lg = 0.9 if p > 0.1 else p
            # Для каждой координаты j в векторе
            if np.random.rand() < 1.0 - CR_lg:
                mutant_vectors[l], [j] = selected_vectors_list[l], [j]

    return mutant_vectors


def generate_offspring(selected_vectors_list, support_vector, g, NP):
    # Вычисляем x_min и x_max для каждой координаты
    x_min = []
    x_max = []
    for i in range(len(selected_vectors_list)):
        x_min_i = []
        x_max_i = []
        for j in range(len(selected_vectors_list[i])):
            x_min_i.append(np.min(selected_vectors_list[i][j]))
            x_max_i.append(np.max(selected_vectors_list[i][j]))
        x_min.append(x_min_i)
        x_max.append(x_max_i)

    # Создаем список для хранения потенциальных потомков
    potential_offspring = []

    # Для каждого вектора в selected_vectors
    for i in range(len(selected_vectors_list)):
        # Для каждой координаты в векторе
        potential_offspring_i = []
        for j in range(len(selected_vectors_list[i])):
            # Вычисляем x_r, e и A
            x_r = support_vector[j]
            e = calculate_e(i, g, NP, len(selected_vectors_list[i]))
            A = calculate_A(-1, 1, x_r, e)

            # Генерируем потенциального потомка
            x_star = generate_potential_offspring(x_r, e, A)

            # Добавляем потенциального потомка в список
            potential_offspring_i.append(x_star)
        potential_offspring.append(potential_offspring_i)

    return potential_offspring


def selection(G, vectors, selected_vectors, mutant_vectors, energy):
    # Для каждого вектора в selected_vectors
    for i in range(len(selected_vectors)):
        vec = selected_vectors
        vec[i] = mutant_vectors[i]
        en = calculate_energy1(G, vectors, vec)

        if en < energy:
            selected_vectors[i] = mutant_vectors[i]

    return selected_vectors


def data():
    n = 12

    # Создаем случайную последовательность n-2 от 0 до n-1
    prufer_sequence = np.random.randint(0, n, n - 2)

    # Восстанавливаем дерево по коду Прюфера
    G1 = nx.from_prufer_sequence(prufer_sequence)

    # Укладываем вершины графа на плоскость методом Фрухтермана-Рейнгольда
    pos = fruchterman_reingold_layout(G1, dim=3)

    # Берем точки на ребрах графа
    edge_points = [pos[node] for node in G1.nodes]
    # [(pos[edge[0]] + pos[edge[1]]) / 2 for edge in G.edges]

    # Смещаем каждую точку на случайный вектор с нормальным распределением
    vectors = []
    for point in edge_points:
        for _ in range(5):  # Создаем три смещенные точки для каждой точки на ребре
            vectors.append(point + np.random.normal(size=3) * 0.05)
    return vectors


def visualize_graphs(vectors, graphs, selected_vectors_list):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")

    # Отображаем все векторы на графике
    for vector in vectors:
        ax.scatter(vector[0], vector[1], vector[2], color="b")
    color_i = [
        "red",
        "blue",
        "green",
        "yellow",
        "orange",
        "lightcoral",
        "darkorange",
        "olive",
        "teal",
        "violet",
        "skyblue",
    ]
    # Отображаем выбранные векторы на графике
    for i in range(len(selected_vectors_list)):
        for vector in selected_vectors_list[i]:
            ax.scatter(vector[0], vector[1], vector[2], color=color_i[i % len(color_i)])

    # Добавляем ребра между вершинами
    for i in range(len(graphs)):
        for edge in graphs[i].edges:
            ax.plot(
                [
                    selected_vectors_list[i][edge[0]][0],
                    selected_vectors_list[i][edge[1]][0],
                ],
                [
                    selected_vectors_list[i][edge[0]][1],
                    selected_vectors_list[i][edge[1]][1],
                ],
                [
                    selected_vectors_list[i][edge[0]][2],
                    selected_vectors_list[i][edge[1]][2],
                ],
                color_i[i % len(color_i)],
            )

    # Добавляем оси
    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    ax.set_zlabel("Z")
    plt.show()


def create_auxiliary_variables(NP):
    max_wedge_energy = [-np.inf] * NP
    max_energy_wedge = [None] * NP
    edge_combination = [0] * NP
    max_vertex_energy = [-np.inf] * NP
    max_edge_energy = [-np.inf] * NP
    max_energy_vertex = [0] * NP
    max_energy_edge = [None] * NP
    return (
        max_wedge_energy,
        max_energy_wedge,
        edge_combination,
        max_vertex_energy,
        max_edge_energy,
        max_energy_vertex,
        max_energy_edge,
    )


def DFS(G, v, visited=None):
    if visited is None:
        visited = set()
    visited.add(v)
    for neighbour in G[v]:
        if neighbour not in visited:
            DFS(G, neighbour, visited)
    return visited


def generate_graphs(vectors, NP):
    graphs = []
    selected_vectors_list = []
    remaining_indices_list = []

    for _ in range(NP):
        indices = np.random.choice(len(vectors), 3, replace=False)
        selected_vectors = [vectors[i] for i in indices]
        all_indices = set(range(len(vectors)))
        remaining_indices = list(all_indices - set(indices))

        distances = distance_matrix(selected_vectors, selected_vectors)
        np.fill_diagonal(distances, np.inf)

        dist_vector_pairs = [
            (np.min(dist), vec) for dist, vec in zip(distances, selected_vectors)
        ]
        dist_vector_pairs.sort(key=lambda x: x[0])
        sorted_vectors = [pair[1] for pair in dist_vector_pairs]
        selected_vectors = sorted_vectors

        distances = distance_matrix(selected_vectors, selected_vectors)
        np.fill_diagonal(distances, np.inf)

        G = nx.Graph()
        for i in range(len(selected_vectors)):
            G.add_node(i)

        visited = set()
        for i in range(len(selected_vectors)):
            nearest_vertex = np.argmin(distances[i])
            if i not in visited:
                G.add_edge(i, nearest_vertex)
                visited.add(nearest_vertex)

        connected = len(DFS(G, 0)) == len(G)
        if not connected:
            components = [list(c) for c in nx.connected_components(G)]
            for i in range(len(components) - 1):
                G.add_edge(components[i][0], components[i + 1][0])

        graphs.append(G)
        selected_vectors_list.append(selected_vectors)
        remaining_indices_list.append(remaining_indices)

    return graphs, selected_vectors_list, remaining_indices_list


if __name__ == "__main__":
    vertex_k = 10  # Замените на ваше значение для коэффициента аппроксимации вершин
    edge_k = 100  # Замените на ваше значение для коэффициента растяжения ребер
    wedge_k = 10  # Замените на ваше значение для коэффициента изгиба клиньев
    NP = 5
    np.random.seed(8)

    vectors = []
    vectors = data()

    graphs, selected_vectors_list, remaining_indices_list = generate_graphs(vectors, NP)
    (
        max_wedge_energy,
        max_energy_wedge,
        edge_combination,
        max_vertex_energy,
        max_edge_energy,
        max_energy_vertex,
        max_energy_edge,
    ) = create_auxiliary_variables(NP)
    # visualize_graphs(vectors, graphs,selected_vectors_list)
    # print(calculate_probabilities(vectors,graphs, selected_vectors_list, vertex_k,edge_k,wedge_k, 1, NP, 1))
    for i in range(2250):
        support_vector = select_support_vector(
            vectors, graphs, selected_vectors_list, vertex_k, edge_k, wedge_k, 1, NP, 1
        )
        potential_offspring = generate_offspring(
            selected_vectors_list, support_vector, 1, NP
        )
        # visualize_graphs(vectors, graphs, potential_offspring)
        mutant_vectors = crossover(selected_vectors_list, potential_offspring, 3)
        for j in range(len(mutant_vectors)):
            energy_mutant = calculate_energy1(graphs[j], vectors, mutant_vectors[j])
            energy_selected = calculate_energy1(
                graphs[j], vectors, selected_vectors_list[j]
            )
            if energy_mutant < energy_selected:
                selected_vectors_list[j] = mutant_vectors[j]
    visualize_graphs(vectors, graphs, mutant_vectors)