import numpy as np
import tensorflow as tf
import argparse
import os
import timeit
import math
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"

parser = argparse.ArgumentParser(description = 'MSCRED encoder-decoder')
parser.add_argument('--batch', type = int, default = 1,
				   help = 'batch_size for data input') #current version supprts batch_size = 1
parser.add_argument('--sensor_n', type = int, default = 30,
				   help = 'number of sensors')
parser.add_argument('--win_size', type = int, default = [10, 30, 60],
				   help = 'window size of each segment')
parser.add_argument('--step_max', type = int, default = 5,
				   help = 'maximum step of ConvLSTM')
parser.add_argument('--learning_rate',  type=float, default= 0.0002,
						help='learning rate')
parser.add_argument('--training_iters',  type = int, default = 5,
						help = 'number of maximum training iterations')
parser.add_argument('--train_start_id',  type = int, default = 11,
						help = 'training start id')
parser.add_argument('--train_end_id',  type = int, default = 800,
						help = 'training end id')
parser.add_argument('--test_start_id',  type = int, default = 800,
						help = 'test start id')
parser.add_argument('--test_end_id',  type = int, default = 2000,
						help = 'test end id')
parser.add_argument('--save_model_step', type = int, default = 1,
						help = 'number of iterations to save model')
parser.add_argument('--model_path', type = str, default = '../MSCRED/',
				   help='path to save models')
parser.add_argument('--raw_data_path', type = str, default = '../data/synthetic_data_with_anomaly-s-1.csv',
				   help='path to load raw data')
parser.add_argument('--matrix_data_path', type = str, default = '../data/matrix_data/',
				   help='matrix data path')
parser.add_argument('--train_test_label', type=int, default = 1,
				   help='train/test label: train (1), test (0)')
parser.add_argument('--GPU_id', type=int, default = 3,
				   help='GPU ID to select')

args = parser.parse_args()
print(args)

# argument settings
batch_size = args.batch
learning_rate =args.learning_rate
training_iters = args.training_iters
sensor_n = args.sensor_n
win_size = args.win_size
step_max= args.step_max

train_start_id = args.train_start_id
train_end_id = args.train_end_id
test_start_id = args.test_start_id
test_end_id = args.test_end_id

raw_data_path = args.raw_data_path
model_path = args.model_path
matrix_data_path = args.matrix_data_path
train_data_path = matrix_data_path + "train_data/"
test_data_path = matrix_data_path + "test_data/"
save_model_step = args.save_model_step

train_test_label = args.train_test_label

scale_n = len(win_size)

# set GPU
#GPU_id = args.GPU_id
#os.environ['CUDA_VISIBLE_DEVICES'] = str(GPU_id)

# 4d data
data_input = tf.placeholder(tf.float32, [step_max, sensor_n, sensor_n, scale_n])

# parameters: adding bias weight get similar performance
conv1_W = tf.Variable(tf.zeros([3, 3, scale_n, 32]), name = "conv1_W")
conv1_W = tf.get_variable("conv1_W", shape = [3, 3, scale_n, 32], initializer=tf.contrib.layers.xavier_initializer())
conv2_W = tf.Variable(tf.zeros([3, 3, 32, 64]), name = "conv2_W")
conv2_W = tf.get_variable("conv2_W", shape = [3, 3, 32, 64], initializer=tf.contrib.layers.xavier_initializer())
conv3_W = tf.Variable(tf.zeros([2, 2, 64, 128]), name = "conv3_W")
conv3_W = tf.get_variable("conv3_W", shape = [2, 2, 64, 128], initializer=tf.contrib.layers.xavier_initializer())
conv4_W = tf.Variable(tf.zeros([2, 2, 128, 256]), name = "conv4_W")
conv4_W = tf.get_variable("conv4_W", shape = [2, 2, 128, 256], initializer=tf.contrib.layers.xavier_initializer())

deconv4_W = tf.Variable(tf.zeros([2, 2, 128, 256]), name = "deconv4_W")
deconv4_W = tf.get_variable("deconv4_W", shape = [2, 2, 128, 256], initializer=tf.contrib.layers.xavier_initializer())
deconv3_W = tf.Variable(tf.zeros([2, 2, 64, 256]), name = "deconv3_W")
deconv3_W = tf.get_variable("deconv3_W", shape = [2, 2, 64, 256], initializer=tf.contrib.layers.xavier_initializer())
deconv2_W = tf.Variable(tf.zeros([3, 3, 32, 128]), name = "deconv2_W")
deconv2_W = tf.get_variable("deconv2_W", shape = [3, 3, 32, 128], initializer=tf.contrib.layers.xavier_initializer())
deconv1_W = tf.Variable(tf.zeros([3, 3, scale_n, 64]), name = "deconv1_W")
deconv1_W = tf.get_variable("deconv1_W", shape = [3, 3, scale_n, 64], initializer=tf.contrib.layers.xavier_initializer())


