text:
"""
@author: Stéphane Clinchant

"""

import math
import random
import warnings

import numpy as np
import tensorflow as tf
try:
    tf.train.AdamOptimizer
except AttributeError:
    import tensorflow.compat.v1 as tf
    tf.disable_v2_behavior()
import sklearn.metrics
#TODO Clean this
# Animesh commented this line out from gcn.gcn_datasets import GCNDataset

#  from gcn.gcn_datasets import GCNDataset

try:
    from . import gcn_datasets
except ImportError:
    import gcn_datasets

from common.trace import traceln

def init_glorot(shape, name=None):
    """Glorot & Bengio (AISTATS 2010) init."""
    init_range = np.sqrt(6.0/(shape[0]+shape[1]))
    initial = tf.random_uniform(shape, minval=-init_range, maxval=init_range, dtype=tf.float32)
    return tf.Variable(initial, name=name)

def init_normal(shape,stddev,name=None):
    initial=tf.random_normal(shape, mean=0.0, stddev=stddev, dtype=np.float32)
    return tf.Variable(initial, name=name)

class MultiGraphNN(object):
    '''
    Abstract Class for a Neural Net learned on a graph list

    '''
    def train_lG(self,session,gcn_graph_train):
        '''
        Train an a list of graph
        :param session:
        :param gcn_graph_train:
        :return:
        '''
        for g in gcn_graph_train:
            self.train(session, g, n_iter=1)

    def test_lG(self, session, gcn_graph_test, verbose=True):
        '''
        Test on a list of Graph
        :param session:
        :param gcn_graph_test:
        :return:
        '''
        acc_tp = np.float64(0.0)
        nb_node_total = np.float64(0.0)
        mean_acc_test = []

        for g in gcn_graph_test:
            acc = self.test(session, g, verbose=False)
            mean_acc_test.append(acc)
            nb_node_total += g.X.shape[0]
            acc_tp += acc * g.X.shape[0]

        g_acc = np.mean(mean_acc_test)
        node_acc = acc_tp / nb_node_total

        if verbose:
            traceln('\t -- Mean Graph Accuracy', '%.4f' % g_acc)
            traceln('\t -- Mean Node  Accuracy', '%.4f' % node_acc)

        return g_acc,node_acc

    def predict_lG(self,session,gcn_graph_predict,verbose=True):
        '''
        Predict for a list of graph
        :param session:
        :param gcn_graph_test:
        :return:
        '''
        lY_pred=[]

        for g in gcn_graph_predict:
            gY_pred = self.predict(session, g, verbose=verbose)
            lY_pred.append(gY_pred)

        return lY_pred

    def predict_prob_lG(self, session, l_gcn_graph, verbose=True):
        '''
                Predict Probabilities for a list of graph
                :param session:
                :param l_gcn_graph:
                :return a list of predictions
                '''
        lY_pred = []

        for g in l_gcn_graph:
            gY_pred = self.prediction_prob(session, g, verbose=verbose)
            lY_pred.append(gY_pred)

        return lY_pred

    def get_nb_params(self):
        total_parameters = 0
        for variable in tf.trainable_variables():
            # shape is an array of tf.Dimension
            shape = variable.get_shape()
            #traceln(shape)
            #traceln(len(shape))
            variable_parameters = 1
            for dim in shape:
                #traceln(dim)
                variable_parameters *= dim.value
            #traceln(variable_parameters)
            total_parameters += variable_parameters
        return total_parameters

    def train_with_validation_set(self,session,graph_train,graph_val,max_iter,eval_iter=10,patience=7,graph_test=None,save_model_path=None):
        '''
        Implements training with a validation set
        The model is trained and accuracy is measure on a validation sets
        In addition, the model can be save and one can perform early stopping thanks to the patience argument

        :param session:
        :param graph_train: the list of graph to train on
        :param graph_val:   the list of graph used for validation
        :param max_iter:  maximum number of epochs
        :param eval_iter: evaluate every eval_iter
        :param patience: stopped training if accuracy is not improved on the validation set after patience_value
        :param graph_test: Optional. If a test set is provided, then accuracy on the test set is reported
        :param save_model_path: checkpoints filename to save the model.
        :return: A Dictionary with training accuracies, validations accuracies and test accuracies if any, and the Wedge parameters
        '''
        best_val_acc=0.0
        wait=0
        stop_training=False
        stopped_iter=max_iter
        train_accuracies=[]
        validation_accuracies=[]
        test_accuracies=[]
        conf_mat=[]

        start_monitoring_val_acc=False

        for i in range(max_iter):
            if stop_training:
                break

            if i % eval_iter == 0:
                traceln('\n -- Epoch ', i,' Patience ', wait)
                _, tr_acc = self.test_lG(session, graph_train, verbose=False)
                traceln(' Train Acc ', '%.4f' % tr_acc)
                train_accuracies.append(tr_acc)

                _, node_acc = self.test_lG(session, graph_val, verbose=False)
                traceln(' --  Valid Acc ', '%.4f' % node_acc)
                validation_accuracies.append(node_acc)

                if save_model_path:
                    save_path = self.saver.save(session, save_model_path, global_step=i)

                if graph_test:
                    test_graph_acc,test_acc = self.test_lG(session, graph_test, verbose=False)
                    traceln(' --   Test Acc ', '%.4f' % test_acc,'  %.4f' % test_graph_acc)
                    test_accuracies.append(test_acc)

                if node_acc > best_val_acc:
                    best_val_acc = node_acc
                    wait = 0
                else:
                    if wait >= patience:
                        stopped_iter = i
                        stop_training = True
                    wait += 1
            else:
                random.shuffle(graph_train)
                for g in graph_train:
                    self.train(session, g, n_iter=1)
        #Final Save
        traceln(' -- Stopped Model Training after : ',stopped_iter)
        traceln(' -- Validation Accuracies : ',['%.4f' % (100*sx) for sx in validation_accuracies])

        #traceln('Final Training Accuracy')
        _,node_train_acc = self.test_lG(session, graph_train)
        traceln(' -- Final Training Accuracy','%.4f' % node_train_acc)

        traceln(' -- Final Valid Acc')
        self.test_lG(session, graph_val)

        R = {}
        R['train_acc']          = train_accuracies
        R['val_acc']            = validation_accuracies
        R['test_acc']           = test_accuracies
        R['stopped_iter']       = stopped_iter
        R['confusion_matrix']   = conf_mat

        #R['W_edge'] =self.get_Wedge(session)
        if graph_test:

            _, final_test_acc = self.test_lG(session, graph_test)
            traceln(' -- Final Test Acc','%.4f' % final_test_acc)
            R['final_test_acc'] = final_test_acc

            val = R['val_acc']
            traceln(' -- Validation scores', val)

            epoch_index = np.argmax(val)
            traceln(' -- Best performance on val set: Epoch', epoch_index,val[epoch_index])
            traceln(' -- Test Performance from val', test_accuracies[epoch_index])

        return R

class EnsembleGraphNN(MultiGraphNN):
    '''
    An ensemble of Graph NN Models
    Construction Outside of class
    '''

    def __init__(self, graph_nn_models):
        self.models = graph_nn_models

    def train_lG(self, session, gcn_graph_train):
        '''
        Train an a list of graph
        :param session:
        :param gcn_graph_train:
        :return:
        '''
        for m in self.models:
            m.train_lG(session, gcn_graph_train)

    def test_lG(self, session, gcn_graph_test, verbose=True):
        '''
        Test on a list of Graph
        :param session:
        :param gcn_graph_test:
        :return:
        '''
        acc_tp = np.float64(0.0)
        nb_node_total = np.float64(0.0)
        mean_acc_test = []

        Y_pred=self.predict_lG(session,gcn_graph_test)
        Y_true =[g.Y for g in gcn_graph_test]
        Y_pred_node = np.vstack(Y_pred)

        node_acc  = sklearn.metrics.accuracy_score(np.argmax(np.vstack(Y_true),axis=1),np.argmax(Y_pred_node,axis=1))

        g_acc =-1
        #node_acc = acc_tp / nb_node_total

        if verbose:
            traceln(' -- Mean Graph Accuracy', '%.4f' % g_acc)
            traceln(' -- Mean Node  Accuracy', '%.4f' % node_acc)

        return g_acc, node_acc

    def predict_lG(self, session, gcn_graph_predict, verbose=True):
        '''
        Predict for a list of graph
        :param session:
        :param gcn_graph_test:
        :return:
        '''
        lY_pred = []

        #Seem Very Slow ... Here
        #I should predict for all graph
        nb_models = float(len(self.models))
        for g in gcn_graph_predict:
            #Average Proba Here
            g_pred=[]
            for m in self.models:
                gY_pred = m.prediction_prob(session, g, verbose=verbose)
                g_pred.append(gY_pred)
                #traceln(gY_pred)
            lY_pred.append(np.sum(g_pred,axis=0)/nb_models)
        return lY_pred

    def train_with_validation_set(self, session, graph_train, graph_val, max_iter, eval_iter=10, patience=7,
                                  graph_test=None, save_model_path=None):
        raise NotImplementedError

class Logit(MultiGraphNN):
    '''
    Logistic Regression for MultiGraph
    '''
    def __init__(self,node_dim,nb_classes,learning_rate=0.1,mu=0.1,node_indim=-1):
        self.node_dim=node_dim
        self.n_classes=nb_classes
        self.learning_rate=learning_rate
        self.activation=tf.nn.relu
        self.mu=mu
        self.optalg = tf.train.AdamOptimizer(self.learning_rate)
        self.stack_instead_add=False
        self.train_Wn0=True

        if node_indim==-1:
            self.node_indim=self.node_dim
        else:
            self.node_indim=node_indim

    def create_model(self):
        '''
        Create the tensorflow graph for the model
        :return:
        '''
        self.nb_node    = tf.placeholder(tf.int32,(), name='nb_node')
        self.node_input = tf.placeholder(tf.float32, [None, self.node_dim], name='X_')
        self.y_input    = tf.placeholder(tf.float32, [None, self.n_classes], name='Y')

        self.Wnode_layers=[]
        self.Bnode_layers=[]

        self.W_classif = tf.Variable(tf.random_uniform((self.node_indim, self.n_classes),
                                                           -1.0 / math.sqrt(self.node_dim),
                                                           1.0 / math.sqrt(self.node_dim)),
                                        name="W_classif",dtype=np.float32)
        self.B_classif = tf.Variable(tf.zeros([self.n_classes]), name='B_classif',dtype=np.float32)

        self.logits =tf.add(tf.matmul(self.node_input,self.W_classif),self.B_classif)
        cross_entropy_source = tf.nn.softmax_cross_entropy_with_logits(logits=self.logits, labels=self.y_input)

        # Global L2 Regulization
        self.loss = tf.reduce_mean(cross_entropy_source)  + self.mu * tf.nn.l2_loss(self.W_classif)

        self.pred = tf.argmax(tf.nn.softmax(self.logits), 1)
        self.correct_prediction = tf.equal(self.pred, tf.argmax(self.y_input, 1))
        self.accuracy = tf.reduce_mean(tf.cast(self.correct_prediction, tf.float32))

        self.grads_and_vars = self.optalg.compute_gradients(self.loss)

        self.train_step = self.optalg.apply_gradients(self.grads_and_vars)

        # Add ops to save and restore all the variables.
        self.init = tf.global_variables_initializer()
        self.saver= tf.train.Saver(max_to_keep=5)

        traceln(' -- Number of Params: ', self.get_nb_params())

    def save_model(self, session, model_filename):
        traceln(" -- Saving Model")
        save_path = self.saver.save(session, model_filename)

    def restore_model(self, session, model_filename):
        self.saver.restore(session, model_filename)
        traceln(" -- Model restored.")

    def train(self,session,graph,verbose=False,n_iter=1):
        '''
        Apply a train operation, ie sgd step for a single graph
        :param session:
        :param graph: a graph from GCN_Dataset
        :param verbose:
        :param n_iter: (default 1) number of steps to perform sgd for this graph
        :return:
        '''
        #TrainEvalSet Here
        for i in range(n_iter):
            #traceln('Train',X.shape,EA.shape)

            feed_batch = {

                self.nb_node: graph.X.shape[0],
                self.node_input: graph.X,
                self.y_input: graph.Y,
            }

            Ops =session.run([self.train_step,self.loss], feed_dict=feed_batch)
            if verbose:
                traceln(' -- Training Loss',Ops[1])

    def test(self,session,graph,verbose=True):
        '''
        Test return the loss and accuracy for the graph passed as argument
        :param session:
        :param graph:
        :param verbose:
        :return:
        '''

        feed_batch = {

            self.nb_node: graph.X.shape[0],
            self.node_input: graph.X,
            self.y_input: graph.Y,
        }

        Ops =session.run([self.loss,self.accuracy], feed_dict=feed_batch)
        if verbose:
            traceln(' -- Test Loss',Ops[0],' Test Accuracy:',Ops[1])
        return Ops[1]

    def predict(self,session,graph,verbose=True):
        '''
        Does the prediction
        :param session:
        :param graph:
        :param verbose:
        :return:
        '''
        feed_batch = {
            self.nb_node: graph.X.shape[0],
            self.node_input: graph.X,
        }
        Ops = session.run([self.pred], feed_dict=feed_batch)
        if verbose:
            traceln(' -- Got Prediction for:',Ops[0].shape)
        return Ops[0]

