#!/usr/bin/env python3
import os.path
import tensorflow as tf
import helper
import warnings
from distutils.version import LooseVersion
import project_tests as tests
import time


# Check TensorFlow Version
assert LooseVersion(tf.__version__) >= LooseVersion('1.0'), 'Please use TensorFlow version 1.0 or newer.  You are using {}'.format(tf.__version__)
print('TensorFlow Version: {}'.format(tf.__version__))

# Check for a GPU
if not tf.test.gpu_device_name():
	warnings.warn('No GPU found. Please use a GPU to train your neural network.')
else:
	print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))



def load_vgg(sess, vgg_path):
	"""
	Load Pretrained VGG Model into TensorFlow.
	:param sess: TensorFlow Session
	:param vgg_path: Path to vgg folder, containing "variables/" and "saved_model.pb"
	:return: Tuple of Tensors from VGG model (image_input, keep_prob, layer3_out, layer4_out, layer7_out)
	"""
	vgg_tag = 'vgg16'
	vgg_input_tensor_name = 'image_input:0'
	vgg_keep_prob_tensor_name = 'keep_prob:0'
	vgg_layer3_out_tensor_name = 'layer3_out:0'
	vgg_layer4_out_tensor_name = 'layer4_out:0'
	vgg_layer7_out_tensor_name = 'layer7_out:0'
	# load graph from file
	tf.saved_model.loader.load(sess, [vgg_tag], vgg_path)
	# assign it to a variable
	graph = tf.get_default_graph()
	# load weights from each layer
	w1  = graph.get_tensor_by_name(vgg_input_tensor_name)
	keep  = graph.get_tensor_by_name(vgg_keep_prob_tensor_name)
	w3  = graph.get_tensor_by_name(vgg_layer3_out_tensor_name)
	w4  = graph.get_tensor_by_name(vgg_layer4_out_tensor_name)
	w7  = graph.get_tensor_by_name(vgg_layer7_out_tensor_name)
	
	
	return w1, keep, w3, w4, w7
	
tests.test_load_vgg(load_vgg, tf)


def layers(vgg_layer3_out, vgg_layer4_out, vgg_layer7_out, num_classes):
	"""
	Create the layers for a fully convolutional network.  Build skip-layers using the vgg layers.
	:param vgg_layer3_out: TF Tensor for VGG Layer 3 output
	:param vgg_layer4_out: TF Tensor for VGG Layer 4 output
	:param vgg_layer7_out: TF Tensor for VGG Layer 7 output
	:param num_classes: Number of classes to classify
	:return: The Tensor for the last layer of output
	"""
	# Feed output of base network into new network
	# Regularize to penalize large weights, avoids overfitting 
	# converting encoder pooling output layers to 1x1 conv 
	vgg_layer7_1x1 = tf.layers.conv2d(	inputs=vgg_layer7_out,\
										filters=num_classes,\
										kernel_size=1,\
										padding='same',\
										kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-3),\
										name='vgg_layer7_1x1')
	# now a deconvlution layer
	# upsample by 2, making it same size as VGG layer 4
	#filter_size = vgg_layer4_out.get_shape().as_list()[-1]
	fcn_layer1 = tf.layers.conv2d_transpose(	inputs=vgg_layer7_1x1,\
												filters=num_classes,\
												kernel_size=4,\
												strides=2,\
												padding='same',\
												kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-3),\
												name='fcn_layer1')
	# converting encoder pooling output layers to 1x1 conv 
	vgg_layer4_1x1 = tf.layers.conv2d(	inputs=vgg_layer4_out,\
										filters=num_classes,\
										kernel_size=1,\
										padding='same',\
										kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-3),\
										name='vgg_layer4_1x1')
	# skip connection pool4 -> layer9
	fcn_layer1_skip = tf.add(fcn_layer1, vgg_layer4_1x1, name='fcn_layer1_skip')
	# upsmaple by 2
	#filter_size = vgg_layer3_out.get_shape().as_list()[-1]
	fcn_layer2 = tf.layers.conv2d_transpose(	inputs=fcn_layer1_skip,\
												filters=num_classes,\
												kernel_size=4,\
												strides=2,\
												padding='same',\
												kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-3),\
												name='fcn_layer2')
	# converting encoder pooling output layers to 1x1 conv 
	vgg_layer3_1x1 = tf.layers.conv2d(	inputs=vgg_layer3_out,\
										filters=num_classes,\
										kernel_size=1,\
										padding='same',\
										kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-3),\
										name='vgg_layer3_1x1')
	# skip connection pool3 -> layer10
	fcn_layer2_skip = tf.add(fcn_layer2, vgg_layer3_1x1, name='fcn_layer2_skip')
	# upsample by 4
	# same size as input image
	output = tf.layers.conv2d_transpose(	inputs=fcn_layer2_skip,\
											filters=num_classes,\
											kernel_size=16,\
											strides=8,\
											padding='same',\
											kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-3),\
											name='output')
	
	
	return output
