from __future__ import division
import numpy as np
import torch
import matplotlib.pyplot as plt
import math
import networkx as nx
from matplotlib.collections import LineCollection
from scipy.spatial import KDTree
import random
import os
import pickle
import sys
import igl

def beam_fea(G_new, adjacency_matrix_flow_, node_dict, boundary_nodes_array, draw):
    lines = get_merged_diff_ga(G_new, adjacency_matrix_flow_)
    lines = np.array(lines, dtype=int)
    needed_nodes = set(lines.flatten())
    old_to_new = {old: new for new, old in enumerate(sorted(needed_nodes))}

    total_force = calculate_global_force_fea(node_dict, lines, boundary_nodes_array, A=1.0e-6, rho=1250.0, g=9.8)
    total_force_dict = {row[0]: row[1:] for row in total_force}
    total_force_new = []
    for old, new in old_to_new.items():
        if old in total_force_dict:
            total_force_new.append(np.append(new, total_force_dict[old]))

    total_force_new = np.array(total_force_new)

    g_coord = np.array([node_dict[old] for old in sorted(needed_nodes)])
    g_num = np.vectorize(old_to_new.get)(lines)

    boundary_nodes_array_fea = [old_to_new[node] for node in boundary_nodes_array if node in old_to_new]

    total_deformation = beam_fea_calculate(g_num, g_coord, total_force_new, boundary_nodes_array_fea, draw)

    return total_deformation

def beam_fea_graph(lines, node_dict, boundary_nodes_array, draw):
    needed_nodes = set(lines.flatten())
    old_to_new = {old: new for new, old in enumerate(sorted(needed_nodes))}
    total_force = calculate_global_force_fea(node_dict, lines, boundary_nodes_array, A=3.14e-6, rho=1250.0, g=9.8)
    total_force_dict = {row[0]: row[1:] for row in total_force}
    total_force_new = []
    for old, new in old_to_new.items():
        if old in total_force_dict:
            total_force_new.append(np.append(new, total_force_dict[old]))

    total_force_new = np.array(total_force_new)

    g_coord = np.array([node_dict[old] for old in sorted(needed_nodes)])
    g_num = np.vectorize(old_to_new.get)(lines)

    boundary_nodes_array_fea = [old_to_new[node] for node in boundary_nodes_array if node in old_to_new]
    total_deformation = beam_fea_calculate(g_num, g_coord, total_force_new, boundary_nodes_array_fea, draw)

    return total_deformation

def calculate_global_stiffness_matrix(coordinates, lines, E=10**6, A=0.000004):
    K_global = np.zeros((3*len(coordinates), 3*len(coordinates)))

    for line in lines:
        point1 = np.array(coordinates[line[0]])
        point2 = np.array(coordinates[line[1]])

        L = np.linalg.norm(point2 - point1) * 0.001
        direction = (point2 - point1) * 0.001 / L
        K_local = E * A / L * np.outer(direction, direction)

        T = np.zeros((3, len(coordinates)*3))
        T[:, line[0]*3:line[0]*3+3] = np.eye(3)
        T[:, line[1]*3:line[1]*3+3] = -np.eye(3)
        K_global += T.T @ K_local @ T

    return K_global

def calculate_global_force_vector(coordinates, lines, boundary_nodes_array, A=0.000004, rho=1000.0, g=9.8):
    F_global = np.zeros(3*len(coordinates))

    for line in lines:
        point1 = np.array(coordinates[line[0]])
        point2 = np.array(coordinates[line[1]])

        L = np.linalg.norm(point2 - point1) * 0.001

        F_self = rho * A * L * g

        F_global[line[0]*3+2] -= F_self / 2
        F_global[line[1]*3+2] -= F_self / 2

    for node in boundary_nodes_array:
        F_global[node*3:node*3+3] = 0

    return F_global


def calculate_displacement(coordinates, lines, boundary_nodes_array, E=10**4, A=0.0016, rho=1250.0, g=9.81):
    K = calculate_global_stiffness_matrix(coordinates, lines, E, A)
    F = calculate_global_force_vector(coordinates, lines, boundary_nodes_array, A, rho, g)

    constrained_dof = []
    for node in boundary_nodes_array:
        constrained_dof.extend([node * 3, node * 3 + 1, node * 3 + 2])

    all_dof = list(range(3 * len(coordinates)))
    free_dof = list(set(all_dof) - set(constrained_dof))

    K_free = K[np.ix_(free_dof, free_dof)]
    F_free = F[free_dof]

    U_free = np.dot(np.linalg.pinv(K_free), F_free)

    U_global = np.zeros(3 * len(coordinates))

    for i in range(len(free_dof)):
        U_global[free_dof[i]] = U_free[i] * 1000

    for dof in constrained_dof:
        U_global[dof] = 0

    U_total = 0
    U_max = 0
    for i in range(len(coordinates)):
        displacement = np.linalg.norm(U_global[i * 3:i * 3 + 3])
        U_total += displacement
        if displacement > U_max:
            U_max = displacement
    return U_max, U_global

def get_merged_diff_ga(G_new, adjacency_matrix_flow_):
    edges_from_matrix = set()
    size = adjacency_matrix_flow_.shape[0]
    for i in range(size):
        for j in range(i+1, size):
            if adjacency_matrix_flow_[i][j] != 0:
                edges_from_matrix.add((i, j))

    edges_from_G_new = set()
    for edge in G_new.edges():
        edges_from_G_new.add(tuple(sorted(edge)))

    diff_edges = edges_from_G_new - edges_from_matrix

    return list(diff_edges)

def calculate_global_force_fea(coordinates, lines, boundary_nodes_array, A=0.0001, rho=1250.0, g=9.8):
    F_global = np.zeros((len(coordinates), 4))

    F_global[:, 0] = np.arange(0, len(coordinates))

    for line in lines:
        point1 = np.array(coordinates[line[0]])
        point2 = np.array(coordinates[line[1]])

        L = np.linalg.norm(point2 - point1) * 0.001

        F_self = rho * A * L * g

        F_global[line[0], 3] -= F_self / 2
        F_global[line[1], 3] -= F_self / 2

    mask = np.ones(len(coordinates), dtype=bool)
    if len(boundary_nodes_array) != 0:
        mask[np.array(boundary_nodes_array)] = False

    F_global = F_global[mask]

    return F_global

def compute_centroid(node_list, coordinates):
    sum_x = sum_y = sum_z = 0
    for node in node_list:
        pos = coordinates[node]
        sum_x += pos[0]
        sum_y += pos[1]
        sum_z += pos[2]
    n = len(node_list)
    return np.array([sum_x / n, sum_y / n, sum_z / n])

def load_data(filename, dataset_path):
    with open(os.path.join(dataset_path, filename + ".pkl"), 'rb') as f:
        new_adjacency_matrix, new_seq = pickle.load(f)
    return new_adjacency_matrix, new_seq

def get_max_index(dataset_path):
    max_index = -1
    for filename in os.listdir(dataset_path):
        if filename.endswith(".pkl"):
            index = int(filename.split('.')[0])
            max_index = max(max_index, index)
    return max_index

def save_data(new_adjacency_matrix, new_seq, dataset_path, env_name, index, i):
    similar_seq = judge_similarity(new_adjacency_matrix, dataset_path)
    if similar_seq:
        return

    new_index = get_max_index(dataset_path) + 1
    with open(os.path.join(dataset_path, str(new_index) + ".pkl"), 'wb') as f:
        pickle.dump((new_adjacency_matrix, new_seq), f)

def judge_similarity(input_adjacency_matrix, dataset_path):
    if not os.path.exists(dataset_path):
        os.makedirs(dataset_path)

    for filename in os.listdir(dataset_path):
        if filename.endswith(".pkl"):
            with open(os.path.join(dataset_path, filename), 'rb') as f:
                new_adjacency_matrix, new_seq = pickle.load(f)
                if new_adjacency_matrix.shape == input_adjacency_matrix.shape \
                        and np.allclose(input_adjacency_matrix, new_adjacency_matrix, rtol=0.02):

                    return new_seq
    return []

def get_new_boundary_nodes(boundary_nodes_array, node_mapping):
    boundary_nodes_array_new = []

    for node_new, node in node_mapping.items():
        if node in boundary_nodes_array:
            boundary_nodes_array_new.append(node_new)

    return boundary_nodes_array_new

def distance_coord(coord1, coord2):
    return math.sqrt((coord1[0] - coord2[0]) ** 2 + (coord1[1] - coord2[1]) ** 2)

