import circadapt
import networkx as nx
import matplotlib.pyplot as plt
from collections import defaultdict
import numpy as np

model = circadapt.VanOsta2024()

components = model._components.keys()

all_objects = []

# Object types
all_object_types = {
    'CiSy': 'ArtVen',
    'CiPu': 'ArtVen',
    'SyArt': 'Artery',
    'SyVen': 'Vein',
    'PuArt': 'Pulmonary Artery',
    'PuVen': 'Pulmonary Vein',
    'Peri': 'Pericardium',
    'Timings': 'Timing Node',
    'PFC': 'Pulmonary Flow Control'
}

for component in components:
    if component in ['Solver', 'General']:
        continue
    all_objects += model[component].objects

    for object in model[component].objects:
        object_split = object.split('.')
        all_object_types[object_split[-1]] = component

all_objects = np.sort(all_objects)


# open plot
plt.figure(figsize=(12, 8))

# Step 1: Build the tree structure
tree = defaultdict(dict)

def add_to_tree(path, tree):
    parts = path.split('.')
    for part in parts[:-1]:
        tree = tree.setdefault(part, {})
    tree[parts[-1]] = {}

for obj in all_objects:
    print(obj)
    add_to_tree(obj, tree)

# Step 2: Flatten the tree for NetworkX
def flatten_tree(tree, parent=None):
    edges = []
    for key, subtree in tree.items():
        if parent:
            edges.append((parent, key))
        edges.extend(flatten_tree(subtree, key))
    return edges

edges = flatten_tree(tree)

# Step 3: Create the graph
G = nx.DiGraph()
G.add_edges_from(edges)

# Ensure "Model" is the root and set layout for a tree
from networkx.drawing.nx_pydot import pydot_layout
pos = pydot_layout(G, prog="dot")  # Use DOT for hierarchical layout

# Draw the tree
nx.draw(G, pos, with_labels=False, node_size=5000, node_color='lightblue', font_size=1, font_weight='bold', edge_color='black')
plt.title("Tree Plot of All Objects", fontsize=20)
plt.axis("off")

# print type
for object_type in all_object_types:
    plt.annotate(object_type, xy=(pos[object_type][0], pos[object_type][1] + 4), ha='center', va='center', fontsize=16)
    plt.annotate(all_object_types[object_type], xy=(pos[object_type][0], pos[object_type][1] - 6), ha='center', va='center', fontsize=12)
plt.annotate('Model', xy=(pos['Model'][0], pos['Model'][1]), ha='center', va='center', fontsize=16, fontweight='bold')


plt.show()