import json
import networkx as nx
import matplotlib.pyplot as plt
from PIL import Image
import matplotlib.image as mpimg
from networkx.drawing.nx_agraph import graphviz_layout

def load_assembly_sequence(filepath):
    with open(filepath, 'r') as file:
        return json.load(file)

def create_assembly_graph(sequence):
    G = nx.DiGraph()
    
    for i, part in enumerate(sequence):
        G.add_node(part, image=f'{part}.svg')
        if i > 0:
            G.add_edge(sequence[i-1], part)
    
    return G

def draw_graph(G):
    pos = graphviz_layout(G, prog='dot')
    fig, ax = plt.subplots(figsize=(12, 8))

    for node in G.nodes(data=True):
        x, y = pos[node[0]]
        img = mpimg.imread(node[1]['image'])
        ax.imshow(img, aspect='auto', extent=(x-0.5, x+0.5, y-0.5, y+0.5), zorder=1)

        ax.text(x, y-0.6, s=node[0], horizontalalignment='center', fontsize=12, zorder=2)

    nx.draw_networkx_edges(G, pos, ax=ax, arrows=True)
    plt.axis('off')
    plt.show()

def main():
    sequence = load_assembly_sequence('path_to_adjacency_matrix.json')
    assembly_graph = create_assembly_graph(sequence)
    
    draw_graph(assembly_graph)

# main()
