#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Import useful packages
from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

# Hide the Configuration and Warnings
import os
#os.environ["TF_CPP_MIN_LOG_LEVEL"] = '3'

import tensorflow as tf
import numpy as np
import pandas as pd
from scipy import sparse
from Models.Network.lib_for_GCN import ResGCN_Model, graph, coarsening
from Models.DatasetAPI.DataLoader import DatasetLoader
from data.gnn.make_dataset import make_dataset
# Model Name
Model = 'Residual_Graph_Convolutional_Neural_Network'

# Clear all the stack and use GPU resources as much as possible
tf.reset_default_graph()
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)

# Your Dataset Location, for example EEG-Motor-Movement-Imagery-Dataset
# The CSV file should be named as training_set.csv, training_label.csv, test_set.csv, and test_label.csv
levels=5
group = '02'
sampRat=250
person=[1,2]
cl=[17, 18, 19, 20, 21, 26, 27, 28, 29, 30, 35, 36, 37, 38, 39]
reName=make_dataset(remake=1,group=group,person=person,path='./data/',sampRat=sampRat)
DIR = './data/gnn/'+reName+'/'
SAVE = 'Saved_Files/' + Model + '/'
if not os.path.exists(SAVE):  # If the SAVE folder doesn't exist, create one
    os.mkdir(SAVE)

# Load the dataset, here it uses one-hot representation for labels
train_data, train_labels, test_data, test_labels = DatasetLoader(DIR=DIR)
train_labels=np.squeeze(train_labels)
test_labels=np.squeeze(test_labels)
# Read the Adjacency matrix
Adjacency_Matrix = np.load(DIR + 'Adjacency_Matrix.npy').astype('float32')
Adjacency_Matrix = sparse.csr_matrix(Adjacency_Matrix)

# This is the coarsen levels, you can definitely change the level to observe the difference
graphs, perm = coarsening.coarsen(Adjacency_Matrix, levels=levels, self_connections=False)
X_train = coarsening.perm_data(train_data, perm)
X_test  = coarsening.perm_data(test_data,  perm)

# Obtain the Graph Laplacian
L = [graph.laplacian(Adjacency_Matrix, normalized=True) for Adjacency_Matrix in graphs]

# Hyper-parameters
params = dict()
params['dir_name']       = Model
params['num_epochs']     = 50
params['batch_size']     = 1024
params['eval_frequency'] = 100

# Building blocks.
params['filter'] = 'chebyshev5'
params['brelu']  = 'b2relu'
params['pool']   = 'mpool1'

# Architecture.
params['F'] = [16, 16, 16, 32, 32, 32, 64, 64, 64, 128, 128, 128, 256, 256, 256, 512, 512, 512]    # Number of graph convolutional filters.
params['K'] = [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3]                               # Polynomial orders.
params['p'] = [1, 1, 2, 1, 1, 2, 1, 1, 2, 1, 1, 2, 1, 1, 2, 1, 1, 2]                               # Pooling sizes.
#params['F'] = params['F'][0:(levels+1)*3]
#params['K'] = params['K'][0:(levels+1)*3]
#params['p'] = params['p'][0:(levels+1)*3]
params['M'] = [3]                                                                                  # Output dimensionality of fully connected layers.

# Optimization.
params['regularization'] = 0.001  # L2 regularization
params['dropout']        = 0.50   # Dropout rate
params['learning_rate']  = 0.001  # Learning rate
params['decay_rate']     = 1      # Learning rate Decay == 1 means no Decay
params['momentum']       = 0      # momentum == 0 means Use Adam Optimizer
params['decay_steps']    = np.shape(train_data)[0] / params['batch_size']

# Train model
f=open(DIR+'/record.txt','a')
f.write('Number of graph convolutional filters.'+str(params['F']))
f.close()
model = ResGCN_Model.cgcnn(L, **params)
accuracy, loss, t_step = model.fit(X_train, train_labels, X_test, test_labels)