# Copyright 2018 Johns Hopkins University. 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.
# ==============================================================================

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from collections import namedtuple

import tensorflow as tf
from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib
from tensorflow.python.ops import parsing_ops


class Pipeline(object):
    def __init__(self, tfrecord_file, feature_map, batch_size=32,
                 num_threads=4, prefetch_buffer_size=1,
                 static_max_length=None, shuffle_buffer_size=10000,
                 shuffle=True, num_epochs=None, one_shot=False):
        self._feature_map = feature_map
        self._batch_size = batch_size
        self._static_max_length = static_max_length

        # Initialize the dataset
        dataset = tf.data.TFRecordDataset(tfrecord_file)

        # Maybe randomize
        if shuffle:
            dataset = dataset.shuffle(shuffle_buffer_size)

        # Maybe repeat
        if num_epochs is None:
            dataset = dataset.repeat()  # repeat indefinitely
        elif num_epochs > 1:
            dataset = dataset.repeat(count=num_epochs)

        dataset = dataset.batch(batch_size)
        dataset = dataset.map(self.parse_example,
                              num_parallel_calls=num_threads)

        # Pre-fetch a batch for faster processing
        dataset = dataset.prefetch(prefetch_buffer_size)

        # Get the iterator
        if one_shot:
            self._iterator = dataset.make_one_shot_iterator()
        else:
            self._iterator = dataset.make_initializable_iterator()
            self._init_op = self._iterator.initializer

        # Get outputs
        self._outputs = self._iterator.get_next()

        # Map to features
        index = 0
        result = {}
        for key in sorted(self._feature_map.keys()):
            result[key] = self._outputs[index]
            index += 1
        self._result = result

    def pad(self, t):
        s = tf.shape(t)
        paddings = [[0, 0], [0, self._static_max_length - s[1]]]
        x = tf.pad(t, paddings, 'CONSTANT', constant_values=0)
        x = tf.reshape(x, [s[0], self._static_max_length])
        assert x.get_shape().as_list()[1] is self._static_max_length
        return x

    def parse_example(self, serialized):
        parsed = parsing_ops.parse_example(serialized, self._feature_map)
        result = []
        for key in sorted(self._feature_map.keys()):
            val = parsed[key]
            if isinstance(val, sparse_tensor_lib.SparseTensor):
                dense_tensor = tf.sparse_tensor_to_dense(val)
                if self._static_max_length is not None:
                    dense_tensor = self.pad(dense_tensor)
                result.append(dense_tensor)
            else:
                result.append(val)
        return tuple(result)

    @property
    def iterator(self):
        return self._iterator

    @property
    def init_op(self):
        return self._init_op

    @property
    def batch(self):
        return self._result


# namedtuple for bucket_info object (used in Pipeline)
# func: a mapping from examples to tf.int64 keys
# pads: a set of tf shapes that correspond to padded examples
bucket_info = namedtuple("bucket_info", "func pads")


def int64_feature(value):
    """ Takes a single int (e.g. 3) and converts it to a tf Feature """
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


def int64_list_feature(sequence):
    """ Sequence of ints (e.g [1,2,3]) to TF feature """
    return tf.train.Feature(int64_list=tf.train.Int64List(value=sequence))

[((112, 14, 112, 52), 'collections.namedtuple', 'namedtuple', ({(112, 25, 112, 38): '"""bucket_info"""', (112, 40, 112, 51): '"""func pads"""'}, {}), "('bucket_info', 'func pads')", False, 'from collections import namedtuple\n'), ((37, 18, 37, 56), 'tensorflow.data.TFRecordDataset', 'tf.data.TFRecordDataset', ({(37, 42, 37, 55): 'tfrecord_file'}, {}), '(tfrecord_file)', True, 'import tensorflow as tf\n'), ((75, 12, 75, 23), 'tensorflow.shape', 'tf.shape', ({(75, 21, 75, 22): 't'}, {}), '(t)', True, 'import tensorflow as tf\n'), ((77, 12, 77, 62), 'tensorflow.pad', 'tf.pad', (), '', True, 'import tensorflow as tf\n'), ((78, 12, 78, 58), 'tensorflow.reshape', 'tf.reshape', ({(78, 23, 78, 24): 'x', (78, 26, 78, 57): '[s[0], self._static_max_length]'}, {}), '(x, [s[0], self._static_max_length])', True, 'import tensorflow as tf\n'), ((83, 17, 83, 73), 'tensorflow.python.ops.parsing_ops.parse_example', 'parsing_ops.parse_example', ({(83, 43, 83, 53): 'serialized', (83, 55, 83, 72): 'self._feature_map'}, {}), '(serialized, self._feature_map)', False, 'from tensorflow.python.ops import parsing_ops\n'), ((117, 39, 117, 72), 'tensorflow.train.Int64List', 'tf.train.Int64List', (), '', True, 'import tensorflow as tf\n'), ((122, 39, 122, 73), 'tensorflow.train.Int64List', 'tf.train.Int64List', (), '', True, 'import tensorflow as tf\n'), ((88, 31, 88, 61), 'tensorflow.sparse_tensor_to_dense', 'tf.sparse_tensor_to_dense', ({(88, 57, 88, 60): 'val'}, {}), '(val)', True, 'import tensorflow as tf\n')]