def update_matrix(matrix, idx1, idx2):
    if matrix[idx1][idx2] < 0 or matrix[idx2][idx1] < 0:
        matrix[idx1][idx2] = -matrix[idx1][idx2]
        matrix[idx2][idx1] = -matrix[idx2][idx1]
    else:
        matrix[idx1][idx2] = 0
        matrix[idx2][idx1] = 0
    return matrix

def update_new_adjacency_matrix(new_adj_matrix, init_adj_matrix, new_state, state, idx1, idx2):
    if init_adj_matrix[state[idx1]][state[idx2]] < 0 and new_adj_matrix[new_state[idx1]][new_state[idx2]] == 0:
        new_adj_matrix[new_state[idx1]][new_state[idx2]] = -init_adj_matrix[state[idx1]][state[idx2]]
        new_adj_matrix[new_state[idx2]][new_state[idx1]] = -init_adj_matrix[state[idx2]][state[idx1]]
    else:
        new_adj_matrix[new_state[idx1]][new_state[idx2]] = init_adj_matrix[state[idx1]][state[idx2]]
        new_adj_matrix[new_state[idx2]][new_state[idx1]] = init_adj_matrix[state[idx2]][state[idx1]]
    return new_adj_matrix

def calculate_angle(node1, node2, node3):
    node1 = np.array(node1)
    node2 = np.array(node2)
    node3 = np.array(node3)
    vec1 = node2 - node1
    vec2 = node3 - node2
    if np.linalg.norm(vec1) == 0 or np.linalg.norm(vec2) == 0:
        cos_angle = 1
    else:
        cos_angle = np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
    if cos_angle > 1:
        angle_ = 0.0
    elif cos_angle < -1:
        angle_ = 3.14159
    else:
        angle_ = np.arccos(cos_angle)
    if np.isnan(angle_):
        angle_ = 3.14159
    return np.degrees(angle_)

def calculate_angle_vector(vector1, vector2):
    unit_vector1 = vector1 / np.linalg.norm(vector1)
    unit_vector2 = vector2 / np.linalg.norm(vector2)
    dot_product = np.dot(unit_vector1, unit_vector2)
    dot_product = np.clip(dot_product, -1.0, 1.0)
    angle = np.arccos(dot_product)
    if np.isnan(angle):
        angle = np.pi
    return angle


def calculate_angle_along_stress(node2, node3, vec1):
    node2 = np.array(node2)
    node3 = np.array(node3)
    vec2 = node3 - node2
    if np.linalg.norm(vec1) == 0 or np.linalg.norm(vec2) == 0:
        cos_angle = 1
    else:
        cos_angle = np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
    if cos_angle > 1:
        angle_ = 0.0
    elif cos_angle < -1:
        angle_ = 3.14159
    else:
        angle_ = np.arccos(cos_angle)
    if np.isnan(angle_):
        angle_ = 3.14159
    angle_ = np.degrees(angle_)
    if angle_ > 90:
        angle_ = 180 - angle_
    return angle_

def draw_graph(env_name, G, node_dict, next_state_, adjacency_matrix_flow_, i=100000, index=100000, output=True, draw=True, subgraph=False, mode='Euler', show='PLA3D', block=0):
    node_colors = []
    if mode == 'Euler':

        sorted_keys = sorted(node_dict.keys())
        node_colors = ["red" if n == next_state_[-1] else "lightblue" for n in sorted_keys]

    if mode == 'Tsp':
        heat_values = [G.nodes[node]['heat'] for node in G.nodes()]
        if np.max(heat_values) == np.min(heat_values):
            norm_heat_values = np.zeros(len(heat_values))
        else:
            norm_heat_values = (heat_values - np.min(heat_values)) / (np.max(heat_values) - np.min(heat_values))
        color_mapping = {0: 'lightblue', 1: 'cyan', 2: 'lightgreen', 3: 'yellow', 4: 'orange', 5: 'red'}
        node_colors = [color_mapping[int(value * 5)] for value in norm_heat_values]

    edge_colors = {}

    if show == 'PLA3D':

        for u, v in G.edges():
            if adjacency_matrix_flow_[u, v] == 0:
                edge_colors[(u, v)] = "black"
            elif adjacency_matrix_flow_[u, v] > 0:
                edge_colors[(u, v)] = "lightgray"
            else:
                edge_colors[(u, v)] = "none"

        fig = plt.figure(figsize=(10, 10), dpi=150)
        ax = fig.add_subplot(111, projection='3d')

        if subgraph:
            for node, coordinates in node_dict.items():
                ax.text(coordinates[0], coordinates[1], coordinates[2], str(node), fontsize=4)

        for edge in G.edges():
            xs, ys, zs = zip(*[(node_dict[node][0], node_dict[node][1], node_dict[node][2]) for node in edge])
            ax.plot(xs, ys, zs, color=edge_colors.get(edge, 'gray'), lw=1)

        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        ax.grid(False)
        ax.xaxis.pane.fill = False
        ax.yaxis.pane.fill = False
        ax.zaxis.pane.fill = False
        ax.xaxis.pane.set_edgecolor('w')
        ax.yaxis.pane.set_edgecolor('w')
        ax.zaxis.pane.set_edgecolor('w')
        ax.set_box_aspect([1, 1, 1])
        ax.axis('off')

    else:

        for u, v in G.edges():
            if adjacency_matrix_flow_[u, v] == 0:
                edge_colors[(u, v)] = "black"
            elif adjacency_matrix_flow_[u, v] > 0:
                edge_colors[(u, v)] = "blue"
            else:
                edge_colors[(u, v)] = "lightgray"

        if mode == 'Tsp':
            temp_dpi = 1200
        else:
            temp_dpi = 200
        fig, ax = plt.subplots(dpi=temp_dpi)
        x_coords = [coordinates[0] for coordinates in node_dict.values()]
        y_coords = [coordinates[1] for coordinates in node_dict.values()]
        colors = [node_colors[node] for node in node_dict.keys()]
        if mode == 'Tsp':
            ax.scatter(x_coords, y_coords, s=0.01, c='none')
        else:
            ax.scatter(x_coords, y_coords, s=2, c=colors)

        if subgraph:
            if mode == 'Tsp':
                for node, coordinates in node_dict.items():
                    ax.text(coordinates[0], coordinates[1], str(node), fontsize=4)
            else:
                for node, coordinates in node_dict.items():
                    ax.text(coordinates[0], coordinates[1], str(node), fontsize=4)

        lines = []
        colors = []
        for edge in G.edges():
            xs, ys = zip(*[(node_dict[node][0], node_dict[node][1]) for node in edge])
            lines.append(list(zip(xs, ys)))
            colors.append(edge_colors.get(edge, 'lightgray'))
        if mode == 'Tsp':
            lc = LineCollection(lines, colors=colors, linewidths=0.35)
        else:
            lc = LineCollection(lines, colors=colors, linewidths=1.4)
        ax.add_collection(lc)
        ax.axis('off')
        ax.set_aspect('equal')

    if output:
        if mode == 'Tsp':
            output_file = f"./figure/{env_name}_{block}_{index}_{i}.png"
        else:
            output_file = f"./figure/{env_name}_{index}_{i}.png"
        plt.savefig(output_file)
    if draw:
        plt.show(block=False)
    plt.close()

def transform_faces(faces, node_mapping):
    reversed_mapping = {v: k for k, v in node_mapping.items()}
    faces_new = []

    for face in faces:
        if all(node in reversed_mapping for node in face):
            new_face = [reversed_mapping[node] for node in face]
            faces_new.append(new_face)
        else:
            missing_nodes = [node for node in face if node not in reversed_mapping]

    faces_new_np = np.array(faces_new, dtype=int)
    return faces_new_np

def lscm_parameterization_libigl(G, faces, new_state_):
    max_node_id = max(G.nodes)
    vertices = np.zeros((max_node_id + 1, 3))
    for n in G.nodes:
        vertices[n] = G.nodes[n]['pos']

    faces = np.array(faces, dtype=np.int32)

    if new_state_[1] != new_state_[2]:
        b = np.array([new_state_[1], new_state_[2]], dtype=np.int32)
    else:
        target_node = new_state_[2]
        min_distance = float('inf')
        nearest_node = None

        for neighbor in G.neighbors(target_node):
            pos_target = np.array(G.nodes[target_node]['pos'])
            pos_neighbor = np.array(G.nodes[neighbor]['pos'])
            distance = np.linalg.norm(pos_target - pos_neighbor)

            if distance < min_distance:
                min_distance = distance
                nearest_node = neighbor

        if nearest_node is None:
            raise ValueError("No neighbors found or no valid nearest node found")

        b = np.array([nearest_node, target_node], dtype=np.int32)

    bc = np.array([[-1.0, 0.0], [0.0, 0.0]])

    success, u = igl.lscm(vertices, faces, b, bc)

    if not success:
        raise ValueError("LSCM parameterization failed")

    return u