def cnn_encoder(input_matrix):
	conv1 = tf.nn.conv2d(
	  input = input_matrix,
	  filter = conv1_W,
	  strides=(1, 1, 1, 1),
	  padding = "SAME")
	conv1 = tf.nn.selu(conv1)

	conv2 = tf.nn.conv2d(
	  input = conv1,
	  filter = conv2_W,
	  strides=(1, 2, 2, 1),
	  padding = "SAME")
	conv2 = tf.nn.selu(conv2)

	conv3 = tf.nn.conv2d(
	  input = conv2,
	  filter = conv3_W,
	  strides=(1, 2, 2, 1),
	  padding = "SAME")
	conv3 = tf.nn.selu(conv3)

	conv4 = tf.nn.conv2d(
	  input = conv3,
	  filter = conv4_W,
	  strides=(1, 2, 2, 1),
	  padding = "SAME")
	conv4 = tf.nn.selu(conv4)

	return  conv1, conv2, conv3, conv4


def conv1_lstm(conv1_out):
	convlstm_layer = tf.contrib.rnn.ConvLSTMCell(
				conv_ndims = 2,
				input_shape = [sensor_n, sensor_n, 32],
				output_channels = 32,
				kernel_shape = [2, 2],
				use_bias = True,
				skip_connection = False,
				forget_bias = 1.0,
				initializers = None,
				name="conv1_lstm_cell")

	outputs, state = tf.nn.dynamic_rnn(convlstm_layer, conv1_out, time_major = False, dtype = conv1_out.dtype)

	# attention based on inner-product between feature representation of last step and other steps
	attention_w = []
	for k in range(step_max):
	  attention_w.append(tf.reduce_sum(tf.multiply(outputs[0][k], outputs[0][-1]))/step_max)
	attention_w = tf.reshape(tf.nn.softmax(tf.stack(attention_w)), [1, step_max])

	outputs = tf.reshape(outputs[0], [step_max, -1])
	outputs = tf.matmul(attention_w, outputs)
	outputs = tf.reshape(outputs, [1, sensor_n, sensor_n, 32])

	return outputs, state[0], attention_w


def conv2_lstm(conv2_out):
	convlstm_layer = tf.contrib.rnn.ConvLSTMCell(
				conv_ndims = 2,
				input_shape = [int(math.ceil(float(sensor_n)/2)), int(math.ceil(float(sensor_n)/2)), 64],
				output_channels = 64,
				kernel_shape = [2, 2],
				use_bias = True,
				skip_connection = False,
				forget_bias = 1.0,
				initializers = None,
				name="conv2_lstm_cell")

	outputs, state = tf.nn.dynamic_rnn(convlstm_layer, conv2_out, time_major = False, dtype = conv2_out.dtype)

	# attention based on inner-product between feature representation of last step and other steps
	attention_w = []
	for k in range(step_max):
	  attention_w.append(tf.reduce_sum(tf.multiply(outputs[0][k], outputs[0][-1]))/step_max)
	attention_w = tf.reshape(tf.nn.softmax(tf.stack(attention_w)), [1, step_max])

	outputs = tf.reshape(outputs[0], [step_max, -1])
	outputs = tf.matmul(attention_w, outputs)
	outputs = tf.reshape(outputs, [1, int(math.ceil(float(sensor_n)/2)), int(math.ceil(float(sensor_n)/2)), 64])

	return outputs, state[0], attention_w


