#! /usr/bin/env python

import random
import tensorflow as tf
from dataset_utils import _dataset_exists, _get_filenames_and_classes, write_label_file, _convert_dataset

#====================================================DEFINE YOUR ARGUMENTS=======================================================================
flags = tf.app.flags

#State your dataset directory
flags.DEFINE_string('training_dataset_dir', None, 'String: Your training dataset directory')
flags.DEFINE_string('validation_dataset_dir', None, 'String: Your validation dataset directory')

# The number of shards to split the dataset into
flags.DEFINE_integer('num_shards', 2, 'Int: Number of shards to split the TFRecord files')

# Seed for repeatability.
flags.DEFINE_integer('random_seed', 0, 'Int: Random seed to use for repeatability.')

#Output filename for the naming the TFRecord file
flags.DEFINE_string('tfrecord_filename', None, 'String: The output filename to name your TFRecord file')

FLAGS = flags.FLAGS

def main():

    #==============================================================CHECKS==========================================================================
    #Check if there is a tfrecord_filename entered
    if not FLAGS.tfrecord_filename:
        raise ValueError('tfrecord_filename is empty. Please state a tfrecord_filename argument.')

    #Check if there is a dataset directory entered
    if not FLAGS.training_dataset_dir:
        raise ValueError('training_dataset_dir is empty. Please state a training_dataset_dir argument.')

    #Check if there is a dataset directory entered
    if not FLAGS.validation_dataset_dir:
        raise ValueError('validation_dataset_dir is empty. Please state a validation_dataset_dir argument.')

    #If the TFRecord files already exist in the directory, then exit without creating the files again
    if _dataset_exists(dataset_dir = FLAGS.training_dataset_dir, _NUM_SHARDS = FLAGS.num_shards, output_filename = FLAGS.tfrecord_filename):
        print('Dataset files already exist. Exiting without re-creating them.')
        return None
    #==============================================================END OF CHECKS===================================================================

    #Get a list of photo_filenames like ['123.jpg', '456.jpg'...] and a list of sorted class names from parsing the subdirectories.
    training_filenames, training_class_names = _get_filenames_and_classes(FLAGS.training_dataset_dir)

    #Get a list of photo_filenames like ['123.jpg', '456.jpg'...] and a list of sorted class names from parsing the subdirectories.
    validation_filenames, validation_class_names = _get_filenames_and_classes(FLAGS.validation_dataset_dir)

    #Refer each of the class name to a specific integer number for predictions later
    training_class_names_to_ids = dict(zip(training_class_names, range(len(training_class_names))))

    #Refer each of the class name to a specific integer number for predictions later
    validation_class_names_to_ids = dict(zip(validation_class_names, range(len(validation_class_names))))

    # First, convert the training and validation sets.
    _convert_dataset('train', training_filenames, training_class_names_to_ids,
                     dataset_dir = FLAGS.training_dataset_dir, tfrecord_filename = FLAGS.tfrecord_filename, _NUM_SHARDS = FLAGS.num_shards)
    _convert_dataset('validation', validation_filenames, validation_class_names_to_ids,
                     dataset_dir = FLAGS.validation_dataset_dir, tfrecord_filename = FLAGS.tfrecord_filename, _NUM_SHARDS = FLAGS.num_shards)

    # Finally, write the labels file:
    labels_to_class_names = dict(zip(list(map(int, training_class_names)), training_class_names))
    write_label_file(labels_to_class_names, FLAGS.training_dataset_dir)

    # Finally, write the labels file:
    labels_to_class_names = dict(zip(list(map(int, validation_class_names)), validation_class_names))
    write_label_file(labels_to_class_names, FLAGS.validation_dataset_dir)

    print('\nFinished converting the %s dataset!' % (FLAGS.tfrecord_filename))

if __name__ == "__main__":
    main()