def reorder_nodes(G, uv_coordinates):
    uv_coordinates = np.array(uv_coordinates)
    indices_sorted = np.lexsort((uv_coordinates[:, 0], uv_coordinates[:, 1]))  # 先x后y
    old_to_new_mapping = {old_index: new_index for new_index, old_index in enumerate(indices_sorted)}
    G_new = nx.relabel_nodes(G, old_to_new_mapping)

    return G_new, old_to_new_mapping

def combine_mappings(new_temp_to_new_mapping, old_to_new_temp_mapping_):
    final_mapping = {}
    for temp_index, new_index in new_temp_to_new_mapping.items():
        if temp_index in old_to_new_temp_mapping_:
            orig_index = old_to_new_temp_mapping_[temp_index]
            final_mapping[new_index] = orig_index
        else:
            print(f"Warning: No original mapping found for temp index {temp_index}")

    return final_mapping

def create_standard_graph(coordinates_, state_, subgraph):
    node1 = coordinates_[state_[-2]]
    node2 = coordinates_[state_[-1]]
    if node1 == node2:
        orig_vector = np.array([1, 0, 0])
    else:
        orig_vector = np.array(node2) - np.array(node1)

    angles = []
    for node in subgraph.nodes:
        if node == state_[-1]:
            continue
        node_coord = coordinates_[node]
        new_vector = np.array(node_coord) - np.array(node2)
        angle = calculate_angle_vector(orig_vector, new_vector)
        vector_magnitude = np.linalg.norm(new_vector)
        angles.append((node, angle, vector_magnitude))

    sorted_nodes = sorted(angles, key=lambda x: (x[1], x[2]))
    node_mapping_ = {0: state_[-1]}
    for index, (orig_node, angle, magnitude) in enumerate(sorted_nodes):
        node_mapping_[index + 1] = orig_node

    return node_mapping_

def align_graph(G, state):
    pos = nx.get_node_attributes(G, 'pos')

    if state[1] == state[2] or state[0] == state[1]:
        min_distance = float('inf')
        closest_node = None
        origin_pos = np.array(pos[state[2]])
        for node, node_pos in pos.items():
            if node != state[2]:
                distance = np.linalg.norm(np.array(node_pos) - origin_pos)
                if distance < min_distance:
                    min_distance = distance
                    closest_node = node
        index1 = closest_node
    else:
        index1 = state[1]

    origin = np.array(pos[state[2]])
    for node in pos:
        pos[node] = np.array(pos[node]) - origin

    target = np.array(pos[index1])
    angle = np.arctan2(target[1], target[0])
    rotation_matrix = np.array([
        [np.cos(angle), np.sin(angle), 0],
        [-np.sin(angle), np.cos(angle), 0],
        [0, 0, 1]
    ])

    for node in pos:
        pos[node] = rotation_matrix @ np.array(pos[node]).reshape(-1, 1)
        pos[node] = pos[node].flatten()

    if pos[index1][0] > 0:
        for node in pos:
            pos[node][0] *= -1

    result = np.array([pos[node] for node in sorted(G.nodes())])
    return result

def find_k_hop_neighbors(G, start_node, k):
    neighbors = {start_node}
    for _ in range(k):
        next_neighbors = set()
        for node in neighbors:
            next_neighbors.update(G.neighbors(node))
        neighbors.update(next_neighbors)
    return neighbors

def create_new_graph(G_orig, coordinates_, edges_, adjacency_matrix_, state_, init_adjacency_matrix_, radius_, state_dim, heat_radius, faces, env_name, index, mode='Tsp', material='PLA3D'):
    k_hop_neighbors = find_k_hop_neighbors(G_orig, state_[-1], radius_)
    subgraph = G_orig.subgraph(k_hop_neighbors)

    if material == 'PLA3D':
        old_to_new_temp_mapping_ = {index: orig_index for index, orig_index in enumerate(subgraph.nodes)}
        G_new_temp = nx.relabel_nodes(subgraph, {v: k for k, v in old_to_new_temp_mapping_.items()}, copy=True)

        faces_new = transform_faces(faces, old_to_new_temp_mapping_)
        new_state_temp = np.zeros_like(state_)
        for ii, old_index in enumerate(state_):
            new_index_ = next((k for k, v in old_to_new_temp_mapping_.items() if v == old_index), None)
            if new_index_ is not None:
                new_state_temp[ii] = new_index_

        uv_coordinates = lscm_parameterization_libigl(G_new_temp, faces_new, new_state_temp)
        G_new_, new_temp_to_new_mapping = reorder_nodes(G_new_temp, uv_coordinates)
        node_mapping_ = combine_mappings(new_temp_to_new_mapping, old_to_new_temp_mapping_)

    elif material == 'CCF':
        old_to_new_temp_mapping_ = {index: orig_index for index, orig_index in enumerate(subgraph.nodes)}
        G_new_temp = nx.relabel_nodes(subgraph, {v: k for k, v in old_to_new_temp_mapping_.items()}, copy=True)
        new_state_temp = np.zeros_like(state_)
        for ii, old_index in enumerate(state_):
            new_index_ = next((k for k, v in old_to_new_temp_mapping_.items() if v == old_index), None)
            if new_index_ is not None:
                new_state_temp[ii] = new_index_

        aligned_positions = align_graph(G_new_temp, new_state_temp)
        G_new_, new_temp_to_new_mapping = reorder_nodes(G_new_temp, aligned_positions)
        node_mapping_ = combine_mappings(new_temp_to_new_mapping, old_to_new_temp_mapping_)

    else:
        node_mapping_ = {index: orig_index for index, orig_index in enumerate(subgraph.nodes)}
        G_new_ = nx.relabel_nodes(subgraph, {v: k for k, v in node_mapping_.items()}, copy=True)

    for new_index, orig_index in node_mapping_.items():
        G_new_.nodes[new_index]['pos'] = coordinates_[orig_index]
        G_new_.nodes[new_index]['heat'] = G_orig.nodes[orig_index]['heat']
        G_new_.nodes[new_index]['p_heat'] = G_orig.nodes[orig_index]['p_heat']
        G_new_.nodes[new_index]['p_p_heat'] = G_orig.nodes[orig_index]['p_p_heat']

    node_dict_ = nx.get_node_attributes(G_new_, 'pos')
    num_nodes = len(node_dict_)
    coords_array = np.zeros((3, num_nodes))
    for index, coord in node_dict_.items():
        coords_array[:, index] = coord

    new_state_ = np.zeros_like(state_)
    for ii, old_index in enumerate(state_):
        new_index_ = next((k for k, v in node_mapping_.items() if v == old_index), None)
        if new_index_ is not None:
            new_state_[ii] = new_index_

    heat_array = np.zeros((1, num_nodes))
    p_heat_array = np.zeros((1, num_nodes))
    p_p_heat_array = np.zeros((1, num_nodes))
    for idx, (node, attrs) in enumerate(G_new_.nodes(data=True)):
        heat_array[0, idx] = attrs['heat']
        p_heat_array[0, idx] = attrs['p_heat']
        p_p_heat_array[0, idx] = attrs['p_p_heat']

    coords_array_repeated = np.vstack((coords_array, heat_array))
    if mode == 'Tsp':

        pre_coords_array = np.vstack((coords_array, p_heat_array))
        pre_pre_coords_array = np.vstack((coords_array, p_p_heat_array))
        coords_array = np.stack((pre_pre_coords_array, pre_coords_array, coords_array_repeated), axis=0)
    else:
        coords_array = np.repeat(coords_array_repeated[np.newaxis, :, :], state_dim, axis=0)

    new_adjacency_matrix_ = np.zeros((len(node_mapping_), len(node_mapping_)), dtype=float)

    for new_index_a, old_index_a in node_mapping_.items():
        for new_index_b, old_index_b in node_mapping_.items():
            new_adjacency_matrix_[new_index_a, new_index_b] = adjacency_matrix_[old_index_a, old_index_b]

    for i in range(-2, -state_dim - 1, -1):
        new_adjacency_matrix_ = update_new_adjacency_matrix(new_adjacency_matrix_, init_adjacency_matrix_, new_state_, state_, i, i + 1)
    new_state_adjacency_ = [new_adjacency_matrix_.copy()]
    for i in range(-state_dim, -1, 1):
        new_adjacency_matrix_ = update_matrix(new_adjacency_matrix_, new_state_[i], new_state_[i + 1])
        new_state_adjacency_.append(new_adjacency_matrix_.copy())

    new_state_adjacency_ = np.array(new_state_adjacency_)

    return G_new_, new_state_, new_adjacency_matrix_, new_state_adjacency_, node_mapping_, node_dict_, coords_array


