import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
import scipy.io as sp


def graph_save(G, dir, i):
    pos = nx.spring_layout(G)
    nx.draw_networkx_nodes(G, pos, cmap=plt.get_cmap('jet'),
                           node_color='b', node_size=300)
    nx.draw_networkx_labels(G, pos)
    nx.draw_networkx_edges(G, pos, edgelist=G.edges(), arrows=False)
    plt.savefig('{}.png'.format(i))
    plt.close()


LENGTH = 6
# missing_dir = 'testgraph_12_,.mat_0_3_miss.txt'
for i in range(0, 11):
    recons_dir = 'out/graphed_{}.mat'.format(i)
    adj_matrix_orig = sp.loadmat(recons_dir)['original_data'].A
    adj_matrix_p = sp.loadmat(recons_dir)['out_data'].A
    graph_orig = nx.convert_matrix.from_numpy_matrix(adj_matrix_orig)
    graph_p = nx.convert_matrix.from_numpy_matrix(adj_matrix_p)
    graph_save(graph_orig, '', i + 1)
    graph_save(graph_p, '', -(i + 1))