import tensorflow as tf
import os.path
import warnings
from distutils.version import LooseVersion
import glob

import helper
import project_tests as tests

#--------------------------
# USER-SPECIFIED DATA
#--------------------------

# Tune these parameters

NUMBER_OF_CLASSES = 2
IMAGE_SHAPE = (160, 576)

EPOCHS = 20
BATCH_SIZE = 1

LEARNING_RATE = 0.0001
DROPOUT = 0.75

# Specify these directory paths

DATA_DIRECTORY = './data'
RUNS_DIRECTORY = './runs'
TRAINING_DATA_DIRECTORY ='./data/data_road/training'
NUMBER_OF_IMAGES = len(glob.glob('./data/data_road/training/calib/*.*'))
VGG_PATH = './data/vgg'

all_training_losses = [] # Used for plotting to visualize if our training is going well given parameters

#--------------------------
# DEPENDENCY CHECK
#--------------------------

# 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()))

#--------------------------
# PLACEHOLDER TENSORS
#--------------------------

correct_label = tf.placeholder(tf.float32, [None, IMAGE_SHAPE[0], IMAGE_SHAPE[1], NUMBER_OF_CLASSES])
learning_rate = tf.placeholder(tf.float32)
keep_prob = tf.placeholder(tf.float32)

#--------------------------
# FUNCTIONS
#--------------------------

def load_vgg(sess, vgg_path):
  """
  Load Pretrained VGG Model into TensorFlow.
  sess: TensorFlow Session
  vgg_path: Path to vgg folder, containing "variables/" and "saved_model.pb"
  return: Tuple of Tensors from VGG model (image_input, keep_prob, layer3, layer4, layer7)
  """
  
  # load the model and weights
  model = tf.saved_model.loader.load(sess, ['vgg16'], vgg_path)

  # Get Tensors to be returned from graph
  graph = tf.get_default_graph()
  image_input = graph.get_tensor_by_name('image_input:0')
  keep_prob = graph.get_tensor_by_name('keep_prob:0')
  layer3 = graph.get_tensor_by_name('layer3_out:0')
  layer4 = graph.get_tensor_by_name('layer4_out:0')
  layer7 = graph.get_tensor_by_name('layer7_out:0')

  return image_input, keep_prob, layer3, layer4, layer7


def conv_1x1(layer, layer_name):
  """ Return the output of a 1x1 convolution of a layer """
  return tf.layers.conv2d(inputs = layer,
                          filters =  NUMBER_OF_CLASSES,
                          kernel_size = (1, 1),
                          strides = (1, 1),
                          name = layer_name)


def upsample(layer, k, s, layer_name):
  """ Return the output of transpose convolution given kernel_size k and strides s """
  return tf.layers.conv2d_transpose(inputs = layer,
                                    filters = NUMBER_OF_CLASSES,
                                    kernel_size = (k, k),
                                    strides = (s, s),
                                    padding = 'same',
                                    name = layer_name)


def layers(vgg_layer3_out, vgg_layer4_out, vgg_layer7_out, num_classes = NUMBER_OF_CLASSES):
  """
  Create the layers for a fully convolutional network.  Build skip-layers using the vgg layers.
  vgg_layerX_out: TF Tensor for VGG Layer X output
  num_classes: Number of classes to classify
  return: The Tensor for the last layer of output
  """

  # Use a shorter variable name for simplicity
  layer3, layer4, layer7 = vgg_layer3_out, vgg_layer4_out, vgg_layer7_out

  # Apply a 1x1 convolution to encoder layers
  layer3x = conv_1x1(layer = layer3, layer_name = "layer3conv1x1")
  layer4x = conv_1x1(layer = layer4, layer_name = "layer4conv1x1")
  layer7x = conv_1x1(layer = layer7, layer_name = "layer7conv1x1")
 
  # Add decoder layers to the network with skip connections and upsampling
  # Note: the kernel size and strides are the same as the example in Udacity Lectures
  #       Semantic Segmentation Scene Understanding Lesson 10-9: FCN-8 - Decoder
  decoderlayer1 = upsample(layer = layer7x, k = 4, s = 2, layer_name = "decoderlayer1")
  decoderlayer2 = tf.add(decoderlayer1, layer4x, name = "decoderlayer2")
  decoderlayer3 = upsample(layer = decoderlayer2, k = 4, s = 2, layer_name = "decoderlayer3")
  decoderlayer4 = tf.add(decoderlayer3, layer3x, name = "decoderlayer4")
  decoderlayer_output = upsample(layer = decoderlayer4, k = 16, s = 8, layer_name = "decoderlayer_output")

  return decoderlayer_output