class EdgeConvNet(MultiGraphNN):
    '''
    Edge-GCN Model for a graph list
    '''

    #Variable ignored by the set_learning_options
    _setter_variables={
        "node_dim":True,"edge_dim":True,"nb_class":True,
        "num_layers":True,"lr":True,"mu":True,
        "node_indim":True,"nconv_edge":True,
        "nb_iter":True,"ratio_train_val":True}

    def __init__(self,node_dim,edge_dim,nb_classes,num_layers=1,learning_rate=0.1,mu=0.1,node_indim=-1,nconv_edge=1,
                 ):
        self.node_dim=node_dim
        self.edge_dim=edge_dim
        self.n_classes=nb_classes
        self.num_layers=num_layers
        self.learning_rate=learning_rate
        self.activation=tf.nn.tanh
        #self.activation=tf.nn.relu
        self.mu=mu
        self.optalg = tf.train.AdamOptimizer(self.learning_rate)
        self.stack_instead_add=False
        self.nconv_edge=nconv_edge
        self.residual_connection=False#deprecated
        self.shared_We = False#deprecated
        self.optim_mode=0 #deprecated
        self.init_fixed=False #ignore --for test purpose
        self.logit_convolve=False#ignore --for test purpose
        self.train_Wn0=True #ignore --for test purpose

        self.dropout_rate_edge_feat= 0.0
        self.dropout_rate_edge = 0.0
        self.dropout_rate_node = 0.0
        self.dropout_rate_H    = 0.0

        self.use_conv_weighted_avg=False
        self.use_edge_mlp=False
        self.edge_mlp_dim = 5
        self.sum_attention=False

        if node_indim==-1:
            self.node_indim=self.node_dim
        else:
            self.node_indim=node_indim

    def set_learning_options(self,dict_model_config):
        """
        Set all learning options that not directly accessible from the constructor

        :param kwargs:
        :return:
        """
        #traceln( -- dict_model_config)
        for attrname,val in dict_model_config.items():
            #We treat the activation function differently as we can not pickle/serialiaze python function
            if attrname=='activation_name':
                if val=='relu':
                    self.activation=tf.nn.relu
                elif val=='tanh':
                    self.activation=tf.nn.tanh
                else:
                    raise Exception('Invalid Activation Function')
            if attrname=='stack_instead_add' or attrname=='stack_convolutions':
                self.stack_instead_add=val
            if attrname not in self._setter_variables:
                try:
                    traceln(' -- set ',attrname,val)
                    setattr(self,attrname,val)
                except AttributeError:
                    warnings.warn("Ignored options for ECN"+attrname+':'+val)

    def fastconvolve(self,Wedge,Bedge,F,S,T,H,nconv,Sshape,nb_edge,dropout_p_edge,dropout_p_edge_feat,
                     stack=True, use_dropout=False,zwe=None,use_weighted_average=False,
                     use_edge_mlp=False,Wedge_mlp=None,Bedge_mlp=None,use_attention=False):
        '''

        :param Wedge: Parameter matrix for edge convolution, with hape (n_conv_edge,edge_dim)
        :param F: The Edge Feature Matrix
        :param S: the Source (Node,Edge) matrix in sparse format
        :param T: the Target (Node,Edge) matrix in sparse format
        :param H: The current node layer
        :param nconv: The numbder of edge convolutions.
        :param Sshape: The shapeof matrix S, and T
        :param nb_edge: The number of edges
        :param stack: whether to concat all the convolutions or add them
        :return: a tensor P of shape ( nconv, node_dim) if stack else P is [node_dim]
        '''

        #F is n_edge time nconv

        #TODO if stack is False we could simply sum,the convolutions and do S diag(sum)T
        #It would be faster

        #Drop convolution individually t
        if use_dropout:
        #if False:
            conv_dropout_ind = tf.nn.dropout(tf.ones([nconv], dtype=tf.float32), 1 - dropout_p_edge_feat)
            ND_conv = tf.diag(conv_dropout_ind)

            FW_ = tf.matmul(F, Wedge, transpose_b=True) + Bedge
            FW  = tf.matmul(FW_,ND_conv)

        elif use_edge_mlp:
            #Wedge mlp is a shared variable across layer which project edge in a lower dim
            FW0 = tf.nn.tanh( tf.matmul(F,Wedge_mlp) +Bedge_mlp )
            traceln(' -- FW0', FW0.get_shape())
            FW = tf.matmul(FW0, Wedge, transpose_b=True) + Bedge
            traceln(' -- FW', FW.get_shape())
        else:
            FW = tf.matmul(F, Wedge, transpose_b=True) + Bedge
            traceln(' -- FW', FW.get_shape())

        self.conv =tf.unstack(FW,axis=1)
        Cops=[]
        alphas=[]

        Tr = tf.SparseTensor(indices=T, values=tf.ones([nb_edge], dtype=tf.float32), dense_shape=[Sshape[1],Sshape[0]])
        Tr = tf.sparse_reorder(Tr)
        TP = tf.sparse_tensor_dense_matmul(Tr,H)

        if use_attention:
            attn_params = va = init_glorot([2, int(self.node_dim)])

        for i, cw in enumerate(self.conv):
            #SD= tf.SparseTensor(indices=S,values=cw,dense_shape=[nb_node,nb_edge])
            #Warning, pay attention to the ordering of edges
            if use_weighted_average:
                cw = zwe[i]*cw
            if use_dropout:
                cwd = tf.nn.dropout(cw, 1.0 -dropout_p_edge)
                SD = tf.SparseTensor(indices=S, values=cwd, dense_shape=Sshape)
            else:
                SD = tf.SparseTensor(indices=S, values=cw, dense_shape=Sshape)
            SD =tf.sparse_reorder(SD)
            #Does this dropout depends on the convolution ?
            #if use_dropout:
            #    SD = tf.nn.dropout(SD, 1.0 - dropout_p_edge)

            Hi =tf.sparse_tensor_dense_matmul(SD,TP)
            Cops.append(Hi)
            if use_attention:
                attn_val = tf.reduce_sum(tf.multiply(attn_params[0], H) + tf.multiply(attn_params[1], Hi), axis=1)
                alphas.append(attn_val)

        if stack is True:
            #If stack we concatenate all the different convolutions
            P=tf.concat(Cops,1)

        elif use_attention:
            alphas_s = tf.stack(alphas, axis=1)
            alphas_l = tf.nn.softmax(tf.nn.leaky_relu(alphas_s))
            #Not Clean to use the dropout for the edge feat
            #Is this dropout necessary Here ? could do without
            alphas_do =tf.nn.dropout(alphas_l,1 - dropout_p_edge_feat)
            wC = [tf.multiply(tf.expand_dims(tf.transpose(alphas_do)[i], 1), C) for i, C in enumerate(Cops)]
            P = tf.add_n(wC)

        else:
            #Else we take the mean
            P=1.0/(tf.cast(nconv,tf.float32))*tf.add_n(Cops)
            #traceln('p_add_n',P.get_shape())
        return P

    @staticmethod
    def logitconvolve_fixed(pY,Yt,A_indegree):
        '''
        Tentative Implement of a fixed logit convolve without taking into account edge features
        '''

        #warning we should test that Yt is column normalized
        pY_Yt = tf.matmul(pY,Yt,transpose_b=True)

        #TODO A is dense but shoudl be sparse ....
        P =tf.matmul(A_indegree,pY_Yt)
        return P

    def create_model_stack_convolutions(self):
        #Create All the Variables
        for i in range(self.num_layers - 1):
            if i == 0:
                Wnli = tf.Variable(
                    tf.random_uniform((self.node_dim * self.nconv_edge + self.node_dim, self.node_indim),
                                      -1.0 / math.sqrt(self.node_indim),
                                      1.0 / math.sqrt(self.node_indim)), name='Wnl', dtype=tf.float32)
            else:

                Wnli = tf.Variable(
                    tf.random_uniform((self.node_indim * self.nconv_edge + self.node_indim, self.node_indim),
                                          -1.0 / math.sqrt(self.node_indim),
                                      1.0 / math.sqrt(self.node_indim)), name='Wnl', dtype=tf.float32)
            traceln(' -- Wnli shape', Wnli.get_shape())
            Bnli = tf.Variable(tf.zeros([self.node_indim]), name='Bnl' + str(i), dtype=tf.float32)

            Weli = init_glorot([int(self.nconv_edge), int(self.edge_dim)], name='Wel_')
            # Weli = tf.Variable(tf.random_normal([int(self.nconv_edge), int(self.edge_dim)], mean=0.0, stddev=1.0),
            #                   dtype=np.float32, name='Wel_')
            Beli = tf.Variable(0.01 * tf.ones([self.nconv_edge]), name='Bel' + str(i), dtype=tf.float32)

            self.Wnode_layers.append(Wnli)
            self.Bnode_layers.append(Bnli)
            self.Wed_layers.append(Weli)
            self.Bed_layers.append(Beli)

        self.train_var.extend((self.Wnode_layers))
        self.train_var.extend((self.Wed_layers))

        self.Hnode_layers = []

        self.W_classif = tf.Variable(
                tf.random_uniform((self.node_indim * self.nconv_edge + self.node_indim, self.n_classes),
                                  -1.0 / math.sqrt(self.node_dim),
                                  1.0 / math.sqrt(self.node_dim)),
                name="W_classif", dtype=np.float32)

        self.B_classif = tf.Variable(tf.zeros([self.n_classes]), name='B_classif', dtype=np.float32)

        self.train_var.append((self.W_classif))
        self.train_var.append((self.B_classif))

        self.node_dropout_ind = tf.nn.dropout(tf.ones([self.nb_node], dtype=tf.float32), 1 - self.dropout_p_node)
        self.ND = tf.diag(self.node_dropout_ind)

        edge_dropout = self.dropout_rate_edge > 0.0 or self.dropout_rate_edge_feat > 0.0
        traceln(' -- Edge Dropout', edge_dropout, self.dropout_rate_edge, self.dropout_rate_edge_feat)
        if self.num_layers == 1:
            self.H = self.activation(tf.add(tf.matmul(self.node_input, self.Wnl0), self.Bnl0))
            self.hidden_layers = [self.H]
            traceln(" -- H shape", self.H.get_shape())

            P = self.fastconvolve(self.Wel0, self.Bel0, self.F, self.Ssparse, self.Tsparse, self.H, self.nconv_edge,
                                  self.Sshape, self.nb_edge,
                                  self.dropout_p_edge, self.dropout_p_edge_feat, stack=self.stack_instead_add,
                                  use_dropout=edge_dropout)

            Hp = tf.concat([self.H, P], 1)
            # Hp= P+self.H

            Hi = self.activation(Hp)
            # Hi_shape = Hi.get_shape()
            # traceln(Hi_shape)
            self.hidden_layers.append(Hi)

        elif self.num_layers > 1:

            if self.dropout_rate_node > 0.0:
                #H0 = self.activation(tf.matmul(self.ND, tf.add(tf.matmul(self.node_input, self.Wnl0), self.Bnl0)))
                H0 = tf.matmul(self.ND, tf.add(tf.matmul(self.node_input, self.Wnl0), self.Bnl0))
            else:
                #H0 = self.activation(tf.add(tf.matmul(self.node_input, self.Wnl0), self.Bnl0))
                H0 = tf.add(tf.matmul(self.node_input, self.Wnl0), self.Bnl0)

            self.Hnode_layers.append(H0)

            # TODO Default to fast convolve but we change update configs, train and test flags
            P = self.fastconvolve(self.Wel0, self.Bel0, self.F, self.Ssparse, self.Tsparse, H0, self.nconv_edge,
                                  self.Sshape, self.nb_edge,
                                  self.dropout_p_edge, self.dropout_p_edge_feat, stack=self.stack_instead_add,
                                  use_dropout=edge_dropout,
                                  )
            Hp = tf.concat([H0, P], 1)

            # TODO add activation Here.
            self.hidden_layers = [self.activation(Hp)]
            #self.hidden_layers = [Hp]

            for i in range(self.num_layers - 1):

                if self.dropout_rate_H > 0.0:
                    Hi_ = tf.nn.dropout(tf.matmul(self.hidden_layers[-1], self.Wnode_layers[i]) + self.Bnode_layers[i],
                                        1 - self.dropout_p_H)
                else:
                    Hi_ = tf.matmul(self.hidden_layers[-1], self.Wnode_layers[i]) + self.Bnode_layers[i]

                if self.residual_connection:
                    Hi_ = tf.add(Hi_, self.Hnode_layers[-1])

                self.Hnode_layers.append(Hi_)

                # traceln('Hi_shape', Hi_.get_shape())
                #traceln('Hi prevous shape', self.hidden_layers[-1].get_shape())

                P = self.fastconvolve(self.Wed_layers[i], self.Bed_layers[i], self.F, self.Ssparse, self.Tsparse, Hi_,
                                      self.nconv_edge, self.Sshape, self.nb_edge,
                                      self.dropout_p_edge, self.dropout_p_edge_feat, stack=self.stack_instead_add,
                                      use_dropout=edge_dropout,
                                      )
                Hp = tf.concat([Hi_, P], 1)
                Hi = self.activation(Hp)
                self.hidden_layers.append(Hi)

    def create_model_sum_convolutions(self):

        #self.Wed_layers.append(Wel0)
        for i in range(self.num_layers-1):

            if i==0:
                # Animesh's code
                Wnli = tf.Variable(tf.random_uniform((2 * self.node_dim, self.node_indim),
                #Wnli = tf.Variable(tf.random_uniform((2 * self.node_indim, self.node_indim),
                                                     -1.0 / math.sqrt(self.node_indim),
                                                     1.0 / math.sqrt(self.node_indim)), name='Wnl',
                                   dtype=tf.float32)
            else:
                Wnli =tf.Variable(tf.random_uniform( (2*self.node_indim, self.node_indim),
                                                               -1.0 / math.sqrt(self.node_indim),
                                                               1.0 / math.sqrt(self.node_indim)),name='Wnl',dtype=tf.float32)

            Bnli = tf.Variable(tf.zeros([self.node_indim]), name='Bnl'+str(i),dtype=tf.float32)

            Weli= init_glorot([int(self.nconv_edge), int(self.edge_dim)], name='Wel_')
            #Weli = tf.Variable(tf.random_normal([int(self.nconv_edge), int(self.edge_dim)], mean=0.0, stddev=1.0),
            #                   dtype=np.float32, name='Wel_')
            Beli = tf.Variable(0.01*tf.ones([self.nconv_edge]), name='Bel'+str(i),dtype=tf.float32)

            self.Wnode_layers.append(Wnli)
            self.Bnode_layers.append(Bnli)
            self.Wed_layers.append  (Weli)
            self.Bed_layers.append(Beli)

        self.train_var.extend((self.Wnode_layers))
        self.train_var.extend((self.Wed_layers))

        self.Hnode_layers=[]

        self.W_classif = tf.Variable(tf.random_uniform((2*self.node_indim, self.n_classes),
                                                           -1.0 / math.sqrt(self.node_dim),
                                                           1.0 / math.sqrt(self.node_dim)),
                                        name="W_classif",dtype=np.float32)
        self.B_classif = tf.Variable(tf.zeros([self.n_classes]), name='B_classif',dtype=np.float32)

        self.train_var.append((self.W_classif))
        self.train_var.append((self.B_classif))

        self.node_dropout_ind = tf.nn.dropout(tf.ones([self.nb_node], dtype=tf.float32), 1 - self.dropout_p_node)
        self.ND = tf.diag(self.node_dropout_ind)

        edge_dropout = self.dropout_rate_edge> 0.0 or self.dropout_rate_edge_feat > 0.0
        traceln(' -- Edge Dropout',edge_dropout, self.dropout_rate_edge,self.dropout_rate_edge_feat)
        if self.num_layers==1:
            self.H = self.activation(tf.add(tf.matmul(self.node_input, self.Wnl0), self.Bnl0))
            self.hidden_layers = [self.H]
            traceln(" -- H shape",self.H.get_shape())

            P = self.fastconvolve(self.Wel0,self.Bel0,self.F,self.Ssparse,self.Tsparse,self.H,self.nconv_edge,self.Sshape,self.nb_edge,
                                  self.dropout_p_edge,self.dropout_p_edge_feat,stack=self.stack_instead_add,use_dropout=edge_dropout,
                                  use_attention=self.sum_attention
                                  )

            Hp = tf.concat([self.H, P], 1)
            Hi=self.activation(Hp)
            self.hidden_layers.append(Hi)

        elif self.num_layers>1:

            if self.dropout_rate_node>0.0:
                H0 = self.activation(tf.matmul(self.ND,tf.add(tf.matmul(self.node_input,self.Wnl0), self.Bnl0)))
            else:
                H0 = self.activation(tf.add(tf.matmul(self.node_input,self.Wnl0),self.Bnl0))

            self.Hnode_layers.append(H0)

            #TODO Default to fast convolve but we change update configs, train and test flags
            P = self.fastconvolve(self.Wel0,self.Bel0, self.F, self.Ssparse, self.Tsparse, H0, self.nconv_edge, self.Sshape,self.nb_edge,
                                  self.dropout_p_edge,self.dropout_p_edge_feat, stack=self.stack_instead_add, use_dropout=edge_dropout,
                                  use_attention=self.sum_attention
                                  )

            if self.use_conv_weighted_avg:
                Hp = self.zH[0] * H0 + P
            else:
                Hp = tf.concat([H0, P], 1)

            #TODO add activation Here.
            #self.hidden_layers = [self.activation(Hp)]
            self.hidden_layers = [Hp]

            for i in range(self.num_layers-1):

                if self.dropout_rate_H > 0.0:
                    Hi_ = tf.nn.dropout(tf.matmul(self.hidden_layers[-1], self.Wnode_layers[i]) + self.Bnode_layers[i], 1-self.dropout_p_H)
                else:
                    Hi_ = tf.matmul(self.hidden_layers[-1], self.Wnode_layers[i]) + self.Bnode_layers[i]

                if self.residual_connection:
                    Hi_= tf.add(Hi_,self.Hnode_layers[-1])

                self.Hnode_layers.append(Hi_)

                #traceln('Hi_shape',Hi_.get_shape())
                # traceln('Hi prevous shape',self.hidden_layers[-1].get_shape())

                P = self.fastconvolve(self.Wed_layers[i],self.Bed_layers[i], self.F, self.Ssparse, self.Tsparse, Hi_, self.nconv_edge,self.Sshape, self.nb_edge,
                                      self.dropout_p_edge,self.dropout_p_edge_feat, stack=self.stack_instead_add, use_dropout=edge_dropout
                                      )

                Hp = tf.concat([Hi_, P], 1)
                Hi = self.activation(Hp)
                self.hidden_layers.append(Hi)

    def create_model(self):
        '''
        Create the tensorflow graph for the model
        :return:
        '''
        self.nb_node = tf.placeholder(tf.int32, (), name='nb_node')
        self.nb_edge = tf.placeholder(tf.int32, (), name='nb_edge')
        self.node_input = tf.placeholder(tf.float32, [None, self.node_dim], name='X_')
        self.y_input = tf.placeholder(tf.float32, [None, self.n_classes], name='Y')
        self.dropout_p_H = tf.placeholder(tf.float32, (), name='dropout_prob_H')
        self.dropout_p_node = tf.placeholder(tf.float32, (), name='dropout_prob_N')
        self.dropout_p_edge = tf.placeholder(tf.float32, (), name='dropout_prob_edges')
        self.dropout_p_edge_feat = tf.placeholder(tf.float32, (), name='dropout_prob_edgefeat')
        self.S = tf.placeholder(tf.float32, name='S')
        self.Ssparse = tf.placeholder(tf.int64, name='Ssparse')  # indices
        self.Sshape = tf.placeholder(tf.int64, name='Sshape')  # indices

        self.T = tf.placeholder(tf.float32, [None, None], name='T')
        self.Tsparse = tf.placeholder(tf.int64, name='Tsparse')

        self.F = tf.placeholder(tf.float32, [None, None], name='F')

        std_dev_in = float(1.0 / float(self.node_dim))

        self.Wnode_layers = []
        self.Bnode_layers = []
        self.Wed_layers = []
        self.Bed_layers = []
        self.zed_layers = []

        self.Wedge_mlp_layers = []
        # Should Project edge as well ...
        self.train_var = []

        # if self.node_indim!=self.node_dim:
        #    Wnl0 = tf.Variable(tf.random_uniform((self.node_dim, self.node_indim),
        #                                                           -1.0 / math.sqrt(self.node_dim),
        #                                                           1.0 / math.sqrt(self.node_dim)),name='Wnl0',dtype=tf.float32)
        #
        self.Wnl0 = tf.Variable(tf.eye(self.node_dim), name='Wnl0', dtype=tf.float32, trainable=self.train_Wn0)

        self.Bnl0 = tf.Variable(tf.zeros([self.node_dim]), name='Bnl0', dtype=tf.float32)

        if self.init_fixed: #For testing Purposes
            self.Wel0 = tf.Variable(100 * tf.ones([int(self.nconv_edge), int(self.edge_dim)]), name='Wel0',
                                    dtype=tf.float32)
        # self.Wel0 =tf.Variable(tf.random_normal([int(self.nconv_edge),int(self.edge_dim)],mean=0.0,stddev=1.0), dtype=np.float32, name='Wel0')
        self.Wel0 = init_glorot([int(self.nconv_edge), int(self.edge_dim)], name='Wel0')
        self.Bel0 = tf.Variable(0.01 * tf.ones([self.nconv_edge]), name='Bel0', dtype=tf.float32)

        traceln(' -- Wel0', self.Wel0.get_shape())
        self.train_var.extend([self.Wnl0, self.Bnl0])
        self.train_var.append(self.Wel0)

        if self.stack_instead_add:
            self.create_model_stack_convolutions()

        else:
            self.create_model_sum_convolutions()

        self.logits = tf.add(tf.matmul(self.hidden_layers[-1], self.W_classif), self.B_classif)
        cross_entropy_source = tf.nn.softmax_cross_entropy_with_logits(logits=self.logits, labels=self.y_input)
        # Global L2 Regulization
        self.loss = tf.reduce_mean(cross_entropy_source) + self.mu * tf.nn.l2_loss(self.W_classif)

        self.predict_proba = tf.nn.softmax(self.logits)
        self.pred = tf.argmax(tf.nn.softmax(self.logits), 1)
        self.correct_prediction = tf.equal(self.pred, tf.argmax(self.y_input, 1))
        self.accuracy = tf.reduce_mean(tf.cast(self.correct_prediction, tf.float32))

        self.grads_and_vars = self.optalg.compute_gradients(self.loss)

        self.gv_Gn = []
        self.train_step = self.optalg.apply_gradients(self.grads_and_vars)

        # Add ops to save and restore all the variables.
        self.init = tf.global_variables_initializer()
        self.saver = tf.train.Saver(max_to_keep=0)

        traceln(' -- Number of Params: ', self.get_nb_params())

    def create_model_old(self):
        '''
        Create the tensorflow graph for the model
        :return:
        '''
        self.nb_node    = tf.placeholder(tf.int32,(), name='nb_node')
        self.nb_edge = tf.placeholder(tf.int32, (), name='nb_edge')
        self.node_input = tf.placeholder(tf.float32, [None, self.node_dim], name='X_')
        self.y_input    = tf.placeholder(tf.float32, [None, self.n_classes], name='Y')
        #self.EA_input   = tf.placeholder(tf.float32, name='EA_input')
        #self.NA_input   = tf.placeholder(tf.float32, name='NA_input')
        self.dropout_p_H    = tf.placeholder(tf.float32,(), name='dropout_prob_H')
        self.dropout_p_node = tf.placeholder(tf.float32, (), name='dropout_prob_N')
        self.dropout_p_edge = tf.placeholder(tf.float32, (), name='dropout_prob_edges')
        self.dropout_p_edge_feat = tf.placeholder(tf.float32, (), name='dropout_prob_edgefeat')
        self.S          = tf.placeholder(tf.float32, name='S')
        self.Ssparse    = tf.placeholder(tf.int64, name='Ssparse') #indices
        self.Sshape     = tf.placeholder(tf.int64, name='Sshape') #indices

        self.T          = tf.placeholder(tf.float32,[None,None], name='T')
        self.Tsparse    = tf.placeholder(tf.int64, name='Tsparse')

        #self.S_indice = tf.placeholder(tf.in, [None, None], name='S')

        self.F          = tf.placeholder(tf.float32,[None,None], name='F')

        #self.NA_indegree = tf.placeholder(tf.float32, name='NA_indegree')

        std_dev_in=float(1.0/ float(self.node_dim))

        self.Wnode_layers=[]
        self.Bnode_layers=[]
        self.Wed_layers=[]
        self.Bed_layers=[]
        #REFACT1 self.zed_layers = []

        #REFACT1 self.Wedge_mlp_layers=[]
        #Should Project edge as well ...
        self.train_var=[]

        #if self.node_indim!=self.node_dim:
        #    Wnl0 = tf.Variable(tf.random_uniform((self.node_dim, self.node_indim),
        #                                                           -1.0 / math.sqrt(self.node_dim),
        #                                                           1.0 / math.sqrt(self.node_dim)),name='Wnl0',dtype=tf.float32)
        #else:
        self.Wnl0 = tf.Variable(tf.eye(self.node_dim),name='Wnl0',dtype=tf.float32,trainable=self.train_Wn0)

        self.Bnl0 = tf.Variable(tf.zeros([self.node_dim]), name='Bnl0',dtype=tf.float32)

        #self.Wel0 =tf.Variable(tf.random_normal([int(self.nconv_edge),int(self.edge_dim)],mean=0.0,stddev=1.0), dtype=np.float32, name='Wel0')
        if self.init_fixed:
            self.Wel0 = tf.Variable(100*tf.ones([int(self.nconv_edge),int(self.edge_dim)]), name='Wel0',dtype=tf.float32)

        elif self.use_edge_mlp:
            self.Wel0 = init_glorot([int(self.nconv_edge), int(self.edge_mlp_dim)], name='Wel0')
        else:
            self.Wel0 = init_glorot([int(self.nconv_edge),int(self.edge_dim)],name='Wel0')

        self.Bel0 = tf.Variable(0.01*tf.ones([self.nconv_edge]), name='Bel0' , dtype=tf.float32)

        #RF self.zel0 = tf.Variable(tf.ones([self.nconv_edge]), name='zel0' , dtype=tf.float32)
        #RF self.zH = tf.Variable(tf.ones([self.num_layers]),name='zH',dtype=tf.float32)

        traceln(' -- Wel0',self.Wel0.get_shape())
        self.train_var.extend([self.Wnl0,self.Bnl0])
        self.train_var.append(self.Wel0)

        #Parameter for convolving the logits
        ''' REFACT1
        if self.logit_convolve:
            #self.Wel_logits = init_glorot([int(self.nconv_edge),int(self.edge_dim)],name='Wel_logit')
            #self.Belg = tf.Variable(tf.zeros( [int(self.nconv_edge)]), name='Belogit' , dtype=tf.float32)
            self.Wel_logits = tf.Variable(tf.zeros([int(1),int(self.edge_dim)]), name='Wlogit0',dtype=tf.float32,trainable=False)
            self.Belg = tf.Variable(tf.ones( [int(1)]), name='Belogit' , dtype=tf.float32)
            #self.logits_Transition = 1.0*tf.Variable(tf.ones([int(self.n_classes) , int(self.n_classes)]), name='logit_Transition')
            self.logits_Transition=init_glorot([int(self.n_classes), int(self.n_classes)], name='Wel_')

        self.Wmlp_edge_0= init_glorot([int(self.edge_dim), int(self.edge_mlp_dim)], name='Wedge_mlp')
        self.Bmlp_edge_0= tf.Variable(tf.ones([self.edge_mlp_dim]),name='Wedge_mlp',dtype=tf.float32)
        '''

        #self.Wed_layers.append(Wel0)
        for i in range(self.num_layers-1):
            if self.stack_instead_add:
                if i==0:
                    Wnli = tf.Variable(
                        tf.random_uniform((self.node_dim * self.nconv_edge + self.node_dim, self.node_indim),
                                          -1.0 / math.sqrt(self.node_indim),
                                          1.0 / math.sqrt(self.node_indim)), name='Wnl', dtype=tf.float32)
                else:

                    Wnli =tf.Variable(tf.random_uniform( (self.node_indim*self.nconv_edge+self.node_indim, self.node_indim),
                                                               -1.0 / math.sqrt(self.node_indim),
                                                               1.0 / math.sqrt(self.node_indim)),name='Wnl',dtype=tf.float32)
                traceln(' -- Wnli shape',Wnli.get_shape())

            elif self.use_conv_weighted_avg:
                Wnli = tf.Variable(
                    tf.random_uniform((self.node_indim, self.node_indim),
                                      -1.0 / math.sqrt(self.node_indim),
                                      1.0 / math.sqrt(self.node_indim)), name='Wnl', dtype=tf.float32)

                #Wnli = tf.eye(self.node_dim,dtype=tf.float32)

                traceln(' -- Wnli shape', Wnli.get_shape())

            else:
                if i==0:
                    Wnli = tf.Variable(tf.random_uniform((2 * self.node_indim, self.node_indim),
                                                         -1.0 / math.sqrt(self.node_indim),
                                                         1.0 / math.sqrt(self.node_indim)), name='Wnl',
                                       dtype=tf.float32)
                else:
                    Wnli =tf.Variable(tf.random_uniform( (2*self.node_indim, self.node_indim),
                                                                   -1.0 / math.sqrt(self.node_indim),
                                                                   1.0 / math.sqrt(self.node_indim)),name='Wnl',dtype=tf.float32)

            Bnli = tf.Variable(tf.zeros([self.node_indim]), name='Bnl'+str(i),dtype=tf.float32)

            #Weli = tf.Variable(tf.ones([int(self.nconv_edge),int(self.edge_dim)],dtype=tf.float32))

            if self.use_edge_mlp:
                #self.Wel0 = init_glorot([int(self.nconv_edge), int(self.edge_mlp_dim)], name='Wel0')
                Weli = init_glorot([int(self.nconv_edge), int(self.edge_mlp_dim)], name='Wel_')
                Beli = tf.Variable(0.01 * tf.ones([self.nconv_edge]), name='Bel' + str(i), dtype=tf.float32)

                #RF Wmlp_edge_i = init_glorot([int(self.edge_dim), int(self.edge_mlp_dim)], name='Wedge_mlp'+str(i))
                #RF Bmlp_edge_i = tf.Variable(tf.ones([self.edge_mlp_dim]), name='Bedge_mlp'+str(i), dtype=tf.float32)
                #RF self.Wedge_mlp_layers.append(Wmlp_edge_i)

            else:
                Weli= init_glorot([int(self.nconv_edge), int(self.edge_dim)], name='Wel_')
                #Weli = tf.Variable(tf.random_normal([int(self.nconv_edge), int(self.edge_dim)], mean=0.0, stddev=1.0),
                #                   dtype=np.float32, name='Wel_')
                Beli = tf.Variable(0.01*tf.ones([self.nconv_edge]), name='Bel'+str(i),dtype=tf.float32)

                #RF Wmlp_edge_i = init_glorot([int(self.edge_dim), int(self.edge_mlp_dim)], name='Wedge_mlp' + str(i))
                #RF Bmlp_edge_i = tf.Variable(tf.ones([self.edge_mlp_dim]), name='Bedge_mlp' + str(i), dtype=tf.float32)
                #RF self.Wedge_mlp_layers.append(Wmlp_edge_i)

            #zeli = tf.Variable(tf.ones([self.nconv_edge]),name='zel'+str(i),dtype=tf.float32)

            self.Wnode_layers.append(Wnli)
            self.Bnode_layers.append(Bnli)
            self.Wed_layers.append  (Weli)
            self.Bed_layers.append(Beli)
            #self.zed_layers.append(zeli)

        self.train_var.extend((self.Wnode_layers))
        self.train_var.extend((self.Wed_layers))

        self.Hnode_layers=[]

        #TODO Do we project the firt layer or not ?
        # Initialize the weights and biases for a simple one full connected network
        if self.stack_instead_add:
            self.W_classif = tf.Variable(tf.random_uniform((self.node_indim*self.nconv_edge+self.node_indim, self.n_classes),
                                                           -1.0 / math.sqrt(self.node_dim),
                                                           1.0 / math.sqrt(self.node_dim)),
                                        name="W_classif",dtype=np.float32)

        elif self.use_conv_weighted_avg:
            self.W_classif = tf.Variable(tf.random_uniform((self.node_indim, self.n_classes),
                                                           -1.0 / math.sqrt(self.node_dim),
                                                           1.0 / math.sqrt(self.node_dim)),
                                         name="W_classif", dtype=np.float32)
        else:
            self.W_classif = tf.Variable(tf.random_uniform((2*self.node_indim, self.n_classes),
                                                           -1.0 / math.sqrt(self.node_dim),
                                                           1.0 / math.sqrt(self.node_dim)),
                                        name="W_classif",dtype=np.float32)
        self.B_classif = tf.Variable(tf.zeros([self.n_classes]), name='B_classif',dtype=np.float32)

        self.train_var.append((self.W_classif))
        self.train_var.append((self.B_classif))

        #Use for true add
        #I = tf.eye(self.nb_node)

        self.node_dropout_ind = tf.nn.dropout(tf.ones([self.nb_node], dtype=tf.float32), 1 - self.dropout_p_node)
        self.ND = tf.diag(self.node_dropout_ind)

        edge_dropout = self.dropout_rate_edge> 0.0 or self.dropout_rate_edge_feat > 0.0
        traceln(' -- Edge Dropout',edge_dropout, self.dropout_rate_edge,self.dropout_rate_edge_feat)
        if self.num_layers==1:
            self.H = self.activation(tf.add(tf.matmul(self.node_input, self.Wnl0), self.Bnl0))
            self.hidden_layers = [self.H]
            traceln("H shape",self.H.get_shape())

            P = self.fastconvolve(self.Wel0,self.Bel0,self.F,self.Ssparse,self.Tsparse,self.H,self.nconv_edge,self.Sshape,self.nb_edge,
                                  self.dropout_p_edge,self.dropout_p_edge_feat,stack=self.stack_instead_add,use_dropout=edge_dropout)

            Hp = tf.concat([self.H, P], 1)
            #Hp= P+self.H

            Hi=self.activation(Hp)
            #Hi_shape = Hi.get_shape()
            #traceln(Hi_shape)
            self.hidden_layers.append(Hi)

        elif self.num_layers>1:

            if self.dropout_rate_node>0.0:
                H0 = self.activation(tf.matmul(self.ND,tf.add(tf.matmul(self.node_input, self.Wnl0), self.Bnl0)))
            else:
                H0 = self.activation(tf.add(tf.matmul(self.node_input,self.Wnl0),self.Bnl0))

            self.Hnode_layers.append(H0)

            #TODO Default to fast convolve but we change update configs, train and test flags
            P = self.fastconvolve(self.Wel0,self.Bel0, self.F, self.Ssparse, self.Tsparse, H0, self.nconv_edge, self.Sshape,self.nb_edge,
                                  self.dropout_p_edge,self.dropout_p_edge_feat, stack=self.stack_instead_add, use_dropout=edge_dropout,
                                  )
                                  #RF zwe=self.zel0,
                                  #RF use_weighted_average=self.use_conv_weighted_avg,
                                  #RF use_edge_mlp=self.use_edge_mlp,
                                  #RFWedge_mlp=self.Wmlp_edge_0,
                                  #RF Bedge_mlp=self.Bmlp_edge_0)

            if self.use_conv_weighted_avg:
                Hp = self.zH[0] * H0 + P
            else:
                Hp = tf.concat([H0, P], 1)

            #TODO add activation Here.
            #self.hidden_layers = [self.activation(Hp)]
            self.hidden_layers = [Hp]

            for i in range(self.num_layers-1):

                if self.dropout_rate_H > 0.0:
                    Hi_ = tf.nn.dropout(tf.matmul(self.hidden_layers[-1], self.Wnode_layers[i]) + self.Bnode_layers[i], 1-self.dropout_p_H)
                else:
                    Hi_ = tf.matmul(self.hidden_layers[-1], self.Wnode_layers[i]) + self.Bnode_layers[i]

                if self.residual_connection:
                    Hi_= tf.add(Hi_,self.Hnode_layers[-1])

                self.Hnode_layers.append(Hi_)

                # traceln(' -- Hi_shape',Hi_.get_shape())
                # traceln(' -- Hi prevous shape',self.hidden_layers[-1].get_shape())

                P = self.fastconvolve(self.Wed_layers[i],self.Bed_layers[i], self.F, self.Ssparse, self.Tsparse, Hi_, self.nconv_edge,self.Sshape, self.nb_edge,
                                      self.dropout_p_edge,self.dropout_p_edge_feat, stack=self.stack_instead_add, use_dropout=edge_dropout,
                                      )
                                      # zwe=self.zed_layers[i],
                                      # use_weighted_average=self.use_conv_weighted_avg,
                                      # use_edge_mlp=self.use_edge_mlp,
                                      # Wedge_mlp=self.Wedge_mlp_layers[i],
                                      #RF Bedge_mlp=self.Bmlp_edge_0)

                if self.use_conv_weighted_avg:
                    Hp = self.zH[i+1]* Hi_ + P

                else:
                    Hp = tf.concat([Hi_, P], 1)

                Hi = self.activation(Hp)

                self.hidden_layers.append(Hi)

        self.logits =tf.add(tf.matmul(self.hidden_layers[-1],self.W_classif),self.B_classif)
        cross_entropy_source = tf.nn.softmax_cross_entropy_with_logits(logits=self.logits, labels=self.y_input)
        # Global L2 Regulization
        self.loss = tf.reduce_mean(cross_entropy_source)  + self.mu * tf.nn.l2_loss(self.W_classif)

        self.pred = tf.argmax(tf.nn.softmax(self.logits), 1)
        self.predict_proba = tf.nn.softmax(self.logits)
        self.correct_prediction = tf.equal(self.pred, tf.argmax(self.y_input, 1))
        self.accuracy = tf.reduce_mean(tf.cast(self.correct_prediction, tf.float32))

        self.grads_and_vars = self.optalg.compute_gradients(self.loss)

        self.gv_Gn=[]
        #TODO Experiment with gradient noise
        #if self.stack_instead_add:
        #    for grad, var in self.grads_and_vars:
        #        traceln(grad,var)
        #        if grad is not None:
        #            self.gv_Gn.append( ( tf.add(grad, tf.random_normal(tf.shape(grad), stddev=0.00001)),var)    )

        #self.gv_Gn = [(tf.add(grad, tf.random_normal(tf.shape(grad), stddev=0.00001)), val) for grad, val in self.grads_and_vars if  is not None]

        self.train_step = self.optalg.apply_gradients(self.grads_and_vars)

        # Add ops to save and restore all the variables.
        self.init = tf.global_variables_initializer()
        self.saver= tf.train.Saver(max_to_keep=0)

        traceln(' -- Number of Params: ', self.get_nb_params())

    def save_model(self, session, model_filename):
        traceln("Saving Model")
        save_path = self.saver.save(session, model_filename)

    def restore_model(self, session, model_filename):
        self.saver.restore(session, model_filename)
        traceln("Model restored.")

    def get_Wedge(self,session):
        '''
        Return the parameters for the Edge Convolutions
        :param session:
        :return:
        '''
        if self.num_layers>1:
            L0=session.run([self.Wel0,self.Wed_layers])
            We0=L0[0]
            list_we=[We0]
            for we in L0[1]:
                list_we.append(we)
            return list_we
        else:
            L0=session.run([self.Wel0])
            We0=L0[0]
            list_we=[We0]
            return list_we

    def train(self,session,graph,verbose=False,n_iter=1):
        '''
        Apply a train operation, ie sgd step for a single graph
        :param session:
        :param graph: a graph from GCN_Dataset
        :param verbose:
        :param n_iter: (default 1) number of steps to perform sgd for this graph
        :return:
        '''
        #TrainEvalSet Here
        for i in range(n_iter):
            #traceln(' -- Train',X.shape,EA.shape)
            #traceln(' -- DropoutEdges',self.dropout_rate_edge)
            feed_batch = {

                self.nb_node: graph.X.shape[0],
                self.nb_edge: graph.F.shape[0],
                self.node_input: graph.X,
                self.Ssparse: np.array(graph.Sind, dtype='int64'),
                self.Sshape: np.array([graph.X.shape[0], graph.F.shape[0]], dtype='int64'),
                self.Tsparse: np.array(graph.Tind, dtype='int64'),
                self.F: graph.F,
                self.y_input: graph.Y,
                self.dropout_p_H: self.dropout_rate_H,
                self.dropout_p_node: self.dropout_rate_node,
                self.dropout_p_edge: self.dropout_rate_edge,
                self.dropout_p_edge_feat: self.dropout_rate_edge_feat,
                #self.NA_indegree:graph.NA_indegree
            }

            Ops =session.run([self.train_step,self.loss], feed_dict=feed_batch)
            if verbose:
                traceln(' -- Training Loss',Ops[1])

    def test(self,session,graph,verbose=True):
        '''
        Test return the loss and accuracy for the graph passed as argument
        :param session:
        :param graph:
        :param verbose:
        :return:
        '''

        feed_batch = {

            self.nb_node: graph.X.shape[0],
            self.nb_edge: graph.F.shape[0],
            self.node_input: graph.X,

            self.Ssparse: np.array(graph.Sind, dtype='int64'),
            self.Sshape: np.array([graph.X.shape[0], graph.F.shape[0]], dtype='int64'),
            self.Tsparse: np.array(graph.Tind, dtype='int64'),

            self.F: graph.F,
            self.y_input: graph.Y,
            self.dropout_p_H: 0.0,
            self.dropout_p_node: 0.0,
            self.dropout_p_edge: 0.0,
            self.dropout_p_edge_feat: 0.0,
            #self.NA_indegree: graph.NA_indegree
        }

        Ops =session.run([self.loss,self.accuracy], feed_dict=feed_batch)
        if verbose:
            traceln(' -- Test Loss',Ops[0],' Test Accuracy:',Ops[1])
        return Ops[1]

    def predict(self,session,graph,verbose=True):
        '''
        Does the prediction
        :param session:
        :param graph:
        :param verbose:
        :return:
        '''
        feed_batch = {
            self.nb_node: graph.X.shape[0],
            self.nb_edge: graph.F.shape[0],
            self.node_input: graph.X,
            # fast_gcn.S: np.asarray(graph.S.todense()).squeeze(),
            # fast_gcn.Ssparse: np.vstack([graph.S.row,graph.S.col]),
            self.Ssparse: np.array(graph.Sind, dtype='int64'),
            self.Sshape: np.array([graph.X.shape[0], graph.F.shape[0]], dtype='int64'),
            self.Tsparse: np.array(graph.Tind, dtype='int64'),
            # fast_gcn.T: np.asarray(graph.T.todense()).squeeze(),
            self.F: graph.F,
            self.dropout_p_H: 0.0,
            self.dropout_p_node: 0.0,
            self.dropout_p_edge: 0.0,
            self.dropout_p_edge_feat: 0.0,
            #self.NA_indegree: graph.NA_indegree
        }
        Ops = session.run([self.pred], feed_dict=feed_batch)
        if verbose:
            traceln(' -- Got Prediction for:',Ops[0].shape)
            print(str(Ops))
        return Ops[0]

    def prediction_prob(self,session,graph,verbose=True):
        '''
        Does the prediction
        :param session:
        :param graph:
        :param verbose:
        :return:
        '''
        feed_batch = {
            self.nb_node: graph.X.shape[0],
            self.nb_edge: graph.F.shape[0],
            self.node_input: graph.X,
            # fast_gcn.S: np.asarray(graph.S.todense()).squeeze(),
            # fast_gcn.Ssparse: np.vstack([graph.S.row,graph.S.col]),
            self.Ssparse: np.array(graph.Sind, dtype='int64'),
            self.Sshape: np.array([graph.X.shape[0], graph.F.shape[0]], dtype='int64'),
            self.Tsparse: np.array(graph.Tind, dtype='int64'),
            # fast_gcn.T: np.asarray(graph.T.todense()).squeeze(),
            self.F: graph.F,
            self.dropout_p_H: 0.0,
            self.dropout_p_node: 0.0,
            self.dropout_p_edge: 0.0,
            self.dropout_p_edge_feat: 0.0,
            #self.NA_indegree: graph.NA_indegree
        }
        Ops = session.run([self.predict_proba], feed_dict=feed_batch)
        if verbose:
            traceln(' -- Got Prediction for:',Ops[0].shape)
        return Ops[0]

    def train_All_lG(self,session,graph_train,graph_val, max_iter, eval_iter = 10, patience = 7, graph_test = None, save_model_path = None):
        '''

        Merge all the graph and train on them

        :param session:
        :param graph_train: the list of graph to train on
        :param graph_val:   the list of graph used for validation
        :param max_iter:  maximum number of epochs
        :param eval_iter: evaluate every eval_iter
        :param patience: stopped training if accuracy is not improved on the validation set after patience_value
        :param graph_test: Optional. If a test set is provided, then accuracy on the test set is reported
        :param save_model_path: checkpoints filename to save the model.
        :return: A Dictionary with training accuracies, validations accuracies and test accuracies if any, and the Wedge parameters
        '''
        best_val_acc = 0.0
        wait = 0
        stop_training = False
        stopped_iter = max_iter
        train_accuracies = []
        validation_accuracies = []
        test_accuracies = []
        conf_mat = []

        start_monitoring_val_acc = False

        # Not Efficient to compute this for

        merged_graph = gcn_datasets.GCNDataset.merge_allgraph(graph_train)

        self.train(session, merged_graph, n_iter=1)

        for i in range(max_iter):
            if stop_training:
                break

            if i % eval_iter == 0:
                traceln('\n -- Epoch', i)
                _, tr_acc = self.test_lG(session, graph_train, verbose=False)
                traceln(' --  Train Acc', '%.4f' % tr_acc)
                train_accuracies.append(tr_acc)

                _, node_acc = self.test_lG(session, graph_val, verbose=False)
                traceln(' --  Valid Acc', '%.4f' % node_acc)
                validation_accuracies.append(node_acc)

                if save_model_path:
                    save_path = self.saver.save(session, save_model_path, global_step=i)

                if graph_test:
                    _, test_acc = self.test_lG(session, graph_test, verbose=False)
                    traceln(' --   Test Acc', '%.4f' % test_acc)
                    test_accuracies.append(test_acc)

                    # Ypred = self.predict_lG(session, graph_test,verbose=False)
                    # Y_true_flat = []
                    # Ypred_flat = []
                    # for graph, ypred in zip(graph_test, Ypred):
                    #    ytrue = np.argmax(graph.Y, axis=1)
                    #    Y_true_flat.extend(ytrue)
                    #    Ypred_flat.extend(ypred)
                    # cm = sklearn.metrics.confusion_matrix(Y_true_flat, Ypred_flat)
                    # conf_mat.append(cm)

                # TODO min_delta
                # if tr_acc>0.99:
                #    start_monitoring_val_acc=True

                if node_acc > best_val_acc:
                    best_val_acc = node_acc
                    wait = 0
                else:
                    if wait >= patience:
                        stopped_iter = i
                        stop_training = True
                    wait += 1
            else:

                self.train(session, merged_graph, n_iter=1)
        # Final Save
        # if save_model_path:
        # save_path = self.saver.save(session, save_model_path, global_step=i)
        # TODO Add the final step
        mean_acc = []
        traceln(' -- Stopped Model Training after ', stopped_iter)
        traceln(' -- Val Accuracies ', validation_accuracies)

        traceln(' -- Final Training Accuracy')
        _, node_train_acc = self.test_lG(session, graph_train)
        traceln(' -- Train Mean Accuracy', '%.4f' % node_train_acc)

        traceln(' -- Final Valid Acc')
        self.test_lG(session, graph_val)

        R = {}
        R['train_acc'] = train_accuracies
        R['val_acc'] = validation_accuracies
        R['test_acc'] = test_accuracies
        R['stopped_iter'] = stopped_iter
        R['confusion_matrix'] = conf_mat
        # R['W_edge'] =self.get_Wedge(session)
        if graph_test:
            _, final_test_acc = self.test_lG(session, graph_test)
            traceln(' -- Final Test Acc', '%.4f' % final_test_acc)
            R['final_test_acc'] = final_test_acc

        return R

