import networkx as nx
import warnings
import os

from pyvis.network import Network

warnings.filterwarnings("ignore")

# Load the GraphML file
file_path = "./sample/graph_chunk_entity_relation.graphml"

assert os.path.exists(file_path), f"File {file_path} does not exist."
G = nx.read_graphml(file_path)

# Convert the graph to node-link data format
dict_graph = nx.node_link_data(G)
print("Number of nodes:", len(dict_graph['nodes']))
print("Number of edges:", len(dict_graph['links']))

# Create a Pyvis network
network = Network(width="100%",
                  height="100vh",
                  notebook=True,
                  bgcolor="#f8fafc",
                  font_color="#1e293b")

# Convert NetworkX graph to Pyvis network
network.from_nx(G)

# Add colors and title to nodes
for node in network.nodes:
    if "description" in node:
        node["title"] = node["description"]

    node['color'] = {'background': '#e0e7ff', 'border': '#6366f1', 'highlight': {'background': '#c7d2fe', 'border': '#4f46e5'}}
    node['font'] = {'size': 14, 'color': '#1e293b'}
    node['shape'] = 'dot'
    node['size'] = 20

# Add title to edges
for edge in network.edges:
    if "description" in edge:
        edge["title"] = edge["description"]

    edge['width'] = 4
    edge['color'] = {'color': '#6366f1', 'highlight': '#4f46e5'}
    edge['font'] = {'size': 12, 'color': '#4b5563', 'face': 'Arial'}

# Save and display the network
filename_out = "knowledge_graph.html"
network.show(filename_out)
print(f"Knowledge graph saved to {filename_out}")