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

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

# Define colors for domains
domain_colors = {
    "Legislation": "red",
    "Healthcare Systems": "blue",
    "Healthcare Policies": "green",
    "Default": "grey"
}

# Load index data
def load_index_data(file_path):
    with open(file_path, "r") as file:
        return json.load(file)

# Load and parse entities
def build_graph(data):
    G = nx.DiGraph()
    for entity_id, entity_info in data["entities"].items():
        label = entity_info.get("label", entity_id)
        domain = entity_info.get("inherits_from", "Default")
        color = domain_colors.get(domain, "grey")  # Set color, defaulting to "grey" if domain is missing
        G.add_node(entity_id, label=label, color=color)

        # Load additional relationships if specified in the entity data
        file_path = entity_info.get("file_path")
        if file_path and os.path.exists(file_path):
            with open(file_path, "r") as f:
                entity_data = json.load(f)
                for rel in entity_data.get("relationships", []):
                    G.add_edge(rel["source"], rel["target"], label=rel["attributes"]["relationship"])
                    
    # Add relationships from index.json
    for relationship in data["relationships"]:
        G.add_edge(relationship["source"], relationship["target"], label=relationship["attributes"].get("relationship", "related_to"))

    return G

# Enhanced visualization
def visualize_graph(G, title="Inferred Contextual Relationships"):
    pos = nx.spring_layout(G)
    plt.figure(figsize=(15, 10))

    # Draw nodes with colors
    node_colors = [G.nodes[node].get("color", "grey") for node in G.nodes]  # Default to "grey" if color is missing
    nx.draw_networkx_nodes(G, pos, node_size=3000, node_color=node_colors, alpha=0.8)

    # Draw labels
    nx.draw_networkx_labels(G, pos, font_size=10, 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=8)

    # Save as PDF and display
    plt.title(title)
    plt.axis("off")
    plt.savefig("graph_visualization.pdf")  # Export as PDF
    plt.show()

# Main execution
if __name__ == "__main__":
    # Load data from 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)