text: # Copyright 2016 The TensorFlow Authors All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  
  """Code for building the input for the prediction model."""
  
  # TODO: Can remove the parts of this code that enable use of state/action.
  # TODO: Do this if you lose overview of the code (its a mess!)
  
  import os
  from collections import OrderedDict
  import numpy as np
  import tensorflow as tf
  
  from tensorflow.python.platform import flags
  from tensorflow.python.platform import gfile
  
  import os, sys
  sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../')))
  import src.utils.utils as utils
  import src.data_loading as dl
  
  FLAGS = flags.FLAGS
  
  # Original image dimensions
  #ORIGINAL_WIDTH = 640
  #ORIGINAL_HEIGHT = 512
  #COLOR_CHAN = 3
  ORIGINAL_WIDTH = 80
  ORIGINAL_HEIGHT = 80
  COLOR_CHAN = 1
  
  # Default image dimensions.
  IMG_WIDTH = 80
  IMG_HEIGHT = 80
  
  # Dimension of the state and action.
  STATE_DIM = 5
  
  
  def build_tfrecord_input(split_string='train', file_nums=[1,2,3,4], training=None, feed_labels=False,
                           return_queue=False, shuffle=True, sequential=False):
    """Create input tfrecord tensors.
  
       Choose between queue output = dict containing images and labels, or
                      queue output = [images, zeros, zeros].
       Default parameters are so that it can be used for training with prediction_train.py.
  
    Args:
      split_string: Only used if FLAGS.custom_data.
                      'train', 'val' or 'test' - use only files within path that have this string in their names.
      file_nums:    Only used if FLAGS.custom_data.
                    A list. Only use files ending in 'i.npy', where i is one of the numbers in file_nums.
      training:     Only used if *not* FLAGS.custom_data. Specifies training or validation data.
      feed_labels:  Ignored if not FLAGS.custom_data OR if FLAGS.use_state.
                       Set this to False for compatibility with prediction_train.py. (For training the model.)
                       If True, the first return value is a dictionary containing images (key: 'image_seq')
                       as well as labels. Else, first return value is the image output of the queue.
      return_queue: Ignored if not FLAGS.custom_data.
                       If True, return the queue itself as a fourth return value
                       (so that the queue size, etc. can be accessed.)
      sequential:   Ignored if not FLAGS.custom_data: If True, use only one thread.
    Flags used:
      FLAGS.data_dir  I believe the tfrecords in there should be one file per sequence.
      FLAGS.custom_data: If True, expect files containing '*train*' and '*val*' within data directory.
                               If False, load first 'train_val_split'-percent of the files if train,
                                else the rest.
      FLAGS.custom_data: to use own bouncing_objects or bouncing_circles datasets.
      FLAGS.sequence_length: OLD: the lenth of sequences in the rf record files
                             NEW, if custom data: Reduce second dimension of the input to this length (discard remainders).
                                                  At most the length of sequences in the tf record files.
      FLAGS.use_state
      IMG_HEIGHT =!= IMG_WIDTH
      ORIGINAL_HEIGHT, ORIGINAL_WIDTH: can have any height to width ratio
      COLOR_CHAN
    Returns:
      list of tensors corresponding to images, actions, and states. The images
      tensor is 5D, batch x time x height x width x channels. The state and
      action tensors are 3D, batch x time x dimension.
      The images values are between 0 and 1.
    Raises:
      RuntimeError: if no files found.
    """
  
    num_threads = min(FLAGS.batch_size, 4)
    if sequential:
        num_threads = 1
  
    if not FLAGS.custom_data:
      filenames = gfile.Glob(os.path.join(FLAGS.data_dir, '*'))
      if not filenames:
        raise RuntimeError('No data files found.')
      index = int(np.floor(FLAGS.train_val_split * len(filenames)))
      if training:
        filenames = filenames[:index]
      else:
        filenames = filenames[index:]
    else:
      filenames = gfile.Glob(os.path.join(FLAGS.data_dir, '*'+split_string+'*.npy'))
              # todo: reduce validation files to only one, for comparability with other models
      filenames = [fn for fn in filenames if np.any([fn.endswith(str(i)+'.npy') for i in file_nums]) ]
      filenames_labels = [fn.replace('.npy', '.npz').replace('_frames_', '_labels_') for fn in filenames]
      for fn in filenames_labels:
          assert fn in gfile.Glob(os.path.join(FLAGS.data_dir, '*'+split_string+'*.npz')), "Label file not found: \n "+fn
    if not filenames:
      raise RuntimeError('No data files found.')
  
  # load the files and create a queue
    if not FLAGS.custom_data:
      filename_queue = tf.train.string_input_producer(filenames, shuffle=True)
      reader = tf.TFRecordReader()
      _, serialized_example = reader.read(filename_queue)
  
    else:
        all_data = None
        all_labels = None
        for fn, fnl in zip(filenames, filenames_labels):
            data = np.load(fn)
            if feed_labels:
                labels = dict(np.load(fnl))
            if data.shape[1] > FLAGS.sequence_length:
                print("prediction_input_custom.py: Need to reduce sequence dimension of data to get target input sequence length.")
                print("\t data.shape was: "+str(data.shape)+"; reduce 2nd dim to length "+str(FLAGS.sequence_length))
                data = data[:, :FLAGS.sequence_length, ...] # reduce to sequence length if needed
                if feed_labels:
                    for key, val in labels.items():
                        labels[key] = val[:, :FLAGS.sequence_length, ...]
            if all_data is None:
                all_data = data
                if feed_labels:
                    all_labels = labels
            else:
                all_data = np.concatenate((all_data, data), axis=0)
                if feed_labels:
                    all_labels = utils.vstack_array_dicts(all_labels, labels, allow_new_keys=False)
  
        shape_arr = [[FLAGS.sequence_length, ORIGINAL_HEIGHT, ORIGINAL_WIDTH, COLOR_CHAN]]
        type_arr = [tf.uint8]
        qnames = ['image_seq']
        if feed_labels:
            all_labels = OrderedDict(all_labels)
            for key, val in all_labels.items():
                assert len(val) == len(all_data)
                #shape_arr.append((FLAGS.sequence_length,)+ val.shape[1:])
                shape_arr.append(val.shape[1:])
                type_arr.append(val.dtype)
                qnames.append(key)
        else:
            all_labels = {}
        if shuffle:
            min_after_dequeue = min(100, int(100*FLAGS.batch_size/10.))
            data_queue = tf.RandomShuffleQueue(capacity=100*FLAGS.batch_size, min_after_dequeue=min_after_dequeue, dtypes=type_arr, shapes=shape_arr,
                                           names=qnames)
        else:
            data_queue = tf.FIFOQueue(capacity=100*FLAGS.batch_size, dtypes=type_arr, shapes=shape_arr,
                                           names=qnames)
        if len(all_data.shape) == 4:
            all_data = all_data[..., np.newaxis]
        assert len(all_data.shape) == 5 and all_data.shape[-1] in [1,2,3,4], "Very weird number of channels found in stored dataset: "+str(data.shape[-1]+". Full shape was: "+str(data.shape))
  
        n_samples = len(all_data)
        data_queue.n_samples = n_samples # uh-oh
        enqueue_op = data_queue.enqueue_many({'image_seq': all_data, **all_labels})
        qr = tf.train.QueueRunner(data_queue, [enqueue_op] * num_threads)
        tf.train.add_queue_runner(qr)
        serialized_example_dict = data_queue.dequeue()
  
  
    image_seq, state_seq, action_seq = [], [], []
  
    for i in range(FLAGS.sequence_length):
      #image_name = 'move/' + str(i) + '/image/encoded'
      #action_name = 'move/' + str(i) + '/commanded_pose/vec_pitch_yaw'
      #state_name = 'move/' + str(i) + '/endeffector/vec_pitch_yaw'
      if not FLAGS.custom_data:
        #image_name, action_name, state_name = get_feature_names_custom(i)
        #assert not FLAGS.use_state, "Dont want to use states in this project."
      #else:
        image_name, action_name, state_name = get_feature_names_finn(i)
        if FLAGS.use_state:
            features = {image_name: tf.FixedLenFeature([1], tf.string),
                        action_name: tf.FixedLenFeature([STATE_DIM], tf.float32),
                        state_name: tf.FixedLenFeature([STATE_DIM], tf.float32)}
        else:
          features = {image_name: tf.FixedLenFeature([1], tf.string)}
          #if not FLAGS.custom_data:
        features = tf.parse_single_example(serialized_example, features=features)
          #else:
          #    features = tf.train.Example()
          #    features.ParseFromString(serialized_example, features=features)
        image_buffer = tf.reshape(features[image_name], shape=[])
        image = tf.image.decode_jpeg(image_buffer, channels=COLOR_CHAN)
        image.set_shape([ORIGINAL_HEIGHT, ORIGINAL_WIDTH, COLOR_CHAN])
  
      else:
          image = serialized_example_dict['image_seq'][i, ...]
          #image = serialized_example[i, ...]
          image = tf.convert_to_tensor(image, np.uint8)
  
      if IMG_HEIGHT != IMG_WIDTH:
        raise ValueError('Unequal height and width unsupported')
  
      crop_size = min(ORIGINAL_HEIGHT, ORIGINAL_WIDTH)
      image = tf.image.resize_image_with_crop_or_pad(image, crop_size, crop_size)
      image = tf.reshape(image, [1, crop_size, crop_size, COLOR_CHAN])
      image = tf.image.resize_bicubic(image, [IMG_HEIGHT, IMG_WIDTH])
      image = tf.cast(image, tf.float32) / 255.0
      image_seq.append(image)
  
      if FLAGS.use_state:
        state = tf.reshape(features[state_name], shape=[1, STATE_DIM])
        state_seq.append(state)
        action = tf.reshape(features[action_name], shape=[1, STATE_DIM])
        action_seq.append(action)
  
    image_seq = tf.concat(image_seq, 0)
  
    if FLAGS.use_state:
      state_seq = tf.concat(state_seq, 0)
      action_seq = tf.concat(action_seq, 0)
      [image_batch, action_batch, state_batch] = tf.train.batch(
          [image_seq, action_seq, state_seq],
          FLAGS.batch_size,
          num_threads=num_threads, #FLAGS.batch_size,
          capacity=100 * FLAGS.batch_size)
      if return_queue:
         return image_batch, action_batch, state_batch, data_queue
      return image_batch, action_batch, state_batch
    else:
      serialized_example_dict['image_seq'] = image_seq
      zeros_batch = tf.zeros([FLAGS.batch_size, FLAGS.sequence_length, STATE_DIM])
      if feed_labels:  # output of the queue will be a dict
          image_label_batch = tf.train.batch(
              { #'image_seq': image_seq,
                  **serialized_example_dict},
              FLAGS.batch_size,
              num_threads=num_threads,  #FLAGS.batch_size,
              capacity=100 * FLAGS.batch_size
          )
      else: # queue output is just the images
          image_label_batch = tf.train.batch(
              [image_seq],
              FLAGS.batch_size,
              num_threads=num_threads,  # FLAGS.batch_size,
              capacity=100 * FLAGS.batch_size)
      if return_queue:
          return image_label_batch, zeros_batch, zeros_batch, data_queue
      return image_label_batch, zeros_batch, zeros_batch
  
  def get_feature_names_finn(i):
      image_name = 'move/' + str(i) + '/image/encoded'
      action_name = 'move/' + str(i) + '/commanded_pose/vec_pitch_yaw'
      state_name = 'move/' + str(i) + '/endeffector/vec_pitch_yaw'
      return image_name, action_name, state_name
  
  def get_feature_names_custom(i):
      image_name = str(i) + '/image_raw'
      return image_name, '', ''
  