import json
import networkx as nx
import matplotlib.pyplot as plt

# Define the path to your index.json file
index_file_path = "graphs/index.json"

# Load the data from index.json
def load_index_data(file_path):
    """Load the index.json file and parse its contents."""
    with open(file_path, "r") as file:
        data = json.load(file)
    return data

def build_graph(data):
    """Builds a directed graph based on entities and relationships."""
    G = nx.DiGraph()

    # Add nodes for each entity
    for entity_id, entity_info in data["entities"].items():
        label = entity_info.get("label", entity_id)
        # Use label separately to avoid duplicate keyword arguments
        G.add_node(entity_id, label=label, **{k: v for k, v in entity_info.items() if k != "label"})

    # Add edges for each relationship
    for relationship in data["relationships"]:
        source = relationship["source"]
        target = relationship["target"]
        relationship_label = relationship["attributes"].get("relationship", "related_to")
        G.add_edge(source, target, label=relationship_label)

    return G

# Visualize the graph using Matplotlib
def visualize_graph(G, title="340B Program - Inferred Contextual Relationships"):
    """Visualizes the graph with nodes and relationships."""
    pos = nx.spring_layout(G)  # Position nodes with a spring layout

    # Draw nodes with labels
    plt.figure(figsize=(15, 10))
    nx.draw_networkx_nodes(G, pos, node_size=3000, node_color="lightblue", alpha=0.7)
    nx.draw_networkx_labels(G, pos, font_size=10, font_color="black", font_weight="bold")

    # Draw edges with labels
    nx.draw_networkx_edges(G, pos, arrowstyle="->", arrowsize=20, edge_color="gray", connectionstyle="arc3,rad=0.1")
    edge_labels = {(u, v): d["label"] for u, v, d in G.edges(data=True)}
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_color="red", font_size=9)

    # Set plot title and display
    plt.title(title)
    plt.axis("off")
    plt.show()

# Main execution
if __name__ == "__main__":
    # Load data from the index.json
    data = load_index_data(index_file_path)

    # Build the graph with entities and relationships
    G = build_graph(data)

    # Visualize the graph
    visualize_graph(G)