import pipmaster as pm

if not pm.is_installed('pyvis'):
    pm.install('pyvis')
if not pm.is_installed('networkx'):
    pm.install('networkx')

import random

import networkx as nx
from pyvis.network import Network

# Load the GraphML file
G = nx.read_graphml('./dickens/graph_chunk_entity_relation.graphml')

# Create a Pyvis network
net = Network(height='100vh', notebook=True)

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


# Add colors and title to nodes
for node in net.nodes:
    node['color'] = f'#{random.randint(0, 0xFFFFFF):06x}'
    if 'description' in node:
        node['title'] = node['description']

# Add title to edges
for edge in net.edges:
    if 'description' in edge:
        edge['title'] = edge['description']

# Save and display the network
net.show('knowledge_graph.html')