from sklearn.metrics import mean_absolute_error
import sys
sys.path.append('../')
from train import *
from args import Args
from GraphCompletion.graph_completion_with_training import graph_show
from GraphCompletion.graph_show import graph_save


def test_completion(x_batch, y_len_unsorted, args, rnn, output, test_batch_size=32, sample_time=1):
    rnn.hidden = rnn.init_hidden(1)
    rnn.eval()
    output.eval()

    # generate graphs
    max_num_node = int(args.max_num_node)
    y_pred = Variable(
        torch.zeros(test_batch_size, max_num_node, args.max_prev_node)).cuda()  # normalized prediction score
    y_pred_long = Variable(torch.zeros(test_batch_size, max_num_node, args.max_prev_node)).cuda()  # discrete prediction
    x_step = Variable(torch.ones(1, 1, args.max_prev_node)).cuda()

    # number_of_missing_nodes = args.number_of_missing_nodes
    number_of_missing_nodes = 2
    for j in range(test_batch_size):
        incomplete_graph_size = y_len_unsorted.data[j] - number_of_missing_nodes
        for i in range(y_len_unsorted.data[j]-1):
            h = rnn(x_step)
            y_pred_step = output(h)
            y_pred[j:j+1, i:i + 1, :] = F.sigmoid(y_pred_step)
            if(i<incomplete_graph_size):
                x_step = (x_batch[j:j+1, i+1:i+2, :]).cuda()
            else:
                x_step = sample_sigmoid(y_pred_step, sample=True, sample_time=sample_time)
            y_pred_long[j:j+1, i:i + 1, :] = x_step
            rnn.hidden = Variable(rnn.hidden.data).cuda()
    y_pred_long_data = y_pred_long.data.long()
    adj_true_list = []
    graph_true_list = []
    adj_incomplete_list = []
    graph_incomplete_list = []
    adj_pred_list = []
    graph_pred_list = []
    for i in range(test_batch_size):
        adj_true = decode_adj(x_batch[i, 1:y_len_unsorted.data[i], :].cpu().numpy())
        adj_true_list.append(adj_true)
        graph_true_list.append(nx.from_numpy_matrix(adj_true))
        adj_incomplete = decode_adj(x_batch[i, 1:y_len_unsorted.data[i] - number_of_missing_nodes, :].cpu().numpy())
        adj_incomplete_list.append(adj_incomplete)
        graph_incomplete_list.append(nx.from_numpy_matrix(adj_incomplete))
        adj_pred = decode_adj(y_pred_long_data[i, 0:y_len_unsorted.data[i]-1, :].cpu().numpy())
        adj_pred_list.append(adj_pred)
        graph_pred_list.append(nx.from_numpy_matrix(adj_pred))
    return adj_true_list, graph_true_list, adj_incomplete_list, graph_incomplete_list, adj_pred_list, graph_pred_list


