# 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 tensorflow.ops.tensor_array_ops."""

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

import numpy as np

from tensorflow.core.protobuf import config_pb2
from tensorflow.python.client import session as session_lib
from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
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 tensor_shape
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 control_flow_util
from tensorflow.python.ops import data_flow_ops
from tensorflow.python.ops import gen_data_flow_ops
from tensorflow.python.ops import gradients_impl
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import tensor_array_grad
from tensorflow.python.ops import tensor_array_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
import tensorflow.python.ops.nn_grad  # pylint: disable=unused-import
from tensorflow.python.platform import test


def _make_converter(tf_dtype):
  def _converter(x):
    if tf_dtype == dtypes.string:
      # In Python3, np.str is unicode, while we always want bytes
      return np.asarray(x).astype("|S")
    x = np.asarray(x).astype(tf_dtype.as_numpy_dtype)
    if tf_dtype.is_complex:
      # Add a non-zero imaginary component to x.
      x -= 1j * x
    return x
  return _converter


def _make_ta(size, name, dtype=dtypes.float32, infer_shape=False):
  return tensor_array_ops.TensorArray(
      dtype=dtype, tensor_array_name=name, size=size, infer_shape=infer_shape)


@test_util.run_all_in_graph_and_eager_modes
@test_util.with_control_flow_v2
class TensorArrayTest(test.TestCase):

  @classmethod
  def setUpClass(cls):
    super(TensorArrayTest, cls).setUpClass()
    cls._workers, _ = test.create_local_cluster(num_workers=3, num_ps=0)

  @classmethod
  def tearDownClass(cls):
    super(TensorArrayTest, cls).tearDownClass()
    session_lib.Session.reset(cls._workers[0].target)

  @test_util.run_in_graph_and_eager_modes
  def testTensorArrayWriteRead(self):
    with self.session(use_gpu=True):
      ta = tensor_array_ops.TensorArray(
          dtype=dtypes.float32,
          tensor_array_name="foo",
          size=3,
          infer_shape=False)

      w0 = ta.write(0, [[4.0, 5.0]])
      w1 = w0.write(1, [[1.0]])
      w2 = w1.write(2, -3.0)

      r0 = w2.read(0)
      r1 = w2.read(1)
      r2 = w2.read(2)

      d0, d1, d2 = self.evaluate([r0, r1, r2])
      self.assertAllEqual([[4.0, 5.0]], d0)
      self.assertAllEqual([[1.0]], d1)
      self.assertAllEqual(-3.0, d2)

  def _testTensorArrayWritePack(self, tf_dtype):
    with self.cached_session(use_gpu=True):
      ta = tensor_array_ops.TensorArray(
          dtype=tf_dtype, tensor_array_name="foo", size=3)

      convert = _make_converter(tf_dtype)

      w0 = ta.write(0, convert([[4.0, 5.0]]))
      w1 = w0.write(1, convert([[6.0, 7.0]]))
      w2 = w1.write(2, convert([[8.0, 9.0]]))

      c0 = w2.stack()

      c0 = self.evaluate(c0)
      self.assertAllEqual(
          convert([[[4.0, 5.0]], [[6.0, 7.0]], [[8.0, 9.0]]]), c0)

  def _testTensorArrayWritePackMaybeLegacy(self):
    self._testTensorArrayWritePack(dtypes.float32)
    self._testTensorArrayWritePack(dtypes.float64)
    self._testTensorArrayWritePack(dtypes.int32)
    self._testTensorArrayWritePack(dtypes.int64)
    self._testTensorArrayWritePack(dtypes.complex64)
    self._testTensorArrayWritePack(dtypes.complex128)
    self._testTensorArrayWritePack(dtypes.string)

  def testTensorArrayWritePack(self):
    self._testTensorArrayWritePackMaybeLegacy()

  def testEmptyTensorArrayPack(self):
    with self.session(use_gpu=True):
      ta = tensor_array_ops.TensorArray(
          dtype=dtypes.float32, tensor_array_name="foo", size=3)

      empty_element = np.zeros((0, 1), dtype=np.float32)
      w0 = ta.write(0, empty_element)
      w1 = w0.write(1, empty_element)
      w2 = w1.write(2, empty_element)

      c0 = w2.stack()

      c0 = self.evaluate(c0)
      self.assertAllEqual([3, 0, 1], c0.shape)

  def _testTensorArrayWriteConcat(self, tf_dtype):
    with self.cached_session(use_gpu=True):
      ta = tensor_array_ops.TensorArray(
          dtype=tf_dtype, tensor_array_name="foo", size=3, infer_shape=False)

      convert = _make_converter(tf_dtype)

      w0 = ta.write(0, convert([[4.0, 5.0], [104.0, 105.0], [204.0, 205.0]]))
      w1 = w0.write(1, convert([[6.0, 7.0], [106.0, 107.0]]))
      w2 = w1.write(2, convert([[8.0, 9.0]]))

      c0 = w2.concat()

      c0 = self.evaluate(c0)
      self.assertAllEqual(
          convert([[4.0, 5.0], [104.0, 105.0], [204.0, 205.0], [6.0, 7.0],
                   [106.0, 107.0], [8.0, 9.0]]), c0)

  @test_util.run_deprecated_v1
  def testTensorArrayWriteConcat(self):
    self._testTensorArrayWriteConcat(dtypes.float32)
    self._testTensorArrayWriteConcat(dtypes.float64)
    self._testTensorArrayWriteConcat(dtypes.int32)
    self._testTensorArrayWriteConcat(dtypes.int64)
    self._testTensorArrayWriteConcat(dtypes.complex64)
    self._testTensorArrayWriteConcat(dtypes.complex128)
    self._testTensorArrayWriteConcat(dtypes.string)

  def _testTensorArrayReadOrPackNotAllValuesAvailableFillsZeros(self):
    with self.cached_session(use_gpu=True):
      ta = tensor_array_ops.TensorArray(
          dtype=dtypes.float32,
          tensor_array_name="foo",
          size=3,
          element_shape=tensor_shape.TensorShape([1, 2]))
      self.assertAllEqual([[0.0, 0.0]], self.evaluate(ta.read(0)))
      self.assertAllEqual([[[0.0, 0.0]], [[4.0, 5.0]], [[0.0, 0.0]]],
                          self.evaluate(ta.write(1, [[4.0, 5.0]]).stack()))
      self.assertAllEqual([[0.0, 0.0], [4.0, 5.0], [0.0, 0.0]],
                          self.evaluate(ta.write(1, [[4.0, 5.0]]).concat()))

  @test_util.run_v1_only("b/122324791")
  def testTensorArrayReadOrPackNotAllValuesAvailableFillsZeros(self):
    self._testTensorArrayReadOrPackNotAllValuesAvailableFillsZeros()

  def _testTensorArrayReadOrPackNotAllValuesAvailableInferShapeFillsZeros(self):
    ta = tensor_array_ops.TensorArray(
        dtype=dtypes.float32,
        tensor_array_name="foo",
        size=3)
    self.assertAllEqual(
        [[0.0, 0.0]], self.evaluate(ta.write(1, [[4.0, 5.0]]).read(0)))
    self.assertAllEqual([[[0.0, 0.0]], [[4.0, 5.0]], [[0.0, 0.0]]],
                        self.evaluate(ta.write(1, [[4.0, 5.0]]).stack()))
    self.assertAllEqual([[0.0, 0.0], [4.0, 5.0], [0.0, 0.0]],
                        self.evaluate(ta.write(1, [[4.0, 5.0]]).concat()))

  @test_util.run_v1_only("b/122324791")
  def testTensorArrayReadOrPackNotAllValuesAvailableInferShapeFillsZeros(self):
    self._testTensorArrayReadOrPackNotAllValuesAvailableInferShapeFillsZeros()

  @test_util.run_v1_only("Uses placeholders")
  def testSkipEagerTensorArrayReadUninitializedInferShapeFillsZeros(self):
    with self.cached_session(use_gpu=True) as sess:
      ta = tensor_array_ops.TensorArray(
          dtype=dtypes.float32,
          tensor_array_name="foo",
          size=3)
      val = array_ops.placeholder(dtypes.float32)
      self.assertAllEqual(
          [[0.0, 0.0]], sess.run(ta.write(1, val).read(0), {val: [[4.0, 5.0]]}))

  def _testTensorArrayUnpackRead(self, tf_dtype):
    with self.cached_session(use_gpu=True):
      convert = _make_converter(tf_dtype)

      ta = _make_ta(3, "foo", dtype=tf_dtype)
      # Unpack a vector into scalars
      w0 = ta.unstack(convert([1.0, 2.0, 3.0]))
      r0 = w0.read(0)
      r1 = w0.read(1)
      r2 = w0.read(2)

      d0, d1, d2 = self.evaluate([r0, r1, r2])
      self.assertAllEqual(convert(1.0), d0)
      self.assertAllEqual(convert(2.0), d1)
      self.assertAllEqual(convert(3.0), d2)

      # Unpack a matrix into vectors
      w1 = ta.unstack(convert([[1.0, 1.1], [2.0, 2.1], [3.0, 3.1]]))
      r0 = w1.read(0)
      r1 = w1.read(1)
      r2 = w1.read(2)

      d0, d1, d2 = self.evaluate([r0, r1, r2])
      self.assertAllEqual(convert([1.0, 1.1]), d0)
      self.assertAllEqual(convert([2.0, 2.1]), d1)
      self.assertAllEqual(convert([3.0, 3.1]), d2)

      # Try unpacking an empty matrix, which should not cause an error.
      w2 = ta.unstack(convert([[], [], []]))
      r0 = w2.read(0)
      r1 = w2.read(1)
      r2 = w2.read(2)

      d0, d1, d2 = self.evaluate([r0, r1, r2])
      self.assertAllEqual(convert([]), d0)
      self.assertAllEqual(convert([]), d1)
      self.assertAllEqual(convert([]), d2)

  def _testTensorArrayUnpackReadMaybeLegacy(self):
    self._testTensorArrayUnpackRead(dtypes.float32)
    self._testTensorArrayUnpackRead(dtypes.float64)
    self._testTensorArrayUnpackRead(dtypes.int32)
    self._testTensorArrayUnpackRead(dtypes.int64)
    self._testTensorArrayUnpackRead(dtypes.complex64)
    self._testTensorArrayUnpackRead(dtypes.complex128)
    self._testTensorArrayUnpackRead(dtypes.string)

  def testTensorArrayUnpackRead(self):
    self._testTensorArrayUnpackReadMaybeLegacy()

  def _testTensorArraySplitRead(self, tf_dtype):
    with self.cached_session(use_gpu=True):
      convert = _make_converter(tf_dtype)

      # Split an empty vector
      ta = _make_ta(3, "foo", dtype=tf_dtype)
      lengths = constant_op.constant([0, 0, 0])
      w0 = ta.split(convert([]), lengths=lengths)
      r0 = w0.read(0)
      r1 = w0.read(1)
      r2 = w0.read(2)

      d0, d1, d2 = self.evaluate([r0, r1, r2])
      self.assertAllEqual(convert([]), d0)
      self.assertAllEqual(convert([]), d1)
      self.assertAllEqual(convert([]), d2)

      # Split a vector
      lengths = constant_op.constant([2, 0, 1])
      w0 = ta.split(convert([1.0, 2.0, 3.0]), lengths=lengths)
      r0 = w0.read(0)
      r1 = w0.read(1)
      r2 = w0.read(2)

      d0, d1, d2 = self.evaluate([r0, r1, r2])
      self.assertAllEqual(convert([1.0, 2.0]), d0)
      self.assertAllEqual(convert([]), d1)
      self.assertAllEqual(convert([3.0]), d2)

      # Split a matrix
      lengths = constant_op.constant([2, 0, 1])
      w0 = ta.split(
          convert([[1.0, 101.0], [2.0, 201.0], [3.0, 301.0]]), lengths=lengths)
      r0 = w0.read(0)
      r1 = w0.read(1)
      r2 = w0.read(2)

      d0, d1, d2 = self.evaluate([r0, r1, r2])
      self.assertAllEqual(convert([[1.0, 101.0], [2.0, 201.0]]), d0)
      self.assertAllEqual(convert([]).reshape(0, 2), d1)
      self.assertAllEqual(convert([[3.0, 301.0]]), d2)

  @test_util.run_deprecated_v1
  def testTensorArraySplitRead(self):
    self._testTensorArraySplitRead(dtypes.float32)
    self._testTensorArraySplitRead(dtypes.float64)
    self._testTensorArraySplitRead(dtypes.int32)
    self._testTensorArraySplitRead(dtypes.int64)
    self._testTensorArraySplitRead(dtypes.complex64)
    self._testTensorArraySplitRead(dtypes.complex128)
    self._testTensorArraySplitRead(dtypes.string)

  @test_util.disable_control_flow_v2("v2 does not support TensorArray.grad.")
  @test_util.run_v1_only("v2 does not support TensorArray.grad.")
  def testSkipEagerTensorGradArrayWriteRead(self):
    with self.session(use_gpu=True) as session:
      ta = tensor_array_ops.TensorArray(
          dtype=dtypes.float32,
          tensor_array_name="foo",
          size=3,
          infer_shape=False)
      g_ta = ta.grad("grad")

      w0 = ta.write(0, [[4.0, 5.0]])
      w1 = w0.write(1, [[1.0]])
      w2 = w1.write(2, -3.0)

      g_w0 = g_ta.write(0, [[5.0, 6.0]])
      g_w1 = g_w0.write(1, [[2.0]])
      g_w2 = g_w1.write(2, -2.0)

      r0 = w2.read(0)
      r1 = w2.read(1)
      r2 = w2.read(2)

      g_r0 = g_w2.read(0)
      g_r1 = g_w2.read(1)
      g_r2 = g_w2.read(2)

      d0, d1, d2, g_d0, g_d1, g_d2 = session.run([r0, r1, r2, g_r0, g_r1, g_r2])
      self.assertAllEqual([[4.0, 5.0]], d0)
      self.assertAllEqual([[1.0]], d1)
      self.assertAllEqual(-3.0, d2)
      self.assertAllEqual([[5.0, 6.0]], g_d0)
      self.assertAllEqual([[2.0]], g_d1)
      self.assertAllEqual(-2.0, g_d2)

  @test_util.run_deprecated_v1
  def testSkipEagerTensorArrayGradGrad(self):
    if not control_flow_util.ENABLE_CONTROL_FLOW_V2:
      self.skipTest("Legacy TensorArray does not support double derivatives.")
    with self.test_session(use_gpu=True) as session:
      x = constant_op.constant(4.0)

      ta = tensor_array_ops.TensorArray(
          dtype=dtypes.float32,
          tensor_array_name="foo",
          size=1,
          infer_shape=False)
      w0 = ta.write(0, x)
      r0 = w0.read(0)
      y = r0 * r0

      g1 = gradients_impl.gradients(ys=[y], xs=[x])
      g2 = gradients_impl.gradients(ys=[g1], xs=[x])
      self.assertAllEqual([2.0], session.run(g2))

  @test_util.disable_control_flow_v2("v2 does not support TensorArray.grad.")
  @test_util.run_v1_only("v2 does not support TensorArray.grad.")
  def testSkipEagerTensorGradArrayDynamicWriteRead(self):
    with self.session(use_gpu=True) as session:
      ta = tensor_array_ops.TensorArray(
          dtype=dtypes.float32,
          tensor_array_name="foo",
          size=0,
          dynamic_size=True,
          infer_shape=False)

      w0 = ta.write(0, [[4.0, 5.0]])
      w1 = w0.write(1, [[1.0]])
      w2 = w1.write(2, -3.0)

      g_ta = w2.grad("grad")  # Get gradient array here so we know the shape

      s = w2.size()
      g_s = g_ta.size()

      g_w0 = g_ta.write(0, [[5.0, 6.0]])
      g_w1 = g_w0.write(1, [[2.0]])
      g_w2 = g_w1.write(2, -2.0)

      r0 = w2.read(0)
      r1 = w2.read(1)
      r2 = w2.read(2)

      g_r0 = g_w2.read(0)
      g_r1 = g_w2.read(1)
      g_r2 = g_w2.read(2)

      d0, d1, d2, g_d0, g_d1, g_d2, vs, g_vs = session.run(
          [r0, r1, r2, g_r0, g_r1, g_r2, s, g_s])
      self.assertAllEqual([[4.0, 5.0]], d0)
      self.assertAllEqual([[1.0]], d1)
      self.assertAllEqual(-3.0, d2)
      self.assertAllEqual([[5.0, 6.0]], g_d0)
      self.assertAllEqual([[2.0]], g_d1)
      self.assertAllEqual(-2.0, g_d2)
      self.assertAllEqual(3, vs)
      self.assertAllEqual(3, g_vs)

  @test_util.disable_control_flow_v2("v2 does not support TensorArray.grad.")
  @test_util.run_v1_only("v2 does not support TensorArray.grad.")
  def testSkipEagerTensorGradAccessTwiceReceiveSameObject(self):
    with self.session(use_gpu=True) as session:
      ta = tensor_array_ops.TensorArray(
          dtype=dtypes.float32, tensor_array_name="foo", size=3)
      g_ta_0 = ta.grad("grad")
      g_ta_1 = ta.grad("grad")

      with ops.control_dependencies([g_ta_0.write(0, [[4.0, 5.0]]).flow]):
        # Write with one gradient handle, read with another copy of it
        r1_0 = g_ta_1.read(0)

      t_g_ta_0, t_g_ta_1, d_r1_0 = session.run(
          [g_ta_0.handle.op, g_ta_1.handle.op, r1_0])
      self.assertAllEqual(t_g_ta_0, t_g_ta_1)
      self.assertAllEqual([[4.0, 5.0]], d_r1_0)

  def testTensorArrayWriteWrongIndexOrDataTypeFails(self):
    with self.session(use_gpu=True):
      ta = _make_ta(3, "foo", dtype=dtypes.float32)
      # Test writing the wrong datatype
      if (control_flow_util.ENABLE_CONTROL_FLOW_V2 and
          not context.executing_eagerly()):
        error_msg = ("Invalid data types; op elements string but list elements "
                     "float")
      else:
        error_msg = (
            "TensorArray dtype is (float|float32) but Op is trying to write "
            "dtype string")
      with self.assertRaisesOpError(error_msg):
        self.evaluate(ta.write(0, "wrong_type_scalar").flow)

      if (control_flow_util.ENABLE_CONTROL_FLOW_V2 and
          not context.executing_eagerly()):
        error_msg = "Trying to modify element -1 in a list with 3 elements."
      else:
        error_msg = "index -1"
      with self.assertRaisesOpError(error_msg):
        self.evaluate(ta.write(-1, 3.0).flow)

      if (control_flow_util.ENABLE_CONTROL_FLOW_V2 and
          not context.executing_eagerly()):
        error_msg = "Trying to modify element 3 in a list with 3 elements"
      else:
        error_msg = ("Tried to write to index 3 but array is not "
                     "resizeable and size is: 3")
      # Test reading from too large an index
      with self.assertRaisesOpError(error_msg):
        self.evaluate(ta.write(3, 3.0).flow)

  def testTensorArrayReadWrongIndexOrDataTypeFails(self):
    with self.session(use_gpu=True):
      ta = _make_ta(3, "foo", dtype=dtypes.float32)

      w0 = ta.write(0, [[4.0, 5.0]])

      # Test reading wrong datatype (only possible when constructing graphs).
      if (not context.executing_eagerly() and
          not control_flow_util.ENABLE_CONTROL_FLOW_V2):
        r0_bad = gen_data_flow_ops.tensor_array_read_v3(
            handle=w0.handle, index=0, dtype=dtypes.float64, flow_in=w0.flow)
        with self.assertRaisesOpError(
            "TensorArray dtype is float but Op requested dtype double."):
          self.evaluate(r0_bad)

      if (control_flow_util.ENABLE_CONTROL_FLOW_V2 and
          not context.executing_eagerly()):
        error_msg = "Trying to access element -1 in a list with 3 elements."
      else:
        error_msg = "index -1"
      # Test reading from a negative index, which is not allowed
      with self.assertRaisesOpError(error_msg):
        self.evaluate(ta.read(-1))

      if (control_flow_util.ENABLE_CONTROL_FLOW_V2 and
          not context.executing_eagerly()):
        error_msg = "Trying to access element 3 in a list with 3 elements."
      else:
        error_msg = "Tried to read from index 3 but array size is: 3"
      # Test reading from too large an index
      with self.assertRaisesOpError(error_msg):
        self.evaluate(ta.read(3))

  @test_util.disable_control_flow_v2("v2 allows multiple writes.")
  @test_util.run_v1_only("v2 allows multiple writes.")
  def testSkipEagerTensorArrayWriteMultipleFails(self):
    with self.session(use_gpu=True):
      ta = tensor_array_ops.TensorArray(
          dtype=dtypes.float32, tensor_array_name="foo", size=3)

      with self.assertRaisesOpError(
          "Could not write to TensorArray index 2 because "
          "it has already been written to."):
        self.evaluate(ta.write(2, 3.0).write(2, 3.0).flow)

  def testTensorArrayConcatIncompatibleShapesFails(self):