
import tensorflow as tf
import numpy as np
import os
import random
import scipy.io as sio
import sys

aType = 'ra'  # attack method 'ra' = rand attack;  'tb' = tar betweenness attack;  'td' = tar degree attack
mat = sio.loadmat(aType + '.mat')   # load training data
num_train_ins = len(mat['data_train'])
num_valid_ins = len(mat['data_x'])
cur_path = os.getcwd()    # get current path
embedding_matrix = np.mat(np.loadtxt(cur_path+'/embedding_matrix.txt'))   # load embedding matrix
shuffle_list = random.sample(range(0, num_train_ins), num_train_ins)   # (re-)shuffle training data


# load training data
def load_batch(batch_size, begin_index):
    A = []
    y = []
    for i in range(batch_size):
        tmpA = mat['data_train'][shuffle_list[i+begin_index], 0]['A'][0, 0].todense()
        tmpy = np.squeeze(mat['data_train'][shuffle_list[i + begin_index], 0]['y'][0, 0])
        A.append(np.expand_dims(tmpA, axis=2))
        y.append(tmpy)
    return A, y


# load cross validation data
def load_batch_val(batch_size, begin_index):
    A = []
    y = []
    for i in range(batch_size):
        tmpA = mat['data_x'][i + begin_index, 0]['A'][0, 0].todense()
        tmpy = np.squeeze(mat['data_x'][i + begin_index, 0]['y'][0, 0])
        A.append(np.expand_dims(tmpA, axis=2))
        y.append(tmpy)
    return A, y


# define functions
def Weight_variable(shape):
    initial = tf.random.truncated_normal(shape, stddev=0.01)
    return tf.Variable(initial)


def Bais_variable(shape):
    initial = tf.constant(0.0, shape=shape)
    return tf.Variable(initial)


def max_pool_2(x):
    return tf.nn.max_pool2d(x, [1, 2, 2, 1], [1, 2, 2, 1], 'SAME')


def conv2d_stride1(x,W):
    return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')


# define the CNN
class CtrlRobustness:
    def __init__(self):
        # layer1
        self.w_conv1_1 = Weight_variable([7, 7, 1, 64])
        self.b_conv1_1 = Bais_variable([64])
        # layer2
        self.w_conv2_1 = Weight_variable([5, 5, 64, 64])
        self.b_conv2_1 = Bais_variable([64])
        # layer3
        self.w_conv3_1 = Weight_variable([3, 3, 64, 128])
        self.b_conv3_1 = Bais_variable([128])
        # layer4
        self.w_conv4_1 = Weight_variable([3, 3, 128, 128])
        self.b_conv4_1 = Bais_variable([128])
        # layer5
        self.w_conv5_1 = Weight_variable([3, 3, 128, 256])
        self.b_conv5_1 = Bais_variable([256])
        # layer6
        self.w_conv6_1 = Weight_variable([3, 3, 256, 256])
        self.b_conv6_1 = Bais_variable([256])
        # layer7
        self.w_conv7_1 = Weight_variable([3, 3, 256, 512])
        self.b_conv7_1 = Bais_variable([512])
        self.w_conv7_2 = Weight_variable([3, 3, 512, 512])
        self.b_conv7_2 = Bais_variable([512])

    def compu_conv(self, input_ma):
        # layer1
        h_conv1_1 = tf.nn.relu(conv2d_stride1(input_ma, self.w_conv1_1) + self.b_conv1_1)
        out_layer1 = max_pool_2(h_conv1_1)
        # layer2
        h_conv2_1 = tf.nn.relu(conv2d_stride1(out_layer1, self.w_conv2_1) + self.b_conv2_1)
        out_layer2 = max_pool_2(h_conv2_1)
        # layer3
        h_conv3_1 = tf.nn.relu(conv2d_stride1(out_layer2, self.w_conv3_1) + self.b_conv3_1)
        out_layer3 = max_pool_2(h_conv3_1)
        # layer4
        h_conv4_1 = tf.nn.relu(conv2d_stride1(out_layer3, self.w_conv4_1) + self.b_conv4_1)
        out_layer4 = max_pool_2(h_conv4_1)
        # layer5
        h_conv5_1 = tf.nn.relu(conv2d_stride1(out_layer4, self.w_conv5_1) + self.b_conv5_1)
        out_layer5 = max_pool_2(h_conv5_1)
        # layer6
        h_conv6_1 = tf.nn.relu(conv2d_stride1(out_layer5, self.w_conv6_1) + self.b_conv6_1)
        out_layer6 = max_pool_2(h_conv6_1)
        # layer7
        h_conv7_1 = tf.nn.relu(conv2d_stride1(out_layer6, self.w_conv7_1) + self.b_conv7_1)
        h_conv7_2 = tf.nn.relu(conv2d_stride1(h_conv7_1, self.w_conv7_2) + self.b_conv7_2)
        out_layer7 = max_pool_2(h_conv7_2)
        return out_layer7