class GraphConvNet(MultiGraphNN):
    '''
    A Deep Standard GCN model for a graph list
    '''
    def __init__(self,node_dim,nb_classes,num_layers=1,learning_rate=0.1,mu=0.1,node_indim=-1,
                 dropout_rate=0.0,dropout_mode=0):
        self.node_dim=node_dim
        self.n_classes=nb_classes
        self.num_layers=num_layers
        self.learning_rate=learning_rate
        self.activation=tf.nn.relu
        self.mu=mu
        self.optalg = tf.train.AdamOptimizer(self.learning_rate)
        self.convolve_last=False
        self.dropout_rate=dropout_rate
        #0 No dropout 1, Node Dropout at input 2 Standard dropout for layer
        # check logit layer
        self.dropout_mode=dropout_mode

        if node_indim==-1:
            self.node_indim=self.node_dim
        else:
            self.node_indim=node_indim

    def create_model(self):
        self.nb_node    = tf.placeholder(tf.int32,(), name='nb_node')
        self.node_input = tf.placeholder(tf.float32, [None, self.node_dim], name='X_')
        self.y_input    = tf.placeholder(tf.float32, [None, self.n_classes], name='Y')
        self.NA_input   = tf.placeholder(tf.float32, name='NA_input') #Normalized Adjacency Matrix Here
        self.dropout_p  = tf.placeholder(tf.float32,(), name='dropout_prob')

        #std_dev_in=float(1.0/ float(self.node_dim))
        self.Wnode_layers=[]
        self.Bnode_layers=[]

        std_dev_input=float(1.0/ float(self.node_dim))
        std_dev_indim=float(1.0/ float(self.node_indim))

        if self.node_indim!=self.node_dim:
            self.Wnode = init_glorot((self.node_dim,self.node_indim),name='Wn0')
            #self.Wnode = init_normal((self.node_dim, self.node_indim),std_dev_input,name='Wn0')
        else:
            self.Wnode = tf.Variable(tf.eye(self.node_dim),name='Wn0',dtype=tf.float32)

        self.Bnode = tf.Variable(tf.zeros([self.node_indim]), name='Bnode',dtype=tf.float32)

        for i in range(self.num_layers-1):
            Wnli =init_glorot((2*self.node_indim, self.node_indim),name='Wnl'+str(i))
            #Wnli = init_normal((self.node_indim, self.node_indim),std_dev_indim, name='Wnl' + str(i))
            self.Wnode_layers.append(Wnli)
            #The GCN do not seem to use a bias term
            #Bnli = tf.Variable(tf.zeros([self.node_indim]), name='Bnl'+str(i),dtype=tf.float32)
            #self.Bnode_layers.append(Bnli)

        self.W_classif = init_glorot((2*self.node_indim, self.n_classes),name="W_classif")
        self.B_classif = tf.Variable(tf.zeros([self.n_classes]), name='B_classif',dtype=np.float32)

        #Input Layer
        #Check the self-loop . Is included in the normalized adjacency matrix
        #Check Residual Connections as weel for deeper models

        #add dropout_placeholder ... to differentiate train and test
        #x = tf.nn.dropout(x, 1 - self.dropout)

        #Dropout some nodes at the input of the graph ?
        #Should I dropout in upper layers as well ?
        #This way this forces the net to infer the node labels from its neighbors only
        #Here I dropout the features, but node the edges ..

        self.node_dropout_ind =  tf.nn.dropout(tf.ones([self.nb_node],dtype=tf.float32),1-self.dropout_p)
        self.ND = tf.diag(self.node_dropout_ind)

        if self.dropout_mode==1:
            #self.H = self.activation(tf.matmul(self.ND,tf.add(tf.matmul(self.node_input, self.Wnode), self.Bnode)))
            P0 = self.activation(tf.matmul(self.ND, tf.add(tf.matmul(self.node_input, self.Wnode), self.Bnode)))
            self.hidden_layers = [self.H]
        else:
            H0 =tf.add(tf.matmul(self.node_input, self.Wnode), self.Bnode)
            P0 =tf.matmul(self.NA_input, H0)  # we should forget the self loop
            H0_ = self.activation(tf.concat([H0, P0], 1))
            self.hidden_layers=[H0_]
            #self.H = self.activation(tf.add(tf.matmul(self.node_input, self.Wnode), self.Bnode))
            #self.hidden_layers = [self.H]

        for i in range(self.num_layers-1):
            if self.dropout_mode==2:
                Hp = tf.nn.dropout(self.hidden_layers[-1],1-self.dropout_p)
                Hi_ = tf.matmul(Hp, self.Wnode_layers[i])
            else:
                Hi_ = tf.matmul(self.hidden_layers[-1], self.Wnode_layers[i])
            P =tf.matmul(self.NA_input, Hi_) #we should forget the self loop
            #Hp = tf.concat([H0, P], 1)
            Hi = self.activation(tf.concat([Hi_,P],1))
            self.hidden_layers.append(Hi)

        #This dropout the logits as in GCN
        if self.dropout_mode==2:
            Hp = tf.nn.dropout(self.hidden_layers[-1], 1 - self.dropout_p)
            self.hidden_layers.append(Hp)

        if self.convolve_last is True:
            logit_0     = tf.add(tf.matmul(self.hidden_layers[-1], self.W_classif), self.B_classif)
            self.logits =  tf.matmul(self.NA_input,logit_0) #No activation function here
        else:
            self.logits =tf.add(tf.matmul(self.hidden_layers[-1],self.W_classif),self.B_classif)

        cross_entropy_source = tf.nn.softmax_cross_entropy_with_logits(logits=self.logits, labels=self.y_input)

        #Global L2 Regulization
        self.loss = tf.reduce_mean(cross_entropy_source) + self.mu * tf.nn.l2_loss(self.W_classif)

        self.correct_prediction = tf.equal(tf.argmax(tf.nn.softmax(self.logits), 1), tf.argmax(self.y_input, 1))
        self.accuracy = tf.reduce_mean(tf.cast(self.correct_prediction, tf.float32))

        self.grads_and_vars = self.optalg.compute_gradients(self.loss)
        self.train_step = self.optalg.apply_gradients(self.grads_and_vars)

        traceln(' -- Number of Params: ', self.get_nb_params())
        # Add ops to save and restore all the variables.
        self.init = tf.global_variables_initializer()

        self.saver = tf.train.Saver()

    def train(self,session,g,n_iter=1,verbose=False):
        #TrainEvalSet Here
        for i in range(n_iter):
            feed_batch={
                        self.nb_node:g.X.shape[0],
                        self.node_input:g.X,
                        self.y_input:g.Y,
                        self.NA_input:g.NA,
                        self.dropout_p:self.dropout_rate
            }
            Ops =session.run([self.train_step,self.loss], feed_dict=feed_batch)
            if verbose:
                traceln(' -- Training Loss',Ops[1])

    def test(self,session,g,verbose=True):
        #TrainEvalSet Here
        feed_batch={
                        self.nb_node:g.X.shape[0],
                        self.node_input:g.X,
                        self.y_input:g.Y,
                        self.NA_input:g.NA,
                        self.dropout_p: 0.0
        }
        Ops =session.run([self.loss,self.accuracy], feed_dict=feed_batch)
        if verbose:
            traceln(' -- Test Loss',Ops[0],' Test Accuracy:',Ops[1])
        return Ops[1]

