# coding=utf-8
# Copyright 2020 The TF-Agents 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
#
#     https://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 example_encoding."""

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

from absl.testing import parameterized
import gin
import numpy as np
import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import
from tf_agents.specs import array_spec
from tf_agents.utils import common
from tf_agents.utils import example_encoding


TYPE_PARAMETERS = (
    ("np.uint8", np.uint8),
    ("np.uint16", np.uint16),
    ("np.int8", np.int8),
    ("np.int16", np.int16),
    ("np.int32", np.int32),
    ("np.int64", np.int64),
    ("np.float16", np.float16),
    ("np.float32", np.float32),
    ("np.float64", np.float64),
    ("python float", float),
    ("python int", int),
)


def example_nested_spec(dtype):
  """Return an example nested array spec."""
  low = -10
  high = 10
  if dtype in (np.uint8, np.uint16):
    low += -low
  return {
      "array_spec_1":
          array_spec.ArraySpec((2, 3), dtype),
      "bounded_spec_1":
          array_spec.BoundedArraySpec((2, 3), dtype, low, high),
      "empty_shape":
          array_spec.BoundedArraySpec((), dtype, low, high),
      "dict_spec": {
          "array_spec_2":
              array_spec.ArraySpec((2, 3), dtype),
          "bounded_spec_2":
              array_spec.BoundedArraySpec((2, 3), dtype, low, high)
      },
      "tuple_spec": (
          array_spec.ArraySpec((2, 3), dtype),
          array_spec.BoundedArraySpec((2, 3), dtype, low, high),
      ),
      "list_spec": [
          array_spec.ArraySpec((2, 3), dtype),
          (array_spec.ArraySpec((2, 3), dtype),
           array_spec.BoundedArraySpec((2, 3), dtype, low, high)),
      ],
  }


class NestExampleEncodeUtilsTest(tf.test.TestCase, parameterized.TestCase):

  @parameterized.named_parameters(*TYPE_PARAMETERS)
  def test_serialize_deserialize(self, dtype):
    spec = example_nested_spec(dtype)
    serializer = example_encoding.get_example_serializer(spec)
    decoder = example_encoding.get_example_decoder(spec)

    sample = array_spec.sample_spec_nest(spec, np.random.RandomState(0))
    example_proto = serializer(sample)

    recovered = self.evaluate(decoder(example_proto))
    tf.nest.map_structure(np.testing.assert_almost_equal, sample, recovered)

  def test_endian_encodings(self):
    spec = {
        "a": array_spec.ArraySpec((2,), np.int16),
        "b": array_spec.ArraySpec((2,), np.int32),
        "c": array_spec.ArraySpec((2,), np.float32),
    }

    serializer = example_encoding.get_example_serializer(spec)
    decoder = example_encoding.get_example_decoder(spec)

    # Little endian encoding.
    le_sample = {
        "a": np.array([100, 25000]).astype("<i2"),
        "b": np.array([-5, 80000000]).astype("<i4"),
        "c": np.array([12.5, np.pi]).astype("<f4")
    }

    example_proto = serializer(le_sample)
    recovered = self.evaluate(decoder(example_proto))
    tf.nest.map_structure(np.testing.assert_almost_equal, le_sample, recovered)

    # Big endian encoding.
    be_sample = {
        "a": np.array([100, 25000]).astype(">i2"),
        "b": np.array([-5, 80000000]).astype(">i4"),
        "c": np.array([12.5, np.pi]).astype(">f4")
    }

    example_proto = serializer(be_sample)
    recovered = self.evaluate(decoder(example_proto))
    tf.nest.map_structure(np.testing.assert_almost_equal, be_sample, recovered)

  def test_shape_validation(self):
    with self.assertRaisesRegexp(ValueError, "is invalid"):
      example_encoding._validate_shape([1, 2, 3, -1])

    with self.assertRaisesRegexp(ValueError, "is invalid"):
      example_encoding._validate_shape([1, None, 3])

    with self.assertRaisesRegexp(ValueError, "is invalid"):
      example_encoding._validate_shape([1, 2.3, 3])

  def test_compress_image(self):
    if not common.has_eager_been_enabled():
      self.skipTest("Image compression only supported in TF2.x")

    gin.parse_config_files_and_bindings([], """
    _get_feature_encoder.compress_image=True
    _get_feature_parser.compress_image=True
    """)
    spec = {
        "image": array_spec.ArraySpec((128, 128, 3), np.uint8),
        "mask": array_spec.ArraySpec((128, 128, 1), np.uint8)
    }
    serializer = example_encoding.get_example_serializer(spec)
    decoder = example_encoding.get_example_decoder(spec)

    sample = {
        "image": 128 * np.ones([128, 128, 3], dtype=np.uint8),
        "mask": 128 * np.ones([128, 128, 1], dtype=np.uint8)
    }
    example_proto = serializer(sample)

    recovered = self.evaluate(decoder(example_proto))
    tf.nest.map_structure(np.testing.assert_almost_equal, sample, recovered)

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