# construct CNN
in_matrix = tf.compat.v1.placeholder(tf.float32)
input_matrix = tf.reshape(in_matrix, [-1, 1000, 1000, 1])
score = tf.compat.v1.placeholder(tf.float32)
network = CtrlRobustness()
out_matr = network.compu_conv(input_matrix)

# get shape of output
shape = out_matr.get_shape()
length = shape[1].value*shape[2].value*shape[3].value
out_o = tf.reshape(out_matr, [-1, length])


# fc = fully connected layers
# fc1
w_fc_1 = Weight_variable([length, 4096])
b_fc_1 = Bais_variable([4096])
h_fc_1 = tf.nn.relu(tf.matmul(out_o, w_fc_1) + b_fc_1)

# fc2
w_fc_2 = Weight_variable([4096, 999])
b_fc_2 = Bais_variable([999])
h_fc_2 = tf.matmul(h_fc_1, w_fc_2) + b_fc_2

# compute the loss value
mse_loss = tf.reduce_sum(tf.square(h_fc_2 - score), 1)
loss = tf.reduce_mean(mse_loss)

# training setting
train_step = tf.compat.v1.train.AdamOptimizer(learning_rate=5e-5, beta1=0.9, beta2=0.999, epsilon=1e-08).minimize(loss)
saver = tf.compat.v1.train.Saver()


#  ----- Main Session -----  #
with tf.compat.v1.Session() as sess:
    sess.run(tf.compat.v1.global_variables_initializer())   # Initialization
    if aType == 'ra':       # random attack
        foldname = '/mdra'
    elif aType == 'tb':     # targeted betweenness attack
        foldname = '/mdtb'
    elif aType == 'td':     # targeted degree attack
        foldname = '/mdtd'

    if len(sys.argv) < 2:   # input check and default paras
        epoch_val =1# 10
    else:
        epoch_val = int(sys.argv[1])

    if len(sys.argv) < 3:
        batch_size =1# 5
    else:
        batch_size = int(sys.argv[2])
    global_loss = float('inf')

    for epoc in range(epoch_val):
        # each epoch:
        for iter in range(num_train_ins//batch_size):
            # each iter:
            A, y = load_batch(batch_size, iter*batch_size)   # load training data
            loss_val = sess.run(loss, feed_dict={input_matrix: A, score: y})   # loss evaluation
            # train the network by a batch of training instances
            train_step.run(feed_dict={input_matrix: A, score: y})
            print(' ..Epoch:(' + str(epoc + 1) + '), Iter:(' + str(iter + 1) + '), Loss Value:(' + str(loss_val) + ')')
        # use validation data to select optimal-trained models
        average_loss = 0
        for iter in range(num_valid_ins//batch_size):
            A, y = load_batch_val(batch_size, iter * batch_size)
            loss_val = sess.run(loss, feed_dict={input_matrix: A, score: y})
            pre_vec = sess.run(h_fc_2, feed_dict={input_matrix: A, score: y})
            average_loss = average_loss + loss_val
        average_loss = average_loss / (num_valid_ins//batch_size)
        if average_loss < global_loss:  # save current model if better
            global_loss = average_loss
            save_path = saver.save(sess, cur_path + foldname + "/model.ckpt")
            print("Model Saved: %s" % save_path)
        print('Validation Average Loss = ' + str(average_loss))