def conv3_lstm(conv3_out):
	convlstm_layer = tf.contrib.rnn.ConvLSTMCell(
				conv_ndims = 2,
				input_shape = [int(math.ceil(float(sensor_n)/4)), int(math.ceil(float(sensor_n)/4)), 128],
				output_channels = 128,
				kernel_shape = [2, 2],
				use_bias = True,
				skip_connection = False,
				forget_bias = 1.0,
				initializers = None,
				name="conv3_lstm_cell")

	outputs, state = tf.nn.dynamic_rnn(convlstm_layer, conv3_out, time_major = False, dtype = conv3_out.dtype)

	attention_w = []
	for k in range(step_max):
	  attention_w.append(tf.reduce_sum(tf.multiply(outputs[0][k], outputs[0][-1]))/step_max)
	attention_w = tf.reshape(tf.nn.softmax(tf.stack(attention_w)), [1, step_max])

	outputs = tf.reshape(outputs[0], [step_max, -1])
	outputs = tf.matmul(attention_w, outputs)
	outputs = tf.reshape(outputs, [1, int(math.ceil(float(sensor_n)/4)), int(math.ceil(float(sensor_n)/4)), 128])

	return outputs, state[0], attention_w


def conv4_lstm(conv4_out):
	convlstm_layer = tf.contrib.rnn.ConvLSTMCell(
				conv_ndims = 2,
				input_shape = [int(math.ceil(float(sensor_n)/8)), int(math.ceil(float(sensor_n)/8)), 256],
				output_channels = 256,
				kernel_shape = [2, 2],
				use_bias = True,
				skip_connection = False,
				forget_bias = 1.0,
				initializers = None,
				name="conv4_lstm_cell")

	#initial_state = convlstm_layer.zero_state(batch_size, dtype = tf.float32)
	outputs, state = tf.nn.dynamic_rnn(convlstm_layer, conv4_out, time_major = False, dtype = conv4_out.dtype)

	attention_w = []
	for k in range(step_max):
	  attention_w.append(tf.reduce_sum(tf.multiply(outputs[0][k], outputs[0][-1]))/step_max)
	attention_w = tf.reshape(tf.nn.softmax(tf.stack(attention_w)), [1, step_max])

	outputs = tf.reshape(outputs[0], [step_max, -1])
	outputs = tf.matmul(attention_w, outputs)
	outputs = tf.reshape(outputs, [1, int(math.ceil(float(sensor_n)/8)), int(math.ceil(float(sensor_n)/8)), 256])

	return outputs, state[0], attention_w


def cnn_decoder(conv1_lstm_out, conv2_lstm_out, conv3_lstm_out, conv4_lstm_out):
	conv1_lstm_out = tf.reshape(conv1_lstm_out, [1, sensor_n, sensor_n, 32])
	conv2_lstm_out = tf.reshape(conv2_lstm_out, [1, int(math.ceil(float(sensor_n)/2)), int(math.ceil(float(sensor_n)/2)), 64])
	conv3_lstm_out = tf.reshape(conv3_lstm_out, [1, int(math.ceil(float(sensor_n)/4)), int(math.ceil(float(sensor_n)/4)), 128])
	conv4_lstm_out = tf.reshape(conv4_lstm_out, [1, int(math.ceil(float(sensor_n)/8)), int(math.ceil(float(sensor_n)/8)), 256])

	deconv4 = tf.nn.conv2d_transpose(
	  value = conv4_lstm_out,
	  filter = deconv4_W,
	  output_shape = [1, int(math.ceil(float(sensor_n)/4)), int(math.ceil(float(sensor_n)/4)), 128],
	  strides = (1, 2, 2, 1),
	  padding = "SAME")
	deconv4 = tf.nn.selu(deconv4)
	deconv4_concat = tf.concat([deconv4, conv3_lstm_out], axis = 3)

	deconv3 = tf.nn.conv2d_transpose(
	  value = deconv4_concat,
	  filter = deconv3_W,
	  output_shape = [1, int(math.ceil(float(sensor_n)/2)), int(math.ceil(float(sensor_n)/2)), 64],
	  strides = (1, 2, 2, 1),
	  padding = "SAME")
	deconv3 = tf.nn.selu(deconv3)
	deconv3_concat = tf.concat([deconv3, conv2_lstm_out], axis = 3)

	deconv2 = tf.nn.conv2d_transpose(
	  value = deconv3_concat,
	  filter = deconv2_W,
	  output_shape = [1, sensor_n, sensor_n, 32],
	  strides = (1, 2, 2, 1),
	  padding = "SAME")
	deconv2 = tf.nn.selu(deconv2)

	deconv2_concat = tf.concat([deconv2, conv1_lstm_out], axis = 3)

	deconv1 = tf.nn.conv2d_transpose(
	  value = deconv2_concat,
	  filter = deconv1_W,
	  output_shape = [1, sensor_n, sensor_n, scale_n],
	  strides = (1, 1, 1, 1),
	  padding = "SAME")
	deconv1 = tf.nn.selu(deconv1)
	deconv1 = tf.reshape(deconv1, [1, sensor_n, sensor_n, scale_n])
	return deconv1


