# Copyright 2015 Google Inc. 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 tensorflow.ops.parsing_ops."""

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

import itertools

import tensorflow.python.platform

import numpy as np
import tensorflow as tf

# Helpers for creating Example objects
example = tf.train.Example
feature = tf.train.Feature
features = lambda d: tf.train.Features(feature=d)
bytes_feature = lambda v: feature(bytes_list=tf.train.BytesList(value=v))
int64_feature = lambda v: feature(int64_list=tf.train.Int64List(value=v))
float_feature = lambda v: feature(float_list=tf.train.FloatList(value=v))


def flatten(list_of_lists):
  """Flatten one level of nesting."""
  return itertools.chain.from_iterable(list_of_lists)


def flatten_values_tensors_or_sparse(tensors_list):
  """Flatten each SparseTensor object into 3 Tensors for session.run()."""
  return list(flatten([[v.indices, v.values, v.shape]
                       if isinstance(v, tf.SparseTensor) else [v]
                       for v in tensors_list]))


def _compare_output_to_expected(
    tester, dict_tensors, expected_tensors, flat_output):
  tester.assertEqual(set(dict_tensors.keys()), set(expected_tensors.keys()))

  i = 0  # Index into the flattened output of session.run()
  for k, v in dict_tensors.items():
    expected_v = expected_tensors[k]
    tf.logging.info("Comparing key: %s", k)
    if isinstance(v, tf.SparseTensor):
      # Three outputs for SparseTensor : indices, values, shape.
      tester.assertEqual([k, len(expected_v)], [k, 3])
      tester.assertAllEqual(expected_v[0], flat_output[i])
      tester.assertAllEqual(expected_v[1], flat_output[i + 1])
      tester.assertAllEqual(expected_v[2], flat_output[i + 2])
      i += 3
    else:
      # One output for standard Tensor.
      tester.assertAllEqual(expected_v, flat_output[i])
      i += 1


class ParseExampleTest(tf.test.TestCase):

  def _test(self, kwargs, expected_values=None, expected_err_re=None):
    with self.test_session() as sess:
      # Pull out some keys to check shape inference
      serialized = kwargs["serialized"]
      dense_keys = kwargs["dense_keys"] if "dense_keys" in kwargs else []
      sparse_keys = kwargs["sparse_keys"] if "sparse_keys" in kwargs else []
      dense_shapes = kwargs["dense_shapes"] if "dense_shapes" in kwargs else []

      # Returns dict w/ Tensors and SparseTensors
      out = tf.parse_example(**kwargs)

      # Check shapes; if serialized is a Tensor we need its size to
      # properly check.
      batch_size = (
          serialized.eval().size if isinstance(serialized, tf.Tensor)
          else np.asarray(serialized).size)
      if dense_shapes:
        self.assertEqual(len(dense_keys), len(dense_shapes))
        for (k, s) in zip(dense_keys, dense_shapes):
          self.assertEqual(
              tuple(out[k].get_shape().as_list()), (batch_size,) + s)
        for k in sparse_keys:
          self.assertEqual(
              tuple(out[k].indices.get_shape().as_list()), (None, 2))
          self.assertEqual(tuple(out[k].values.get_shape().as_list()), (None,))
          self.assertEqual(tuple(out[k].shape.get_shape().as_list()), (2,))

      # Check values
      result = flatten_values_tensors_or_sparse(out.values())  # flatten values
      if expected_err_re is None:
        tf_result = sess.run(result)
        _compare_output_to_expected(self, out, expected_values, tf_result)
      else:
        with self.assertRaisesOpError(expected_err_re):
          sess.run(result)

  def testEmptySerializedWithAllDefaults(self):
    cname = "c:has_a_tricky_name"
    dense_keys = ["a", "b", cname]
    dense_shapes = [(1, 3), (3, 3), (2,)]
    dense_types = [tf.int64, tf.string, tf.float32]
    dense_defaults = {
        "a": [0, 42, 0],
        "b": np.random.rand(3, 3).astype(bytes),
        cname: np.random.rand(2).astype(np.float32),
    }

    expected_st_a = (  # indices, values, shape
        np.empty((0, 2), dtype=np.int64),  # indices
        np.empty((0,), dtype=np.int64),  # sp_a is DT_INT64
        np.array([2, 0], dtype=np.int64))  # batch == 2, max_elems = 0

    expected_output = {
        "st_a": expected_st_a,
        "a": np.array(2 * [[dense_defaults["a"]]]),
        "b": np.array(2 * [dense_defaults["b"]]),
        cname: np.array(2 * [dense_defaults[cname]]),
    }

    self._test(
        {
            "names": np.empty((0,), dtype=bytes),
            # empty serialized input Examples
            "serialized": tf.convert_to_tensor(["", ""]),
            "dense_defaults": dense_defaults,
            "sparse_keys": ["st_a"],
            "sparse_types": [tf.int64],
            "dense_keys": dense_keys,
            "dense_types": dense_types,
            "dense_shapes": dense_shapes
        }, expected_output)

  def testEmptySerializedWithoutDefaultsShouldFail(self):
    dense_shapes = [(1, 3), (3, 3), (2,)]
    dense_defaults = {
        "a": [0, 42, 0],
        "b": np.random.rand(3, 3).astype(bytes),
        # Feature "c" is missing, since there's gaps it will cause failure.
    }
    self._test(
        {
            "serialized": ["", ""],  # empty serialized input Examples
            "names": ["in1", "in2"],
            "dense_defaults": dense_defaults,
            "sparse_keys": ["st_a"],
            "sparse_types": [tf.int64],
            "dense_keys": ["a", "b", "c"],
            "dense_types": [tf.int64, tf.string, tf.float32],
            "dense_shapes": dense_shapes
        },
        expected_err_re="Name: in1, Feature: c is required")

  def testDenseNotMatchingShapeShouldFail(self):
    dense_shapes = [(1, 3)]
    dense_defaults = {
        # no default!
    }

    original = [
        example(features=features({
            "a": float_feature([1, 1, 3]),
        })),
        example(features=features({
            "a": float_feature([-1, -1]),
        }))
    ]

    names = ["passing", "failing"]
    serialized = [m.SerializeToString() for m in original]

    self._test(
        {
            "serialized": tf.convert_to_tensor(serialized),
            "names": names,
            "dense_defaults": dense_defaults,
            "dense_keys": ["a"],
            "dense_types": [tf.float32],
            "dense_shapes": dense_shapes,
        },
        expected_err_re="Name: failing, Key: a.  Number of float values")

  def testDenseDefaultNoShapeShouldFail(self):
    original = [