# Copyright 2018 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 `tf.data.Optional`."""
import functools

from absl.testing import parameterized
import numpy as np

from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.data.ops import optional_ops
from tensorflow.python.data.util import structure
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.framework import combinations
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.framework import tensor_shape
from tensorflow.python.framework import tensor_spec
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test


def _optional_spec_test_combinations():
  # pylint: disable=g-long-lambda
  cases = [
      ("Dense", lambda: c
onstant_op.constant(37.0),
       tensor_spec.TensorSpec([], dtypes.float32)),
      ("Sparse", lambda: sparse_tensor.SparseTensor(
          indices=[[0, 1]],
          values=constant_op.constant([0], dtype=dtypes.int32),
          dense_shape=[10, 10]),
       sparse_tensor.SparseTensorSpec([10, 10], dtypes.int32)),
      ("Nest", lambda: {
          "a": constant_op.constant(37.0),
          "b": (constant_op.constant(["Foo"]), 
constant_op.constant("Bar"))
      }, {
          "a":
              tensor_spec.TensorSpec([], dtypes.float32),
          "b": (
              tensor_spec.TensorSpec([1], dtypes.string),
              tensor_spec.TensorSpec([], dtypes.string),
          )
      }),
      ("Optional", lambda: optional_ops.Optional.from_value(37.0),
       optional_ops.OptionalSpec(tensor_spec.TensorSpec([], dtypes.float32))),
  ]

  def reduce_fn(x, y):
    name, value_fn, expected_structure = y
    return x + combinations.combine(
        tf_value_fn=combinations.NamedObject(name, value_fn),
        expected_value_structure=expected_structure)

  return functools.reduce(reduce_fn, cases, [])


def _get_next_as_optional_test_combinations():
  # pylint: disable=g-long-lambda
  cases = [
      ("Dense", np.array([1, 2, 3], dtype=np.int32),
       lambda: constant_op.constant([4, 5, 6], dtype=dtypes.int32), True),
      ("Sparse",
       sparse_tensor.SparseTensorValue(
           indices=[[0, 0], [1, 1]],
           values=np.array([-1., 1.], dtype=np.float32),
           dense_shape=[2, 2]),
       lambda: sparse_tensor.SparseTensor(
           indices=[[0, 1], [1, 0]], values=[37.0, 42.0], dense_shape=[2, 2]),
       False),
      ("Nest", {
          "a":
              np.array([1, 2, 3], dtype=np.int32),
          "b":
              sparse_tensor.SparseTensorValue(
                  indices=[[0, 0], [1, 1]],
                  values=np.array([-1., 1.], dtype=np.float32),
                  dense_shape=[2, 2])
      }, lambda: {
          "a":
              constant_op.constant([4, 5, 6], dtype=dtypes.int32),
          "b":
              sparse_tensor.SparseTensor(
                  indices=[[0, 1], [1, 0]],
                  values=[37.0, 42.0],
                  dense_shape=[2, 2])
      }, False),
  ]

  def reduce_fn(x, y):
    name, value, value_fn, gpu_compatible = y
    return x + combinations.combine(
        np_value=value,
        tf_value_fn=combinations.NamedObject(name, value_fn),
        gpu_compatible=gpu_compatible)

  return functools.reduce(reduce_fn, cases, [])


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

  @combinations.generate(test_base.default_test_combinations())
  def testFromValue(self):
    opt = optional_ops.Optional.from_value(constant_op.constant(37.0))
    self.assertTrue(self.evaluate(opt.has_value()))
    self.assertEqual(37.0, self.evaluate(opt.get_value()))

  @combinations.generate(test_base.default_test_combinations())
  def testFromStructuredValue(self):
    opt = optional_ops.Optional.from_value({
        "a": constant_op.constant(37.0),
        "b": (constant_op.constant(["Foo"]), constant_op.constant("Bar"))
    })
    self.assertTrue(self.evaluate(opt.has_value()))
    self.assertEqual({
        "a": 37.0,
        "b": ([b"Foo"], b"Bar")
    }, self.evaluate(opt.get_value()))

  @combinations.generate(test_base.default_test_combinations())
  def testFromSparseTensor(self):
    st_0 = sparse_tensor.SparseTensorValue(
        indices=np.array([[0]]),
        values=np.array([0], dtype=np.int64),
        dense_shape=np.array([1]))
    st_1 = sparse_tensor.SparseTensorValue(
        indices=np.array([[0, 0], [1, 1]]),
        values=np.array([-1., 1.], dtype=np.float32),
        dense_shape=np.array([2, 2]))
    opt = optional_ops.Optional.from_value((st_0, st_1))
    self.assertTrue(self.evaluate(opt.has_value()))
    val_0, val_1 = opt.get_value()
    for expected, actual in [(st_0, val_0), (st_1, val_1)]:
      self.assertAllEqual(expected.indices, self.evaluate(actual.indices))
      self.assertAllEqual(expected.values, self.evaluate(actual.values))
      self.assertAllEqual(expected.dense_shape,
                          self.evaluate(actual.dense_shape))

  @combinations.generate(test_base.default_test_combinations())
  def testFromNone(self):
    value_structure = tensor_spec.TensorSpec([], dtypes.float32)
    opt = optional_ops.Optional.empty(value_structure)
    self.assertTrue(opt.element_spec.is_compatible_with(value_structure))
    self.assertFalse(
        opt.element_spec.is_compatible_with(
            tensor_spec.TensorSpec([1], dtypes.float32)))
    self.assertFalse(
        opt.element_spec.is_compatible_with(
            tensor_spec.TensorSpec([], dtypes.int32)))
    self.assertFalse(self.evaluate(opt.has_value()))
    with self.assertRaises(errors.InvalidArgumentError):
      self.evaluate(opt.get_value())

  @combinations.generate(test_base.default_test_combinations())
  def testAddN(self):
    devices = ["/cpu:0"]
    if test_util.is_gpu_available():
      devices.append("/gpu:0")
    for device in devices:
      with ops.device(device):
        # With value
        opt1 = optional_ops.Optional.from_value((1.0, 2.0))
        opt2 = optional_ops.Optional.from_value((3.0, 4.0))

        add_tensor = math_ops.add_n(
            [opt1._variant_tensor, opt2._variant_tensor])
        add_opt = optional_ops._OptionalImpl(add_tensor, opt1.element_spec)
        self.assertAllEqual(self.evaluate(add_opt.get_value()), (4.0, 6.0))

        # Without value
        opt_none1 = optional_ops.Optional.empty(opt1.element_spec)
        opt_none2 = optional_ops.Optional.empty(opt2.element_spec)
        add_tensor = math_ops.add_n(
            [opt_none1._variant_tensor, opt_none2._variant_tensor])
        add_opt = optional_ops._OptionalImpl(add_tensor, opt_none1.element_spec)
        self.assertFalse(self.evaluate(add_opt.has_value()))

  @combinations.generate(test_base.default_test_combinations())
  def testNestedAddN(self):
    devices = ["/cpu:0"]
    if test_util.is_gpu_available():
      devices.append("/gpu:0")
    for device in devices:
      with ops.device(device):
        opt1 = optional_ops.Optional.from_value([1, 2.0])
        opt2 = optional_ops.Optional.from_value([3, 4.0])
        opt3 = optional_ops.Optional.fr