tests.test_layers(layers)


def optimize(nn_last_layer, correct_label, learning_rate, num_classes):
	"""
	Build the TensorFLow loss and optimizer operations.
	:param nn_last_layer: TF Tensor of the last layer in the neural network
	:param correct_label: TF Placeholder for the correct label image
	:param learning_rate: TF Placeholder for the learning rate
	:param num_classes: Number of classes to classify
	:return: Tuple of (logits, train_op, cross_entropy_loss)
	"""
	# Resize Logits, flatten
	logits = tf.reshape(nn_last_layer, (-1, num_classes))
	correct_label_flatten = tf.reshape(correct_label, (-1, num_classes))
	# calculate loss and take mean
	cross_entropy_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=correct_label_flatten))
	# Optimizer 
	opt = tf.train.AdamOptimizer(learning_rate).minimize(cross_entropy_loss)
	return logits, opt, cross_entropy_loss
	
tests.test_optimize(optimize)


def train_nn(sess, epochs, batch_size, get_batches_fn, train_op, cross_entropy_loss, input_image,
			 correct_label, keep_prob, learning_rate):
	"""
	Train neural network and print out the loss during training.
	:param sess: TF Session
	:param epochs: Number of epochs
	:param batch_size: Batch size
	:param get_batches_fn: Function to get batches of training data.  Call using get_batches_fn(batch_size)
	:param train_op: TF Operation to train the neural network
	:param cross_entropy_loss: TF Tensor for the amount of loss
	:param input_image: TF Placeholder for input images
	:param correct_label: TF Placeholder for label images
	:param keep_prob: TF Placeholder for dropout keep probability
	:param learning_rate: TF Placeholder for learning rate
	"""
	
	keep_prob_value = 0.75
	learning_rate_value = 0.0001
	
	for epoch in range(epochs):
		total_loss = 0
		batch = 0
		for image, label in get_batches_fn(batch_size):
			# training
			loss, optimizer = sess.run([cross_entropy_loss, train_op], feed_dict={input_image: image, correct_label: label, keep_prob: keep_prob_value, learning_rate: learning_rate_value}) 
			print('Epoch {} - Batch {} - Loss {}'.format(epoch + 1, batch+1, loss))
			total_loss += loss;
			batch += 1
		
		print(" ")	
		print("EPOCH {} ...".format(epoch + 1))
		print("Loss = {:.3f}".format(total_loss))
		print()
	
tests.test_train_nn(train_nn)


def run():
	num_classes = 2
	image_shape = (160, 576)  # KITTI dataset uses 160x576 images
	data_dir = './data'
	runs_dir = './runs'
	tests.test_for_kitti_dataset(data_dir)
	
	

	# Download pretrained vgg model
	helper.maybe_download_pretrained_vgg(data_dir)

	# OPTIONAL: Train and Inference on the cityscapes dataset instead of the Kitti dataset.
	# You'll need a GPU with at least 10 teraFLOPS to train on.
	#  https://www.cityscapes-dataset.com/
	
	# Model parameters
	EPOCHS = 110
	BATCH_SIZE = 16
		
	# create tensors 
	correct_label = tf.placeholder(tf.float32, [None, image_shape[0], image_shape[1], num_classes])
	learning_rate = tf.placeholder(tf.float32)
	keep_prob = tf.placeholder(tf.float32)
	

	with tf.Session() as sess:
		
		
		
		# Path to vgg model
		vgg_path = os.path.join(data_dir, 'vgg')
		# Create function to get batches
		get_batches_fn = helper.gen_batch_function(os.path.join(data_dir, 'data_road/training'), image_shape)
		
		# OPTIONAL: Augment Images for better results
		#  https://datascience.stackexchange.com/questions/5224/how-to-prepare-augment-images-for-neural-network

		# Build NN using load_vgg, layers, and optimize function
		image_input, keep_prob, layer3, layer4, layer7 = load_vgg(sess, vgg_path)
		model_result = layers(layer3, layer4, layer7, num_classes) 
		
		logits, opt, cross_entropy_loss = optimize(model_result, correct_label, learning_rate, num_classes)

		# Create saver
		saver = tf.train.Saver()
		
		# Initialize all variables
		sess.run( tf.global_variables_initializer() )
		sess.run( tf.local_variables_initializer() )
		# Train NN using the train_nn function
		train_nn(sess, EPOCHS, BATCH_SIZE, get_batches_fn, opt, cross_entropy_loss, image_input, correct_label, keep_prob, learning_rate)
		
		# TODO: Save inference data using helper.save_inference_samples
		helper.save_inference_samples(runs_dir, data_dir, sess, image_shape, logits, keep_prob, image_input)

		# OPTIONAL: Apply the trained model to a video

		#save model
		print('Saving model...')
		save_time = str(int(time.time()))
		save_path = saver.save(sess, './models/model_'+save_time+'/model_'+save_time)

if __name__ == '__main__':
	run()
