# -*- coding: utf-8 -*-
#
# Copyright 2017-2019 Spotify AB.
#
# 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, division, print_function

import logging
from typing import Tuple, Union, Dict, Iterator  # noqa: F401

import six
import numpy as np  # noqa: F401
import pandas as pd
import tensorflow as tf
from spotify_tensorflow.tf_schema_utils import parse_schema_file, schema_to_feature_spec
from tensorflow_metadata.proto.v0.schema_pb2 import Schema  # noqa: F401


logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


class Datasets(object):

    @staticmethod
    def _assert_eager(endpoint):
        assert tf.executing_eagerly(), "Eager execution is required for a %s endpoint! " \
                                       "Add this add the begining of your main:\n\nimport " \
                                       "tensorflow as tf\ntf.enable_eager_execution()\n\n" % \
                                       endpoint

    @classmethod
    def parse_schema(cls, schema_path):
        # type: (str) -> Tuple[Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]], Schema]  # noqa: E501
        """
        Returns TensorFlow Feature Spec and parsed tf.metadata Schema for given tf.metadata Schema.

        :param schema_path: tf.metadata Schema path
        """
        schema = parse_schema_file(schema_path)
        return schema_to_feature_spec(schema), schema

    @classmethod
    def parse_schema_from_stats(cls, stats_path):
        # type: (str) -> Tuple[Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]], Schema]  # noqa: E501
        """
        Returns TensorFlow Feature Spec and parsed tf.metadata Schema for given tf.metadata
        DatasetFeatureStatisticsList.

        :param stats_path: tf.metadata DatasetFeatureStatisticsList path
        """
        import tensorflow_data_validation as tfdv
        stats = tfdv.load_statistics(stats_path)
        schema = tfdv.infer_schema(stats)
        return schema_to_feature_spec(schema), schema

    @classmethod
    def examples_via_schema(cls,
                            file_pattern,  # type: str
                            schema_path,  # type: str
                            compression_type=None,  # type: str
                            batch_size=128,  # type: int
                            shuffle=True,  # type: bool
                            num_epochs=1,  # type: int
                            shuffle_buffer_size=10000,  # type: int
                            shuffle_seed=None,  # type: int
                            prefetch_buffer_size=1,  # type: int
                            reader_num_threads=1,  # type: int
                            parser_num_threads=2,  # type: int
                            sloppy_ordering=False,  # type: bool
                            drop_final_batch=False  # type: bool
                            ):
        # type: (...) -> tf.data.Dataset
        """Get `Dataset` of parsed `Example` protos.

        :param file_pattern: List of files or patterns of file paths containing
                             `Example` records. See `tf.gfile.Glob` for pattern rules
        :param schema_path: tf.metadata Schema path
        :param compression_type: TFRecord compression type, see `tf.data.TFRecordDataset` doc
        :param batch_size: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param shuffle: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param num_epochs: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param shuffle_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param shuffle_seed: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param prefetch_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param reader_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param parser_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param sloppy_ordering: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param drop_final_batch: see `tensorflow.contrib.data.make_batched_features_dataset` doc

        :return `Dataset`, which holds results of the parsing of `Example` protos
        """
        return cls._examples(file_pattern,
                             schema_path=schema_path,
                             compression_type=compression_type,
                             batch_size=batch_size,
                             shuffle=shuffle,
                             num_epochs=num_epochs,
                             shuffle_buffer_size=shuffle_buffer_size,
                             shuffle_seed=shuffle_seed,
                             prefetch_buffer_size=prefetch_buffer_size,
                             reader_num_threads=reader_num_threads,
                             parser_num_threads=parser_num_threads,
                             sloppy_ordering=sloppy_ordering,
                             drop_final_batch=drop_final_batch)

    @classmethod
    def examples_via_feature_spec(cls,
                                  file_pattern,  # type: str
                                  feature_spec,  # type: Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]]  # noqa: E501
                                  compression_type=None,  # type: str
                                  batch_size=128,  # type: int
                                  shuffle=True,  # type: bool
                                  num_epochs=1,  # type: int
                                  shuffle_buffer_size=10000,  # type: int
                                  shuffle_seed=None,  # type: int
                                  prefetch_buffer_size=1,  # type: int
                                  reader_num_threads=1,  # type: int
                                  parser_num_threads=2,  # type: int
                                  sloppy_ordering=False,  # type: bool
                                  drop_final_batch=False  # type: bool
                                  ):
        # type: (...) -> tf.data.Dataset
        """Get `Dataset` of parsed `Example` protos.

        :param file_pattern: List of files or patterns of file paths containing
                             `Example` records. See `tf.gfile.Glob` for pattern rules
        :param feature_spec: TensorFlow feature spec
        :param compression_type: TFRecord compression type, see `tf.data.TFRecordDataset` doc
        :param batch_size: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param shuffle: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param num_epochs: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param shuffle_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param shuffle_seed: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param prefetch_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param reader_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param parser_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param sloppy_ordering: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param drop_final_batch: see `tensorflow.contrib.data.make_batched_features_dataset` doc

        :return `Dataset`, which holds results of the parsing of `Example` protos
        """
        return cls._examples(file_pattern,
                             feature_spec=feature_spec,
                             compression_type=compression_type,
                             batch_size=batch_size,
                             shuffle=shuffle,
                             num_epochs=num_epochs,
                             shuffle_buffer_size=shuffle_buffer_size,
                             shuffle_seed=shuffle_seed,
                             prefetch_buffer_size=prefetch_buffer_size,
                             reader_num_threads=reader_num_threads,
                             parser_num_threads=parser_num_threads,
                             sloppy_ordering=sloppy_ordering,
                             drop_final_batch=drop_final_batch)

    @classmethod
    def _examples(cls,
                  file_pattern,  # type: str
                  schema_path=None,  # type: str
                  feature_spec=None,  # type: Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]]  # noqa: E501
                  compression_type=None,  # type: str
                  batch_size=128,  # type: int
                  shuffle=True,  # type: bool
                  num_epochs=1,  # type: int
                  shuffle_buffer_size=10000,  # type: int
                  shuffle_seed=None,  # type: int
                  prefetch_buffer_size=1,  # type: int
                  reader_num_threads=1,  # type: int
                  parser_num_threads=2,  # type: int
                  sloppy_ordering=False,  # type: bool
                  drop_final_batch=False  # type: bool
                  ):
        # type: (...) -> tf.data.Dataset

        if schema_path:
            feature_spec, _ = cls.parse_schema(schema_path)

        logger.debug("Will parse features from: `%s`, using features spec: %s",
                     file_pattern,
                     str(feature_spec))

        from tensorflow.contrib.data import make_batched_features_dataset
        reader_args = [compression_type] if compression_type else None
        dataset = make_batched_features_dataset(file_pattern,
                                                batch_size=batch_size,
                                                features=feature_spec,
                                                reader_args=reader_args,
                                                num_epochs=num_epochs,
                                                shuffle=shuffle,
                                                shuffle_buffer_size=shuffle_buffer_size,
                                                shuffle_seed=shuffle_seed,
                                                prefetch_buffer_size=prefetch_buffer_size,
                                                reader_num_threads=reader_num_threads,
                                                parser_num_threads=parser_num_threads,
                                                sloppy_ordering=sloppy_ordering,
                                                drop_final_batch=drop_final_batch)
        return dataset

    class __DictionaryEndpoint(object):

        @classmethod
        def examples_via_schema(cls,
                                file_pattern,  # type: str
                                schema_path,  # type: str
                                default_value=0,  # type: float
                                batch_size=128,  # type: int
                                compression_type=None,  # type: str
                                shuffle=True,  # type: bool
                                num_epochs=1,  # type: int
                                shuffle_buffer_size=10000,  # type: int
                                shuffle_seed=42,  # type: int
                                prefetch_buffer_size=1,  # type: int
                                reader_num_threads=1,  # type: int
                                parser_num_threads=2,  # type: int
                                sloppy_ordering=False,  # type: bool
                                drop_final_batch=False  # type: bool
                                ):
            # type: (...) -> Iterator[Dict[str, np.ndarray]]
            """
            Read a TF dataset and load it into a dictionary of NumPy Arrays.

            :param file_pattern: List of files or patterns of file paths containing
                                 `Example` records. See `tf.gfile.Glob` for pattern rules
            :param default_value: Value used if a sparse feature is missing.
            :param schema_path: tf.metadata Schema path
            :param compression_type: TFRecord compression type, see `tf.data.TFRecordDataset` doc
            :param batch_size: batch size, set to the size of the dataset to read all data at once
            :param shuffle: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param num_epochs: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param shuffle_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`
                                        doc
            :param shuffle_seed: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param prefetch_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`
                                         doc
            :param reader_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`
                                       doc
            :param parser_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`
                                       doc
            :param sloppy_ordering: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param drop_final_batch: see `tensorflow.contrib.data.make_batched_features_dataset` doc

            :return Dictionary of NumPy arrays
            """
            return cls._examples(file_pattern=file_pattern,
                                 schema_path=schema_path,
                                 default_value=default_value,
                                 batch_size=batch_size,
                                 compression_type=compression_type,
                                 shuffle=shuffle,
                                 num_epochs=num_epochs,
                                 shuffle_buffer_size=shuffle_buffer_size,
                                 shuffle_seed=shuffle_seed,
                                 prefetch_buffer_size=prefetch_buffer_size,
                                 reader_num_threads=reader_num_threads,
                                 parser_num_threads=parser_num_threads,
                                 sloppy_ordering=sloppy_ordering,
                                 drop_final_batch=drop_final_batch)

        @classmethod
        def examples_via_feature_spec(cls,
                                      file_pattern,  # type: str
                                      feature_spec,  # type: Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]]  # noqa: E501
                                      default_value=0,  # type: float
                                      batch_size=128,  # type: int
                                      compression_type=None,  # type: str
                                      shuffle=True,  # type: bool
                                      num_epochs=1,  # type: int
                                      shuffle_buffer_size=10000,  # type: int
                                      shuffle_seed=42,  # type: int
                                      prefetch_buffer_size=1,  # type: int
                                      reader_num_threads=1,  # type: int
                                      parser_num_threads=2,  # type: int
                                      sloppy_ordering=False,  # type: bool
                                      drop_final_batch=False  # type: bool
                                      ):
            # type: (...) -> Iterator[Dict[str, np.ndarray]]
            """
            Read a TF dataset and load it into a dictionary of NumPy Arrays.

            :param file_pattern: List of files or patterns of file paths containing
                                 `Example` records. See `tf.gfile.Glob` for pattern rules
            :param feature_spec: TensorFlow feature spec
            :param default_value: Value used if a sparse feature is missing.
            :param compression_type: TFRecord compression type, see `tf.data.TFRecordDataset` doc
            :param batch_size: batch size, set to the size of the dataset to read all data at once
            :param shuffle: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param num_epochs: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param shuffle_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`
                                        doc
            :param shuffle_seed: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param prefetch_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`
                                         doc
            :param reader_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`
                                       doc
            :param parser_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`
                                       doc
            :param sloppy_ordering: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param drop_final_batch: see `tensorflow.contrib.data.make_batched_features_dataset` doc

            :return Dictionary of NumPy arrays
            """
            return cls._examples(file_pattern=file_pattern,
                                 feature_spec=feature_spec,
                                 default_value=default_value,
                                 batch_size=batch_size,
                                 compression_type=compression_type,
                                 shuffle=shuffle,
                                 num_epochs=num_epochs,
                                 shuffle_buffer_size=shuffle_buffer_size,
                                 shuffle_seed=shuffle_seed,
                                 prefetch_buffer_size=prefetch_buffer_size,
                                 reader_num_threads=reader_num_threads,
                                 parser_num_threads=parser_num_threads,
                                 sloppy_ordering=sloppy_ordering,
                                 drop_final_batch=drop_final_batch)

        @classmethod
        def _examples(cls,
                      file_pattern,  # type: str
                      schema_path=None,  # type: str
                      feature_spec=None,  # type: Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]]  # noqa: E501
                      default_value=0,  # type: float
                      compression_type=None,  # type: str
                      batch_size=128,  # type: int
                      shuffle=True,  # type: bool
                      num_epochs=1,  # type: int
                      shuffle_buffer_size=10000,  # type: int
                      shuffle_seed=None,  # type: int
                      prefetch_buffer_size=1,  # type: int
                      reader_num_threads=1,  # type: int
                      parser_num_threads=2,  # type: int
                      sloppy_ordering=False,  # type: bool
                      drop_final_batch=False  # type: bool
                      ):
            # type: (...) -> Iterator[Dict[str, np.ndarray]]
            Datasets._assert_eager("Dictionary")

            def get_numpy(tensor):
                if isinstance(tensor, tf.Tensor):
                    return tensor.numpy()
                elif isinstance(tensor, tf.SparseTensor):
                    # If it's a SparseTensor, which is the representation of VarLenFeature and
                    # SparseFeature, we convert it to dense representation, and further is it's
                    # a scalar, we reshape to to a vector

                    shape = tensor.dense_shape.numpy()
                    # first element is batch size
                    if shape[1] == 0:
                        # this feature is not defined for any of the examples in the batch
                        return np.repeat(default_value, shape[0])

                    numpy_dense = tf.sparse_tensor_to_dense(tensor,
                                                            default_value=default_value).numpy()
                    if shape[1] == 1:
                        # this is scalar feature, reshape to a vector
                        return numpy_dense.reshape(shape[0])
                    else:
                        return numpy_dense
                else:
                    raise ValueError("This type %s is not supported!", type(tensor).__name__)

            dataset = Datasets._examples(file_pattern=file_pattern,
                                         schema_path=schema_path,
                                         feature_spec=feature_spec,
                                         compression_type=compression_type,
                                         batch_size=batch_size,
                                         shuffle=shuffle,
                                         num_epochs=num_epochs,
                                         shuffle_buffer_size=shuffle_buffer_size,
                                         shuffle_seed=shuffle_seed,
                                         prefetch_buffer_size=prefetch_buffer_size,
                                         reader_num_threads=reader_num_threads,
                                         parser_num_threads=parser_num_threads,
                                         sloppy_ordering=sloppy_ordering,
                                         drop_final_batch=drop_final_batch)
            for batch in dataset:
                yield {name: get_numpy(eager_tensor) for name, eager_tensor in six.iteritems(batch)}

    dict = __DictionaryEndpoint()

    class __DataFrameEndpoint(object):

        @classmethod
        def examples_via_schema(cls,
                                file_pattern,  # type: str
                                schema_path,  # type: str
                                default_value=0,  # type: float
                                batch_size=128,  # type: int
                                compression_type=None,  # type: str
                                shuffle=True,  # type: bool
                                num_epochs=1,  # type: int
                                shuffle_buffer_size=10000,  # type: int
                                shuffle_seed=42,  # type: int
                                prefetch_buffer_size=1,  # type: int
                                reader_num_threads=1,  # type: int
                                parser_num_threads=2,  # type: int
                                sloppy_ordering=False,  # type: bool
                                drop_final_batch=False  # type: bool
                                ):
            # type: (...) -> Iterator[pd.DataFrame]
            """
            Read a TF dataset in batches, each one yields a Pandas DataFrame.

            :param file_pattern: List of files or patterns of file paths containing
                                 `Example` records. See `tf.gfile.Glob` for pattern rules
            :param schema_path: tf.metadata Schema path
            :param default_value: Value used if a sparse feature is missing.
            :param batch_size: batch size, set to the size of the dataset to read all data at once
            :param compression_type: TFRecord compression type, see `tf.data.TFRecordDataset` doc
            :param shuffle: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param num_epochs: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param shuffle_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`
                                        doc
            :param shuffle_seed: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param prefetch_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`
                                         doc
            :param reader_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`
                                       doc
            :param parser_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`
                                       doc
            :param sloppy_ordering: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param drop_final_batch: see `tensorflow.contrib.data.make_batched_features_dataset` doc

            :return A Python Generator, yielding batches of data in a Pandas DataFrame
            """
            return cls._examples(file_pattern=file_pattern,
                                 schema_path=schema_path,
                                 default_value=default_value,
                                 batch_size=batch_size,
                                 compression_type=compression_type,
                                 shuffle=shuffle,
                                 num_epochs=num_epochs,
                                 shuffle_buffer_size=shuffle_buffer_size,
                                 shuffle_seed=shuffle_seed,
                                 prefetch_buffer_size=prefetch_buffer_size,
                                 reader_num_threads=reader_num_threads,
                                 parser_num_threads=parser_num_threads,
                                 sloppy_ordering=sloppy_ordering,
                                 drop_final_batch=drop_final_batch)

        @classmethod
        def examples_via_feature_spec(cls,
                                      file_pattern,  # type: str
                                      feature_spec,  # type: Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]]  # noqa: E501
                                      default_value=0,  # type: float
                                      batch_size=128,  # type: int
                                      compression_type=None,  # type: str
                                      shuffle=True,  # type: bool
                                      num_epochs=1,  # type: int
                                      shuffle_buffer_size=10000,  # type: int
                                      shuffle_seed=42,  # type: int
                                      prefetch_buffer_size=1,  # type: int
                                      reader_num_threads=1,  # type: int
                                      parser_num_threads=2,  # type: int
                                      sloppy_ordering=False,  # type: bool
                                      drop_final_batch=False  # type: bool
                                      ):
            # type: (...) -> Iterator[pd.DataFrame]
            """
            Read a TF dataset in batches, each one yields a Pandas DataFrame.

            :param file_pattern: List of files or patterns of file paths containing
                                 `Example` records. See `tf.gfile.Glob` for pattern rules
            :param feature_spec: TensorFlow feature spec
            :param default_value: Value used if a sparse feature is missing.
            :param batch_size: batch size, set to the size of the dataset to read all data at once
            :param compression_type: TFRecord compression type, see `tf.data.TFRecordDataset` doc
            :param shuffle: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param num_epochs: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param shuffle_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`
                                        doc
            :param shuffle_seed: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param prefetch_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`
                                         doc
            :param reader_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`
                                       doc
            :param parser_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`
                                       doc
            :param sloppy_ordering: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param drop_final_batch: see `tensorflow.contrib.data.make_batched_features_dataset` doc

            :return A Python Generator, yielding batches of data in a Pandas DataFrame
            """
            return cls._examples(file_pattern=file_pattern,
                                 feature_spec=feature_spec,
                                 default_value=default_value,
                                 batch_size=batch_size,
                                 compression_type=compression_type,
                                 shuffle=shuffle,
                                 num_epochs=num_epochs,
                                 shuffle_buffer_size=shuffle_buffer_size,
                                 shuffle_seed=shuffle_seed,
                                 prefetch_buffer_size=prefetch_buffer_size,
                                 reader_num_threads=reader_num_threads,
                                 parser_num_threads=parser_num_threads,
                                 sloppy_ordering=sloppy_ordering,
                                 drop_final_batch=drop_final_batch)

        @classmethod
        def _examples(cls,
                      file_pattern,  # type: str
                      schema_path=None,  # type: str
                      feature_spec=None,  # type: Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]]  # noqa: E501
                      default_value=0,  # type: float
                      compression_type=None,  # type: str
                      batch_size=128,  # type: int
                      shuffle=True,  # type: bool
                      num_epochs=1,  # type: int
                      shuffle_buffer_size=10000,  # type: int
                      shuffle_seed=None,  # type: int
                      prefetch_buffer_size=1,  # type: int
                      reader_num_threads=1,  # type: int
                      parser_num_threads=2,  # type: int
                      sloppy_ordering=False,  # type: bool
                      drop_final_batch=False  # type: bool
                      ):
            # type: (...) -> Iterator[pd.DataFrame]
            Datasets._assert_eager("DataFrame")
            dataset = Datasets.dict._examples(file_pattern=file_pattern,
                                              schema_path=schema_path,
                                              default_value=default_value,
                                              feature_spec=feature_spec,
                                              compression_type=compression_type,
                                              batch_size=batch_size,
                                              shuffle=shuffle,
                                              num_epochs=num_epochs,
                                              shuffle_buffer_size=shuffle_buffer_size,
                                              shuffle_seed=shuffle_seed,
                                              prefetch_buffer_size=prefetch_buffer_size,
                                              reader_num_threads=reader_num_threads,
                                              parser_num_threads=parser_num_threads,
                                              sloppy_ordering=sloppy_ordering,
                                              drop_final_batch=drop_final_batch)
            for d in dataset:
                yield pd.DataFrame(data=d)

    dataframe = __DataFrameEndpoint()

