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()

    relationship_columns = ['father', 'sibling', 'spouse', 'mother', 'child']

    # G.add_nodes_from(df["itemLabel"])
    inf_labels = ["residence", "occupation", "class"]
    for index, row in df.iterrows():
        if row["itemLabel"] not in G.nodes():
            G.add_node(row["itemLabel"])
            information = ""
            for i in inf_labels:
                if pd.notna(row[i]):
                    information+=f"{i[:-5]}: {row[i]}\n"
            G.nodes[row["itemLabel"]]["attr"]=information

        for index, row in df.iterrows():
            main_entity = row['itemLabel']
            for relationship in relationship_columns:
                if pd.notna(row[relationship]) and not G.has_edge(main_entity, row[relationship]) and not G.has_edge(row[relationship], main_entity):
                    # if pd.notna(row[relationship]):

                    G.add_edge(row[relationship], main_entity, label=str(relationship)+ "_of")

    for x in G.nodes():
        if "attr" not in  G.nodes[x]:
            G.nodes[x]["attr"]=""

    # 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="1500px", height="1000px", bgcolor="#FFFFFF", font_color="black", directed=True)
    net.from_nx(G)

    for node in net.nodes:
        node["title"] = G.nodes[node["id"]]["attr"]
        node["value"] = len(G[node["id"]])
        node["font"]={"size": 20}

    for edge in net.edges:
        # Set edge title to the relationship from your graph
        relationship = G[edge["from"]][edge["to"]].get("label", "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>"""