def create_current_state(lst, state_dim, init_adjacency_matrix, mode='Euler', material='CCF'):
    arr = np.zeros(state_dim, dtype=int)
    for i in range(min(3, len(lst))):
        arr[-(i + 1)] = lst[-(i + 1)]

    if mode == 'Tsp' or material == 'PLA3D':
        if len(lst) > 1 and init_adjacency_matrix[lst[-1]][lst[-2]] == 0:
            for ii in range(len(arr)):
                arr[ii] = lst[-1]
    return arr


def create_adj_matrix(adjacency_matrix, beam_seq, mode='Euler'):
    for i in range(len(beam_seq) - 1):
        if adjacency_matrix[beam_seq[i]][beam_seq[i + 1]] < 0:
            adjacency_matrix[beam_seq[i]][beam_seq[i + 1]] = -adjacency_matrix[beam_seq[i]][beam_seq[i + 1]]
            adjacency_matrix[beam_seq[i + 1]][beam_seq[i]] = -adjacency_matrix[beam_seq[i + 1]][beam_seq[i]]
        else:
            adjacency_matrix[beam_seq[i]][beam_seq[i + 1]] = 0
            adjacency_matrix[beam_seq[i + 1]][beam_seq[i]] = 0

        if mode == 'Tsp' and len(beam_seq) > 1:
            if adjacency_matrix[beam_seq[i]][beam_seq[i - 1]] == 0:
                adjacency_matrix[beam_seq[i]][beam_seq[i - 1]] = 1000
                adjacency_matrix[beam_seq[i - 1]][beam_seq[i]] = 1000

    return adjacency_matrix


def o2n_seq_trans(old_seq, node_mapping):
    new_seq = old_seq.copy()
    for ii, old_index in enumerate(old_seq):
        new_index_ = next((k for k, v in node_mapping.items() if v == old_index), None)
        if new_index_ is not None:
            new_seq[ii] = new_index_
    return new_seq


def n2o_seq_trans(new_seq, node_mapping):
    old_seq = new_seq.copy()
    for ii, new_index in enumerate(new_seq):
        old_seq[ii] = node_mapping[new_index]
    return old_seq


def remove_duplicates_and_keep_max(input_list):
    input_list.sort(key=lambda x: x[1], reverse=True)

    result_dict = {}
    for item in input_list:
        key = tuple(item[0])
        if key not in result_dict:
            result_dict[key] = item

    result = list(result_dict.values())
    return result

def adjacency_matrix_to_graph(matrix):
    graph = nx.from_numpy_array(np.array(matrix), create_using=nx.Graph())
    return graph


def calculate_maximum_common_subgraph_similarity(matrix1, matrix2):
    graph1 = adjacency_matrix_to_graph(matrix1)
    graph2 = adjacency_matrix_to_graph(matrix2)

    gm = nx.algorithms.isomorphism.GraphMatcher(graph1, graph2)
    mcs = max(gm.subgraph_isomorphisms_iter(), key=lambda x: len(x))

    similarity = len(mcs) / min(len(graph1.nodes), len(graph2.nodes))
    return similarity


def find_most_similar_graph(target_matrix, other_matrices):
    max_similarity = -1
    most_similar_index = None

    for i, matrix in enumerate(other_matrices):
        similarity = calculate_maximum_common_subgraph_similarity(target_matrix, matrix)

        if similarity > max_similarity:
            max_similarity = similarity
            most_similar_index = i

    return most_similar_index, max_similarity

def dfs_cycles(G, source, target, visited, depth, max_depth, cycles):
    visited[source] = True

    if depth == max_depth:
        if target in G[source]:
            cycles.append(tuple(visited.keys()))
    else:
        for neighbor in G[source]:
            if neighbor not in visited:
                dfs_cycles(G, neighbor, target, visited.copy(), depth + 1, max_depth, cycles)


def find_cycles_of_length(G, length):
    cycles = []

    for node in G.nodes:
        dfs_cycles(G, node, node, {}, 0, length - 1, cycles)

    cycles = list(set(cycles))

    for i in range(len(cycles)):
        cycle = cycles[i]
        if len(cycle) == 3:
            A, B, C = cycle
            pos_A = np.array(G.nodes[A]['pos'])
            pos_B = np.array(G.nodes[B]['pos'])
            pos_C = np.array(G.nodes[C]['pos'])
            AB = pos_B - pos_A
            AC = pos_C - pos_A
            cross_product = np.cross(AB, AC)
            if cross_product[2] < 0:
                cycles[i] = (A, C, B)


    return cycles

def remove_duplicate_cycles(cycles):
    unique_cycle_sets = set()
    unique_cycles = []

    for cycle in cycles:
        cycle_set = frozenset(cycle)
        if cycle_set not in unique_cycle_sets:
            unique_cycle_sets.add(cycle_set)
            unique_cycles.append(cycle)

    return unique_cycles


def remove_nested_cycles(cycles_length_3, cycles_length_4):
    nested_cycles = set()

    for cycle3_a in cycles_length_3:
        for cycle3_b in cycles_length_3:
            if cycle3_a == cycle3_b:
                continue
            for cycle4 in cycles_length_4:
                if set(cycle3_a).issubset(cycle4) and set(cycle3_b).issubset(cycle4):
                    nested_cycles.add(cycle4)
                    break

    filtered_cycles_length_4 = [cycle4 for cycle4 in cycles_length_4 if cycle4 not in nested_cycles]

    return cycles_length_3 + filtered_cycles_length_4

def find_boundary_edges(G):
    cycles = find_faces(G)
    edge_count = {}
    for cycle in cycles:
        for i in range(len(cycle)):
            edge = frozenset((cycle[i], cycle[(i + 1) % len(cycle)]))
            if edge in edge_count:
                edge_count[edge] += 1
            else:
                edge_count[edge] = 1

    boundary_nodes = set()
    boundary_edges_list = []
    for edge, count in edge_count.items():
        if count == 1:
            boundary_nodes.update(edge)
            boundary_edges_list.append(list(edge))
    boundary_edges = np.array(boundary_edges_list)
    boundary_nodes = np.array(list(boundary_nodes))

    return boundary_edges, boundary_nodes

def find_faces(G):
    cycles_length_3 = find_cycles_of_length(G, 3)
    cycles_length_4 = find_cycles_of_length(G, 4)
    unique_cycles_length_3 = remove_duplicate_cycles(cycles_length_3)
    unique_cycles_length_4 = remove_duplicate_cycles(cycles_length_4)
    cycles = remove_nested_cycles(unique_cycles_length_3, unique_cycles_length_4)
    return cycles

def find_boundary_nodes(G):
    cycles = find_faces(G)
    edge_count = {}
    for cycle in cycles:
        for i in range(len(cycle)):
            edge = frozenset((cycle[i], cycle[(i + 1) % len(cycle)]))
            if edge in edge_count:
                edge_count[edge] += 1
            else:
                edge_count[edge] = 1

    boundary_nodes = set()
    for edge, count in edge_count.items():
        if count == 1:
            boundary_nodes.update(edge)

    single_degree_nodes = {node for node in G.nodes() if G.degree(node) == 1}
    boundary_nodes.update(single_degree_nodes)

    return boundary_nodes

def find_fix_nodes(G, G_new, adjacency_matrix, new_adjacency_matrix, node_mapping):
    boundary_nodes_new = list(G_new.nodes())
    new_counts = count_zeros_in_adjacent_nodes(G_new, new_adjacency_matrix, boundary_nodes_new)
    old_counts = count_zeros_in_adjacent_nodes(G, adjacency_matrix, [node_mapping[node] for node in boundary_nodes_new])
    difference_nodes = [node for node in boundary_nodes_new if new_counts[node] != old_counts[node_mapping[node]]]
    return difference_nodes


def count_zeros_in_adjacent_nodes(G_new, matrix, rows):
    counts = {}
    for row in rows:
        adjacent_nodes = list(G_new.adj[row])
        count = 0
        for node in adjacent_nodes:
            if matrix[row, node] == 0:
                count += 1
        counts[row] = count
    return counts