class EdgeLogit(Logit):
    '''
    Logistic Regression for MultiGraph
    '''

    def train(self,session,graph,verbose=False,n_iter=1):
        '''
        Apply a train operation, ie sgd step for a single graph
        :param session:
        :param graph: a graph from GCN_Dataset
        :param verbose:
        :param n_iter: (default 1) number of steps to perform sgd for this graph
        :return:
        '''
        #TrainEvalSet Here
        for i in range(n_iter):
            #traceln(' -- Train',X.shape,EA.shape)

            nb_edge =graph.E.shape[0]
            half_edge  =nb_edge/2
            feed_batch = {

                self.nb_node: graph.EC.shape[0], #here we pass the number of edges
                self.node_input: graph.EC,
                self.y_input: graph.Yedge,

                #self.nb_node: half_edge, #here we pass the number of edges
                #self.node_input: graph.F[:half_edge],
                #self.y_input: graph.Yedge[:half_edge],
            }

            Ops =session.run([self.train_step,self.loss], feed_dict=feed_batch)
            if verbose:
                traceln(' -- Training Loss',Ops[1])

    def test(self,session,graph,verbose=True):
        '''
        Test return the loss and accuracy for the graph passed as argument
        :param session:
        :param graph:
        :param verbose:
        :return:
        '''

        nb_edge = graph.E.shape[0]
        half_edge = nb_edge / 2

        feed_batch = {

            self.nb_node: graph.EC.shape[0],
            self.node_input: graph.EC,
            self.y_input: graph.Yedge,

            #self.nb_node: half_edge,  # here we pass the number of edges
            #self.node_input: graph.F[:half_edge],
            #self.y_input: graph.Yedge[:half_edge],
        }

        Ops =session.run([self.loss,self.accuracy], feed_dict=feed_batch)
        if verbose:
            traceln(' -- Test Loss',Ops[0],' Test Accuracy:',Ops[1])
        return Ops[1]

    def predict(self,session,graph,verbose=True):
        '''
        Does the prediction
        :param session:
        :param graph:
        :param verbose:
        :return:
        '''

        nb_edge = graph.E.shape[0]
        half_edge = nb_edge / 2

        feed_batch = {
            self.nb_node: graph.EC.shape[0],
            self.node_input: graph.EC,

            #self.nb_node: half_edge,  # here we pass the number of edges
            #self.node_input: graph.F[:half_edge],
            #self.y_input: graph.Yedge[:, half_edge],
        }
        Ops = session.run([self.pred], feed_dict=feed_batch)
        if verbose:
            traceln(' -- Got Prediction for:',Ops[0].shape)
        return Ops[0]

