# 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 nested structure coding."""

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

import collections

from google.protobuf import text_format
from tensorflow.core.protobuf import struct_pb2
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.platform import test
from tensorflow.python.saved_model import nested_structure_coder


class NestedStructureTest(test.TestCase):

  def setUp(self):
    self._coder = nested_structure_coder.StructureCoder()

  def testEncodeDecodeList(self):
    structure = [1.5, 2.5, 3.0]
    self.assertTrue(self._coder.can_encode(structure))
    encoded = self._coder.encode_structure(structure)
    expected = struct_pb2.StructuredValue()
    expected.list_value.values.add().float64_value = 1.5
    expected.list_value.values.add().float64_value = 2.5
    expected.list_value.values.add().float64_value = 3.0
    self.assertEqual(expected, encoded)
    decoded = self._coder.decode_proto(encoded)
    self.assertEqual(structure, decoded)

  def testEncodeDecodeTuple(self):
    structure = ("hello", [3, (2, 1)])
    self.assertTrue(self._coder.can_encode(structure))
    encoded = self._coder.encode_structure(structure)
    expected = struct_pb2.StructuredValue()
    expected.tuple_value.values.add().string_value = "hello"
    list_value = expected.tuple_value.values.add().list_value
    list_value.values.add().int64_value = 3
    tuple_value = list_value.values.add().tuple_value
    tuple_value.values.add().int64_value = 2
    tuple_value.values.add().int64_value = 1
    self.assertEqual(expected, encoded)
    decoded = self._coder.decode_proto(encoded)
    self.assertEqual(structure, decoded)

  def testEncodeDecodeDict(self):
    structure = dict(a=3, b=[7, 2.5])
    self.assertTrue(self._coder.can_encode(structure))
    encoded = self._coder.encode_structure(structure)
    expected = struct_pb2.StructuredValue()
    expected.dict_value.fields["a"].int64_value = 3
    list_value = expected.dict_value.fields["b"].list_value
    list_value.values.add().int64_value = 7
    list_value.values.add().float64_value = 2.5
    self.assertEqual(expected, encoded)
    decoded = self._coder.decode_proto(encoded)
    self.assertIsInstance(decoded["a"], int)
    self.assertEqual(structure, decoded)

  def testEncodeDecodeTensorShape(self):
    structure = [tensor_shape.TensorShape([1, 2, 3]), "hello"]
    self.assertTrue(self._coder.can_encode(structure))
    encoded = self._coder.encode_structure(structure)
    expected = struct_pb2.StructuredValue()
    expected_list = expected.list_value
    expected_tensor_shape = expected_list.values.add().tensor_shape_value
    expected_tensor_shape.dim.add().size = 1
    expected_tensor_shape.dim.add().size = 2
    expected_tensor_shape.dim.add().size = 3
    expected_tensor_shape = expected_list.values.add().string_value = "hello"
    self.assertEqual(expected, encoded)
    decoded = self._coder.decode_proto(encoded)
    self.assertEqual(structure, decoded)

  def testEncodeDecodeNamedTuple(self):
    named_tuple_type = collections.namedtuple("NamedTuple", ["x", "y"])
    named_tuple = named_tuple_type(x=[1, 2], y="hello")
    self.assertTrue(self._coder.can_encode(named_tuple))
    encoded = self._coder.encode_structure(named_tuple)
    expected = struct_pb2.StructuredValue()
    expected_named_tuple = expected.named_tuple_value
    expected_named_tuple.name = "NamedTuple"
    key_value_pair = expected_named_tuple.values.add()
    key_value_pair.key = "x"
    list_value = key_value_pair.value.list_value
    list_value.values.add().int64_value = 1
    list_value.values.add().int64_value = 2
    key_value_pair = expected_named_tuple.values.add()
    key_value_pair.key = "y"
    key_value_pair.value.string_value = "hello"
    self.assertEqual(expected, encoded)
    decoded = self._coder.decode_proto(encoded)
    self.assertEqual(named_tuple._asdict(), decoded._asdict())
    self.assertEqual(named_tuple.__class__.__name__, decoded.__class__.__name__)

  def testNone(self):
    structure = [1.0, None]
    self.assertTrue(self._coder.can_encode(structure))
    encoded = self._coder.encode_structure(structure)
    expected = struct_pb2.StructuredValue()
    expected.list_value.values.add().float64_value = 1.0
    expected.list_value.values.add().none_value.CopyFrom(struct_pb2.NoneValue())
    self.assertEqual(expected, encoded)
    decoded = self._coder.decode_proto(encoded)
    self.assertEqual(structure, decoded)

  def testBool(self):
    structure = [False]
    self.assertTrue(self._coder.can_encode(structure))
    encoded = self._coder.encode_structure(structure)
    expected = struct_pb2.StructuredValue()
    expected.list_value.values.add().bool_value = False
    self.assertEqual(expected, encoded)
    decoded = self._coder.decode_proto(encoded)
    self.assertEqual(structure, decoded)

  def testEmptyStructures(self):
    structure = [list(), dict(), tuple()]
    self.assertTrue(self._coder.can_encode(structure))
    encoded = self._coder.encode_structure(structure)
    expected = struct_pb2.StructuredValue()
    expected.list_value.values.add().list_value.CopyFrom(struct_pb2.ListValue())
    expected.list_value.values.add().dict_value.CopyFrom(struct_pb2.DictValue())
    expected.list_value.values.add().tuple_value.CopyFrom(
        struct_pb2.TupleValue())
    self.assertEqual(expected, encoded)
    decoded = self._coder.decode_proto(encoded)
    self.assertEqual(structure, decoded)

  def testDtype(self):
    structure = [dtypes.int64]
    self.assertTrue(self._coder.can_encode(structure))
    encoded = self._coder.encode_structure(structure)
    expected = struct_pb2.StructuredValue()
    list_value = expected.list_value.values.add()
    list_value.tensor_dtype_value = dtypes.int64.as_datatype_enum
    self.assertEqual(expected, encoded)
    decoded = self._coder.decode_proto(encoded)
    self.assertEqual(structure, decoded)

  def testEncodeDecodeTensorSpec(self):
    structure = [tensor_spec.TensorSpec([1, 2, 3], dtypes.int64, "hello")]
    self.assertTrue(self._coder.can_encode(structure))
    encoded = self._coder.encode_structure(structure)
    expected = struct_pb2.StructuredValue()
    expected_list = expected.list_value
    expected_tensor_spec = expected_list.values.add().tensor_spec_value
    expected_tensor_spec.shape.dim.add().size = 1
    expected_tensor_spec.shape.dim.add().size = 2
    expected_tensor_spec.shape.dim.add().size = 3
    expected_tensor_spec.name = "hello"
    expected_tensor_spec.dtype = dtypes.int64.as_datatype_enum
    self.assertEqual(expected, encoded)
    decoded = self._coder.decode_proto(encoded)
    self.assertEqual(structure, decoded)

  def testEncodeDecodeTensorSpecWithNoName(self):
    structure = [tensor_spec.TensorSpec([1, 2, 3], dtypes.int64)]
    self.assertTrue(self._coder.can_encode(structure))
    encoded = self._coder.encode_structure(structure)
    expected = struct_pb2.StructuredValue()
    expected_list = expected.list_value
    expected_tensor_spec = expected_list.values.add().tensor_spec_value
    expected_tensor_spec.shape.dim.add().size = 1
    expected_tensor_spec.shape.dim.add().size = 2
    expected_tensor_spec.shape.dim.add().size = 3
    expected_tensor_spec.name = ""
    expected_tensor_spec.dtype = dtypes.int64.as_datatype_enum
    self.assertEqual(expected, encoded)
    decoded = self._coder.decode_proto(encoded)
    self.assertEqual(structure, decoded)

  def testEncodeDecodeRaggedTensorSpec(self):
    structure = [ragged_tensor.RaggedTensorSpec(
        [1, 2, 3], dtypes.int64, 2, dtypes.int32)]
    self.assertTrue(self._coder.can_encode(structure))
    encoded = self._coder.encode_structure(structure)
    expected_pbtxt = r"""
      list_value {
        values {
          type_spec_value {
            type_spec_class: RAGGED_TENSOR_SPEC
            type_spec_class_name: 'RaggedTensorSpec'
            type_state {
              tuple_value {
                # spec._shape
                values {
                  tensor_shape_value {
                    dim { size: 1 }
                    dim { size: 2 }
                    dim { size: 3 }
                  }
                }
                # spec._dtype
                values { tensor_dtype_value: DT_INT64 }
                # spec._ragged_rank
                values { int64_value: 2 }
                # spec._row_splits_dtype
                values { tensor_dtype_value: DT_INT32 }
              }
            }
          }
        }
      }
    """
    expected = struct_pb2.StructuredValue()
    text_format.Parse(expected_pbtxt, expected)
    self.assertEqual(expected, encoded)
    decoded = self._coder.decode_proto(encoded)
    self.assertEqual(structure, decoded)

  def testEncodeDecodeSparseTensorSpec(self):
    structure = [sparse_tensor.SparseTensorSpec([10, 20], dtypes.float32)]
    self.assertTrue(self._coder.can_encode(structure))
    encoded = self._coder.encode_structure(structure)
    expected_pbtxt = r"""
      list_value {
        values {
          type_spec_value {
            type_spec_class: SPARSE_TENSOR_SPEC
            type_spec_class_name: 'SparseTensorSpec'
            type_state {
              tuple_value {
                # spec._shape
                values {
                  tensor_shape_value {
                    dim { size: 10 }
                    dim { size: 20 }
                  }
                }
                # spec._dtype
                values { tensor_dtype_value: DT_FLOAT }
              }
            }
          }
        }
      }
    """
    expected = struct_pb2.StructuredValue()
    text_format.Parse(expected_pbtxt, expected)
    self.assertEqual(expected, encoded)
    decoded = self._coder.decode_proto(encoded)
    self.assertEqual(structure, decoded)

  def testDecodeUnknownTensorSpec(self):
    encoded = struct_pb2.StructuredValue()
    encoded.type_spec_value.type_spec_class = 0
    encoded.type_spec_value.type_spec_class_name = "FutureTensorSpec"
    with self.assertRaisesRegexp(
        ValueError, "The type 'FutureTensorSpec' is not supported"):
      self._coder.decode_proto(encoded)

  def testEncodeDataSetSpec(self):
    structure = [dataset_ops.DatasetSpec(
        {"rt": ragged_tensor.RaggedTensorSpec([10, None], dtypes.int32),
         "st": sparse_tensor.SparseTensorSpec([10, 20], dtypes.float32),
         "t": tensor_spec.TensorSpec([10, 8], dtypes.string)})]
    self.assertTrue(self._coder.can_encode(structure))
    encoded = self._coder.encode_structure(structure)
    decoded = self._coder.decode_proto(encoded)
    self.assertEqual(structure, decoded)

  def testNotEncodable(self):

    class NotEncodable(object):
      pass

    self.assertFalse(self._coder.can_encode([NotEncodable()]))


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