import time
import os
import numpy as np
import tensorflow as tf
import sys
import cv2

slim = tf.contrib.slim
sys.path.append(os.getcwd())
from nets import model as model
from matplotlib import pyplot as plt
from utils.tf_records import read_tfrecord_and_decode_into_image_annotation_pair_tensors

tf.app.flags.DEFINE_string('model_type', 'refinenet', '')
tf.app.flags.DEFINE_string('test_data_path', 'data/test.tfrecords', '')
tf.app.flags.DEFINE_string('gpu_list', '1', '')
tf.app.flags.DEFINE_integer('num_classes', 2, '')
tf.app.flags.DEFINE_string('checkpoint_path', 'checkpoints/', '')
tf.app.flags.DEFINE_string('result_path', 'result/', '')
tf.app.flags.DEFINE_integer('test_size', 384, '')

FLAGS = tf.app.flags.FLAGS

def main(argv=None):
	os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list

	if not os.path.exists(FLAGS.result_path):
		os.makedirs(FLAGS.result_path)

	filename_queue = tf.train.string_input_producer([FLAGS.test_data_path], num_epochs=1)
	image, annotation = read_tfrecord_and_decode_into_image_annotation_pair_tensors(filename_queue)

	image_batch_tensor = tf.expand_dims(image, axis=0)
	annotation_batch_tensor = tf.expand_dims(annotation, axis=0)

	input_image_shape = tf.shape(image_batch_tensor)
	image_height_width = input_image_shape[1:3]
	image_height_width_float = tf.to_float(image_height_width)
	image_height_width_multiple = tf.to_int32(tf.round(image_height_width_float / 32) * 32)

	image_batch_tensor = tf.image.resize_images(image_batch_tensor, image_height_width_multiple)

	global_step = tf.get_variable('global_step', [], initializer=tf.constant_initializer(0), trainable=False)
	logits = model.model(FLAGS.model_type, image_batch_tensor, is_training=False)
	pred = tf.argmax(logits, dimension=3)
	pred = tf.expand_dims(pred, 3)
	pred = tf.image.resize_bilinear(images=pred, size=image_height_width)
	annotation_batch_tensor = tf.image.resize_bilinear(images=annotation_batch_tensor, size=image_height_width)
	annotation_batch_tensor = tf.div(annotation_batch_tensor, 255)

	pred = tf.reshape(pred, [-1, ])
	gt = tf.reshape(annotation_batch_tensor, [-1, ])

	acc, acc_update_op = tf.contrib.metrics.streaming_accuracy(pred, gt)
	miou, miou_update_op = tf.contrib.metrics.streaming_mean_iou(pred, gt, num_classes=FLAGS.num_classes)

	with tf.get_default_graph().as_default():
		global_vars_init_op = tf.global_variables_initializer()
		local_vars_init_op = tf.local_variables_initializer()
		init = tf.group(local_vars_init_op, global_vars_init_op)

		variable_averages = tf.train.ExponentialMovingAverage(0.997, global_step)
		saver = tf.train.Saver(variable_averages.variables_to_restore())

		gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=1.0)
		config = tf.ConfigProto(allow_soft_placement=True,
		                        log_device_placement=False,
		                        gpu_options=gpu_options)
		config.gpu_options.allow_growth = True

		with tf.Session(config=config) as sess:
			sess.run(init)
			ckpt_state = tf.train.get_checkpoint_state(FLAGS.checkpoint_path)
			model_path = os.path.join(FLAGS.checkpoint_path, os.path.basename(ckpt_state.model_checkpoint_path))
			print('Restore from {}'.format(model_path))
			saver.restore(sess, model_path)

			coord = tf.train.Coordinator()
			threads = tf.train.start_queue_runners(coord=coord)

			for i in range(150):
				start = time.time()
				image_np, annotation_np, pred_np, tmp_acc, tmp_miou = sess.run(
					[image, annotation, pred, acc_update_op, miou_update_op])
				_diff_time = time.time() - start
				print('{}: cost {:.0f}ms').format(i, _diff_time * 1000)
				# upsampled_predictions = pred_np.squeeze()
				# plt.subplot(131)
				# plt.imshow(image_np)
				# plt.subplot(132)
				# plt.imshow(annotation_np.squeeze(), cmap='gray')
				# plt.subplot(133)
				# plt.imshow(np.reshape(pred_np, (annotation_np.shape[0], annotation_np.shape[1])).squeeze(), cmap='gray')
				# plt.savefig(os.path.join(FLAGS.result_path, str(i) + '.png'))
				prediction = np.reshape(pred_np, (annotation_np.shape[0], annotation_np.shape[1])).squeeze() * 255
				cv2.imwrite(os.path.join(FLAGS.result_path, str(i) + '.png'), prediction)
			print('Test Finished !')

	coord.request_stop()
	coord.join(threads)


if __name__ == '__main__':
	tf.app.run()