# Copyright 2017 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.
# ==============================================================================
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

import threading
import time

import numpy as np

from tensorflow.contrib import lookup
from tensorflow.contrib.eager.python import datasets
from tensorflow.python.data.experimental.ops import threadpool
from tensorflow.python.data.experimental.ops import unique
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.eager import test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import script_ops
from tensorflow.python.training import checkpoint_management
from tensorflow.python.training.tracking import util as trackable_utils


class IteratorTest(test.TestCase):

  def testBasic(self):
    got = []
    for t in datasets.Iterator(dataset_ops.Dataset.range(4)):
      got.append(t.numpy())
    self.assertAllEqual([0, 1, 2, 3], got)

  def testBasicOneShotIterator(self):
    got = []
    for t in dataset_ops.Dataset.range(4).make_one_shot_iterator():
      got.append(t.numpy())
    self.assertAllEqual([0, 1, 2, 3], got)

  def testBasicImplicitIterator(self):
    got = []
    for t in dataset_ops.Dataset.range(4):
      got.append(t.numpy())
    self.assertAllEqual([0, 1, 2, 3], got)

  def testGetNext(self):
    iterator = datasets.Iterator(dataset_ops.Dataset.range(4))
    self.assertEqual(0, iterator.get_next().numpy())
    self.assertEqual(1, iterator.get_next().numpy())
    self.assertEqual(2, iterator.get_next().numpy())
    self.assertEqual(3, iterator.get_next().numpy())
    with self.assertRaises(errors.OutOfRangeError):
      iterator.get_next()

  def testGetNextOneShotIterator(self):
    iterator = dataset_ops.Dataset.range(4).make_one_shot_iterator()
    self.assertEqual(0, iterator.get_next().numpy())
    self.assertEqual(1, iterator.get_next().numpy())
    self.assertEqual(2, iterator.get_next().numpy())
    self.assertEqual(3, iterator.get_next().numpy())
    with self.assertRaises(errors.OutOfRangeError):
      iterator.get_next()

  def testMultipleIteratorsOnTheSameDataset(self):
    ds = dataset_ops.Dataset.range(4)
    it1 = datasets.Iterator(ds)
    it2 = datasets.Iterator(ds)
    got = [x.numpy() for x in it1]
    self.assertAllEqual([0, 1, 2, 3], got)

    got = [x.numpy() for x in it2]
    self.assertAllEqual([0, 1, 2, 3], got)

  def testNestedOutputs(self):
    ds = dataset_ops.Dataset.zip(
        (dataset_ops.Dataset.range(4),
         dataset_ops.Dataset.zip(
             (dataset_ops.Dataset.range(4), dataset_ops.Dataset.range(4)))))
    total = 0
    # The Iterator will return a nested structure of Tensor objects.
    # Some funkiness to compare against simple integers.
    for (i, x) in enumerate(datasets.Iterator(ds)):
      want = (i, (i, i))
      got = (x[0].numpy(), (x[1][0].numpy(), x[1][1].numpy()))
      self.assertEqual(got, want)
      total += 1
    self.assertEqual(4, total)

  def testMapAndFilter(self):

    def even(x):
      return math_ops.equal(math_ops.mod(x, 2), 0)

    it = datasets.Iterator(
        dataset_ops.Dataset.range(8).map(math_ops.square).filter(even))
    got = [x.numpy() for x in it]
    self.assertAllEqual([0, 4, 16, 36], got)

  def testMapCaptureLookupTable(self):
    default_val = -1
    keys = constant_op.constant(['brain', 'salad', 'surgery'])
    values = constant_op.constant([0, 1, 2], dtypes.int64)
    table = lookup.HashTable(
        lookup.KeyValueTensorInitializer(keys, values), default_val)
    dataset = dataset_ops.Dataset.from_tensor_slices(
        ['brain', 'salad', 'surgery'])
    dataset = dataset.map(table.lookup)
    it = datasets.Iterator(dataset)
    got = [x.numpy() for x in it]
    self.assertAllEqual([0, 1, 2], got)

  def testMultipleIteratorsOnADatasetThatUsesFunctions(self):
    ds = dataset_ops.Dataset.from_tensor_slices([1, 2, 3, 4, 5,
                                                 6]).map(math_ops.square)

    got1 = [x.numpy() for x in datasets.Iterator(ds)]
    self.assertAllEqual([1, 4, 9, 16, 25, 36], got1)
    got2 = [x.numpy() for x in datasets.Iterator(ds)]
    self.assertAllEqual(got1, got2)

  def assertSparseValuesEqual(self, a, b):
    self.assertAllEqual(a.indices, b.indices)
    self.assertAllEqual(a.values, b.values)
    self.assertAllEqual(a.dense_shape, b.dense_shape)

  def testSparseTensorElements(self):
    components = (sparse_tensor.SparseTensorValue(
        indices=np.array([[0, 0], [1, 0], [2, 0]]),
        values=np.array([0, 0, 0]),
        dense_shape=np.array([3, 1])),
                  sparse_tensor.SparseTensorValue(
                      indices=np.array([[0, 0], [1, 1], [2, 2]]),
                      values=np.array([1, 2, 3]),
                      dense_shape=np.array([3, 3])))

    expected = [
        (sparse_tensor.SparseTensorValue(
            indices=np.array([[0]]),
            values=np.array([0]),
            dense_shape=np.array([1])),
         sparse_tensor.SparseTensorValue(
             indices=np.array([[0]]),
             values=np.array([1]),
             dense_shape=np.array([3]))),
        (sparse_tensor.SparseTensorValue(
            indices=np.array([[0]]),
            values=np.array([0]),
            dense_shape=np.array([1])),
         sparse_tensor.SparseTensorValue(
             indices=np.array([[1]]),
             values=np.array([2]),
             dense_shape=np.array([3]))),
        (sparse_tensor.SparseTensorValue(
            indices=np.array([[0]]),
            values=np.array([0]),
            dense_shape=np.array([1])),
         sparse_tensor.SparseTensorValue(
             indices=np.array([[2]]),
             values=np.array([3]),
             dense_shape=np.array([3]))),
    ]

    for i, result in enumerate(
        datasets.Iterator(dataset_ops.Dataset.from_tensor_slices(components))):
      self.assertSparseValuesEqual(expected[i][0], result[0])
      self.assertSparseValuesEqual(expected[i][1], result[1])

  def testPyFunc(self):

    def my_map(inp):
      return [[x + 1 for x in inp]]

    ds = dataset_ops.Dataset.range(4).map(
        lambda x: script_ops.py_func(my_map, [[x]], dtypes.int64))
    got = [x.numpy() for x in datasets.Iterator(ds)]
    self.assertAllEqual([[1], [2], [3], [4]], got)

  def testTensorsPlacedOnDevice(self):
    ds = dataset_ops.Dataset.from_tensors([0., 1.])
    with ops.device(test.gpu_device_name()):
      x = datasets.Iterator(ds).next()
      x = math_ops.add(x, x)
    self.assertAllEqual([0., 2.], x.numpy())

  def testGpuTensor(self):
    ds = dataset_ops.Dataset.from_tensors([0., 1.])
    with ops.device(test.gpu_device_name()):
      for x in ds:
        y = math_ops.add(x, x)
    self.assertAllEqual([0., 2.], y.numpy())

  def testOverrideThreadPool(self):

    def get_thread_id(_):
      # Python creates a dummy thread object to represent the current
      # thread when called from an "alien" thread (such as a
      # `PrivateThreadPool` thread in this case). It does not include
      # the TensorFlow-given display name, but it has a unique
      # identifier that maps one-to-one with the underlying OS thread.
      return np.array(threading.current_thread().ident).astype(np.int64)

    for num_threads in [1, 2, 4, 8, 16]:

      dataset = (
          dataset_ops.Dataset.range(1000).map(
              lambda x: script_ops.py_func(get_thread_id, [x], dtypes.int64),
              num_parallel_calls=32).apply(unique.unique()))

      dataset = threadpool.override_threadpool(
          dataset,
          threadpool.PrivateThreadPool(
              num_threads, display_name='private_thread_pool_%d' % num_threads))

      thread_ids = []
      for next_element in datasets.Iterator(dataset):
        thread_ids.append(next_element)
      self.assertEqual(len(thread_ids), len(set(thread_ids)))
      self.assertGreater(len(thread_ids), 0)
      # NOTE(mrry): We don't control the thread pool scheduling, and
      # so cannot guarantee that all of the threads in the pool will
      # perform work.
      self.assertLessEqual(len(thread_ids), num_threads)

  def testSaveRestore(self):
    checkpoint_directory = self.get_temp_dir()
    checkpoint_prefix = os.path.join(checkpoint_directory, 'ckpt')
    dataset = dataset_ops.Dataset.from_tensor_slices(
        [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
    dataset = dataset.map(math_ops.square).batch(2)
    # TODO(b/138399725): Re-enable default optimizations.
    options = dataset_ops.Options()
    options.experimental_optimization.apply_default_optimizations = False
    dataset = dataset.with_options(options)
    iterator = datasets.Iterator(dataset)
    checkpoint = trackable_utils.Checkpoint(iterator=iterator)
    self.assertAllEqual([1, 4], iterator.get_next().numpy())
    save_path = checkpoint.save(checkpoint_prefix)
    self.assertAllEqual([9, 16], iterator.get_next().numpy())
    self.assertAllEqual([25, 36], iterator.get_next().numpy())
    checkpoint.restore(save_path)
    self.assertAllEqual([9, 16], iterator.get_next().numpy())
    self.assertAllEqual([25, 36], iterator.get_next().numpy())

  def testSaveRestoreMultipleIterator(self):
    checkpoint_directory = self.get_temp_dir()
    checkpoint_prefix = os.path.join(checkpoint_directory, 'ckpt')
    dataset = dataset_ops.Dataset.from_tensor_slices(
        [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
    dataset = dataset.map(math_ops.square).batch(2)
    # TODO(b/138399725): Re-enable default optimizations.
    options = dataset_ops.Options()
    options.experimental_optimization.apply_default_optimizations = False
    dataset = dataset.with_options(options)
    iterator_1 = datasets.Iterator(dataset)
    iterator_2 = datasets.Iterator(dataset)
    dataset_2 = dataset_ops.Dataset.range(10)
    iterator_3 = datasets.Iterator(dataset_2)

    checkpoint = trackable_utils.Checkpoint(
        iterator_1=iterator_1, iterator_2=iterator_2, iterator_3=iterator_3)
    self.assertAllEqual([1, 4], iterator_1.get_next().numpy())
    self.assertEqual(0, iterator_3.get_next().numpy())
    self.assertEqual(1, iterator_3.get_next().numpy())
    self.assertEqual(2, iterator_3.get_next().numpy())

    save_path = checkpoint.save(checkpoint_prefix)
    self.assertAllEqual([1, 4], iterator_2.get_next().numpy())
    self.assertAllEqual([9, 16], iterator_2.get_next().numpy())
    self.assertEqual(3, iterator_3.get_next().numpy())
    checkpoint.restore(save_path)
    self.assertAllEqual([9, 16], iterator_1.get_next().numpy())
    self.assertAllEqual([1, 4], iterator_2.get_next().numpy())
    self.assertEqual(3, iterator_3.get_next().numpy())

  def testRestoreExhaustedIterator(self):
    checkpoint_directory = self.get_temp_dir()
    checkpoint_prefix = os.path.join(checkpoint_directory, 'ckpt')
    dataset = dataset_ops.Dataset.range(3)
    iterator = datasets.Iterator(dataset)

    checkpoint = trackable_utils.Checkpoint(iterator=iterator)
    self.assertEqual(0, iterator.get_next().numpy())
    self.assertEqual(1, iterator.get_next().numpy())
    save_path = checkpoint.save(checkpoint_prefix)
    self.assertEqual(2, iterator.get_next().numpy())
    checkpoint.restore(save_path)
    self.assertEqual(2, iterator.get_next().numpy())

  def testRestoreInReconstructedIterator(self):
    checkpoint_directory = self.get_temp_dir()
    checkpoint_prefix = os.path.join(checkpoint_directory, 'ckpt')
    dataset = dataset_ops.Dataset.range(10)
    for i in range(5):
      iterator = datasets.Iterator(dataset)
      checkpoint = trackable_utils.Checkpoint(iterator=iterator)
      checkpoint.restore(
          checkpoint_management.latest_checkpoint(checkpoint_directory))
      for j in range(2):
        self.assertEqual(i * 2 + j, iterator.get_next().numpy())
      checkpoint.save(file_prefix=checkpoint_prefix)


class DatasetConstructorBenchmark(test.Benchmark):

  def benchmarkSliceRepeatBatchEager(self):
    input_size = 10000
    batch_size = 100
    num_epochs = 100

    input_data = np.random.randn(input_size)

    dataset = (
        dataset_ops.Dataset.from_tensor_slices(input_data).repeat(
            num_epochs).batch(batch_size))
    iterator = datasets.Iterator(dataset)

    ends = [time.time()]
    for _ in iterator:
      ends.append(time.time())

    deltas = np.ediff1d(ends)
    median_wall_time = np.median(deltas)
    print('Slice/repeat/batch eager input size: %d batch size: %d Median wall '
          'time per element: %f' % (input_size, batch_size, median_wall_time))
    self.report_benchmark(
        iters=len(deltas),
        wall_time=median_wall_time,
        name='benchmark_slice_repeat_batch_eager_input_%d_batch_%d' %
        (input_size, batch_size))

  def benchmarkSliceBatchCacheRepeatCallable(self):
    input_size = 10000
    batch_size = 100
    num_epochs = 100

    input_data = np.random.randn(input_size)

    dataset = (
        dataset_ops.Dataset.from_tensor_slices(input_data).batch(
            batch_size).cache().repeat(num_epochs))
    iterator = datasets.Iterator(dataset)

    ends = [time.time()]
    for _ in iterator:
      ends.append(time.time())

    deltas = np.ediff1d(ends)
    median_wall_time = np.median(deltas)
    print('Slice/batch/cache/repeat eager input size: %d batch size: %d Median '
          'wall time per element: %f' %
          (input_size, batch_size, median_wall_time))
    self.report_benchmark(
        iters=len(deltas),
        wall_time=median_wall_time,
        name='benchmark_slice_batch_cache_repeat_eager_input_%d_batch_%d' %
        (input_size, batch_size))


if __name__ == '__main__':
  test.main()