def find_smallest_z_set(G):
    z_values = {}
    for node in G.nodes():
        z_values[node] = G.nodes[node]['pos'][-1]
    nodes_with_z_smaller_than_one = {node for node, z_value in z_values.items() if z_value < 1}

    return nodes_with_z_smaller_than_one

def load_checkpoint():
    checkpoint = torch.load('checkpoint/model2_0.pt')
    adj_matrix = checkpoint['adj_matrix']
    print(adj_matrix)
    return adj_matrix


def calculate_distance(pos1, pos2):
    pos1_np, pos2_np = np.array(pos1), np.array(pos2)
    return np.linalg.norm(pos1_np - pos2_np)


def create_heat_field(beam_seq, G_orig, heat_radius, mode='Tsp'):
    nodes = G_orig.nodes
    for node in nodes:
        nodes[node]['heat'] = 0
        nodes[node]['p_heat'] = 0
        nodes[node]['p_p_heat'] = 0

    if mode == 'Tsp':
        ini_temp = 0
        node_positions = [nodes[node]['pos'] for node in nodes]
        kdtree = KDTree(node_positions)
        for i, center_node in enumerate(beam_seq[-8:], start=1):
            for node in nodes:
                nodes[node]['heat'] = max(nodes[node]['heat'] * 0.52, 0)

            center_node_pos = nodes[center_node]['pos']
            indices = kdtree.query_ball_point(center_node_pos, heat_radius + 0.01)

            for index in indices:
                node = list(nodes)[index]
                distance_ = calculate_distance(center_node_pos, nodes[node]['pos'])
                heat_value = heat_radius * (1 - (distance_ / heat_radius) ** 0.8) + ini_temp
                nodes[node]['heat'] = min(nodes[node]['heat'] + heat_value, heat_radius)

            if len(beam_seq) < 3:
                for node in nodes:
                    nodes[node]['p_heat'] = 0
                    nodes[node]['p_p_heat'] = 0

            if i == len(beam_seq[-8:]) - 2:
                for node in nodes:
                    nodes[node]['p_p_heat'] = nodes[node]['heat']
            elif i == len(beam_seq[-8:]) - 1:
                for node in nodes:
                    nodes[node]['p_heat'] = nodes[node]['heat']

    return G_orig

def update_heat_field(next_state, coords_array, heat_radius, mode='Tsp'):
    next_coords_array = np.copy(coords_array)
    if mode == 'Tsp':
        new_heat_values = np.maximum(coords_array[3] * 0.52, 0)
        next_coords_array[3] = new_heat_values
        center_node_idx = next_state
        center_node_coords = next_coords_array[:3, center_node_idx]
        tree = KDTree(next_coords_array[:3].T)
        in_radius_indices = tree.query_ball_point(center_node_coords, heat_radius + 0.01)
        distances = np.sqrt(np.sum((next_coords_array[:3, in_radius_indices] - center_node_coords.reshape(3, 1)) ** 2, axis=0))
        new_heat_values = np.maximum(heat_radius * (1 - (distances / heat_radius) ** 0.8), 0)
        next_coords_array[3, in_radius_indices] += new_heat_values
        next_coords_array[3, in_radius_indices] = np.minimum(next_coords_array[3, in_radius_indices], heat_radius)

    return next_coords_array