conv1_out, conv2_out, conv3_out, conv4_out = cnn_encoder(data_input)
conv1_out = tf.reshape(conv1_out, [-1, step_max, sensor_n, sensor_n, 32])
conv2_out = tf.reshape(conv2_out, [-1, step_max, int(math.ceil(float(sensor_n)/2)), int(math.ceil(float(sensor_n)/2)), 64])
conv3_out = tf.reshape(conv3_out, [-1, step_max, int(math.ceil(float(sensor_n)/4)), int(math.ceil(float(sensor_n)/4)), 128])
conv4_out = tf.reshape(conv4_out, [-1, step_max, int(math.ceil(float(sensor_n)/8)), int(math.ceil(float(sensor_n)/8)), 256])

conv1_lstm_attention_out, conv1_lstm_last_out, atten_weight_1 = conv1_lstm(conv1_out)
conv2_lstm_attention_out, conv2_lstm_last_out, atten_weight_2 = conv2_lstm(conv2_out)
conv3_lstm_attention_out, conv3_lstm_last_out, atten_weight_3 = conv3_lstm(conv3_out)
conv4_lstm_attention_out, conv4_lstm_last_out, atten_weight_4 = conv4_lstm(conv4_out)

deconv_out = cnn_decoder(conv1_lstm_attention_out, conv2_lstm_attention_out, conv3_lstm_attention_out, conv4_lstm_attention_out)

# loss function: reconstruction error of last step matrix
loss = tf.reduce_mean(tf.square(data_input[-1] - deconv_out))
optimizer = tf.train.AdamOptimizer(learning_rate = learning_rate).minimize(loss)

init = tf.global_variables_initializer()
saver = tf.train.Saver(max_to_keep = 10)

starting_iter = 0
if train_test_label == 1: # model training
	with tf.Session(config=tf.ConfigProto(inter_op_parallelism_threads=80, intra_op_parallelism_threads=80)) as sess:
		if starting_iter == 0:
			sess.run(init)
		else:
			saver.restore(sess, model_path + str(starting_iter-1) + ".ckpt")
		start = timeit.default_timer()
		for idx in range(starting_iter, starting_iter+training_iters):
			print ("training iteration " + str(idx) + "...")
			for train_id in range(train_start_id, train_end_id):
				matrix_data_path = train_data_path + "train_data_" + str(train_id) + ".npy"
				matrix_gt = np.load(matrix_data_path)
				matrix_gt = np.transpose(matrix_gt, (0, 2, 3, 1))

				feed_dict = {data_input: np.asarray(matrix_gt)}
				a, loss_value = sess.run([optimizer, loss], feed_dict)

			if idx % save_model_step == 0:
				if not os.path.exists(model_path):
					os.makedirs(model_path)
				save_path = saver.save(sess, model_path + str(idx) + ".ckpt")
				print ("mse of last train data: " + str(loss_value))

		stop = timeit.default_timer()

		print(str(training_iters) + " iterations training finish, use time: " + str(stop - start))

if train_test_label == 0: # model test

	with tf.Session(config=tf.ConfigProto(inter_op_parallelism_threads=20, intra_op_parallelism_threads=20)) as sess:
		restore_idx = 4 # setting restore_idx of learned model
		saver.restore(sess, model_path + str(restore_idx) + ".ckpt")

		reconstructed_data_path = matrix_data_path + "reconstructed_data/"
		if not os.path.exists(reconstructed_data_path):
			os.makedirs(reconstructed_data_path)

		print ("model test: generate recontrucuted matrices"+ "...")
		loss_value_ave = 0
		for test_id in range(test_start_id, test_end_id):
			matrix_data_path = test_data_path + 'test_data_' + str(test_id) + ".npy"
			matrix_gt = np.load(matrix_data_path)
			matrix_gt = np.transpose(matrix_gt, (0, 2, 3, 1))

			feed_dict = {data_input: np.asarray(matrix_gt)}
			reconstructed_matrix = sess.run([deconv_out], feed_dict)

			path_temp = os.path.join(reconstructed_data_path, 'reconstructed_data_' + str(test_id) + ".npy")
			np.save(path_temp, np.asarray(reconstructed_matrix))

		print ("reconstructed matrices generation finish.")