def optimize(nn_last_layer, correct_label, learning_rate, num_classes = NUMBER_OF_CLASSES):
  """
  Build the TensorFLow loss and optimizer operations.
  nn_last_layer: TF Tensor of the last layer in the neural network
  correct_label: TF Placeholder for the correct label image
  learning_rate: TF Placeholder for the learning rate
  num_classes: Number of classes to classify
  return: Tuple of (logits, train_op, cross_entropy_loss)
  """
  # Reshape 4D tensors to 2D, each row represents a pixel, each column a class
  logits = tf.reshape(nn_last_layer, (-1, num_classes))
  class_labels = tf.reshape(correct_label, (-1, num_classes))

  # The cross_entropy_loss is the cost which we are trying to minimize to yield higher accuracy
  cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits = logits, labels = class_labels)
  cross_entropy_loss = tf.reduce_mean(cross_entropy)

  # The model implements this operation to find the weights/parameters that would yield correct pixel labels
  train_op = tf.train.AdamOptimizer(learning_rate).minimize(cross_entropy_loss)

  return logits, train_op, cross_entropy_loss


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.
  sess: TF Session
  epochs: Number of epochs
  batch_size: Batch size
  get_batches_fn: Function to get batches of training data.  Call using get_batches_fn(batch_size)
  train_op: TF Operation to train the neural network
  cross_entropy_loss: TF Tensor for the amount of loss
  input_image: TF Placeholder for input images
  correct_label: TF Placeholder for label images
  keep_prob: TF Placeholder for dropout keep probability
  learning_rate: TF Placeholder for learning rate
  """

  for epoch in range(EPOCHS):
    
    losses, i = [], 0
    
    for images, labels in get_batches_fn(BATCH_SIZE):
        
      i += 1
    
      feed = { input_image: images,
               correct_label: labels,
               keep_prob: DROPOUT,
               learning_rate: LEARNING_RATE }
        
      _, partial_loss = sess.run([train_op, cross_entropy_loss], feed_dict = feed)
      
      print("---> iteration: ", i, " partial loss:", partial_loss)
      losses.append(partial_loss)
          
    training_loss = sum(losses) / len(losses)
    all_training_losses.append(training_loss)
    
    print("------------------")
    print("epoch: ", epoch + 1, " of ", EPOCHS, "training loss: ", training_loss)
    print("------------------")
    

def run_tests():
  tests.test_layers(layers)
  tests.test_optimize(optimize)
  tests.test_for_kitti_dataset(DATA_DIRECTORY)
  tests.test_train_nn(train_nn)


def run():
  """ Run a train a model and save output images resulting from the test image fed on the trained model """

  # Get vgg model if we can't find it where it should be
  helper.maybe_download_pretrained_vgg(DATA_DIRECTORY)

  # A function to get batches
  get_batches_fn = helper.gen_batch_function(TRAINING_DATA_DIRECTORY, IMAGE_SHAPE)
  
  with tf.Session() as session:
        
    # Returns the three layers, keep probability and input layer from the vgg architecture
    image_input, keep_prob, layer3, layer4, layer7 = load_vgg(session, VGG_PATH)

    # The resulting network architecture from adding a decoder on top of the given vgg model
    model_output = layers(layer3, layer4, layer7, NUMBER_OF_CLASSES)

    # Returns the output logits, training operation and cost operation to be used
    # - logits: each row represents a pixel, each column a class
    # - train_op: function used to get the right parameters to the model to correctly label the pixels
    # - cross_entropy_loss: function outputting the cost which we are minimizing, lower cost should yield higher accuracy
    logits, train_op, cross_entropy_loss = optimize(model_output, correct_label, learning_rate, NUMBER_OF_CLASSES)
    
    # Initialize all variables
    session.run(tf.global_variables_initializer())
    session.run(tf.local_variables_initializer())

    # Train the neural network
    train_nn(session, EPOCHS, BATCH_SIZE, get_batches_fn, 
             train_op, cross_entropy_loss, image_input,
             correct_label, keep_prob, learning_rate)

    # Run the model with the test images and save each painted output image (roads painted green)
    helper.save_inference_samples(RUNS_DIRECTORY, DATA_DIRECTORY, session, IMAGE_SHAPE, logits, keep_prob, image_input)

#--------------------------
# MAIN
#--------------------------
if __name__ == "__main__":
  run_tests()
  run() # Run a train a model and save output images resulting from the test image fed on the trained model
  print(all_training_losses)
  

[((54, 16, 54, 101), 'tensorflow.placeholder', 'tf.placeholder', ({(54, 31, 54, 41): 'tf.float32', (54, 43, 54, 100): '[None, IMAGE_SHAPE[0], IMAGE_SHAPE[1], NUMBER_OF_CLASSES]'}, {}), '(tf.float32, [None, IMAGE_SHAPE[0], IMAGE_SHAPE[1],\n    NUMBER_OF_CLASSES])', True, 'import tensorflow as tf\n'), ((55, 16, 55, 42), 'tensorflow.placeholder', 'tf.placeholder', ({(55, 31, 55, 41): 'tf.float32'}, {}), '(tf.float32)', True, 'import tensorflow as tf\n'), ((56, 12, 56, 38), 'tensorflow.placeholder', 'tf.placeholder', ({(56, 27, 56, 37): 'tf.float32'}, {}), '(tf.float32)', True, 'import tensorflow as tf\n'), ((30, 23, 30, 71), 'glob.glob', 'glob.glob', ({(30, 33, 30, 70): '"""./data/data_road/training/calib/*.*"""'}, {}), "('./data/data_road/training/calib/*.*')", False, 'import glob\n'), ((40, 7, 40, 35), 'distutils.version.LooseVersion', 'LooseVersion', ({(40, 20, 40, 34): 'tf.__version__'}, {}), '(tf.__version__)', False, 'from distutils.version import LooseVersion\n'), ((40, 39, 40, 58), 'distutils.version.LooseVersion', 'LooseVersion', ({(40, 52, 40, 57): '"""1.0"""'}, {}), "('1.0')", False, 'from distutils.version import LooseVersion\n'), ((45, 7, 45, 32), 'tensorflow.test.gpu_device_name', 'tf.test.gpu_device_name', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((46, 2, 46, 79), 'warnings.warn', 'warnings.warn', ({(46, 16, 46, 78): '"""No GPU found. Please use a GPU to train your neural network."""'}, {}), "('No GPU found. Please use a GPU to train your neural network.')", False, 'import warnings\n'), ((71, 10, 71, 63), 'tensorflow.saved_model.loader.load', 'tf.saved_model.loader.load', ({(71, 37, 71, 41): 'sess', (71, 43, 71, 52): "['vgg16']", (71, 54, 71, 62): 'vgg_path'}, {}), "(sess, ['vgg16'], vgg_path)", True, 'import tensorflow as tf\n'), ((74, 10, 74, 32), 'tensorflow.get_default_graph', 'tf.get_default_graph', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((86, 9, 90, 44), 'tensorflow.layers.conv2d', 'tf.layers.conv2d', (), '', True, 'import tensorflow as tf\n'), ((95, 9, 100, 54), 'tensorflow.layers.conv2d_transpose', 'tf.layers.conv2d_transpose', (), '', True, 'import tensorflow as tf\n'), ((123, 18, 123, 72), 'tensorflow.add', 'tf.add', (), '', True, 'import tensorflow as tf\n'), ((125, 18, 125, 72), 'tensorflow.add', 'tf.add', (), '', True, 'import tensorflow as tf\n'), ((141, 11, 141, 55), 'tensorflow.reshape', 'tf.reshape', ({(141, 22, 141, 35): 'nn_last_layer', (141, 37, 141, 54): '(-1, num_classes)'}, {}), '(nn_last_layer, (-1, num_classes))', True, 'import tensorflow as tf\n'), ((142, 17, 142, 61), 'tensorflow.reshape', 'tf.reshape', ({(142, 28, 142, 41): 'correct_label', (142, 43, 142, 60): '(-1, num_classes)'}, {}), '(correct_label, (-1, num_classes))', True, 'import tensorflow as tf\n'), ((145, 18, 145, 97), 'tensorflow.nn.softmax_cross_entropy_with_logits', 'tf.nn.softmax_cross_entropy_with_logits', (), '', True, 'import tensorflow as tf\n'), ((146, 23, 146, 52), 'tensorflow.reduce_mean', 'tf.reduce_mean', ({(146, 38, 146, 51): 'cross_entropy'}, {}), '(cross_entropy)', True, 'import tensorflow as tf\n'), ((198, 2, 198, 27), 'project_tests.test_layers', 'tests.test_layers', ({(198, 20, 198, 26): 'layers'}, {}), '(layers)', True, 'import project_tests as tests\n'), ((199, 2, 199, 31), 'project_tests.test_optimize', 'tests.test_optimize', ({(199, 22, 199, 30): 'optimize'}, {}), '(optimize)', True, 'import project_tests as tests\n'), ((200, 2, 200, 46), 'project_tests.test_for_kitti_dataset', 'tests.test_for_kitti_dataset', ({(200, 31, 200, 45): 'DATA_DIRECTORY'}, {}), '(DATA_DIRECTORY)', True, 'import project_tests as tests\n'), ((201, 2, 201, 31), 'project_tests.test_train_nn', 'tests.test_train_nn', ({(201, 22, 201, 30): 'train_nn'}, {}), '(train_nn)', True, 'import project_tests as tests\n'), ((208, 2, 208, 54), 'helper.maybe_download_pretrained_vgg', 'helper.maybe_download_pretrained_vgg', ({(208, 39, 208, 53): 'DATA_DIRECTORY'}, {}), '(DATA_DIRECTORY)', False, 'import helper\n'), ((211, 19, 211, 82), 'helper.gen_batch_function', 'helper.gen_batch_function', ({(211, 45, 211, 68): 'TRAINING_DATA_DIRECTORY', (211, 70, 211, 81): 'IMAGE_SHAPE'}, {}), '(TRAINING_DATA_DIRECTORY, IMAGE_SHAPE)', False, 'import helper\n'), ((213, 7, 213, 19), 'tensorflow.Session', 'tf.Session', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((237, 4, 237, 119), 'helper.save_inference_samples', 'helper.save_inference_samples', ({(237, 34, 237, 48): 'RUNS_DIRECTORY', (237, 50, 237, 64): 'DATA_DIRECTORY', (237, 66, 237, 73): 'session', (237, 75, 237, 86): 'IMAGE_SHAPE', (237, 88, 237, 94): 'logits', (237, 96, 237, 105): 'keep_prob', (237, 107, 237, 118): 'image_input'}, {}), '(RUNS_DIRECTORY, DATA_DIRECTORY, session,\n    IMAGE_SHAPE, logits, keep_prob, image_input)', False, 'import helper\n'), ((48, 40, 48, 65), 'tensorflow.test.gpu_device_name', 'tf.test.gpu_device_name', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((149, 13, 149, 50), 'tensorflow.train.AdamOptimizer', 'tf.train.AdamOptimizer', ({(149, 36, 149, 49): 'learning_rate'}, {}), '(learning_rate)', True, 'import tensorflow as tf\n'), ((228, 16, 228, 49), 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((229, 16, 229, 48), 'tensorflow.local_variables_initializer', 'tf.local_variables_initializer', ({}, {}), '()', True, 'import tensorflow as tf\n')]