# Copyright 2015 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.subscribe."""

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

import collections

from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import subscribe
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import script_ops
from tensorflow.python.ops import sparse_ops
from tensorflow.python.ops import tensor_array_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import googletest


class SubscribeTest(test_util.TensorFlowTestCase):

  def _ExpectSubscribedIdentities(self, container):
    """Convenience function to test a container of subscribed identities."""
    self.assertTrue(
        all(subscribe._is_subscribed_identity(x) for x in container))

  def testSideEffect(self):
    a = constant_op.constant(1)
    b = constant_op.constant(1)
    c = math_ops.add(a, b)
    with ops.control_dependencies([c]):
      d = constant_op.constant(42)
    n = math_ops.negative(c)

    shared = []

    def sub(t):
      shared.append(t)
      return t

    c0 = c
    self.assertTrue(c0.op in d.op.control_inputs)
    c = subscribe.subscribe(c,
                            lambda t: script_ops.py_func(sub, [t], [t.dtype]))
    # Verify that control dependencies are correctly moved to the subscription.
    self.assertFalse(c0.op in d.op.control_inputs)
    self.assertTrue(c.op in d.op.control_inputs)

    with self.test_session() as sess:
      c_out = sess.run([c])
      n_out = sess.run([n])
      d_out = sess.run([d])

    self.assertEqual(n_out, [-2])
    self.assertEqual(c_out, [2])
    self.assertEqual(d_out, [42])
    self.assertEqual(shared, [2, 2, 2])

  def testSupportedTypes(self):
    """Confirm that supported types are correctly detected and handled."""

    a = constant_op.constant(1)
    b = constant_op.constant(1)
    c = math_ops.add(a, b)

    def sub(t):
      return t

    # Tuples.
    subscribed = subscribe.subscribe(
        (a, b), lambda t: script_ops.py_func(sub, [t], [t.dtype]))
    self.assertIsInstance(subscribed, tuple)
    self._ExpectSubscribedIdentities(subscribed)

    # Lists.
    subscribed = subscribe.subscribe(
        [a, b], lambda t: script_ops.py_func(sub, [t], [t.dtype]))
    self.assertIsInstance(subscribed, list)
    self._ExpectSubscribedIdentities(subscribed)

    # Dictionaries.
    subscribed = subscribe.subscribe({
        'first': a,
        'second': b
    }, lambda t: script_ops.py_func(sub, [t], [t.dtype]))
    self.assertIsInstance(subscribed, dict)
    self._ExpectSubscribedIdentities(subscribed.values())

    # Namedtuples.
    # pylint: disable=invalid-name
    TensorPair = collections.namedtuple('TensorPair', ['first', 'second'])
    # pylint: enable=invalid-name
    pair = TensorPair(a, b)
    subscribed = subscribe.subscribe(
        pair, lambda t: script_ops.py_func(sub, [t], [t.dtype]))
    self.assertIsInstance(subscribed, TensorPair)
    self._ExpectSubscribedIdentities(subscribed)

    # Expect an exception to be raised for unsupported types.
    with self.assertRaisesRegexp(TypeError, 'has invalid type'):
      subscribe.subscribe(c.name,
                          lambda t: script_ops.py_func(sub, [t], [t.dtype]))

  def testCaching(self):
    """Confirm caching of control output is recalculated between calls."""
    a = constant_op.constant(1)
    b = constant_op.constant(2)
    with ops.control_dependencies([a]):
      c = constant_op.constant(42)

    shared = {}

    def sub(t):
      shared[t] = shared.get(t, 0) + 1