# coding=utf-8
# Copyright 2018 The TensorFlow Datasets Authors.
#
# 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.

"""Test utilities."""

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

import collections
import contextlib
import os
import tempfile

import tensorflow as tf

from tensorflow_datasets.core import dataset_utils
from tensorflow_datasets.core import file_format_adapter


@contextlib.contextmanager
def tmp_dir(dirname=None):
  tmp = make_tmp_dir(dirname)
  yield tmp
  rm_tmp_dir(tmp)


def make_tmp_dir(dirname=None):
  if dirname and not tf.gfile.Exists(dirname):
    tf.gfile.MakeDirs(dirname)
  return tempfile.mkdtemp(dir=dirname)


def rm_tmp_dir(dirname):
  tf.gfile.DeleteRecursively(dirname)


class FeatureExpectation(
    collections.namedtuple('_FeatureExpectation',
                           ['name', 'feature', 'value', 'expected'])):
  pass


def features_encode_decode(specs_dict, sample):
  """Runs the full pipeline: encode > write > tmp files > read > decode."""
  # Encode sample
  encoded_sample = specs_dict.encode_sample(sample)

  with tmp_dir() as tmp_dir_:
    tmp_filename = os.path.join(tmp_dir_, 'tmp.tfrecord')

    # Read/write the file
    file_adapter = file_format_adapter.TFRecordExampleAdapter(
        specs_dict.get_specs())
    file_adapter.write_from_generator(
        generator_fn=lambda: [encoded_sample],
        output_files=[tmp_filename],
    )
    dataset = file_adapter.dataset_from_filename(tmp_filename)

    # Decode the sample
    dataset = dataset.map(specs_dict.decode_sample)

    for el in dataset_utils.iterate_over_dataset(dataset):
      return el
