import os

import tensorflow as tf

from distort_image import (distort_image_for_train, distort_image_for_eval)

BATCH_SIZE = 128
MIN_FRACTION_OF_EXAMPLES_IN_QUEUE = 0.4
NUM_THREADS = 16


def inputs(dataset, batch_size=BATCH_SIZE, train=True, eval_data=False):
    """Constructs inputs from a dataset.

    Args:
        dataset: Instance of the dataset to use.
        batch_size: Number of data per batch.
        train: Boolean indiciating if one wants to train or evaluate. Default:
          True.
        eval_data: Boolean indicating if one should use the train or eval data
          set. Default: False.

    Returns:
        data_batch: 4D tensor of [batch_size, height, width, depth] size.
        label_batch: 1D tensor of [batch_size] size.
    """

    # When training, we want to apply a different distortion as if we are
    # evaluating.
    if train:
        distort = distort_image_for_train
        shuffle = True
    else:
        distort = distort_image_for_eval
        shuffle = False

    # Choose the right dataset.
    if not eval_data:
        filenames = dataset.train_filenames
        num_examples_per_epoch = dataset.num_examples_per_epoch_for_train
    else:
        filenames = dataset.eval_filenames
        num_examples_per_epoch = dataset.num_examples_per_epoch_for_eval

    # Construct the inputs.
    return _inputs(filenames, dataset.read, distort, num_examples_per_epoch,
                   batch_size, shuffle)


def _inputs(filenames, read, distort, num_examples_per_epoch, batch_size,
            shuffle):

    """Constructs inputs from data files using the read operation.

    Args:
        filenames: The filenames of the data batches in the data directory.
        read: Reader operation that returns a single record.
        distort: Distort operation on the data of the example.
        num_examples_per_epoch: Number of examples per epoch.
        batch_size: Number of data per batch.
        shuffle: Boolean indicating whether to use a shuffling queue.

    Returns:
        data_batch: 4D tensor of [batch_size, height, width, depth] size.
        label_batch: 1D tensor of [batch_size] size.
    """

    # Create a queue that produces the filenames to read.
    filename_queue = tf.train.string_input_producer(filenames)

    # Read examples from files in the filename queue.
    record = read(filename_queue)

    # Distort the data.
    record = distort(record)

    min_queue_examples = int(num_examples_per_epoch *
                             MIN_FRACTION_OF_EXAMPLES_IN_QUEUE)
    capacity = min_queue_examples + 3 * batch_size

    print('Filling queue with {} examples before starting. This can take a '
          'few minutes.'.format(min_queue_examples))

    # Create a queue that shuffles the examples, and then read batch_size
    # data + labels from the example queue.
    if shuffle:
        data_batch, label_batch = tf.train.shuffle_batch(
            [record.data, record.label],
            batch_size=batch_size,
            num_threads=NUM_THREADS,
            capacity=capacity,
            min_after_dequeue=min_queue_examples)
    else:
        data_batch, label_batch = tf.train.batch(
            [record.data, record.label],
            batch_size=batch_size,
            num_threads=NUM_THREADS,
            capacity=capacity)

    return data_batch, tf.reshape(label_batch, [batch_size])

@@ -201,16 +201,37 @@
 dataset,
+ distort_inputs=True,
  batch_s
@@ -245,22 +245,52 @@
 CH_SIZE,
- train
+%0A           num_epochs=None, shuffle
 =True, e
@@ -419,59 +419,281 @@
     
-batch_size: Number of data per batch.%0A        train
+distort_inputs: Boolean whether to distort the inputs (optional).%0A        batch_size: Number of data per batch (optional).%0A        num_epochs: Number indicating the maximal number of epochs iterations%0A          before raising an OutOfRange error (optional).%0A        shuffle
 : Bo
@@ -730,35 +730,26 @@
  to 
-train or evaluate. Default:
+shuffle the inputs
 %0A   
@@ -755,20 +755,26 @@
         
-True
+(optional)
 .%0A      
@@ -863,24 +863,19 @@
  set
-. Default: False
+ (optional)
 .%0A%0A 
@@ -1037,37 +1037,44 @@
 hen 
-training, w
+shuffling on
 e want
+s
  to 
+also 
 apply a 
 diff
@@ -1073,17 +1073,14 @@
 y a 
-different
+random
  dis
@@ -1090,38 +1090,8 @@
 tion
- as if we are%0A    # evaluating
 .%0A  
@@ -1095,21 +1095,23 @@
 %0A    if 
-train
+shuffle
 :%0A      
@@ -1150,31 +1150,8 @@
 ain%0A
-        shuffle = True%0A
     
@@ -1200,32 +1200,8 @@
 eval
-%0A        shuffle = False
 %0A%0A  
@@ -1508,339 +1508,71 @@
  # C
-onstruct the inputs.%0A    return _inputs(filenames, dataset.read, distort, num_examples_per_epoch,%0A                   batch_size, shuffle)%0A%0A%0Adef _inputs(filenames, read, distort, num_examples_per_epoch, batch_size,%0A            shuffle):%0A%0A    %22%22%22Constructs inputs from data files using the read operation.%0A%0A    Args:%0A  
+reate a queue that produces the filenames to read.%0A
     
-  
 filename
 s: T
@@ -1571,639 +1571,87 @@
 name
-s: The filenames of the data batches in the data directory.%0A        read: Reader operation that returns a single record.%0A        distort: Distort operation on the data of the example.%0A        num_examples_per_epoch: Number of examples per epoch.%0A        batch_size: Number of data per batch.%0A        shuffle: Boolean indicating whether to use a shuffling queue.%0A%0A    Returns:%0A        data_batch: 4D tensor of %5Bbatch_size, height, width, depth%5D size.%0A        label_batch: 1D tensor of %5Bbatch_size%5D size.%0A    %22%22%22%0A%0A    # Create a queue that produces the filenames to read.%0A    filename_queue = tf.train.string_input_producer(filenames
+_queue = tf.train.string_input_producer(%0A        filenames, num_epochs, shuffle
 )%0A%0A 
@@ -1716,16 +1716,24 @@
 ecord = 
+dataset.
 read(fil
@@ -1770,16 +1770,43 @@
 e data.%0A
+    if distort_inputs:%0A    
     reco

clean up inputs code