def create_new_graph_with_grid_nodes_block(G, divisions=3, interval=0.1, max_nodes=4000):
    pos = nx.get_node_attributes(G, 'pos')

    min_x, max_x = float('inf'), float('-inf')
    min_y, max_y = float('inf'), float('-inf')

    for _, position in pos.items():
        x, y, _ = position
        min_x = min(min_x, x)
        max_x = max(max_x, x)
        min_y = min(min_y, y)
        max_y = max(max_y, y)

    width = max_x - min_x
    height = max_y - min_y

    grid_width = width / divisions
    grid_height = height / divisions

    G_new_list = [nx.Graph() for _ in range(divisions * divisions)]

    new_node_id = 0
    for x in np.arange(min_x, max_x + interval, interval):
        for y in np.arange(min_y, max_y + interval, interval):
            if x > max_x or y > max_y:
                continue
            grid_x = int((x - min_x) // grid_width)
            grid_y = int((y - min_y) // grid_height)

            grid_x = min(divisions - 1, grid_x)
            grid_y = min(divisions - 1, grid_y)
            grid_index = grid_y * divisions + grid_x

            G_new_list[grid_index].add_node(new_node_id, pos=(x, y, 0))
            new_node_id += 1

    G_new_list_final = []
    for G_new in G_new_list:
        num_nodes = G_new.number_of_nodes()
        if num_nodes > max_nodes:
            G_sub_list = create_new_graph_with_grid_nodes_block(G_new, divisions=2, interval=interval, max_nodes=4000)
            G_new_list_final.extend(G_sub_list)
        else:
            G_new_list_final.append(G_new)

    return G_new_list_final

def traverse_matrix_euler(matrix, row_index):
    if not any(matrix[row_index, :]):
        return
    columns = np.where(matrix[row_index, :] != 0)[0]
    matrix[row_index, :] = 0
    matrix[:, row_index] = 0

    for col_index in columns:
        traverse_matrix_euler(matrix, col_index)

def recursive_search(G, node, node_dict):
    neighbors = [n for n in G.neighbors(node) if node_dict[n] == 0]
    for neighbor in neighbors:
        node_dict[neighbor] = 1
        recursive_search(G, neighbor, node_dict)

def angle_between(v1, v2):
    dot_product = np.dot(v1, v2)
    norms = np.linalg.norm(v1) * np.linalg.norm(v2)
    cos_theta = dot_product / norms
    cos_theta = np.clip(cos_theta, -1.0, 1.0)

    angle = np.arccos(cos_theta)
    return np.degrees(angle)


def is_within_sector(v_base, v1, v2):
    angle_v1 = angle_between(v_base, v1)
    angle_v2 = angle_between(v_base, v2)

    if angle_v2 >= angle_v1 - 0.01:
        return False

    cross_v1 = np.cross(v_base, v1)
    cross_v2 = np.cross(v_base, v2)
    if angle_between(cross_v1, cross_v2) < 89:
        return True
    else:
        return False

def anti_self_locking_subgraph(G, G_new, adjacency_matrix_, node_mapping_, state_, beam_seq, new_adjacency_matrix_,
                               new_state_, coordinates, heat_radius, init_adjacency_matrix_, rays, train_mode, mode='Tsp', material='CCF'):

    optional_action = np.array([])
    lifting = False
    lines = get_merged_diff(init_adjacency_matrix_, adjacency_matrix_)
    if mode == 'Tsp':
        optional_action = np.where(np.array(adjacency_matrix_[state_[-1]]) != 0)[0]
        row_col_array = np.concatenate(np.argwhere(adjacency_matrix_ > 0).T)
        optional_action = np.setdiff1d(optional_action, row_col_array)

        if optional_action.size == 0:
            lifting = True
            candidate_node = np.array(list(G.nodes()))
            unique_nodes = np.setdiff1d(candidate_node, row_col_array)
            distances = [nx.shortest_path_length(G, state_[-1], node) for node in unique_nodes]
            sorted_nodes = [node for _, node in sorted(zip(distances, unique_nodes), reverse=True)]

            for index in sorted_nodes:
                connectivity = True
                optional_action_jump = np.where(np.array(adjacency_matrix_[index]) != 0)[0]
                optional_action_jump = np.setdiff1d(optional_action_jump, row_col_array)
                if len(optional_action_jump) == 2:
                    angle = calculate_angle(coordinates[optional_action_jump[0]], coordinates[index], coordinates[optional_action_jump[1]])
                    if angle < 1:
                        connectivity = False

                if connectivity:
                    optional_action = np.append(optional_action, index)
                    break

        new_optional_action_ = np.zeros_like(optional_action)
        if lifting:
            new_optional_action_ = optional_action
        else:
            for ii, old_index in enumerate(optional_action):
                new_index_ = next((k for k, v in node_mapping_.items() if v == old_index), None)
                if new_index_ is not None:
                    new_optional_action_[ii] = new_index_

        optional_action = new_optional_action_

    if mode == 'Euler':
        optional_action = np.where(np.array(adjacency_matrix_[state_[-1]]) != 0)[0]

        if material == 'CCF':
            for row_index in optional_action:
                adjacency_matrix_next = adjacency_matrix_.copy()

                if adjacency_matrix_next[state_[-1]][row_index] < 0 or adjacency_matrix_next[row_index][state_[-1]] < 0:
                    adjacency_matrix_next[state_[-1]][row_index] = -adjacency_matrix_next[state_[-1]][row_index]
                    adjacency_matrix_next[row_index][state_[-1]] = -adjacency_matrix_next[row_index][state_[-1]]
                else:
                    adjacency_matrix_next[state_[-1]][row_index] = 0.0
                    adjacency_matrix_next[row_index][state_[-1]] = 0.0

                connectivity = True
                if not connectivity:
                    optional_action = optional_action[optional_action != row_index]

        if material == 'PLA3D':
            collision_check = False
            optional_action_temp = optional_action.copy()
            for action in optional_action:
                temp_lines = lines.copy()
                if coordinates[state_[-1]][2] > coordinates[action][2]:
                    temp_lines.append((state_[-1], action))

                if collision_check:
                    collision, _, _ = collision_check_simulation(coordinates, state_[-1], action, rays, temp_lines, norm_output=False)
                else:
                    collision = False

                if collision:
                    optional_action_temp = optional_action_temp[optional_action_temp != action]

            optional_action = optional_action_temp
            if optional_action.size == 0:
                lifting = True
                row_col_array = np.concatenate(np.argwhere(adjacency_matrix_ > 0).T)
                unique_nodes = np.unique(row_col_array)
                distances = [nx.shortest_path_length(G, state_[-1], node) for node in unique_nodes]
                sorted_nodes = [node for _, node in sorted(zip(distances, unique_nodes))]
                for row_index in sorted_nodes:
                    append_node = False
                    optional_action_jump = np.where(np.array(adjacency_matrix_[row_index]) != 0)[0]
                    for index in optional_action_jump:

                        temp_lines = lines.copy()
                        if coordinates[row_index][2] > coordinates[index][2]:
                            temp_lines.append((row_index, index))

                        if collision_check:
                            collision, _, _ = collision_check_simulation(coordinates, row_index, index, rays, temp_lines, norm_output=False)
                        else:
                            collision = False

                        if row_index in beam_seq and not collision:
                            append_node = True
                            break

                    if append_node:
                        optional_action = np.append(optional_action, row_index)
                        break

        new_optional_action_ = np.zeros_like(optional_action)
        if lifting:
            new_optional_action_ = optional_action
        else:
            for ii, old_index in enumerate(optional_action):
                new_index_ = next((k for k, v in node_mapping_.items() if v == old_index), None)
                if new_index_ is not None:
                    new_optional_action_[ii] = new_index_

        optional_action = new_optional_action_

    return optional_action, lifting, lines


def aabb_tree(nodes, lines, coordinates):
    x_min = min(nodes[0][0], nodes[-1][0]) - 180
    x_max = max(nodes[0][0], nodes[-1][0]) + 180

    y_min = min(nodes[0][1], nodes[-1][1]) - 180
    y_max = max(nodes[0][1], nodes[-1][1]) + 180

    z_min = min(nodes[0][2], nodes[-1][2]) - 30
    z_max = max(nodes[0][2], nodes[-1][2]) + 180

    filtered_lines = []

    for line in lines:
        x1, y1, z1 = coordinates[line[0]]
        x2, y2, z2 = coordinates[line[1]]

        if (x_min <= x1 <= x_max and y_min <= y1 <= y_max and z_min <= z1 <= z_max and
                x_min <= x2 <= x_max and y_min <= y2 <= y_max and z_min <= z2 <= z_max):
            filtered_lines.append(line)

    return filtered_lines

def collision_check_simulation(coordinates, current, action, rays, lines, interval=5, norm_output=True, specified_vector=np.array([0., 0., 1.])):
    current_coords = np.array(coordinates[current])
    action_coords = np.array(coordinates[action])
    distance = np.linalg.norm(current_coords - action_coords)
    if distance <= interval:
        nodes = [current_coords, action_coords]
    else:
        num_of_new_nodes = int(distance / interval)
        new_nodes = [current_coords + i * interval * (action_coords - current_coords) / distance for i in range(1, num_of_new_nodes)]
        nodes = [current_coords] + new_nodes + [action_coords]

    lines = aabb_tree(nodes, lines, coordinates)

    radius = 32
    height_cylinder = 150
    height_cone = 41

    node_rays = [[] for _ in nodes]
    node_index = 0
    for node in nodes:
        for ray_dir in rays:
            cone_center = node - height_cone * ray_dir
            cylinder_center = node - (height_cylinder + height_cone) * ray_dir
            collision = False
            for line in lines:
                line_nodes = sample_line_nodes(line, coordinates, step=5)
                for P in line_nodes:
                    if if_in_cone(P, cone_center, ray_dir, radius, height_cone):
                        collision = True
                        break
                    if if_in_cylinder(P, cylinder_center, ray_dir, radius, height_cylinder):
                        collision = True
                        break

                if collision:
                    break

            if not collision:
                node_rays[node_index].append(ray_dir)

        node_index = node_index + 1

    all_rays = [ray for rays in node_rays for ray in rays]
    duplicate_rays = []
    threshold = 1e-6

    for i, ray_i in enumerate(all_rays):
        if any(np.linalg.norm(dup - ray_i) < threshold for dup in duplicate_rays):
            continue

        count = sum(1 for ray_j in all_rays if np.linalg.norm(ray_i - ray_j) < threshold)
        if count == len(node_rays):
            duplicate_rays.append(ray_i)

    final_collision = True
    normal = None
    min_angle = 0
    if len(duplicate_rays) != 0:
        final_collision = False
        if norm_output:
            neg_duplicate_rays = [-1 * ray for ray in duplicate_rays]
            angles = [np.arccos(np.clip(np.dot(ray, specified_vector) / (np.linalg.norm(ray) * np.linalg.norm(specified_vector)), -1, 1)) for ray in neg_duplicate_rays]

            min_angle_index = np.argmin(angles)
            normal = neg_duplicate_rays[min_angle_index]
            normal = normal.tolist()
            min_angle = angles[min_angle_index]

    return final_collision, normal, min_angle

def generate_rays(degrees):
    rays = []
    init_vector = np.array([1, 0, 0])

    for elev in range(-35, -65, -13):
        rot_mat_y = np.array([
            [math.cos(math.radians(elev)), 0, -math.sin(math.radians(elev))],
            [0, 1, 0],
            [math.sin(math.radians(elev)), 0, math.cos(math.radians(elev))]
        ])

        curr_vector = np.matmul(rot_mat_y, init_vector)

        for azim in range(0, 360, 30):
            rot_mat_z = np.array([
                [math.cos(math.radians(azim)), -math.sin(math.radians(azim)), 0],
                [math.sin(math.radians(azim)), math.cos(math.radians(azim)), 0],
                [0, 0, 1]
            ])
            ray_dir = np.matmul(rot_mat_z, curr_vector)
            rays.append(ray_dir)

    rays.append(np.array([0, 0, -1]))
    return rays

def sample_line_nodes(line, coordinates, step=2):
    start, end = line
    start_coords = np.array(coordinates[start])
    end_coords = np.array(coordinates[end])
    length = np.linalg.norm(end_coords - start_coords)
    nodes_count = int(length / step)
    new_nodes_coords = [start_coords + (end_coords - start_coords) * i / nodes_count for i in range(1, nodes_count)]
    nodes = [start_coords] + new_nodes_coords + [end_coords]
    return nodes


def get_merged_diff(init_adjacency_matrix_, adjacency_matrix_):
    diff = np.where(init_adjacency_matrix_ != adjacency_matrix_)
    lines = list(zip(diff[0], diff[1]))
    merged_lines = set()
    for line in lines:
        line = tuple(sorted(line))
        merged_lines.add(line)

    return list(merged_lines)

def collision_checking(coordinates, adjacency_matrix_, vector_z, current, action, optional_action, beam_seq, rays):
    collision = False
    optional_action_lo = optional_action[optional_action != action]
    for other_action in optional_action_lo:
        vector1 = np.array(coordinates[action]) - np.array(coordinates[current])
        vector2 = np.array(coordinates[other_action]) - np.array(coordinates[current])
        if is_within_sector(vector_z, vector1, vector2):
            collision = True


    if len(beam_seq) > 1:
        optional_action_right = np.where(np.array(adjacency_matrix_[action]) != 0)[0]
        optional_action_ro = optional_action_right[optional_action_right != current]
        for other_action in optional_action_ro:
            vector1 = np.array(coordinates[current]) - np.array(coordinates[action])
            vector2 = np.array(coordinates[other_action]) - np.array(coordinates[action])
            if is_within_sector(vector_z, vector1, vector2):
                collision = True

    return collision


def if_in_cylinder(P, cylinder_center, ray_dir, radius, height):
    ray_dir = ray_dir / np.linalg.norm(ray_dir)
    P_prime = P - cylinder_center
    distance_along_ray = np.dot(P_prime, ray_dir)
    if distance_along_ray < 0 or distance_along_ray > height - 0.01:
        return False

    P_proj = distance_along_ray * ray_dir
    distance_to_ray = np.linalg.norm(P_proj - P_prime)

    return distance_to_ray < radius - 0.01


def if_in_cone(P, node, ray_dir, radius, height):
    ray_dir = ray_dir / np.linalg.norm(ray_dir)
    P_prime = P - node
    distance_along_ray = np.dot(P_prime, ray_dir)

    if distance_along_ray < 0 or distance_along_ray > height - 0.01:
        return False

    P_proj = distance_along_ray * ray_dir
    distance_to_ray = np.linalg.norm(P_proj - P_prime)

    cone_radius_at_P = (height - distance_along_ray) * radius / height
    return distance_to_ray < cone_radius_at_P - 0.01

def choose_start_nodes(G, boundary_nodes_array, material):
    if material == 'PLA3D':
        selected_node = np.random.choice(boundary_nodes_array)
        neighbors = list(G.neighbors(selected_node))
        boundary_neighbors = [node for node in neighbors if node in boundary_nodes_array]
        selected_boundary_neighbor = np.random.choice(boundary_neighbors)
        start_nodes = [selected_node, selected_boundary_neighbor]
    elif material == 'CCF':
        available_nodes = list(set(G.nodes) - set(boundary_nodes_array))
        random_node = random.choice(available_nodes)
        start_nodes = [random_node]
    else:
        start_nodes = [0]

    return start_nodes

def normalize_vector(v):
    norm = np.linalg.norm(v)
    if norm == 0:
        return v
    return v / norm

def calculate_mid_vector(v1, v2):
    v1_norm = normalize_vector(v1)
    v2_norm = normalize_vector(v2)
    return normalize_vector(v1_norm + v2_norm)

def calculate_offset_point(point, mid_vector, offset_distance):
    return point + mid_vector * offset_distance

def generate_beam_sequence(env_name, existing_coordinates):
    parts = env_name.split('-')
    common_name = '-'.join(parts[:-1]) + '_wireframe'
    common_name_2 = parts[-1]

    best_path_file = f'data/{parts[0]}/{common_name}.txt'
    node_coordinates_file = f'data/{parts[0]}/Node{common_name}.txt'
    second_node_coordinates_file = f'data/{parts[0]}/Node{parts[0]}-{common_name_2}.txt'

    def extract_best_path(file_path):
        with open(file_path, 'r') as file:
            for line in file:
                if line.startswith("Best path:"):
                    best_path_str = line.split(":")[1].strip()
                    best_path_ = eval(best_path_str)
                    return best_path_
        return None

    def extract_coordinates(file_path):
        coordinates = {}
        with open(file_path, 'r') as file:
            for i, line in enumerate(file):
                coords = [float(x) for x in line.split()]
                coordinates[i] = coords
        return coordinates

    def compare_coordinates(best_path_, node_coordinates_, existing_coordinates_):
        beam_seq_ = []
        close_nodes = []

        for index in best_path_:
            if index in node_coordinates_:
                x, y, z = node_coordinates_[index]
                for key, value in existing_coordinates_.items():
                    if all(abs(a - b) <= 0.01 for a, b in zip(value, [x, y, z])):
                        close_nodes.append(key)

        if close_nodes:
            beam_seq_.append([close_nodes, 0.0])

        return beam_seq_

    def add_min_z_coordinate_node(beam_seq_, second_node_coordinates_, existing_coordinates_):
        if beam_seq_:
            min_z_node = min(second_node_coordinates_, key=lambda k: second_node_coordinates_[k][2])
            min_z_coords = second_node_coordinates_[min_z_node]

            corresponding_node = next((key for key, value in existing_coordinates_.items() if all(abs(a - b) <= 0.01 for a, b in zip(value, min_z_coords))), None)
            if corresponding_node is not None:
                beam_seq_[0][0].append(corresponding_node)

    best_path = extract_best_path(best_path_file)
    node_coordinates = extract_coordinates(node_coordinates_file)
    second_node_coordinates = extract_coordinates(second_node_coordinates_file)
    beam_seq = compare_coordinates(best_path, node_coordinates, existing_coordinates)
    add_min_z_coordinate_node(beam_seq, second_node_coordinates, existing_coordinates)

    return beam_seq


def calculate_total_max_deformation(node_file, path_file):
    def read_coordinates(filename):
        coordinates_ = {}
        with open(filename, 'r') as file:
            for index, line in enumerate(file):
                x, y, z = map(float, line.split())
                coordinates_[index] = [x, y, z]
        return coordinates_

    def read_path(filename):
        path_ = []
        with open(filename, 'r') as file:
            for line in file:
                if line.startswith("Best path:"):
                    path_ = line.split(':')[1].strip().strip('[]').split(', ')
                    path_ = list(map(int, path_))
                    break
        return [(path_[i], path_[i + 1]) for i in range(len(path_) - 1)]

    coordinates = read_coordinates(node_file)
    path = read_path(path_file)

    boundary_nodes_array = np.array([index for index, coord in coordinates.items() if coord[2] < 1])
    lines = []
    max_deformation = 0
    lines_max_deformation = []
    index = 0
    max_index = 0
    for edge in path:
        if edge not in lines:
            lines.append(edge)
            current_lines = np.array(lines)
            current_deformation = beam_fea_graph(current_lines, coordinates, boundary_nodes_array, draw=False)

            if current_deformation > max_deformation:
                max_deformation = current_deformation
                lines_max_deformation = list(lines)
                max_index = index

        index += 1

    lines_max_deformation_array = np.array(lines_max_deformation)
    beam_fea_graph(lines_max_deformation_array, coordinates, boundary_nodes_array, draw=False)

    return max_deformation, max_index


def save_average_to_file(return_list, env_name, savept):

    average = round(sum(return_list[:100]) / 100, 3)
    file_suffix = "P" if savept else "I"
    file_name = f"checkpoint/100_average_{env_name}_{file_suffix}.txt"

    with open(file_name, "a") as file:
        file.write(f"{average}\n")

def update_progress(beam_seq, total_nodes, bar_length=50):
    unique_nodes = len(set(node for path in beam_seq for node in path[0]))

    progress = unique_nodes / total_nodes
    arrow = int(round(progress * bar_length - 1)) if unique_nodes != total_nodes else bar_length
    spaces = bar_length - arrow
    progress_bar = 'Progress: [{0}{1}] {2}%'.format('>' * arrow, ' ' * spaces, round(progress * 100, 2))

    sys.stdout.write("\r" + progress_bar)
    sys.stdout.flush()


def update_progress_post(ii, total, bar_length=50):
    """Update the progress bar based on the iteration index and total length."""
    progress = ii / total

    arrow = int(round(progress * bar_length - 1)) if ii < total else bar_length
    spaces = bar_length - arrow

    progress_bar = 'Progress: [{0}{1}] {2}%'.format('>' * arrow, ' ' * spaces, round(progress * 100, 2))

    sys.stdout.write("\r" + progress_bar)
    sys.stdout.flush()


def final_progress(bar_length=50):
    """Function to display the progress bar at 100% completion and then move to a new line."""
    progress_bar = 'Progress: [{0}] {1}%'.format('>' * bar_length, 100.0)
    print("\r" + progress_bar)

def find_most_similar_checkpoint(current_adj_matrix, checkpoint_dir="checkpoint"):
    max_similarity = -1
    most_similar_checkpoint_path = None
    normalized_matrix = normalize_matrix(current_adj_matrix)

    for filename in os.listdir(checkpoint_dir):
        if filename.endswith(".pt"):
            filepath = os.path.join(checkpoint_dir, filename)
            checkpoint = torch.load(filepath, map_location=torch.device('cpu'))
            adj_matrix = checkpoint['adj_matrix']

            if adj_matrix.shape != normalized_matrix.shape:
                larger_shape = tuple(max(s1, s2) for s1, s2 in zip(adj_matrix.shape, normalized_matrix.shape))
                adj_matrix = adjust_and_pad_matrix(adj_matrix, larger_shape)
                normalized_matrix = adjust_and_pad_matrix(normalized_matrix, larger_shape)

            similarity = calculate_similarity(normalized_matrix, adj_matrix)

            if similarity > max_similarity:
                max_similarity = similarity
                most_similar_checkpoint_path = filepath

    return most_similar_checkpoint_path, max_similarity

def normalize_matrix(matrix):
    max_val = matrix.max()
    if max_val > 0:
        normalized_matrix = matrix / max_val
    else:
        normalized_matrix = matrix
    return normalized_matrix

def calculate_similarity(matrix_a, matrix_b):
    if not isinstance(matrix_a, torch.Tensor):
        matrix_a = torch.tensor(matrix_a, dtype=torch.float32)
    if not isinstance(matrix_b, torch.Tensor):
        matrix_b = torch.tensor(matrix_b, dtype=torch.float32)

    diff = torch.norm(matrix_a - matrix_b, p='fro') * 0.2
    similarity = 1 / (1 + diff)
    return similarity.item()

def adjust_and_pad_matrix(matrix, target_shape):
    padding = [(0, max(0, t - c)) for t, c in zip(target_shape, matrix.shape)]

    if isinstance(matrix, np.ndarray):
        return np.pad(matrix, padding, mode='constant', constant_values=0)
    elif isinstance(matrix, torch.Tensor):
        pad = (0, padding[2][1], 0, padding[1][1], 0, padding[0][1])
        return torch.nn.functional.pad(matrix, pad, "constant", 0)
    else:
        raise TypeError("Unsupported type. The input must be either a numpy array or a torch tensor.")

def extract_faces_from_obj(env_name):
    file_name = f"data/{env_name}.obj"
    faces = []
    try:
        with open(file_name, 'r') as file:
            for line in file:
                if line.startswith('f '):
                    face_indices = line.strip().split()[1:]
                    face_indices = [int(index) - 1 for index in face_indices]
                    faces.append(face_indices)

        faces_np = np.array(faces, dtype=int)
        return faces_np
    except FileNotFoundError:
        print(f"file {file_name} not found。")
        return None
    except Exception as e:
        print(f"read eroor：{e}")
        return None

def pad_matrix(matrix, target_dim):
    current_depth, current_height, current_width = matrix.shape
    padding_height = max(0, target_dim - current_height)
    padding_width = max(0, target_dim - current_width)
    new_height = current_height + padding_height
    new_width = current_width + padding_width
    new_shape = (current_depth, new_height, new_width)
    padded_matrix = np.zeros(new_shape, dtype=matrix.dtype)
    padded_matrix[:, :current_height, :current_width] = matrix
    return padded_matrix

def compute_edges_length(V, Edge):
    diff = V[Edge[:, 1], :] - V[Edge[:, 0], :]
    return np.sqrt(np.sum(diff ** 2, axis=1))

def beam_fea_calculate(g_num, g_coord, total_force_new, boundary_nodes_array_fea, draw, index=0):
    E = 2636
    rho = 1250
    A = 3.14e-6
    G = 1419
    r = 0.001

    G = G * 2.22
    E = E + 700

    V = g_coord * 0.001
    Edge = g_num

    le = compute_edges_length(V, Edge)

    F = np.zeros(6 * V.shape[0])
    for force in total_force_new:
        F[int(force[0]) * 6 + 2] = force[-1]
    F = F * 0.001

    Iz = np.pi * r ** 4 / 4
    Iy = np.pi * r ** 4 / 4
    J = np.pi * r ** 4 / 2

    K = np.zeros((6 * V.shape[0], 6 * V.shape[0]))
    for i in range(Edge.shape[0]):
        v1 = Edge[i, 0]
        v2 = Edge[i, 1]
        DOF = np.hstack((np.arange(6 * v1, 6 * (v1 + 1)), np.arange(6 * v2, 6 * (v2 + 1))))
        kk = np.zeros((12, 12))

        kk[0, 0] = E * A / le[i]
        kk[6, 0] = -kk[0, 0]
        kk[0, 6] = -E * A / le[i]
        kk[6, 6] = -kk[0, 6]
        kk[1, 1] = 12 * E * Iz / le[i] ** 3
        kk[7, 1] = -kk[1, 1]
        kk[1, 5] = 6 * E * Iz / le[i] ** 2
        kk[7, 5] = -kk[1, 5]
        kk[1, 7] = -12 * E * Iz / le[i] ** 3
        kk[7, 7] = -kk[1, 7]
        kk[1, 11] = 6 * E * Iz / le[i] ** 2
        kk[7, 11] = -kk[1, 11]
        kk[2, 2] = 12 * E * Iy / le[i] ** 3
        kk[8, 2] = -kk[2, 2]
        kk[2, 4] = -6 * E * Iy / le[i] ** 2
        kk[8, 4] = -kk[2, 4]
        kk[2, 8] = -12 * E * Iy / le[i] ** 3
        kk[8, 8] = -kk[2, 8]
        kk[2, 10] = -6 * E * Iy / le[i] ** 2
        kk[8, 10] = -kk[2, 10]
        kk[3, 3] = G * J / le[i]
        kk[9, 3] = -kk[3, 3]
        kk[3, 9] = -G * J / le[i]
        kk[9, 9] = -kk[3, 9]
        kk[4, 2] = -6 * E * Iy / le[i] ** 2
        kk[10, 2] = kk[4, 2]
        kk[4, 4] = 4 * E * Iy / le[i]
        kk[10, 4] = 2 * E * Iy / le[i]
        kk[4, 8] = 6 * E * Iy / le[i] ** 2
        kk[10, 8] = kk[4, 8]
        kk[4, 10] = 2 * E * Iy / le[i]
        kk[10, 10] = 4 * E * Iy / le[i]
        kk[5, 1] = 6 * E * Iz / le[i] ** 2
        kk[11, 1] = kk[5, 1]
        kk[5, 5] = 4 * E * Iz / le[i]
        kk[11, 5] = 2 * E * Iz / le[i]
        kk[5, 7] = -6 * E * Iz / le[i] ** 2
        kk[11, 7] = kk[5, 7]
        kk[5, 11] = 2 * E * Iz / le[i]
        kk[11, 11] = 4 * E * Iz / le[i]

        l = (V[v2, 0] - V[v1, 0]) / le[i]
        m = (V[v2, 1] - V[v1, 1]) / le[i]
        n = (V[v2, 2] - V[v1, 2]) / le[i]
        D = np.sqrt(l ** 2 + m ** 2)

        if D == 0:
            if n > 0:
                R = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]])
            else:
                R = np.array([[0, 0, -1], [0, 1, 0], [1, 0, 0]])
        else:
            R = np.array([[l, m, n], [-m / D, l / D, 0], [-l * n / D, -m * n / D, D]])

        RR = np.zeros((12, 12))
        RR[:3, :3] = R
        RR[3:6, 3:6] = R
        RR[6:9, 6:9] = R
        RR[9:, 9:] = R

        K[np.ix_(DOF, DOF)] += RR.T @ kk @ RR

    U = np.zeros(6 * V.shape[0])
    fixed_node = boundary_nodes_array_fea
    fixed_free_dof = np.unique(np.hstack([np.arange(6 * node, 6 * (node + 1)) for node in fixed_node]))
    all_free_dof = np.arange(6 * V.shape[0])
    free_dof = np.setdiff1d(all_free_dof, fixed_free_dof)
    U[free_dof] = np.linalg.solve(K[np.ix_(free_dof, free_dof)], F[free_dof])
    U = U.reshape(-1, 6)
    s = U[:, :3]
    VV = V + s

    if draw:
        distances = np.linalg.norm(V - VV, axis=1)

        color_factor = 100 * distances / 8
        color_factor = np.clip(color_factor, 0, 1)

        fig = plt.figure(dpi=500)
        ax = fig.add_subplot(111, projection='3d')

        for i in range(Edge.shape[0]):
            ax.plot(*V[Edge[i, :], :].T, color='grey', linewidth=1.4, alpha=0.9)

        for i in range(V.shape[0]):
            ax.scatter(*V[i], color='grey', s=2, alpha=0.9)

        ax.axis('off')

        output_dir = 'FEA_simu/ori'
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        file_name = f"{output_dir}/plot_{index}.png"

        elev = 30
        azim = 0
        ax.view_init(elev=elev, azim=azim)

        ax.set_xlim([0, 0.1])
        ax.set_ylim([0, 0.1])
        ax.set_zlim([0, 0.1])
        plt.savefig(file_name)

        plt.clf()

    max_deformation_length = 0
    for i in range(VV.shape[0]):
        deformation_length = np.linalg.norm(V[i] - VV[i]) * 100
        if deformation_length > max_deformation_length:
            max_deformation_length = deformation_length

    return max_deformation_length