# -*- coding: utf-8 -*-
#
# Copyright 2017-2019 Spotify AB.
#
# 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, division, print_function

import logging
from typing import Tuple, Union, Dict, Iterator  # noqa: F401

import six
import numpy as np  # noqa: F401
import pandas as pd
import tensorflow as tf
from spotify_tensorflow.tf_schema_utils import parse_schema_file, schema_to_feature_spec
from tensorflow_metadata.proto.v0.schema_pb2 import Schema  # noqa: F401


logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


class Datasets(object):

    @staticmethod
    def _assert_eager(endpoint):
        assert tf.executing_eagerly(), "Eager execution is required for a %s endpoint! " \
                                       "Add this add the begining of your main:\n\nimport " \
                                       "tensorflow as tf\ntf.enable_eager_execution()\n\n" % \
                                       endpoint

    @classmethod
    def parse_schema(cls, schema_path):
        # type: (str) -> Tuple[Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]], Schema]  # noqa: E501
        """
        Returns TensorFlow Feature Spec and parsed tf.metadata Schema for given tf.metadata Schema.

        :param schema_path: tf.metadata Schema path
        """
        schema = parse_schema_file(schema_path)
        return schema_to_feature_spec(schema), schema

    @classmethod
    def parse_schema_from_stats(cls, stats_path):
        # type: (str) -> Tuple[Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]], Schema]  # noqa: E501
        """
        Returns TensorFlow Feature Spec and parsed tf.metadata Schema for given tf.metadata
        DatasetFeatureStatisticsList.

        :param stats_path: tf.metadata DatasetFeatureStatisticsList path
        """
        import tensorflow_data_validation as tfdv
        stats = tfdv.load_statistics(stats_path)
        schema = tfdv.infer_schema(stats)
        return schema_to_feature_spec(schema), schema

    @classmethod
    def examples_via_schema(cls,
                            file_pattern,  # type: str
                            schema_path,  # type: str
                            compression_type=None,  # type: str
                            batch_size=128,  # type: int
                            shuffle=True,  # type: bool
                            num_epochs=1,  # type: int
                            shuffle_buffer_size=10000,  # type: int
                            shuffle_seed=None,  # type: int
                            prefetch_buffer_size=1,  # type: int
                            reader_num_threads=1,  # type: int
                            parser_num_threads=2,  # type: int
                            sloppy_ordering=False,  # type: bool
                            drop_final_batch=False  # type: bool
                            ):
        # type: (...) -> tf.data.Dataset
        """Get `Dataset` of parsed `Example` protos.

        :param file_pattern: List of files or patterns of file paths containing
                             `Example` records. See `tf.gfile.Glob` for pattern rules
        :param schema_path: tf.metadata Schema path
        :param compression_type: TFRecord compression type, see `tf.data.TFRecordDataset` doc
        :param batch_size: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param shuffle: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param num_epochs: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param shuffle_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param shuffle_seed: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param prefetch_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param reader_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param parser_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param sloppy_ordering: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param drop_final_batch: see `tensorflow.contrib.data.make_batched_features_dataset` doc

        :return `Dataset`, which holds results of the parsing of `Example` protos
        """
        return cls._examples(file_pattern,
                             schema_path=schema_path,
                             compression_type=compression_type,
                             batch_size=batch_size,
                             shuffle=shuffle,
                             num_epochs=num_epochs,
                             shuffle_buffer_size=shuffle_buffer_size,
                             shuffle_seed=shuffle_seed,
                             prefetch_buffer_size=prefetch_buffer_size,
                             reader_num_threads=reader_num_threads,
                             parser_num_threads=parser_num_threads,
                             sloppy_ordering=sloppy_ordering,
                             drop_final_batch=drop_final_batch)

    @classmethod
    def examples_via_feature_spec(cls,
                                  file_pattern,  # type: str
                                  feature_spec,  # type: Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]]  # noqa: E501
                                  compression_type=None,  # type: str
                                  batch_size=128,  # type: int
                                  shuffle=True,  # type: bool
                                  num_epochs=1,  # type: int
                                  shuffle_buffer_size=10000,  # type: int
                                  shuffle_seed=None,  # type: int
                                  prefetch_buffer_size=1,  # type: int
                                  reader_num_threads=1,  # type: int
                                  parser_num_threads=2,  # type: int
                                  sloppy_ordering=False,  # type: bool
                                  drop_final_batch=False  # type: bool
                                  ):
        # type: (...) -> tf.data.Dataset
        """Get `Dataset` of parsed `Example` protos.

        :param file_pattern: List of files or patterns of file paths containing
                             `Example` records. See `tf.gfile.Glob` for pattern rules
        :param feature_spec: TensorFlow feature spec
        :param compression_type: TFRecord compression type, see `tf.data.TFRecordDataset` doc
        :param batch_size: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param shuffle: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param num_epochs: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param shuffle_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param shuffle_seed: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param prefetch_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param reader_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param parser_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param sloppy_ordering: see `tensorflow.contrib.data.make_batched_features_dataset` doc
        :param drop_final_batch: see `tensorflow.contrib.data.make_batched_features_dataset` doc

        :return `Dataset`, which holds results of the parsing of `Example` protos
        """
        return cls._examples(file_pattern,
                             feature_spec=feature_spec,
                             compression_type=compression_type,
                             batch_size=batch_size,
                             shuffle=shuffle,
                             num_epochs=num_epochs,
                             shuffle_buffer_size=shuffle_buffer_size,
                             shuffle_seed=shuffle_seed,
                             prefetch_buffer_size=prefetch_buffer_size,
                             reader_num_threads=reader_num_threads,
                             parser_num_threads=parser_num_threads,
                             sloppy_ordering=sloppy_ordering,
                             drop_final_batch=drop_final_batch)

    @classmethod
    def _examples(cls,
                  file_pattern,  # type: str
                  schema_path=None,  # type: str
                  feature_spec=None,  # type: Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]]  # noqa: E501
                  compression_type=None,  # type: str
                  batch_size=128,  # type: int
                  shuffle=True,  # type: bool
                  num_epochs=1,  # type: int
                  shuffle_buffer_size=10000,  # type: int
                  shuffle_seed=None,  # type: int
                  prefetch_buffer_size=1,  # type: int
                  reader_num_threads=1,  # type: int
                  parser_num_threads=2,  # type: int
                  sloppy_ordering=False,  # type: bool
                  drop_final_batch=False  # type: bool
                  ):
        # type: (...) -> tf.data.Dataset

        if schema_path:
            feature_spec, _ = cls.parse_schema(schema_path)

        logger.debug("Will parse features from: `%s`, using features spec: %s",
                     file_pattern,
                     str(feature_spec))

        from tensorflow.contrib.data import make_batched_features_dataset
        reader_args = [compression_type] if compression_type else None
        dataset = make_batched_features_dataset(file_pattern,
                                                batch_size=batch_size,
                                                features=feature_spec,
                                                reader_args=reader_args,
                                                num_epochs=num_epochs,
                                                shuffle=shuffle,
                                                shuffle_buffer_size=shuffle_buffer_size,
                                                shuffle_seed=shuffle_seed,
                                                prefetch_buffer_size=prefetch_buffer_size,
                                                reader_num_threads=reader_num_threads,
                                                parser_num_threads=parser_num_threads,
                                                sloppy_ordering=sloppy_ordering,
                                                drop_final_batch=drop_final_batch)
        return dataset

    class __DictionaryEndpoint(object):

        @classmethod
        def examples_via_schema(cls,
                                file_pattern,  # type: str
                                schema_path,  # type: str
                                default_value=0,  # type: float
                                batch_size=128,  # type: int
                                compression_type=None,  # type: str
                                shuffle=True,  # type: bool
                                num_epochs=1,  # type: int
                                shuffle_buffer_size=10000,  # type: int
                                shuffle_seed=42,  # type: int
                                prefetch_buffer_size=1,  # type: int
                                reader_num_threads=1,  # type: int
                                parser_num_threads=2,  # type: int
                                sloppy_ordering=False,  # type: bool
                                drop_final_batch=False  # type: bool
                                ):
            # type: (...) -> Iterator[Dict[str, np.ndarray]]
            """
            Read a TF dataset and load it into a dictionary of NumPy Arrays.

            :param file_pattern: List of files or patterns of file paths containing
                                 `Example` records. See `tf.gfile.Glob` for pattern rules
            :param default_value: Value used if a sparse feature is missing.
            :param schema_path: tf.metadata Schema path
            :param compression_type: TFRecord compression type, see `tf.data.TFRecordDataset` doc
            :param batch_size: batch size, set to the size of the dataset to read all data at once
            :param shuffle: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param num_epochs: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param shuffle_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`
                                        doc
            :param shuffle_seed: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param prefetch_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`
                                         doc
            :param reader_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`
                                       doc
            :param parser_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`
                                       doc
            :param sloppy_ordering: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param drop_final_batch: see `tensorflow.contrib.data.make_batched_features_dataset` doc

            :return Dictionary of NumPy arrays
            """
            return cls._examples(file_pattern=file_pattern,
                                 schema_path=schema_path,
                                 default_value=default_value,
                                 batch_size=batch_size,
                                 compression_type=compression_type,
                                 shuffle=shuffle,
                                 num_epochs=num_epochs,
                                 shuffle_buffer_size=shuffle_buffer_size,
                                 shuffle_seed=shuffle_seed,
                                 prefetch_buffer_size=prefetch_buffer_size,
                                 reader_num_threads=reader_num_threads,
                                 parser_num_threads=parser_num_threads,
                                 sloppy_ordering=sloppy_ordering,
                                 drop_final_batch=drop_final_batch)

        @classmethod
        def examples_via_feature_spec(cls,
                                      file_pattern,  # type: str
                                      feature_spec,  # type: Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]]  # noqa: E501
                                      default_value=0,  # type: float
                                      batch_size=128,  # type: int
                                      compression_type=None,  # type: str
                                      shuffle=True,  # type: bool
                                      num_epochs=1,  # type: int
                                      shuffle_buffer_size=10000,  # type: int
                                      shuffle_seed=42,  # type: int
                                      prefetch_buffer_size=1,  # type: int
                                      reader_num_threads=1,  # type: int
                                      parser_num_threads=2,  # type: int
                                      sloppy_ordering=False,  # type: bool
                                      drop_final_batch=False  # type: bool
                                      ):
            # type: (...) -> Iterator[Dict[str, np.ndarray]]
            """
            Read a TF dataset and load it into a dictionary of NumPy Arrays.

            :param file_pattern: List of files or patterns of file paths containing
                                 `Example` records. See `tf.gfile.Glob` for pattern rules
            :param feature_spec: TensorFlow feature spec
            :param default_value: Value used if a sparse feature is missing.
            :param compression_type: TFRecord compression type, see `tf.data.TFRecordDataset` doc
            :param batch_size: batch size, set to the size of the dataset to read all data at once
            :param shuffle: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param num_epochs: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param shuffle_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`
                                        doc
            :param shuffle_seed: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param prefetch_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`
                                         doc
            :param reader_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`
                                       doc
            :param parser_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`
                                       doc
            :param sloppy_ordering: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param drop_final_batch: see `tensorflow.contrib.data.make_batched_features_dataset` doc

            :return Dictionary of NumPy arrays
            """
            return cls._examples(file_pattern=file_pattern,
                                 feature_spec=feature_spec,
                                 default_value=default_value,
                                 batch_size=batch_size,
                                 compression_type=compression_type,
                                 shuffle=shuffle,
                                 num_epochs=num_epochs,
                                 shuffle_buffer_size=shuffle_buffer_size,
                                 shuffle_seed=shuffle_seed,
                                 prefetch_buffer_size=prefetch_buffer_size,
                                 reader_num_threads=reader_num_threads,
                                 parser_num_threads=parser_num_threads,
                                 sloppy_ordering=sloppy_ordering,
                                 drop_final_batch=drop_final_batch)

        @classmethod
        def _examples(cls,
                      file_pattern,  # type: str
                      schema_path=None,  # type: str
                      feature_spec=None,  # type: Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]]  # noqa: E501
                      default_value=0,  # type: float
                      compression_type=None,  # type: str
                      batch_size=128,  # type: int
                      shuffle=True,  # type: bool
                      num_epochs=1,  # type: int
                      shuffle_buffer_size=10000,  # type: int
                      shuffle_seed=None,  # type: int
                      prefetch_buffer_size=1,  # type: int
                      reader_num_threads=1,  # type: int
                      parser_num_threads=2,  # type: int
                      sloppy_ordering=False,  # type: bool
                      drop_final_batch=False  # type: bool
                      ):
            # type: (...) -> Iterator[Dict[str, np.ndarray]]
            Datasets._assert_eager("Dictionary")

            def get_numpy(tensor):
                if isinstance(tensor, tf.Tensor):
                    return tensor.numpy()
                elif isinstance(tensor, tf.SparseTensor):
                    # If it's a SparseTensor, which is the representation of VarLenFeature and
                    # SparseFeature, we convert it to dense representation, and further is it's
                    # a scalar, we reshape to to a vector

                    shape = tensor.dense_shape.numpy()
                    # first element is batch size
                    if shape[1] == 0:
                        # this feature is not defined for any of the examples in the batch
                        return np.repeat(default_value, shape[0])

                    numpy_dense = tf.sparse_tensor_to_dense(tensor,
                                                            default_value=default_value).numpy()
                    if shape[1] == 1:
                        # this is scalar feature, reshape to a vector
                        return numpy_dense.reshape(shape[0])
                    else:
                        return numpy_dense
                else:
                    raise ValueError("This type %s is not supported!", type(tensor).__name__)

            dataset = Datasets._examples(file_pattern=file_pattern,
                                         schema_path=schema_path,
                                         feature_spec=feature_spec,
                                         compression_type=compression_type,
                                         batch_size=batch_size,
                                         shuffle=shuffle,
                                         num_epochs=num_epochs,
                                         shuffle_buffer_size=shuffle_buffer_size,
                                         shuffle_seed=shuffle_seed,
                                         prefetch_buffer_size=prefetch_buffer_size,
                                         reader_num_threads=reader_num_threads,
                                         parser_num_threads=parser_num_threads,
                                         sloppy_ordering=sloppy_ordering,
                                         drop_final_batch=drop_final_batch)
            for batch in dataset:
                yield {name: get_numpy(eager_tensor) for name, eager_tensor in six.iteritems(batch)}

    dict = __DictionaryEndpoint()

    class __DataFrameEndpoint(object):

        @classmethod
        def examples_via_schema(cls,
                                file_pattern,  # type: str
                                schema_path,  # type: str
                                default_value=0,  # type: float
                                batch_size=128,  # type: int
                                compression_type=None,  # type: str
                                shuffle=True,  # type: bool
                                num_epochs=1,  # type: int
                                shuffle_buffer_size=10000,  # type: int
                                shuffle_seed=42,  # type: int
                                prefetch_buffer_size=1,  # type: int
                                reader_num_threads=1,  # type: int
                                parser_num_threads=2,  # type: int
                                sloppy_ordering=False,  # type: bool
                                drop_final_batch=False  # type: bool
                                ):
            # type: (...) -> Iterator[pd.DataFrame]
            """
            Read a TF dataset in batches, each one yields a Pandas DataFrame.

            :param file_pattern: List of files or patterns of file paths containing
                                 `Example` records. See `tf.gfile.Glob` for pattern rules
            :param schema_path: tf.metadata Schema path
            :param default_value: Value used if a sparse feature is missing.
            :param batch_size: batch size, set to the size of the dataset to read all data at once
            :param compression_type: TFRecord compression type, see `tf.data.TFRecordDataset` doc
            :param shuffle: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param num_epochs: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param shuffle_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`
                                        doc
            :param shuffle_seed: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param prefetch_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`
                                         doc
            :param reader_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`
                                       doc
            :param parser_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`
                                       doc
            :param sloppy_ordering: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param drop_final_batch: see `tensorflow.contrib.data.make_batched_features_dataset` doc

            :return A Python Generator, yielding batches of data in a Pandas DataFrame
            """
            return cls._examples(file_pattern=file_pattern,
                                 schema_path=schema_path,
                                 default_value=default_value,
                                 batch_size=batch_size,
                                 compression_type=compression_type,
                                 shuffle=shuffle,
                                 num_epochs=num_epochs,
                                 shuffle_buffer_size=shuffle_buffer_size,
                                 shuffle_seed=shuffle_seed,
                                 prefetch_buffer_size=prefetch_buffer_size,
                                 reader_num_threads=reader_num_threads,
                                 parser_num_threads=parser_num_threads,
                                 sloppy_ordering=sloppy_ordering,
                                 drop_final_batch=drop_final_batch)

        @classmethod
        def examples_via_feature_spec(cls,
                                      file_pattern,  # type: str
                                      feature_spec,  # type: Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]]  # noqa: E501
                                      default_value=0,  # type: float
                                      batch_size=128,  # type: int
                                      compression_type=None,  # type: str
                                      shuffle=True,  # type: bool
                                      num_epochs=1,  # type: int
                                      shuffle_buffer_size=10000,  # type: int
                                      shuffle_seed=42,  # type: int
                                      prefetch_buffer_size=1,  # type: int
                                      reader_num_threads=1,  # type: int
                                      parser_num_threads=2,  # type: int
                                      sloppy_ordering=False,  # type: bool
                                      drop_final_batch=False  # type: bool
                                      ):
            # type: (...) -> Iterator[pd.DataFrame]
            """
            Read a TF dataset in batches, each one yields a Pandas DataFrame.

            :param file_pattern: List of files or patterns of file paths containing
                                 `Example` records. See `tf.gfile.Glob` for pattern rules
            :param feature_spec: TensorFlow feature spec
            :param default_value: Value used if a sparse feature is missing.
            :param batch_size: batch size, set to the size of the dataset to read all data at once
            :param compression_type: TFRecord compression type, see `tf.data.TFRecordDataset` doc
            :param shuffle: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param num_epochs: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param shuffle_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`
                                        doc
            :param shuffle_seed: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param prefetch_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`
                                         doc
            :param reader_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`
                                       doc
            :param parser_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`
                                       doc
            :param sloppy_ordering: see `tensorflow.contrib.data.make_batched_features_dataset` doc
            :param drop_final_batch: see `tensorflow.contrib.data.make_batched_features_dataset` doc

            :return A Python Generator, yielding batches of data in a Pandas DataFrame
            """
            return cls._examples(file_pattern=file_pattern,
                                 feature_spec=feature_spec,
                                 default_value=default_value,
                                 batch_size=batch_size,
                                 compression_type=compression_type,
                                 shuffle=shuffle,
                                 num_epochs=num_epochs,
                                 shuffle_buffer_size=shuffle_buffer_size,
                                 shuffle_seed=shuffle_seed,
                                 prefetch_buffer_size=prefetch_buffer_size,
                                 reader_num_threads=reader_num_threads,
                                 parser_num_threads=parser_num_threads,
                                 sloppy_ordering=sloppy_ordering,
                                 drop_final_batch=drop_final_batch)

        @classmethod
        def _examples(cls,
                      file_pattern,  # type: str
                      schema_path=None,  # type: str
                      feature_spec=None,  # type: Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]]  # noqa: E501
                      default_value=0,  # type: float
                      compression_type=None,  # type: str
                      batch_size=128,  # type: int
                      shuffle=True,  # type: bool
                      num_epochs=1,  # type: int
                      shuffle_buffer_size=10000,  # type: int
                      shuffle_seed=None,  # type: int
                      prefetch_buffer_size=1,  # type: int
                      reader_num_threads=1,  # type: int
                      parser_num_threads=2,  # type: int
                      sloppy_ordering=False,  # type: bool
                      drop_final_batch=False  # type: bool
                      ):
            # type: (...) -> Iterator[pd.DataFrame]
            Datasets._assert_eager("DataFrame")
            dataset = Datasets.dict._examples(file_pattern=file_pattern,
                                              schema_path=schema_path,
                                              default_value=default_value,
                                              feature_spec=feature_spec,
                                              compression_type=compression_type,
                                              batch_size=batch_size,
                                              shuffle=shuffle,
                                              num_epochs=num_epochs,
                                              shuffle_buffer_size=shuffle_buffer_size,
                                              shuffle_seed=shuffle_seed,
                                              prefetch_buffer_size=prefetch_buffer_size,
                                              reader_num_threads=reader_num_threads,
                                              parser_num_threads=parser_num_threads,
                                              sloppy_ordering=sloppy_ordering,
                                              drop_final_batch=drop_final_batch)
            for d in dataset:
                yield pd.DataFrame(data=d)

    dataframe = __DataFrameEndpoint()

