from __future__ import print_function
import numpy as np
import networkx as nx
import node2vec
import random
from config import *
from evaluation import *
from model import *
from utils import *
import tensorflow as tf
import math
import time
import os

tf.app.flags.DEFINE_string('datasets', 'citeseer', 'datasets descriptions')
tf.app.flags.DEFINE_integer('dimensions', 128, 'embedding dimensions')
tf.app.flags.DEFINE_integer('walk_length', 80, 'walk length')
tf.app.flags.DEFINE_integer('num_walks', 10, 'number of walks')
tf.app.flags.DEFINE_integer('window_size', 10, 'window size')
tf.app.flags.DEFINE_float('p', 1.0, 'p value')
tf.app.flags.DEFINE_float('q', 1.0, 'q value')
tf.app.flags.DEFINE_boolean('weighted', False, 'weighted edges')
tf.app.flags.DEFINE_boolean('directed', True, 'undirected edges')
tf.app.flags.DEFINE_float('alpha',  1.0, 'alpha in FakeWeight')
tf.app.flags.DEFINE_float('reg', 10.0, 'reg in config')
tf.app.flags.DEFINE_string('train_type', 'class', 'classification or link prediction')
tf.app.flags.DEFINE_integer('iters', 12000, 'max_iters')
tf.app.flags.DEFINE_string('target_neighbors', 'WAN', 'the way to construct target neighbors')
tf.app.flags.DEFINE_float('a',  1.0, 'alpha in config')
#set struct in config

os.environ['CUDA_VISIBLE_DEVICES'] = '1'
config_tf = tf.ConfigProto()
config_tf.gpu_options.allow_growth = True


def generate_graph_context_all_pairs(path, window_size):
    # generating graph context pairs
    all_pairs = []
    for k in range(len(path)):
        for i in range(len(path[k])):
            for j in range(i - window_size, i + window_size + 1):
                if i == j or j < 0 or j >= len(path[k]):
                    continue
                else:
                    all_pairs.append([path[k][i], path[k][j]])

    return np.array(all_pairs, dtype=np.int32)


def graph_context_batch_iter(all_pairs, batch_size):
    while True:
        start_idx = np.random.randint(0, len(all_pairs) - batch_size)
        batch_idx = np.array(range(start_idx, start_idx + batch_size))
        batch_idx = np.random.permutation(batch_idx)
        batch = np.zeros(batch_size, dtype=np.int32)
        labels = np.zeros((batch_size, 1), dtype=np.int32)
        batch[:] = all_pairs[batch_idx, 0]
        labels[:, 0] = all_pairs[batch_idx, 1]
        yield batch, labels


def construct_traget_neighbors(nx_G, X, FLAGS, mode='WAN'):
    # construct target neighbor feature matrix
    X_target = np.zeros(X.shape)
    nodes = nx_G.nodes()

    if mode == 'OWN':
        # autoencoder for reconstructing itself
        return X
    elif mode == 'EMN':
        # autoencoder for reconstructing Elementwise Median Neighbor
        for node in nodes:
            neighbors = list(nx_G.neighbors(node))
            if len(neighbors) == 0:
                X_target[node] = X[node]
            else:
                temp = np.array(X[node])
                for n in neighbors:
                    if FLAGS.weighted:
                        # weighted sum
                        temp = np.vstack((temp, X[n] * nx_G[node][n]['weight']))
                    else:
                        temp = np.vstack((temp, X[n]))
                temp = np.median(temp, axis=0)
                X_target[node] = temp
        return X_target
    elif mode == 'WAN':
        # autoencoder for reconstructing Weighted Average Neighbor
        for node in nodes:
            # put myself ( maybe can be deleted, because second-order will count itself)
            temp = np.array(X[node])

            alpha = 0.1
            first_order = 1
            neighbors = list(nx_G.neighbors(node))
            neighbors_length = len(neighbors)
            if neighbors_length != 0:
                for n in neighbors:
                    if FLAGS.weighted:
                        # weighted sum
                        temp = np.vstack((temp, X[n] * nx_G[node][n]['weight']))
                    else:
                        temp = np.vstack((temp, X[n]))
                        set_n = set(nx_G.successors(n))
                        set_n = set_n.intersection(set(neighbors))
                        second_order = len(set_n) / neighbors_length
                        w = alpha * first_order + (1 - alpha) * second_order
                        nx_G[node][n]['weight'] = w

                temp = np.mean(temp, axis=0)
            X_target[node] = temp
        return X_target


# ratio is the test set size divide by the whole graph
def make_small_graph(G, ratio):
    test_size = int(G.number_of_edges() * ratio)

    edge_list=list(G.edges)
    np.random.shuffle(edge_list)
    test_set = edge_list[:test_size]
    train_set = edge_list[test_size:]

    non_edge_list=list(nx.non_edges(G))    
    np.random.shuffle(non_edge_list)
    test_negative_set = non_edge_list[:test_size]

    return train_set, test_set, test_negative_set


