from pyvis.network import Network
from graph_builder import load_graph

def visualize_graph(G):
    net = Network(
        directed=True,
        notebook=False   
    )

    for node in G.nodes():
        net.add_node(node)

    for u, v, data in G.edges(data=True):
        net.add_edge(
            u,
            v,
            label=data.get("relation", "")
        )

    output_path = "kg.html"

    net.write_html(output_path)

    print(f"Graph saved to {output_path}")
    
#test
if __name__ == "__main__":

    G = load_graph()

    visualize_graph(G)