if __name__ == '__main__':
    # All necessary arguments are defined in args.py
    args = Args()
    os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda)
    # check if necessary directories exist
    if not os.path.isdir(args.model_save_path):
        os.makedirs(args.model_save_path)
    if not os.path.isdir(args.graph_save_path):
        os.makedirs(args.graph_save_path)
    if not os.path.isdir(args.figure_save_path):
        os.makedirs(args.figure_save_path)
    if not os.path.isdir(args.timing_save_path):
        os.makedirs(args.timing_save_path)
    if not os.path.isdir(args.figure_prediction_save_path):
        os.makedirs(args.figure_prediction_save_path)
    if not os.path.isdir(args.nll_save_path):
        os.makedirs(args.nll_save_path)

    graphs = create_graphs.create(args)

    # split datasets
    random.seed(123)
    shuffle(graphs)
    graphs_len = len(graphs)

    args.max_num_node = max([graphs[i].number_of_nodes() for i in range(len(graphs))])
    max_num_edge = max([graphs[i].number_of_edges() for i in range(len(graphs))])
    min_num_edge = min([graphs[i].number_of_edges() for i in range(len(graphs))])

    # show graphs statistics
    print('total graph num: {}'.format(len(graphs)))
    print('max number node: {}'.format(args.max_num_node))
    print('max/min number edge: {}; {}'.format(max_num_edge, min_num_edge))
    print('max previous node: {}'.format(args.max_prev_node))

    if 'nobfs' in args.note:
        args.max_prev_node = args.max_num_node-1
    dataset = Graph_sequence_sampler_pytorch_nobfs_for_completion(graphs,
                                                 max_num_node=args.max_num_node)
    sample_strategy = torch.utils.data.sampler.WeightedRandomSampler([1.0 / len(dataset) for i in range(len(dataset))],
                                                                     num_samples=args.batch_size * args.batch_ratio,
                                                                     replacement=True)
    dataset_loader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, num_workers=args.num_workers,
                                                 sampler=sample_strategy)

    rnn_GraphRNN = GRU_plain(input_size=args.max_prev_node, embedding_size=args.embedding_size_rnn,
                    hidden_size=args.hidden_size_rnn, num_layers=args.num_layers, has_input=True,
                    has_output=False).cuda()
    output_GraphRNN = MLP_plain(h_size=args.hidden_size_rnn, embedding_size=args.embedding_size_output,
                       y_size=args.max_prev_node).cuda()
    fname = args.model_save_path + args.fname_GraphRNN + 'lstm_' + str(args.load_epoch) + '.dat'
    rnn_GraphRNN.load_state_dict(torch.load(fname))
    fname = args.model_save_path + args.fname_GraphRNN + 'output_' + str(args.load_epoch) + '.dat'
    output_GraphRNN.load_state_dict(torch.load(fname))

    # ******************************************************************
    rnn_GraphCompletion = GRU_plain(input_size=args.max_prev_node, embedding_size=args.embedding_size_rnn,
                    hidden_size=args.hidden_size_rnn, num_layers=args.num_layers, has_input=True,
                    has_output=False).cuda()
    output_GraphCompletion = MLP_plain(h_size=args.hidden_size_rnn, embedding_size=args.embedding_size_output,
                       y_size=args.max_prev_node).cuda()

    fname = args.model_save_path + args.fname + 'lstm_' + args.graph_completion_string + str(args.epochs) + '.dat'
    rnn_GraphCompletion.load_state_dict(torch.load(fname))
    fname = args.model_save_path + args.fname + 'output_' + args.graph_completion_string + str(args.epochs) + '.dat'
    output_GraphCompletion.load_state_dict(torch.load(fname))
    # ******************************************************************
    MAE_GraphRNN = []
    MAE_GraphCompletion = []
    for batch_idx, data in enumerate(dataset_loader):
        if True:
            x_unsorted = data['x'].float()
            y_unsorted = data['y'].float()
            y_len_unsorted = data['len']
            # *********************************
            adj_true_list, graph_true_list, adj_incomplete_list, graph_incomplete_list, adj_pred_list, graph_pred_list\
                = test_completion(x_unsorted, y_len_unsorted, args, rnn_GraphRNN, output_GraphRNN)
            adj_true_list, graph_true_list, adj_incomplete_list, graph_incomplete_list, adj_pred_list_completion_model\
                , graph_pred_list_completion_model \
                = test_completion(x_unsorted, y_len_unsorted, args, rnn_GraphCompletion, output_GraphCompletion)
            mae = np.sum(np.absolute((adj_pred_list[0].astype("float") - adj_true_list[0].astype("float"))))
            # print("adj_true: ")
            # print(adj_true_list[0])
            # graph_show(nx.from_numpy_matrix(adj_true_list[0]),"adj_true" )
            # graph_show(nx.from_numpy_matrix(adj_incomplete_list[0]), "adj_incomplete")
            # print("my error")
            # print(mae)
            # print(mean_absolute_error(adj_pred_list[0], adj_true_list[0]))
            # print("adj_pred_list:")
            # print(adj_pred_list[0])
            # graph_show(nx.from_numpy_matrix(adj_pred_list[0]), "pred_true")

            for i in range(len(graph_true_list)):
                # graph_save(graph_true_list[i], graph_incomplete_list[i],
                #            graph_pred_list[i], graph_pred_list_completion_model[i], i, args.graph_save_path)
                mae = mean_absolute_error(adj_pred_list[i], adj_true_list[i])
                MAE_GraphRNN.append(mae)
                mae = mean_absolute_error(adj_pred_list_completion_model[i], adj_true_list[i])
                MAE_GraphCompletion.append(mae)
            # *********************************
            # print(MAE_GraphCompletion)
            # G_pred_step = test_mlp(x_unsorted, y_len_unsorted, epoch, args, rnn_for_graph_completion,
            #                        output_for_graph_completion)
            # nx.write_gpickle(G_pred_step, "completed_graphs_with_training.dat")
    print("MAE_GraphRNN:")
    print(np.mean(MAE_GraphRNN))
    # print(MAE_GraphRNN)
    print("MAE_GraphCompletion:")
    print(np.mean(MAE_GraphCompletion))