#TODO Benchmark on Snake, GCN, ECN, graphAttNet vs Cora
#TODO Refactorize Code
#TODO Add L2 Regularization
#TODO Stack or Add Convolution -> Reduce the size
# Force the attention to preserve the node information i.e alpha'= 0.8 I +0.2 alpha
# Force doing attention only for the logit ?
# with a two layer
# Branching factor --> Attention
# Logit Layer and Attention
# There is one diff with ECN the feature for the edges are dynamically calculated
# whereas for GAT they are conditionned on the currect Node features
# 0.88
# Do a dot product attention or something different ...
# Change Initialization of the attention vector
# with Attn vector equal  [x;0] the attention keeps the source features and do not propagate ...
# Should reduce Nb of parameters
# This notion of edges is completely arbritrary
# We could at all nodes in the graph to see whether there are some depencies, no ?, interesting exp

class GraphAttNet(MultiGraphNN):
    '''
    Graph Attention Network
    '''

    # Variable ignored by the set_learning_options
    _setter_variables = {
        "node_dim": True, "edge_dim": True, "nb_class": True,
        "num_layers": True, "lr": True,
        "node_indim": True, "nb_attention": True,
        "nb_iter": True, "ratio_train_val": True}

    def __init__(self,node_dim,nb_classes,num_layers=1,learning_rate=0.1,node_indim=-1,nb_attention=3
                 ):
        self.node_dim=node_dim

        self.n_classes=nb_classes
        self.num_layers=num_layers
        self.learning_rate=learning_rate
        self.activation=tf.nn.elu
        #self.activation=tf.nn.relu

        self.optalg = tf.train.AdamOptimizer(self.learning_rate)
        self.stack_instead_add=False
        self.residual_connection=False#deprecated
        self.mu=0.0
        self.dropout_rate_node = 0.0
        self.dropout_rate_attention = 0.0
        self.nb_attention=nb_attention
        self.distinguish_node_from_neighbor=False
        self.original_model=False
        self.attn_type=0
        self.dense_model=False

        if node_indim==-1:
            self.node_indim=self.node_dim
        else:
            self.node_indim=node_indim

    #TODO GENERIC Could be move in MultigraphNN
    def set_learning_options(self,dict_model_config):
        """
        Set all learning options that not directly accessible from the constructor

        :param kwargs:
        :return:
        """
        traceln(dict_model_config)
        for attrname,val in dict_model_config.items():
            #We treat the activation function differently as we can not pickle/serialiaze python function
            if attrname=='activation_name':
                if val=='relu':
                    self.activation=tf.nn.relu
                elif val=='tanh':
                    self.activation=tf.nn.tanh
                else:
                    raise Exception('Invalid Activation Function')
            if attrname=='stack_instead_add' or attrname=='stack_convolutions':
                self.stack_instead_add=val
            if attrname not in self._setter_variables:
                try:
                    traceln(' -- set',attrname,val)
                    setattr(self,attrname,val)
                except AttributeError:
                    warnings.warn("Ignored options for ECN"+attrname+':'+val)

    def dense_graph_attention_layer(self,H,W,A,nb_node,dropout_attention,dropout_node,use_dropout=False):
        '''
        Implement a dense attention layer where every node is connected to everybody
        :param A:
        :param H:
        :param W:
        :param dropout_attention:
        :param dropout_node:
        :param use_dropout:
        :return:
        '''

        '''
        for all i,j aHi + bHj
        repmat all first column contains H1 second columns H2  etc
        diag may be a special case

        '''
        with tf.name_scope('graph_att_dense_attn'):

            P = tf.matmul(H, W)
            Aij_forward = tf.expand_dims(A[0], 0)  # attention vector for forward edge and backward edge
            Aij_backward = tf.expand_dims(A[1], 0)  # Here we assume it is the same on contrary to the paper
            # Compute the attention weight for target node, ie a . Whj if j is the target node
            att_target_node = tf.matmul(P, Aij_backward,transpose_b=True)
            # Compute the attention weight for the source node, ie a . Whi if j is the target node
            att_source_node = tf.matmul(P, Aij_forward, transpose_b=True)

            Asrc_vect = tf.tile(att_source_node,[nb_node,1])
            Asrc = tf.reshape(Asrc_vect,[nb_node,nb_node])

            Atgt_vect = tf.tile(att_target_node, [nb_node,1])
            Atgt = tf.reshape(Atgt_vect, [nb_node, nb_node])

            Att = tf.nn.leaky_relu(Asrc+Atgt)
            #Att = tf.nn.leaky_relu(Asrc)
            alphas = tf.nn.softmax(Att)

            # dropout is done after the softmax
            if use_dropout:
                traceln(' -- ... using dropout for attention layer')
                alphasD = tf.nn.dropout(alphas, 1.0 - dropout_attention)
                P_D = tf.nn.dropout(P, 1.0 - dropout_node)
                alphasP = tf.matmul(alphasD, P_D)
                return alphasD, alphasP
            else:
                # We compute the features given by the attentive neighborhood
                alphasP = tf.matmul(alphas, P)
                return alphas, alphasP

    #TODO Change the transpose of the A parameter
    def simple_graph_attention_layer(self,H,W,A,S,T,Adjind,Sshape,nb_edge,
                                     dropout_attention,dropout_node,
                                     use_dropout=False,add_self_loop=False,attn_type=0):
        '''
        :param H: The current node feature
        :param W: The node projection for this layer
        :param A: The attention weight vector: a
        :param S: The source edge matrix indices
        :param T: The target edge matrix indices
        :param Adjind: The adjcency matrix indices
        :param Sshape: Shape of S
        :param nb_edge: Number of edge
        :param dropout_attention: dropout_rate for the attention
        :param use_dropout: wether to use dropout
        :param add_self_loop: wether to add edge (i,i)
        :return: alphas,nH
            where alphas is the attention-based adjancency matrix alpha[i,j] correspond to alpha_ij
            nH correspond to the new features for this layer ie alphas(H,W)
        '''

        with tf.name_scope('graph_att_net_attn'):
            # This has shape (nb_node,in_dim) and correspond to the project W.h in the paper
            P=tf.matmul(H,W)
            #traceln(P.get_shape())

            #This has shape #shape,(nb_edge,nb_node)
            #This sparse tensor contains target nodes for edges.
            #The indices are [edge_idx,node_target_index]
            Tr = tf.SparseTensor(indices=T, values=tf.ones([nb_edge], dtype=tf.float32),
                                 dense_shape=[Sshape[1], Sshape[0]])
            Tr = tf.sparse_reorder(Tr) # reorder so that sparse operations work correctly
            # This tensor has shape(nb_edge,in_dim) and contains the node target projection, ie Wh
            TP = tf.sparse_tensor_dense_matmul(Tr, P,name='TP')

            # This has shape #shape,(nb_node,nb_edge)
            # This sparse tensor contains source nodes for edges.
            # The indices are [node_source_index,edge_idx]
            SD = tf.SparseTensor(indices=S, values=tf.ones([nb_edge],dtype=tf.float32), dense_shape=Sshape)
            SD = tf.sparse_reorder(SD) #shape,(nb_edge,nb_node)
            # This tensor has shape(nb_edge,in_dim) and contains the node source projection, ie Wh
            SP = tf.sparse_tensor_dense_matmul(tf.sparse_transpose(SD), P,name='SP') #shape(nb_edge,in_dim)
            #traceln(' -- SP', SP.get_shape())

            #Deprecated
            if attn_type==1:
                #Mutlitplication Attn Module
                Aij_forward = A  # attention vector for forward edge and backward edge
                Aij_backward = A  # Here we assume it is the same on contrary to the paper
                # Compute the attention weight for target node, ie a . Whj if j is the target node
                att_target_node = tf.multiply(TP,  Aij_forward[0])
                # Compute the attention weight for the source node, ie a . Whi if j is the target node
                att_source_node = tf.multiply(SP, Aij_backward[0])

                # The attention values for the edge ij is the sum of attention of node i and j
                # Attn( node_i, node_j) = Sum_k (a_k)^2 Hik Hjk Is this what we want ?
                att_source_target_node = tf.reduce_sum( tf.multiply(att_source_node,att_target_node),axis=1)
                attn_values = tf.nn.leaky_relu( att_source_target_node)
            #
            elif attn_type==2:
                #Inspired by learning to rank approach on w(x+-x-)
                # Attn( node_i, node_j) = Sum_k (a_k)  (Hik- Hjk) Is this what we want ?
                att_source_target_node = tf.reduce_sum( tf.multiply(SP-TP,A[0]),axis=1)
                attn_values = tf.nn.leaky_relu( att_source_target_node)

            else:
                Aij_forward=tf.expand_dims(A[0],0)  # attention vector for forward edge and backward edge
                Aij_backward=tf.expand_dims(A[1],0) # Here we assume it is the same on contrary to the paper
                # Compute the attention weight for target node, ie a . Whj if j is the target node
                att_target_node  =tf.matmul(TP,Aij_backward,transpose_b=True)
                # Compute the attention weight for the source node, ie a . Whi if j is the target node
                att_source_node = tf.matmul(SP,Aij_forward,transpose_b=True)

                # The attention values for the edge ij is the sum of attention of node i and j
                attn_values = tf.nn.leaky_relu(tf.squeeze(att_target_node) + tf.squeeze(att_source_node))

            # From that we build a sparse adjacency matrix containing the correct values
            # which we then feed to a sparse softmax
            AttAdj = tf.SparseTensor(indices=Adjind, values=attn_values, dense_shape=[Sshape[0], Sshape[0]])
            AttAdj = tf.sparse_reorder(AttAdj)

            #Note very efficient to do this, we should add the loop in the preprocessing
            if add_self_loop:
                node_indices=tf.range(Sshape[0])
                #Sparse Idendity
                Aij_forward = tf.expand_dims(A[0], 0)
                id_indices = tf.stack([node_indices, node_indices], axis=1)
                val =tf.squeeze(tf.matmul(P,Aij_forward,transpose_b=True))
                spI = tf.SparseTensor(indices=id_indices,values=2.0*val,dense_shape=[Sshape[0], Sshape[0]])

                AttAdj_I = tf.sparse_add(AttAdj,spI)
                alphas = tf.sparse_softmax(AttAdj_I)

            else:
                alphas = tf.sparse_softmax(AttAdj)

            #dropout is done after the softmax
            if use_dropout:
                traceln(' -- ... using dropout for attention layer')
                alphasD = tf.SparseTensor(indices=alphas.indices,values=tf.nn.dropout(alphas.values, 1.0 - dropout_attention),dense_shape=alphas.dense_shape)
                P_D =tf.nn.dropout(P,1.0-dropout_node)
                alphasP = tf.sparse_tensor_dense_matmul(alphasD, P_D)
                return alphasD, alphasP
            else:
                #We compute the features given by the attentive neighborhood
                alphasP = tf.sparse_tensor_dense_matmul(alphas,P)
                return alphas,alphasP

    def _create_original_model(self):
        std_dev_in = float(1.0 / float(self.node_dim))
        self.use_dropout = self.dropout_rate_attention > 0 or self.dropout_rate_node > 0
        self.hidden_layer = []
        attns0 = []

        # Define the First Layer from the Node Input
        for a in range(self.nb_attention):
            # H0 = Maybe do a common H0 and have different attention parameters
            # Change the attention, maybe ?
            # Do multiplicative
            # How to add edges here
            # Just softmax makes a differences
            # I could stack [current_node,representation; edge_features;] and do a dot product on that
            Wa = init_glorot([int(self.node_dim), int(self.node_indim)], name='Wa0' + str(a))
            va = init_glorot([2, int(self.node_indim)], name='va0' + str(a))

            if self.distinguish_node_from_neighbor:
                H0 = tf.matmul(self.node_input, Wa)
                attns0.append(H0)

            _, nH = self.simple_graph_attention_layer(self.node_input, Wa, va, self.Ssparse, self.Tsparse, self.Aind,
                                                      self.Sshape, self.nb_edge, self.dropout_p_attn,
                                                      self.dropout_p_node,
                                                      use_dropout=self.use_dropout, add_self_loop=True)
            attns0.append(nH)
        self.hidden_layer.append(
            self.activation(tf.concat(attns0, axis=-1)))  # Now dims should be indim*self.nb_attention

        # Define Intermediate Layers
        for i in range(1, self.num_layers):
            attns = []
            for a in range(self.nb_attention):
                if self.distinguish_node_from_neighbor:
                    Wia = init_glorot(
                        [int(self.node_indim * self.nb_attention + self.node_indim), int(self.node_indim)],
                        name='Wa' + str(i) + '_' + str(a))
                else:
                    Wia = init_glorot([int(self.node_indim * self.nb_attention), int(self.node_indim)],
                                      name='Wa' + str(i) + '_' + str(a))

                via = init_glorot([2, int(self.node_indim)], name='va' + str(i) + '_' + str(a))
                _, nH = self.simple_graph_attention_layer(self.hidden_layer[-1], Wia, via, self.Ssparse, self.Tsparse,
                                                          self.Aind,
                                                          self.Sshape, self.nb_edge, self.dropout_p_attn,
                                                          self.dropout_p_node,
                                                          use_dropout=self.use_dropout, add_self_loop=True)
                attns.append(nH)

            self.hidden_layer.append(self.activation(tf.concat(attns, axis=-1)))

        # Define Logit Layer
        out = []
        for i in range(self.nb_attention):
        #for i in range(1):
            logits_a = init_glorot([int(self.node_indim * self.nb_attention), int(self.n_classes)],
                                   name='Logita' + '_' + str(a))
            via = init_glorot([2, int(self.n_classes)], name='LogitA' + '_' + str(a))
            _, nL = self.simple_graph_attention_layer(self.hidden_layer[-1], logits_a, via, self.Ssparse, self.Tsparse,
                                                      self.Aind,
                                                      self.Sshape, self.nb_edge, self.dropout_p_attn,
                                                      self.dropout_p_node,
                                                      use_dropout=self.use_dropout, add_self_loop=True)
            out.append(nL)

        self.logits = tf.add_n(out) / self.nb_attention
        #self.logits = out[0]

    def _create_nodedistint_model(self):
        '''
        Create a model the separe node distinct models
        :return:
        '''

        std_dev_in = float(1.0 / float(self.node_dim))
        self.use_dropout = self.dropout_rate_attention > 0 or self.dropout_rate_node > 0
        self.hidden_layer = []
        attns0 = []

        # Define the First Layer from the Node Input
        Wa = tf.eye(int(self.node_dim), name='I0')
        H0 = tf.matmul(self.node_input, Wa)
        attns0.append(H0)

        I = tf.Variable(tf.eye(self.node_dim), trainable=False)
        for a in range(self.nb_attention):
            # H0 = Maybe do a common H0 and have different attention parameters
            # Change the attention, maybe ?
            # Do multiplicative
            # How to add edges here
            # Just softmax makes a differences
            # I could stack [current_node,representation; edge_features;] and do a dot product on that
            va = init_glorot([2, int(self.node_dim)], name='va0' + str(a))

            _, nH = self.simple_graph_attention_layer(H0, I, va, self.Ssparse, self.Tsparse, self.Aind,
                                                      self.Sshape, self.nb_edge, self.dropout_p_attn,
                                                      self.dropout_p_node,
                                                      use_dropout=self.use_dropout, add_self_loop=False,attn_type=self.attn_type)
            attns0.append(nH)
        self.hidden_layer.append(
            self.activation(tf.concat(attns0, axis=-1)))  # Now dims should be indim*self.nb_attention

        # Define Intermediate Layers
        for i in range(1, self.num_layers):
            attns = []
            if i == 1:
                previous_layer_dim =int(self.node_dim * self.nb_attention + self.node_dim)
                Wia = init_glorot([previous_layer_dim, int(self.node_indim)],
                    name='Wa' + str(i) + '_' + str(a))
            else:
                previous_layer_dim = int(self.node_indim * self.nb_attention + self.node_indim)
                Wia = init_glorot( [previous_layer_dim, int(self.node_indim)], name='Wa' + str(i) + '_' + str(a))

            Hi = tf.matmul(self.hidden_layer[-1], Wia)
            attns.append(Hi)
            Ia = tf.Variable(tf.eye(self.node_indim), trainable=False)

            for a in range(self.nb_attention):
                via = init_glorot([2, int(self.node_indim)], name='va' + str(i) + '_' + str(a))
                _, nH = self.simple_graph_attention_layer(Hi, Ia, via, self.Ssparse, self.Tsparse,
                                                          self.Aind,
                                                          self.Sshape, self.nb_edge, self.dropout_p_attn,
                                                          self.dropout_p_node,
                                                          use_dropout=self.use_dropout, add_self_loop=False,attn_type=self.attn_type)
                attns.append(nH)

            self.hidden_layer.append(self.activation(tf.concat(attns, axis=-1)))

        # Define Logit Layer
        #TODO Add Attention on Logit Layer
        #It would not cost too much to add an attn mecha once I get the logits
        #If x,y are indicated in the node feature then we can implicitly find the type of edges that we are using ...

        if self.num_layers>1:
            logits_a = init_glorot([int(self.node_indim * self.nb_attention+self.node_indim), int(self.n_classes)],
                                   name='Logita' + '_' + str(a))
        else:
            logits_a = init_glorot([int(self.node_dim * self.nb_attention + self.node_dim), int(self.n_classes)],
                                   name='Logita' + '_' + str(a))
        Bc = tf.ones([int(self.n_classes)], name='LogitA' + '_' + str(a))
        # self.logits = tf.add_n(out) / self.nb_attention
        self.logits = tf.matmul(self.hidden_layer[-1],logits_a) +Bc

    def _create_densegraph_model(self):
        '''
        Create a model the separe node distinct models
        :return:
        '''

        std_dev_in = float(1.0 / float(self.node_dim))
        self.use_dropout = self.dropout_rate_attention > 0 or self.dropout_rate_node > 0
        self.hidden_layer = []
        attns0 = []

        # Define the First Layer from the Node Input
        Wa = tf.eye(int(self.node_dim), name='I0')
        H0 = tf.matmul(self.node_input, Wa)
        attns0.append(H0)

        I = tf.Variable(tf.eye(self.node_dim), trainable=False)
        for a in range(self.nb_attention):
            # H0 = Maybe do a common H0 and have different attention parameters
            # Change the attention, maybe ?
            # Do multiplicative
            # How to add edges here
            # Just softmax makes a differences
            # I could stack [current_node,representation; edge_features;] and do a dot product on that
            va = init_glorot([2, int(self.node_dim)], name='va0' + str(a))

            _, nH = self.dense_graph_attention_layer(H0, I, va, self.nb_node, self.dropout_p_attn,
                                                      self.dropout_p_node,
                                                      use_dropout=self.use_dropout)
            attns0.append(nH)
        self.hidden_layer.append(
            self.activation(tf.concat(attns0, axis=-1)))  # Now dims should be indim*self.nb_attention

        # Define Intermediate Layers
        for i in range(1, self.num_layers):
            attns = []
            if i == 1:
                previous_layer_dim =int(self.node_dim * self.nb_attention + self.node_dim)
                Wia = init_glorot([previous_layer_dim, int(self.node_indim)],
                    name='Wa' + str(i) + '_' + str(a))
            else:
                previous_layer_dim = int(self.node_indim * self.nb_attention + self.node_indim)
                Wia = init_glorot( [previous_layer_dim, int(self.node_indim)], name='Wa' + str(i) + '_' + str(a))

            Hi = tf.matmul(self.hidden_layer[-1], Wia)
            attns.append(Hi)
            Ia = tf.Variable(tf.eye(self.node_indim), trainable=False)

            for a in range(self.nb_attention):
                via = init_glorot([2, int(self.node_indim)], name='va' + str(i) + '_' + str(a))
                _, nH = self.dense_graph_attention_layer(Hi, Ia, via, self.nb_node,
                                                          self.dropout_p_attn,
                                                          self.dropout_p_node,
                                                          use_dropout=self.use_dropout)
                attns.append(nH)

            self.hidden_layer.append(self.activation(tf.concat(attns, axis=-1)))

        # Define Logit Layer
        #TODO Add Attention on Logit Layer
        #It would not cost too much to add an attn mecha once I get the logits
        #If x,y are indicated in the node feature then we can implicitly find the type of edges that we are using ...

        if self.num_layers>1:
            logits_a = init_glorot([int(self.node_indim * self.nb_attention+self.node_indim), int(self.n_classes)],
                                   name='Logita' + '_' + str(a))
        else:
            logits_a = init_glorot([int(self.node_dim * self.nb_attention + self.node_dim), int(self.n_classes)],
                                   name='Logita' + '_' + str(a))
        Bc = tf.ones([int(self.n_classes)], name='LogitA' + '_' + str(a))
        # self.logits = tf.add_n(out) / self.nb_attention
        self.logits = tf.matmul(self.hidden_layer[-1],logits_a) +Bc

    def create_model(self):
        '''
        Create the tensorflow graph for the model
        :return:
        '''
        self.nb_node    = tf.placeholder(tf.int32,(), name='nb_node')
        self.nb_edge = tf.placeholder(tf.int32, (), name='nb_edge')
        self.node_input = tf.placeholder(tf.float32, [None, self.node_dim], name='X_')
        self.y_input    = tf.placeholder(tf.float32, [None, self.n_classes], name='Y')
        #self.dropout_p_H    = tf.placeholder(tf.float32,(), name='dropout_prob_H')
        self.dropout_p_node = tf.placeholder(tf.float32, (), name='dropout_prob_N')
        self.dropout_p_attn = tf.placeholder(tf.float32, (), name='dropout_prob_edges')
        self.S          = tf.placeholder(tf.float32, name='S')
        self.Ssparse    = tf.placeholder(tf.int64, name='Ssparse') #indices
        self.Sshape     = tf.placeholder(tf.int64, name='Sshape') #indices
        self.Aind       =tf.placeholder(tf.int64, name='Sshape') #Adjacency indices
        self.T          = tf.placeholder(tf.float32,[None,None], name='T')
        self.Tsparse    = tf.placeholder(tf.int64, name='Tsparse')

        #self.S_indice = tf.placeholder(tf.in, [None, None], name='S')
        #self.F          = tf.placeholder(tf.float32,[None,None], name='F')

        if self.original_model:
            self._create_original_model()

        elif self.dense_model:
            self._create_densegraph_model()
        else:
            self._create_nodedistint_model()

        cross_entropy_source = tf.nn.softmax_cross_entropy_with_logits(logits=self.logits, labels=self.y_input)
        self.loss = tf.reduce_mean(cross_entropy_source)
        self.predict_proba = tf.nn.softmax(self.logits)
        self.pred = tf.argmax(tf.nn.softmax(self.logits), 1)
        self.correct_prediction = tf.equal(self.pred, tf.argmax(self.y_input, 1))
        self.accuracy = tf.reduce_mean(tf.cast(self.correct_prediction, tf.float32))

        self.grads_and_vars = self.optalg.compute_gradients(self.loss)
        self.train_step = self.optalg.apply_gradients(self.grads_and_vars)

        # Add ops to save and restore all the variables.
        self.init = tf.global_variables_initializer()
        self.saver= tf.train.Saver(max_to_keep=0)

        traceln(' -- Number of Params: ', self.get_nb_params())

    #TODO Move in MultigraphNN
    def save_model(self, session, model_filename):
        traceln("Saving Model")
        save_path = self.saver.save(session, model_filename)

    def restore_model(self, session, model_filename):
        self.saver.restore(session, model_filename)
        traceln("Model restored.")

    def train(self,session,graph,verbose=False,n_iter=1):
        '''
        Apply a train operation, ie sgd step for a single graph
        :param session:
        :param graph: a graph from GCN_Dataset
        :param verbose:
        :param n_iter: (default 1) number of steps to perform sgd for this graph
        :return:
        '''
        #TrainEvalSet Here
        for i in range(n_iter):
            #traceln(' -- Train',X.shape,EA.shape)
            #traceln(' -- DropoutEdges',self.dropout_rate_edge)
            Aind = np.array(np.stack([graph.Sind[:, 0], graph.Tind[:, 1]], axis=-1), dtype='int64')
            feed_batch = {

                self.nb_node: graph.X.shape[0],
                self.nb_edge: graph.F.shape[0],
                self.node_input: graph.X,
                self.Ssparse: np.array(graph.Sind, dtype='int64'),
                self.Sshape: np.array([graph.X.shape[0], graph.F.shape[0]], dtype='int64'),
                self.Tsparse: np.array(graph.Tind, dtype='int64'),
                #self.F: graph.F,
                self.Aind: Aind,
                self.y_input: graph.Y,
                #self.dropout_p_H: self.dropout_rate_H,
                self.dropout_p_node: self.dropout_rate_node,
                self.dropout_p_attn: self.dropout_rate_attention,
            }

            Ops =session.run([self.train_step,self.loss], feed_dict=feed_batch)
            if verbose:
                traceln(' -- Training Loss',Ops[1])

    def test(self,session,graph,verbose=True):
        '''
        Test return the loss and accuracy for the graph passed as argument
        :param session:
        :param graph:
        :param verbose:
        :return:
        '''
        Aind = np.array(np.stack([graph.Sind[:, 0], graph.Tind[:, 1]], axis=-1), dtype='int64')
        feed_batch = {

            self.nb_node: graph.X.shape[0],
            self.nb_edge: graph.F.shape[0],
            self.node_input: graph.X,

            self.Ssparse: np.array(graph.Sind, dtype='int64'),
            self.Sshape: np.array([graph.X.shape[0], graph.F.shape[0]], dtype='int64'),
            self.Tsparse: np.array(graph.Tind, dtype='int64'),
            self.Aind: Aind,
            #self.F: graph.F,
            self.y_input: graph.Y,
            #self.dropout_p_H: 0.0,
            self.dropout_p_node: 0.0,
            self.dropout_p_attn: 0.0,
            #self.dropout_p_edge_feat: 0.0,
            #self.NA_indegree: graph.NA_indegree
        }

        Ops =session.run([self.loss,self.accuracy], feed_dict=feed_batch)
        if verbose:
            traceln(' -- Test Loss',Ops[0],' Test Accuracy:',Ops[1])
        return Ops[1]

    def predict(self,session,graph,verbose=True):
        '''
        Does the prediction
        :param session:
        :param graph:
        :param verbose:
        :return:
        '''
        Aind = np.array(np.stack([graph.Sind[:, 0], graph.Tind[:, 1]], axis=-1), dtype='int64')
        feed_batch = {
            self.nb_node: graph.X.shape[0],
            self.nb_edge: graph.F.shape[0],
            self.node_input: graph.X,
            # fast_gcn.S: np.asarray(graph.S.todense()).squeeze(),
            # fast_gcn.Ssparse: np.vstack([graph.S.row,graph.S.col]),
            self.Ssparse: np.array(graph.Sind, dtype='int64'),
            self.Sshape: np.array([graph.X.shape[0], graph.F.shape[0]], dtype='int64'),
            self.Tsparse: np.array(graph.Tind, dtype='int64'),
            # fast_gcn.T: np.asarray(graph.T.todense()).squeeze(),
            #self.F: graph.F,
            self.Aind: Aind,
            self.dropout_p_node: 0.0,
            self.dropout_p_attn: 0.0,
        }
        Ops = session.run([self.pred], feed_dict=feed_batch)
        if verbose:
            traceln(' -- Got Prediction for:',Ops[0].shape)
        return Ops[0]

    def prediction_prob(self, session, graph, verbose=True):
        '''
        Does the prediction
        :param session:
        :param graph:
        :param verbose:
        :return:
        '''
        Aind = np.array(np.stack([graph.Sind[:, 0], graph.Tind[:, 1]], axis=-1), dtype='int64')
        feed_batch = {
            self.nb_node: graph.X.shape[0],
            self.nb_edge: graph.F.shape[0],
            self.node_input: graph.X,
            # fast_gcn.S: np.asarray(graph.S.todense()).squeeze(),
            # fast_gcn.Ssparse: np.vstack([graph.S.row,graph.S.col]),
            self.Ssparse: np.array(graph.Sind, dtype='int64'),
            self.Sshape: np.array([graph.X.shape[0], graph.F.shape[0]], dtype='int64'),
            self.Tsparse: np.array(graph.Tind, dtype='int64'),
            # fast_gcn.T: np.asarray(graph.T.todense()).squeeze(),
            # self.F: graph.F,
            self.Aind: Aind,
            self.dropout_p_node: 0.0,
            self.dropout_p_attn: 0.0,
        }
        Ops = session.run([self.predict_proba], feed_dict=feed_batch)
        if verbose:
            traceln(' -- Got Prediction for:', Ops[0].shape)
        return Ops[0]

    #TODO Move that MultiGraphNN
    def train_All_lG(self,session,graph_train,graph_val, max_iter, eval_iter = 10, patience = 7, graph_test = None, save_model_path = None):
        '''

        Merge all the graph and train on them

        :param session:
        :param graph_train: the list of graph to train on
        :param graph_val:   the list of graph used for validation
        :param max_iter:  maximum number of epochs
        :param eval_iter: evaluate every eval_iter
        :param patience: stopped training if accuracy is not improved on the validation set after patience_value
        :param graph_test: Optional. If a test set is provided, then accuracy on the test set is reported
        :param save_model_path: checkpoints filename to save the model.
        :return: A Dictionary with training accuracies, validations accuracies and test accuracies if any, and the Wedge parameters
        '''
        best_val_acc = 0.0
        wait = 0
        stop_training = False
        stopped_iter = max_iter
        train_accuracies = []
        validation_accuracies = []
        test_accuracies = []
        conf_mat = []

        start_monitoring_val_acc = False

        # Not Efficient to compute this for

        merged_graph = gcn_datasets.GCNDataset.merge_allgraph(graph_train)

        self.train(session, merged_graph, n_iter=1)

        for i in range(max_iter):
            if stop_training:
                break

            if i % eval_iter == 0:
                traceln('\n -- Epoch', i)
                _, tr_acc = self.test_lG(session, graph_train, verbose=False)
                traceln(' --  Train Acc', '%.4f' % tr_acc)
                train_accuracies.append(tr_acc)

                _, node_acc = self.test_lG(session, graph_val, verbose=False)
                traceln(' --  Valid Acc', '%.4f' % node_acc)
                validation_accuracies.append(node_acc)

                if save_model_path:
                    save_path = self.saver.save(session, save_model_path, global_step=i)

                if graph_test:
                    _, test_acc = self.test_lG(session, graph_test, verbose=False)
                    traceln(' --   Test Acc', '%.4f' % test_acc)
                    test_accuracies.append(test_acc)

                # TODO min_delta
                # if tr_acc>0.99:
                #    start_monitoring_val_acc=True

                if node_acc > best_val_acc:
                    best_val_acc = node_acc
                    wait = 0
                else:
                    if wait >= patience:
                        stopped_iter = i
                        stop_training = True
                    wait += 1
            else:

                self.train(session, merged_graph, n_iter=1)
        # Final Save
        # if save_model_path:
        # save_path = self.saver.save(session, save_model_path, global_step=i)
        # TODO Add the final step
        mean_acc = []
        traceln(' -- Stopped Model Training after', stopped_iter)
        traceln(' -- Val Accuracies', validation_accuracies)

        traceln(' -- Final Training Accuracy')
        _, node_train_acc = self.test_lG(session, graph_train)
        traceln(' -- Train Mean Accuracy', '%.4f' % node_train_acc)

        traceln(' -- Final Valid Acc')
        self.test_lG(session, graph_val)

        R = {}
        R['train_acc'] = train_accuracies
        R['val_acc'] = validation_accuracies
        R['test_acc'] = test_accuracies
        R['stopped_iter'] = stopped_iter
        R['confusion_matrix'] = conf_mat
        # R['W_edge'] =self.get_Wedge(session)
        if graph_test:
            _, final_test_acc = self.test_lG(session, graph_test)
            traceln(' -- Final Test Acc', '%.4f' % final_test_acc)
            R['final_test_acc'] = final_test_acc

        return R