network.py

import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt
from pyvis.network import Network

from data import df


def analysis():
    G = nx.DiGraph()

    # Define the columns to consider for relationships
    relationship_columns = ['father', 'sibling', 'spouse', 'mother', 'child']

    G.add_nodes_from(df["itemLabel"])

    for index, row in df.iterrows():
        main_entity = row['itemLabel']
        for relationship in relationship_columns:
            if pd.notna(row[relationship]):
                G.add_edge(main_entity, row[relationship], relationship=str(relationship), label=relationship)

    plt.figure(figsize=(50, 20))  # Set the size of the plot
    # pos = nx.kamada_kawai_layout(G)  # Compute the positions of the nodes

    # # Draw the nodes and edges with labels
    # nx.draw(G, pos, with_labels=True, node_size=20, node_color='skyblue', font_size=10, font_color='black', font_weight='bold', alpha=0.7)

    # # Draw edge labels (the relationships)
    # edge_labels = nx.get_edge_attributes(G, 'relationship')
    # nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_color='red')

    # plt.title('Relationship Graph with Labels')
    # plt.axis('off')  # Hide the axes for clarity
    # plt.show()


    net = Network(width="1200px", height="800px", bgcolor="#222222", font_color="white", directed=True)
    net.from_nx(G)

    for node in net.nodes:
        node["title"] = node["id"]
        node["value"] = len(G[node["id"]])


    for edge in net.edges:
        # Set edge title to the relationship from your graph
        relationship = G[edge["from"]][edge["to"]].get("relationship", "Unknown")
        edge["title"] = relationship  # This will show as a tooltip
        edge["color"] = "blue"
        edge["width"] = 2 if relationship != "Unknown" else 1

    html = net.generate_html()
    #need to remove ' from HTML
    html = html.replace("'", "\"")
    
    return f"""<iframe style="width: 100%; height: 800px;margin:0 auto" name="result" allow="midi; geolocation; microphone; camera; 
    display-capture; encrypted-media;" sandbox="allow-modals allow-forms 
    allow-scripts allow-same-origin allow-popups 
    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
    allowpaymentrequest="" frameborder="0" srcdoc='{html}'></iframe>"""