# 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.
# ==============================================================================
"""Tests for the experimental input pipeline ops."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl.testing import parameterized
import numpy as np

from tensorflow.contrib.data.python.ops import batching
from tensorflow.contrib.data.python.ops import grouping
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import sparse_ops
from tensorflow.python.platform import test


class WindowDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):

  def _structuredDataset(self, structure, shape, dtype):
    if structure is None:
      return dataset_ops.Dataset.from_tensors(
          array_ops.zeros(shape, dtype=dtype))
    else:
      return dataset_ops.Dataset.zip(
          tuple([
              self._structuredDataset(substructure, shape, dtype)
              for substructure in structure
          ]))

  def _structuredElement(self, structure, shape, dtype):
    if structure is None:
      return array_ops.zeros(shape, dtype=dtype)
    else:
      return tuple([
          self._structuredElement(substructure, shape, dtype)
          for substructure in structure
      ])

  def _assertEqual(self, xs, ys):
    self.assertEqual(type(xs), type(ys))
    if isinstance(xs, tuple) and isinstance(ys, tuple):
      self.assertEqual(len(xs), len(ys))
      for x, y in zip(xs, ys):
        self._assertEqual(x, y)
    elif isinstance(xs, np.ndarray) and isinstance(ys, np.ndarray):
      self.assertAllEqual(xs, ys)
    else:
      self.assertEqual(xs, ys)

  @parameterized.named_parameters(
      ("1", None, np.int32([]), dtypes.bool),
      ("2", None, np.int32([]), dtypes.int32),
      ("3", None, np.int32([]), dtypes.float32),
      ("4", None, np.int32([]), dtypes.string),
      ("5", None, np.int32([2]), dtypes.int32),
      ("6", None, np.int32([2, 2]), dtypes.int32),
      ("7", (None, None, None), np.int32([]), dtypes.int32),
      ("8", (None, (None, None)), np.int32([]), dtypes.int32),
  )
  def testWindowDatasetFlatMap(self, structure, shape, dtype):
    """Tests windowing by chaining it with flat map.

    Args:
      structure: the input structure
      shape: the input shape
      dtype: the input data type
    """

    def fn(*args):
      if len(args) == 1 and not isinstance(args[0], tuple):
        return args[0]
      return dataset_ops.Dataset.zip(
          tuple([fn(*arg) if isinstance(arg, tuple) else arg for arg in args]))

    dataset = self._structuredDataset(structure, shape, dtype).repeat(5).apply(
        grouping.window_dataset(5)).flat_map(fn)
    get_next = dataset.make_one_shot_iterator().get_next()
    with self.cached_session() as sess:
      expected = sess.run(self._structuredElement(structure, shape, dtype))
      for _ in range(5):
        actual = sess.run(get_next)
        self._assertEqual(expected, actual)

  @parameterized.named_parameters(
      ("1", None, np.int32([]), dtypes.bool),
      ("2", None, np.int32([]), dtypes.int32),
      ("3", None, np.int32([]), dtypes.float32),
      ("4", None, np.int32([]), dtypes.string),
      ("5", None, np.int32([2]), dtypes.int32),
      ("6", None, np.int32([2, 2]), dtypes.int32),
      ("7", (None, None, None), np.int32([]), dtypes.int32),
      ("8", (None, (None, None)), np.int32([]), dtypes.int32),
  )
  def testWindowDatasetBatchDense(self, structure, shape, dtype):
    """Tests batching of dense tensor windows.

    Args:
      structure: the input structure
      shape: the input shape
      dtype: the input data type
    """

    def fn(*args):
      if len(args) == 1 and not isinstance(args[0], tuple):
        return batching.batch_window(args[0])

      return tuple([
          fn(*arg) if isinstance(arg, tuple) else batching.batch_window(arg)
          for arg in args
      ])

    dataset = self._structuredDataset(structure, shape, dtype).repeat(5).apply(
        grouping.window_dataset(5)).apply(grouping._map_x_dataset(fn))
    get_next = dataset.make_one_shot_iterator().get_next()
    with self.cached_session() as sess:
      expected = sess.run(
          self._structuredElement(structure, np.concatenate(
              ([5], shape), axis=0), dtype))
      actual = sess.run(get_next)
      self._assertEqual(expected, actual)

  @parameterized.named_parameters(
      ("1", np.int32([])),
      ("2", np.int32([1])),
      ("3", np.int32([1, 2, 3])),
  )
  def testWindowDatasetBatchDenseDynamicShape(self, shape):
    """Tests batching of dynamically shaped dense tensor windows.

    Args:
      shape: the input shape
    """

    shape_t = array_ops.placeholder(dtypes.int32)
    dataset = dataset_ops.Dataset.from_tensors(
        array_ops.zeros(shape_t)).repeat(5).apply(
            grouping.window_dataset(5)).apply(
                grouping._map_x_dataset(batching.batch_window))
    iterator = dataset.make_initializable_iterator()
    init_op = iterator.initializer
    get_next = iterator.get_next()
    with self.cached_session() as sess:
      sess.run(init_op, {shape_t: shape})
      expected = sess.run(
          self._structuredElement(None, np.concatenate(([5], shape), axis=0),
                                  dtypes.int32))
      actual = sess.run(get_next)
      self._assertEqual(expected, actual)

  def _make_dense_to_sparse_fn(self, is_scalar):

    def dense_to_sparse_scalar(tensor):
      indices = [[]]
      values = array_ops.expand_dims(tensor, 0)
      shape = []
      return sparse_tensor.SparseTensorValue(indices, values, shape)

    def dense_to_sparse_non_scalar(tensor):
      indices = array_ops.where(array_ops.ones_like(tensor, dtype=dtypes.bool))
      values = array_ops.gather_nd(tensor, indices)
      shape = array_ops.shape(tensor, out_type=dtypes.int64)
      return sparse_tensor.SparseTensorValue(indices, values, shape)

    if is_scalar:
      return dense_to_sparse_scalar
    return dense_to_sparse_non_scalar

  def _structuredSparseDataset(self, structure, shape, dtype):
    dense_to_sparse = self._make_dense_to_sparse_fn(len(shape) == 0)  # pylint: disable=g-explicit-length-test
    if structure is None:
      return dataset_ops.Dataset.from_tensors(
          dense_to_sparse(array_ops.zeros(shape, dtype=dtype)))
    else:
      return dataset_ops.Dataset.zip(
          tuple([
              self._structuredSparseDataset(substructure, shape, dtype)
              for substructure in structure
          ]))

  def _structuredSparseElement(self, structure, shape, dtype):
    dense_to_sparse = self._make_dense_to_sparse_fn(len(shape) == 0)  # pylint: disable=g-explicit-length-test
    if structure is None:
      return dense_to_sparse(array_ops.zeros(shape, dtype=dtype))
    else:
      return tuple([
          self._structuredSparseElement(substructure, shape, dtype)
          for substructure in structure
      ])

  @parameterized.named_parameters(
      ("1", None, np.int32([]), dtypes.bool),
      ("2", None, np.int32([]), dtypes.int32),
      ("3", None, np.int32([]), dtypes.float32),
      ("4", None, np.int32([]), dtypes.string),
      ("5", None, np.int32([2]), dtypes.int32),
      ("6", None, np.int32([2, 2]), dtypes.int32),
      ("7", (None, None, None), np.int32([]), dtypes.int32),
      ("8", (None, (None, None)), np.int32([]), dtypes.int32),
  )
  def testWindowDatasetBatchSparse(self, structure, shape, dtype):
    """Tests batching of sparse tensor windows.

    Args:
      structure: the input structure
      shape: the input shape
      dtype: the input data type
    """

    def fn(*args):
      if len(args) == 1 and not isinstance(args[0], tuple):
        return batching.batch_window(args[0])

      return tuple([
          fn(*arg) if isinstance(arg, tuple) else batching.batch_window(arg)
          for arg in args
      ])

    dataset = self._structuredSparseDataset(
        structure, shape, dtype).repeat(5).apply(
            grouping.window_dataset(5)).apply(grouping._map_x_dataset(fn))
    get_next = dataset.make_one_shot_iterator().get_next()
    with self.cached_session() as sess:
      expected = sess.run(
          self._structuredSparseElement(structure,
                                        np.concatenate(([5], shape), axis=0),
                                        dtype))
      actual = sess.run(get_next)
      self._assertEqual(expected, actual)

  @parameterized.named_parameters(
      ("1", np.int32([])),
      ("2", np.int32([1])),
      ("3", np.int32([1, 2, 3])),
  )
  def testWindowDatasetBatchSparseDynamicShape(self, shape):
    """Tests batching of dynamically shaped sparse tensor windows.

    Args:
      shape: the input shape
    """

    shape_t = array_ops.placeholder(dtypes.int32)
    dataset = dataset_ops.Dataset.from_tensors(array_ops.zeros(shape_t)).map(
        self._make_dense_to_sparse_fn(len(shape) == 0)).repeat(5).apply(  # pylint: disable=g-explicit-length-test
            grouping.window_dataset(5)).apply(
                grouping._map_x_dataset(batching.batch_window))
    iterator = dataset.make_initializable_iterator()
    init_op = iterator.initializer
    get_next = iterator.get_next()
    with self.cached_session() as sess:
      sess.run(init_op, {shape_t: shape})
      expected = sess.run(
          self._structuredSparseElement(None,
                                        np.concatenate(([5], shape), axis=0),
                                        dtypes.int32))
      actual = sess.run(get_next)
      self._assertEqual(expected, actual)

  def _structuredRaggedDataset(self, structure, shapes, dtype):

    if structure is None:
      return dataset_ops.Dataset.from_tensor_slices(shapes).map(
          lambda shape: array_ops.zeros(shape, dtype=dtype))
    else:
      return dataset_ops.Dataset.zip(
          tuple([
              self._structuredRaggedDataset(substructure, shapes, dtype)
              for substructure in structure
          ]))

  @parameterized.named_parameters(
      ("1", None, np.int32([[1], [2], [3]]), dtypes.bool, [-1]),
      ("2", None, np.int32([[1], [2], [3]]), dtypes.int32, [-1]),
      ("3", None, np.int32([[1], [2], [3]]), dtypes.float32, [-1]),
      ("4", None, np.int32([[1], [2], [3]]), dtypes.string, [-1]),
      ("5", None, np.int32([[1, 3], [2, 2], [3, 1]]), dtypes.int32, [-1, -1]),
      ("6", None, np.int32([[3, 1, 3], [1, 3, 1]]), dtypes.int32, [-1, -1, -1]),
      ("7", (None, None, None), np.int32([[1], [2], [3]]), dtypes.int32, [-1]),
      ("8", (None,
             (None, None)), np.int32([[1], [2], [3]]), dtypes.int32, [-1]),
      ("9", None, np.int32([[1], [2], [3]]), dtypes.int32, [-1]),
      ("10", None, np.int32([[1], [2], [3]]), dtypes.int32, np.int32([10])),
  )
  def testWindowDatasetPaddedBatchDense(self, structure, shapes, dtype,
                                        padded_shape):
    """Tests padded batching of dense tensor windows.

    Args:
      structure: the input structure
      shapes: the input shapes
      dtype: the input data type
      padded_shape: the shape to pad the output to
    """

    def fn(*args):
      if len(args) == 1 and not isinstance(args[0], tuple):
        return batching.padded_batch_window(args[0], padded_shape)

      return tuple([
          fn(*arg) if isinstance(arg, tuple) else batching.padded_batch_window(
              arg, padded_shape) for arg in args
      ])

    dataset = self._structuredRaggedDataset(structure, shapes, dtype).apply(
        grouping.window_dataset(len(shapes))).apply(
            grouping._map_x_dataset(fn))
    get_next = dataset.make_one_shot_iterator().get_next()
    with self.cached_session() as sess:
      expected_shape = np.maximum(np.amax(shapes, axis=0), padded_shape)
      expected = sess.run(
          self._structuredElement(
              structure,
              np.concatenate((np.int32([len(shapes)]), expected_shape)), dtype))
      actual = sess.run(get_next)
      self._assertEqual(expected, actual)

  @parameterized.named_parameters(
      ("1", np.int32([[1], [2], [3]]), [-1]),
      ("2", np.int32([[1, 3], [2, 2], [3, 1]]), [-1, -1]),
      ("3", np.int32([[3, 1, 3], [1, 3, 1]]), [-1, -1, -1]),
  )
  def testWindowDatasetPaddedBatchDenseDynamicShape(self, shapes, padded_shape):
    """Tests padded batching of dynamically shaped dense tensor windows.

    Args:
      shapes: the input shapes
      padded_shape: the shape to pad the output to
    """

    shapes_t = array_ops.placeholder(dtypes.int32)
    dataset = dataset_ops.Dataset.from_tensor_slices(shapes_t).map(
        lambda shape: array_ops.zeros(shape, dtype=dtypes.int32)).apply(
            grouping.window_dataset(len(shapes))).apply(
                grouping._map_x_dataset(
                    lambda x: batching.padded_batch_window(x, padded_shape)))
    iterator = dataset.make_initializable_iterator()
    init_op = iterator.initializer
    get_next = iterator.get_next()
    with self.cached_session() as sess:
      sess.run(init_op, {shapes_t: shapes})
      expected_shape = np.maximum(np.amax(shapes, axis=0), padded_shape)
      expected = sess.run(
          self._structuredElement(
              None, np.concatenate((np.int32([len(shapes)]), expected_shape)),
              dtypes.int32))
      actual = sess.run(get_next)
      self._assertEqual(expected, actual)

  @parameterized.named_parameters(
      ("1", np.int32([[1]]), np.int32([0])),
      ("2", np.int32([[10], [20]]), np.int32([15])),
  )
  def testWindowDatasetPaddedBatchDenseInvalid(self, shapes, padded_shape):
    """Tests invalid padded batching of dense tensor windows.

    Args:
      shapes: the input shapes
      padded_shape: the shape to pad the output to
    """

    dataset = dataset_ops.Dataset.from_tensor_slices(shapes).map(
        lambda shape: array_ops.zeros(shape, dtype=dtypes.int32)).apply(
            grouping.window_dataset(len(shapes))).apply(
                grouping._map_x_dataset(
                    lambda x: batching.padded_batch_window(x, padded_shape)))
    get_next = dataset.make_one_shot_iterator().get_next()
    with self.cached_session() as sess:
      with self.assertRaises(errors.InvalidArgumentError):
        sess.run(get_next)

  def _structuredRaggedSparseDataset(self, structure, shapes, dtype):

    def map_fn(shape):
      dense_to_sparse = self._make_dense_to_sparse_fn(False)
      return dense_to_sparse(array_ops.zeros(shape, dtype=dtype))

    if structure is None:
      return dataset_ops.Dataset.from_tensor_slices(shapes).map(map_fn)
    else:
      return dataset_ops.Dataset.zip(
          tuple([
              self._structuredRaggedSparseDataset(substructure, shapes, dtype)
              for substructure in structure
          ]))

  def _structuredRaggedSparseElement(self, structure, shapes, dtype,
                                     padded_shape):
    if structure is None:
      dense_shape = np.maximum(np.amax(shapes, axis=0), padded_shape)
      values = []
      for shape in shapes:
        dense_to_sparse = self._make_dense_to_sparse_fn(len(shape) == 0)  # pylint: disable=g-explicit-length-test
        sparse = dense_to_sparse(array_ops.zeros(shape, dtype=dtype))
        padded_sparse = sparse_tensor.SparseTensor(sparse.indices,
                                                   sparse.values, dense_shape)
        reshaped_sparse = sparse_ops.sparse_reshape(
            padded_sparse,
            array_ops.concat([np.array([1], dtype=np.int64), dense_shape], 0))
        values.append(reshaped_sparse)
      return sparse_ops.sparse_concat(0, values)
    else:
      return tuple([
          self._structuredRaggedSparseElement(substructure, shapes, dtype,
                                              padded_shape)
          for substructure in structure
      ])

  @parameterized.named_parameters(
      ("1", None, np.int64([[1], [2], [3]]), dtypes.bool, [-1]),
      ("2", None, np.int64([[1], [2], [3]]), dtypes.int32, [-1]),
      ("3", None, np.int64([[1], [2], [3]]), dtypes.float32, [-1]),
      ("4", None, np.int64([[1], [2], [3]]), dtypes.string, [-1]),
      ("5", None, np.int64([[1, 3], [2, 2], [3, 1]]), dtypes.int32, [-1, -1]),
      ("6", None, np.int64([[1, 3, 1], [3, 1, 3]]), dtypes.int32, [-1, -1, -1]),
      ("7", (None, None, None), np.int64([[1], [2], [3]]), dtypes.int32, [-1]),
      ("8", (None,
             (None, None)), np.int64([[1], [2], [3]]), dtypes.int32, [-1]),
      ("9", None, np.int64([[1], [2], [3]]), dtypes.int32, [-1]),
      ("10", None, np.int64([[1], [2], [3]]), dtypes.int32, np.int64([10])),
  )
  def testWindowDatasetPaddedBatchSparse(self, structure, shapes, dtype,
                                         padded_shape):
    """Tests padded batching of sparse tensor windows.

    Args:
      structure: the input structure
      shapes: the input shapes
      dtype: the input data type
      padded_shape: the shape to pad the output to
    """

    def fn(*args):
      if len(args) == 1 and not isinstance(args[0], tuple):
        return batching.padded_batch_window(args[0], padded_shape)

      return tuple([
          fn(*arg) if isinstance(arg, tuple) else batching.padded_batch_window(
              arg, padded_shape) for arg in args
      ])

    dataset = self._structuredRaggedSparseDataset(
        structure, shapes, dtype).apply(grouping.window_dataset(
            len(shapes))).apply(grouping._map_x_dataset(fn))
    get_next = dataset.make_one_shot_iterator().get_next()
    with self.cached_session() as sess:
      expected = sess.run(
          self._structuredRaggedSparseElement(structure, shapes, dtype,
                                              padded_shape))
      actual = sess.run(get_next)
      self._assertEqual(expected, actual)

  @parameterized.named_parameters(
      ("1", np.int64([[1], [2], [3]]), [-1]),
      ("2", np.int64([[1, 3], [2, 2], [3, 1]]), [-1, -1]),
      ("3", np.int64([[3, 1, 3], [1, 3, 1]]), [-1, -1, -1]),
  )
  def testWindowDatasetPaddedBatchSparseDynamicShape(self, shapes,
                                                     padded_shape):
    """Tests padded batching of dynamically shaped sparse tensor windows.

    Args:
      shapes: the input shapes
      padded_shape: the shape to pad the output to
    """

    shapes_t = array_ops.placeholder(dtypes.int32)
    dataset = dataset_ops.Dataset.from_tensor_slices(shapes_t).map(
        lambda shape: array_ops.zeros(shape, dtype=dtypes.int32)).map(
            self._make_dense_to_sparse_fn(False)
        ).apply(grouping.window_dataset(len(shapes))).apply(
            grouping._map_x_dataset(
                lambda x: batching.padded_batch_window(x, padded_shape)))
    iterator = dataset.make_initializable_iterator()
    init_op = iterator.initializer
    get_next = iterator.get_next()
    with self.cached_session() as sess:
      sess.run(init_op, {shapes_t: shapes})
      expected = sess.run(
          self._structuredRaggedSparseElement(None, shapes, dtypes.int32,
                                              padded_shape))
      actual = sess.run(get_next)
      self._assertEqual(expected, actual)

  @parameterized.named_parameters(
      ("1", np.int64([[1]]), [0]),
      ("2", np.int64([[10], [20]]), [15]),
  )
  def testWindowDatasetPaddedBatchSparseInvalid(self, shapes, padded_shape):
    """Tests invalid padded batching of sparse tensor windows.

    Args:
      shapes: the input shapes
      padded_shape: the shape to pad the output to
    """

    dataset = dataset_ops.Dataset.from_tensor_slices(shapes).map(
        lambda shape: array_ops.zeros(shape, dtype=dtypes.int32)).map(
            self._make_dense_to_sparse_fn(False)
        ).apply(grouping.window_dataset(len(shapes))).apply(
            grouping._map_x_dataset(
                lambda x: batching.padded_batch_window(x, padded_shape)))
    get_next = dataset.make_one_shot_iterator().get_next()
    with self.cached_session() as sess:
      with self.assertRaises(errors.InvalidArgumentError):
        sess.run(get_next)


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