
import timeit
import tensorflow as tf
import numpy as np
import os
import scipy.io as sio

start = timeit.default_timer()

aType = 'ra'   # attack method 'ra' = rand attack;  'tb' = tar betweenness attack;  'td' = tar degree attack
vname = 'test'
mat = sio.loadmat(vname + '.mat')
num_test_ins = len(mat[vname])


# Load testing instances
def load_mat(idx):
    A = []
    y = []
    tmpA = mat[vname][idx, 0]['A'][0, 0].todense()
    tmpy = np.squeeze(mat[vname][idx, 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.truncated_normal(shape, stddev=0.005)
    return tf.Variable(initial)


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


def max_pool_2(x):
    return tf.nn.max_pool(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.placeholder(tf.float32)
input_matrix = tf.reshape(in_matrix, [-1, 1000, 1000, 1])
score = tf.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.train.AdamOptimizer(learning_rate=5e-5, beta1=0.9, beta2=0.999, epsilon=1e-08).minimize(loss)
saver = tf.train.Saver()


#  ----- Main Session -----  #
with tf.Session() as sess:
    if aType == 'ra':       # random attack
        foldname = '/mdra'
    elif aType == 'tb':     # targeted betweenness attack
        foldname = '/mdtb'
    elif aType == 'td':     # targeted degree attack
        foldname = '/mdtd'
    # reload trained model
    cur_path = os.getcwd()  # Get current path
    saver.restore(sess, cur_path + foldname + "/model.ckpt")
    pv = []
    for i in range(num_test_ins):
        in_ma, in_score = load_mat(i)
        print(i)
        pre_vec = sess.run(h_fc_2, feed_dict={input_matrix: in_ma})
        if i == 0:
            pv = abs(pre_vec)
        else:
            pv = np.concatenate((pv, abs(pre_vec)), axis=0)
        # save results ('pred_' = 'predicted data')
    pr_add = 'pred_' + aType + '.mat'
    sio.savemat(pr_add, {'pv': pv})   # row[i] of 'pv' == predicted curve of test_sample[i]
    print('pv_' + aType + ' saved')


stop = timeit.default_timer()
print('Time: ', stop - start)


# total run time = 67.94s for 160 attacks
# average run time = 0.42s
# Intel(R) Core(TM) i7-9750H CPU @ 2.60GHz 2.59GHz
# Installed memory (RAM): 16.0GB (15.8 usable)
# Windows 10 Home 64-bit Operating System