# -*- coding: utf-8 -*-
"""
Created on Thu Sep  6 19:05:22 2018

@author: admin
"""
import os
path=os.path.dirname(os.path.abspath(__file__))
import time
import tensorflow as tf
import numpy as np
from utils import *
from models import GCN, MLP

# Set random seed
seed = 123
np.random.seed(seed)
tf.set_random_seed(seed)


import sys
name=sys.argv[3]
times=sys.argv[1]
num=sys.argv[2]


labelpath=sys.argv[4]
dataset=sys.argv[5]

errlog=open(os.path.join(labelpath,'errlog.txt'),'a')
sys.stderr=errlog
sys.stdout=errlog

# Settings
flags = tf.app.flags
FLAGS = flags.FLAGS
flags.DEFINE_string('dataset', dataset, 'Dataset string.')  # 'cora', 'citeseer', 'pubmed'
flags.DEFINE_string('model', 'gcn', 'Model string.')  # 'gcn', 'gcn_cheby', 'dense'
flags.DEFINE_float('learning_rate', 0.01, 'Initial learning rate.')
flags.DEFINE_integer('epochs', 200, 'Number of epochs to train.')
flags.DEFINE_integer('hidden1', 16, 'Number of units in hidden layer 1.')
flags.DEFINE_float('dropout', 0.5, 'Dropout rate (1 - keep probability).')
flags.DEFINE_float('weight_decay', 5e-4, 'Weight for L2 loss on embedding matrix.')
flags.DEFINE_integer('early_stopping', 200, 'Tolerance for early stopping (# of epochs).')
flags.DEFINE_integer('max_degree', 3, 'Maximum Chebyshev polynomial degree.')

# Load data
adj, features, y_train, y_val, y_test, train_mask, val_mask, test_mask = load_data(FLAGS.dataset)

# Some preprocessing
features = preprocess_features(features)
if FLAGS.model == 'gcn':
    support = [preprocess_adj(adj)]
    num_supports = 1
    model_func = GCN
elif FLAGS.model == 'gcn_cheby':
    support = chebyshev_polynomials(adj, FLAGS.max_degree)
    num_supports = 1 + FLAGS.max_degree
    model_func = GCN
elif FLAGS.model == 'dense':
    support = [preprocess_adj(adj)]  # Not used
    num_supports = 1
    model_func = MLP
else:
    raise ValueError('Invalid argument for model: ' + str(FLAGS.model))

# Define placeholders
placeholders = {
    'support': [tf.sparse_placeholder(tf.float32) for _ in range(num_supports)],
    'features': tf.sparse_placeholder(tf.float32, shape=tf.constant(features[2], dtype=tf.int64)),
    'labels': tf.placeholder(tf.float32, shape=(None, y_train.shape[1])),
    'labels_mask': tf.placeholder(tf.int32),
    'dropout': tf.placeholder_with_default(0., shape=()),
    'num_features_nonzero': tf.placeholder(tf.int32)  # helper variable for sparse dropout
}

# Create model
model = model_func(placeholders, input_dim=features[2][1], logging=True)
saver=tf.train.Saver()
# Initialize session
sess = tf.Session()


# Define model evaluation function
def evaluate(features, support, labels, mask, placeholders):
    t_test = time.time()
    feed_dict_val = construct_feed_dict(features, support, labels, mask, placeholders)
    outs_val = sess.run([model.loss, model.accuracy], feed_dict=feed_dict_val)
    return outs_val[0], outs_val[1], (time.time() - t_test)


# Init variables
sess.run(tf.global_variables_initializer())

cost_val = []


queue_acc=0
queue_time=0
queue_ckpt=os.path.join(labelpath,'best_model.ckpt')

# Train model
for epoch in range(FLAGS.epochs):

    t = time.time()
    # Construct feed dictionary
    feed_dict = construct_feed_dict(features, support, y_train, train_mask, placeholders)
    feed_dict.update({placeholders['dropout']: FLAGS.dropout})

    # Training step
    outs = sess.run([model.opt_op, model.loss, model.accuracy], feed_dict=feed_dict)

    # Validation
    cost, acc, duration = evaluate(features, support, y_val, val_mask, placeholders)
    cost_val.append(cost)

    # Print results
    '''
    print("Epoch:", '%04d' % (epoch + 1), "train_loss=", "{:.5f}".format(outs[1]),
          "train_acc=", "{:.5f}".format(outs[2]), "val_loss=", "{:.5f}".format(cost),
          "val_acc=", "{:.5f}".format(acc), "time=", "{:.5f}".format(time.time() - t))
    '''
    if epoch > FLAGS.early_stopping and cost_val[-1] > np.mean(cost_val[-(FLAGS.early_stopping+1):-1]):
        print("Early stopping...")
        break

    if epoch>0:
        if acc > queue_acc:
            try:
                os.remove(queue_ckpt)
            except:
                pass
            saver.save(sess, queue_ckpt)
            queue_acc = acc
            queue_time = 0
            best_epoch = epoch
        elif acc == queue_acc:
            queue_time += 1
            if np.random.random() < (1 / (queue_time + 1)):
                try:
                    os.remove(queue_ckpt)
                except:
                    pass
                saver.save(sess, queue_ckpt)
                best_epoch = epoch
                # saver.save()
        
    
    
    

# print("Optimization Finished!")
# Testing
#choose the final model
print('best epoch is ',best_epoch)

saver.restore(sess,queue_ckpt)
#evaluate the final model
train_cost, train_acc, train_duration = evaluate(features, support, y_train, train_mask, placeholders)
valid_cost, valid_acc, valid_duration = evaluate(features, support, y_val, val_mask, placeholders)
test_cost, test_acc, test_duration = evaluate(features, support, y_test, test_mask, placeholders)
    



final_train_acc=train_acc
final_valid_acc=valid_acc
final_test_acc=test_acc



resultFile=open(os.path.join(labelpath,name),'a')
resultFile.write(str(final_train_acc)+' '+str(final_valid_acc)+' '+str(final_test_acc)+'\n')

resultFile.close()

errlog.close()
