import networkx as nx
import matplotlib.pyplot as plt
import matplotlib.pylab as pl
import matplotlib.gridspec as gridspec
from matplotlib.pyplot import figure


def graph_save(main_graphs, incomplete_graphs, completed_graphs, completed_graphs_with_training,
               name, path):
    # Create 2x2 sub plots
    gs = gridspec.GridSpec(2, 2)
    pl.figure()
    figure(num=None, figsize=(10, 10), dpi=80, facecolor='w', edgecolor='k')
    # main_graphs
    ax = pl.subplot(gs[0, 0])  # row 0, col 0
    pos = nx.spring_layout(main_graphs)
    nx.draw_networkx_nodes(main_graphs, pos, cmap=plt.get_cmap('jet'),
                           node_color='b', node_size=500)
    nx.draw_networkx_labels(main_graphs, pos)
    nx.draw_networkx_edges(main_graphs, pos, edgelist=main_graphs.edges(), arrows=False)
    ax.set_yticklabels([])
    ax.set_xticklabels([])
    pl.title('Main')
    pl.plot()
    # incomplete_graphs
    ax = pl.subplot(gs[0, 1])  # row 0, col 0
    pos = nx.spring_layout(incomplete_graphs)
    nx.draw_networkx_nodes(incomplete_graphs, pos, cmap=plt.get_cmap('jet'),
                           node_color='b', node_size=500)
    nx.draw_networkx_labels(incomplete_graphs, pos)
    nx.draw_networkx_edges(incomplete_graphs, pos, edgelist=incomplete_graphs.edges(), arrows=False)
    ax.set_yticklabels([])
    ax.set_xticklabels([])
    pl.title('Incomplete')
    pl.plot()
    # completed_graphs
    ax = pl.subplot(gs[1, 0])  # row 0, col 0
    pos = nx.spring_layout(completed_graphs)
    nx.draw_networkx_nodes(completed_graphs, pos, cmap=plt.get_cmap('jet'),
                           node_color='b', node_size=500)
    nx.draw_networkx_labels(completed_graphs, pos)
    nx.draw_networkx_edges(completed_graphs, pos, edgelist=completed_graphs.edges(), arrows=False)
    ax.set_yticklabels([])
    ax.set_xticklabels([])
    pl.title('Completed without Training')
    pl.plot()
    # completed_graphs_with_training
    ax = pl.subplot(gs[1, 1])  # row 0, col 0
    pos = nx.spring_layout(completed_graphs_with_training)
    nx.draw_networkx_nodes(completed_graphs_with_training, pos, cmap=plt.get_cmap('jet'),
                           node_color='b', node_size=500)
    nx.draw_networkx_labels(completed_graphs_with_training, pos)
    nx.draw_networkx_edges(completed_graphs_with_training, pos,
                           edgelist=completed_graphs_with_training.edges(), arrows=False)
    ax.set_yticklabels([])
    ax.set_xticklabels([])
    pl.title('Completed with Training')
    pl.plot()

    # pl.show()
    plt.savefig(path + str(name + 1) + ".png", format="PNG")
    plt.close()


def graph_show(G):
    pos = nx.spring_layout(G)
    nx.draw_networkx_nodes(G, pos, cmap=plt.get_cmap('jet'),
                           node_color='b', node_size=500)
    nx.draw_networkx_labels(G, pos)
    nx.draw_networkx_edges(G, pos, edgelist=G.edges(), arrows=False)
    plt.show()


if __name__ == '__main__':
    name = "small_grid_graphs"
    path = './results/' + name + '/'
    main_graphs = nx.read_gpickle(path + "main_graphs.dat")
    incomplete_graphs = nx.read_gpickle(path + "incomplete_graphs.dat")
    completed_graphs = nx.read_gpickle(path + "completed_graphs.dat")
    completed_graphs_with_training = nx.read_gpickle(path + "completed_graphs_with_training.dat")
    for i in range(len(main_graphs)):
        graph_save(main_graphs[i], incomplete_graphs[i],
               completed_graphs[i], completed_graphs_with_training[i], i, path)
# i = 2
# graph_save(main_graphs[i], incomplete_graphs[i],
#                completed_graphs[i], i, path)