# Simple precedence graph generator. Loads assembly tiers pickle file, loads into networkX graph and plots it.
from src.graph_utilities.graph_builder import *
import networkx as nx
import matplotlib.pyplot as plt
import pickle
import pandas as pd


def main():
    product_name = 'centrifugal_pump'
    with open('../out/assembly_tiers/' + product_name + '.pickle', 'rb') as handle:
        disassembly_tiers = pickle.load(handle)
    assembly_tiers = list(reversed(disassembly_tiers))

    subset_sizes = create_subset_sizes(assembly_tiers)
    G = create_multilayered_graph(*subset_sizes)
    mapping = create_mapping_of_names_to_nodes(assembly_tiers)
    G = nx.relabel_nodes(G, mapping)
    pos = nx.multipartite_layout(G, subset_key='layer')
    plt.figure(figsize=(8, 8))
    nx.draw(G, pos, with_labels=True)
    plt.axis("equal")
    plt.show()

    pickle_path = '../out/precedence_graph/' + product_name + '.pickle'
    with open(pickle_path, 'wb') as handle:
        pickle.dump(G, handle, protocol=pickle.HIGHEST_PROTOCOL)

if __name__ == '__main__':
    main()