def main():
    FLAGS = tf.app.flags.FLAGS
    print('alpha in FakeWeight: %.4f, alpha in config: %.4f, reg: %.4f, p in node2vec: %.4f' %(FLAGS.alpha, FLAGS.a, FLAGS.reg, FLAGS.p))
    print(FLAGS.datasets + '  ' + FLAGS.train_type)
    print(FLAGS.target_neighbors)
    inputEdgeFile = 'graph/' + FLAGS.datasets + '.edgelist'
    inputFeatureFile = 'graph/' + FLAGS.datasets + '.feature.mat'
    inputLabelFile = 'graph/' + FLAGS.datasets + '.label'
    outputEmbedFile = 'graph/' + FLAGS.datasets + '.embed'
    window_size = FLAGS.window_size

    # Read graph
    nx_G = read_graph(FLAGS, inputEdgeFile)

    if FLAGS.train_type == 'class':
        pass
    elif FLAGS.train_type == 'link':
        train_set, test_set, test_negative_set = make_small_graph(nx_G, 0.5)
        nx_G.remove_edges_from(test_set)

    # Read features
    print('reading features...')
    X = read_feature_mat(inputFeatureFile)
    _ , feaDims = X.shape

    # construct traget neighbors and build FakeWeight
    X_target = construct_traget_neighbors(nx_G, X, FLAGS, mode=FLAGS.target_neighbors)

    # Perform random walks to generate graph context
    G = node2vec.Graph(nx_G, FLAGS.directed, FLAGS.p, FLAGS.q)
    G.preprocess_transition_probs()
    walks = G.simulate_walks(FLAGS.num_walks, FLAGS.walk_length)

    print('generating graph context pairs...')
    start_time = time.time()
    all_pairs = generate_graph_context_all_pairs(walks, window_size)
    end_time = time.time()
    print('time consumed for constructing graph context: %.2f' % (end_time - start_time))

    nodes = nx_G.nodes()

    # Total number nodes
    N = len(nodes)
    dims = FLAGS.dimensions

    config = Config()
    config.struct[0] = feaDims
    config.struct[-1] = FLAGS.dimensions
    config.reg = FLAGS.reg
    config.max_iters = FLAGS.iters
    config.alpha = FLAGS.a
    model = Model(config, N, dims, X_target)

    init = tf.global_variables_initializer()
    sess = tf.Session(config=config_tf)
    sess.run(init)

    batch_size = config.batch_size
    max_iters = config.max_iters

    idx = 0
    print_every_k_iterations = 500
    start = time.time()

    loss_sg = 0
    loss_ae = 0

    for iter_cnt in range(max_iters):
        idx += 1

        batch_index, batch_labels = next(graph_context_batch_iter(all_pairs, batch_size))

        # train for autoencoder model
        start_idx = np.random.randint(0, N - batch_size)
        batch_idx = np.array(range(start_idx, start_idx + batch_size))
        batch_idx = np.random.permutation(batch_idx)
        batch_X = X[batch_idx]
        feed_dict = {model.X: batch_X, model.inputs: batch_idx}
        _, loss_ae_value = sess.run([model.train_opt_ae, model.loss_ae], feed_dict=feed_dict)
        loss_ae += loss_ae_value

        # train for skip-gram model
        batch_X = X[batch_index]
        feed_dict = {model.X: batch_X, model.labels: batch_labels}
        _, loss_sg_value = sess.run([model.train_opt_sg, model.loss_sg], feed_dict=feed_dict)
        loss_sg += loss_sg_value

        if idx % print_every_k_iterations == 0:
            end = time.time()
            print('iterations: %d' % idx + ', time elapsed: %.2f, ' % (end - start), end='')
            total_loss = loss_sg / idx + loss_ae / idx
            print('loss: %.2f, ' % total_loss, end='')
            embedding_result = sess.run(model.Y, feed_dict={model.X: X})

            if FLAGS.train_type == 'class':
                y = read_label(inputLabelFile)
                macro_f1, micro_f1 = multiclass_node_classification_eval(embedding_result, y, 0.7)
                print('[macro_f1 = %.4f, micro_f1 = %.4f]' % (macro_f1, micro_f1))
            else:
                roc_score, ap_score = get_roc_score(test_set, test_negative_set, embedding_result)
                print('[roc = %.4f, ap = %.4f]' % (roc_score, ap_score))

    print('optimization finished...')
    embedding_result = sess.run(model.Y, feed_dict={model.X: X})
    if FLAGS.train_type == 'class':
        y = read_label(inputLabelFile)
        for i in range(5):
            ratio = 1.0 - float(i + 1.0) / 10
            print('repeat 10 times for node classification with random split at ratio %.2f' %(ratio))
            node_classification_F1(embedding_result, y, ratio)

    if FLAGS.train_type == 'link':
        roc_score, ap_score = get_roc_score(test_set, test_negative_set, embedding_result)
        print('for link prediction:')
        print('roc = ' + str(roc_score))
        print('ap = ' + str(ap_score))

    # print('saving embedding result...')
    # write_embedding(embedding_result, outputEmbedFile)


if __name__ == '__main__':
    main()