# -*- coding: utf-8 -*- # # Copyright 2017-2019 Spotify AB. # # 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. # # noqa: F401 # noqa: F401 # noqa: F401 # type: (str) -> Tuple[Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]], Schema]  # noqa: E501 Returns TensorFlow Feature Spec and parsed tf.metadata Schema for given tf.metadata Schema.          :param schema_path: tf.metadata Schema path # type: (str) -> Tuple[Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]], Schema]  # noqa: E501 Returns TensorFlow Feature Spec and parsed tf.metadata Schema for given tf.metadata         DatasetFeatureStatisticsList.          :param stats_path: tf.metadata DatasetFeatureStatisticsList path # type: str # type: str # type: str # type: int # type: bool # type: int # type: int # type: int # type: int # type: int # type: int # type: bool # type: bool # type: (...) -> tf.data.Dataset Get `Dataset` of parsed `Example` protos.          :param file_pattern: List of files or patterns of file paths containing                              `Example` records. See `tf.gfile.Glob` for pattern rules         :param schema_path: tf.metadata Schema path         :param compression_type: TFRecord compression type, see `tf.data.TFRecordDataset` doc         :param batch_size: see `tensorflow.contrib.data.make_batched_features_dataset` doc         :param shuffle: see `tensorflow.contrib.data.make_batched_features_dataset` doc         :param num_epochs: see `tensorflow.contrib.data.make_batched_features_dataset` doc         :param shuffle_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset` doc         :param shuffle_seed: see `tensorflow.contrib.data.make_batched_features_dataset` doc         :param prefetch_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset` doc         :param reader_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset` doc         :param parser_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset` doc         :param sloppy_ordering: see `tensorflow.contrib.data.make_batched_features_dataset` doc         :param drop_final_batch: see `tensorflow.contrib.data.make_batched_features_dataset` doc          :return `Dataset`, which holds results of the parsing of `Example` protos # type: str # type: Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]]  # noqa: E501 # type: str # type: int # type: bool # type: int # type: int # type: int # type: int # type: int # type: int # type: bool # type: bool # type: (...) -> tf.data.Dataset Get `Dataset` of parsed `Example` protos.          :param file_pattern: List of files or patterns of file paths containing                              `Example` records. See `tf.gfile.Glob` for pattern rules         :param feature_spec: TensorFlow feature spec         :param compression_type: TFRecord compression type, see `tf.data.TFRecordDataset` doc         :param batch_size: see `tensorflow.contrib.data.make_batched_features_dataset` doc         :param shuffle: see `tensorflow.contrib.data.make_batched_features_dataset` doc         :param num_epochs: see `tensorflow.contrib.data.make_batched_features_dataset` doc         :param shuffle_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset` doc         :param shuffle_seed: see `tensorflow.contrib.data.make_batched_features_dataset` doc         :param prefetch_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset` doc         :param reader_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset` doc         :param parser_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset` doc         :param sloppy_ordering: see `tensorflow.contrib.data.make_batched_features_dataset` doc         :param drop_final_batch: see `tensorflow.contrib.data.make_batched_features_dataset` doc          :return `Dataset`, which holds results of the parsing of `Example` protos # type: str # type: str # type: Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]]  # noqa: E501 # type: str # type: int # type: bool # type: int # type: int # type: int # type: int # type: int # type: int # type: bool # type: bool # type: (...) -> tf.data.Dataset # type: str # type: str # type: float # type: int # type: str # type: bool # type: int # type: int # type: int # type: int # type: int # type: int # type: bool # type: bool # type: (...) -> Iterator[Dict[str, np.ndarray]] Read a TF dataset and load it into a dictionary of NumPy Arrays.              :param file_pattern: List of files or patterns of file paths containing                                  `Example` records. See `tf.gfile.Glob` for pattern rules             :param default_value: Value used if a sparse feature is missing.             :param schema_path: tf.metadata Schema path             :param compression_type: TFRecord compression type, see `tf.data.TFRecordDataset` doc             :param batch_size: batch size, set to the size of the dataset to read all data at once             :param shuffle: see `tensorflow.contrib.data.make_batched_features_dataset` doc             :param num_epochs: see `tensorflow.contrib.data.make_batched_features_dataset` doc             :param shuffle_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`                                         doc             :param shuffle_seed: see `tensorflow.contrib.data.make_batched_features_dataset` doc             :param prefetch_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`                                          doc             :param reader_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`                                        doc             :param parser_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`                                        doc             :param sloppy_ordering: see `tensorflow.contrib.data.make_batched_features_dataset` doc             :param drop_final_batch: see `tensorflow.contrib.data.make_batched_features_dataset` doc              :return Dictionary of NumPy arrays # type: str # type: Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]]  # noqa: E501 # type: float # type: int # type: str # type: bool # type: int # type: int # type: int # type: int # type: int # type: int # type: bool # type: bool # type: (...) -> Iterator[Dict[str, np.ndarray]] Read a TF dataset and load it into a dictionary of NumPy Arrays.              :param file_pattern: List of files or patterns of file paths containing                                  `Example` records. See `tf.gfile.Glob` for pattern rules             :param feature_spec: TensorFlow feature spec             :param default_value: Value used if a sparse feature is missing.             :param compression_type: TFRecord compression type, see `tf.data.TFRecordDataset` doc             :param batch_size: batch size, set to the size of the dataset to read all data at once             :param shuffle: see `tensorflow.contrib.data.make_batched_features_dataset` doc             :param num_epochs: see `tensorflow.contrib.data.make_batched_features_dataset` doc             :param shuffle_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`                                         doc             :param shuffle_seed: see `tensorflow.contrib.data.make_batched_features_dataset` doc             :param prefetch_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`                                          doc             :param reader_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`                                        doc             :param parser_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`                                        doc             :param sloppy_ordering: see `tensorflow.contrib.data.make_batched_features_dataset` doc             :param drop_final_batch: see `tensorflow.contrib.data.make_batched_features_dataset` doc              :return Dictionary of NumPy arrays # type: str # type: str # type: Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]]  # noqa: E501 # type: float # type: str # type: int # type: bool # type: int # type: int # type: int # type: int # type: int # type: int # type: bool # type: bool # type: (...) -> Iterator[Dict[str, np.ndarray]] # If it's a SparseTensor, which is the representation of VarLenFeature and # SparseFeature, we convert it to dense representation, and further is it's # a scalar, we reshape to to a vector # first element is batch size # this feature is not defined for any of the examples in the batch # this is scalar feature, reshape to a vector # type: str # type: str # type: float # type: int # type: str # type: bool # type: int # type: int # type: int # type: int # type: int # type: int # type: bool # type: bool # type: (...) -> Iterator[pd.DataFrame] Read a TF dataset in batches, each one yields a Pandas DataFrame.              :param file_pattern: List of files or patterns of file paths containing                                  `Example` records. See `tf.gfile.Glob` for pattern rules             :param schema_path: tf.metadata Schema path             :param default_value: Value used if a sparse feature is missing.             :param batch_size: batch size, set to the size of the dataset to read all data at once             :param compression_type: TFRecord compression type, see `tf.data.TFRecordDataset` doc             :param shuffle: see `tensorflow.contrib.data.make_batched_features_dataset` doc             :param num_epochs: see `tensorflow.contrib.data.make_batched_features_dataset` doc             :param shuffle_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`                                         doc             :param shuffle_seed: see `tensorflow.contrib.data.make_batched_features_dataset` doc             :param prefetch_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`                                          doc             :param reader_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`                                        doc             :param parser_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`                                        doc             :param sloppy_ordering: see `tensorflow.contrib.data.make_batched_features_dataset` doc             :param drop_final_batch: see `tensorflow.contrib.data.make_batched_features_dataset` doc              :return A Python Generator, yielding batches of data in a Pandas DataFrame # type: str # type: Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]]  # noqa: E501 # type: float # type: int # type: str # type: bool # type: int # type: int # type: int # type: int # type: int # type: int # type: bool # type: bool # type: (...) -> Iterator[pd.DataFrame] Read a TF dataset in batches, each one yields a Pandas DataFrame.              :param file_pattern: List of files or patterns of file paths containing                                  `Example` records. See `tf.gfile.Glob` for pattern rules             :param feature_spec: TensorFlow feature spec             :param default_value: Value used if a sparse feature is missing.             :param batch_size: batch size, set to the size of the dataset to read all data at once             :param compression_type: TFRecord compression type, see `tf.data.TFRecordDataset` doc             :param shuffle: see `tensorflow.contrib.data.make_batched_features_dataset` doc             :param num_epochs: see `tensorflow.contrib.data.make_batched_features_dataset` doc             :param shuffle_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`                                         doc             :param shuffle_seed: see `tensorflow.contrib.data.make_batched_features_dataset` doc             :param prefetch_buffer_size: see `tensorflow.contrib.data.make_batched_features_dataset`                                          doc             :param reader_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`                                        doc             :param parser_num_threads: see `tensorflow.contrib.data.make_batched_features_dataset`                                        doc             :param sloppy_ordering: see `tensorflow.contrib.data.make_batched_features_dataset` doc             :param drop_final_batch: see `tensorflow.contrib.data.make_batched_features_dataset` doc              :return A Python Generator, yielding batches of data in a Pandas DataFrame # type: str # type: str # type: Dict[str, Union[tf.FixedLenFeature, tf.VarLenFeature, tf.SparseFeature]]  # noqa: E501 # type: float # type: str # type: int # type: bool # type: int # type: int # type: int # type: int # type: int # type: int # type: bool # type: bool # type: (...) -> Iterator[pd.DataFrame]