# 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.python.framework.ops."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tensorflow.python.framework import device as pydev
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import test_ops
from tensorflow.python.framework import test_util
from tensorflow.python.framework import versions
from tensorflow.python.ops import common_shapes
from tensorflow.python.ops import constant_op
# Import gradients to register _IndexedSlicesToTensor.
import tensorflow.python.ops.gradients  # pylint: disable=unused-import
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
from tensorflow.python.platform import googletest


class TensorTest(test_util.TensorFlowTestCase):

  def testShape(self):
    op = ops.Operation(ops._NodeDef("noop", "myop"), ops.Graph(),
                       [], [dtypes.float32])
    t = op.outputs[0]
    self.assertEqual(tensor_shape.unknown_shape(), t.get_shape())
    t.set_shape([1, 2, 3])
    self.assertEqual([1, 2, 3], t.get_shape())

  def testIterable(self):
    op = ops.Operation(
        ops._NodeDef("noop", "myop"), ops.Graph(), [], [dtypes.float32])
    t = op.outputs[0]
    self.assertTrue(isinstance(t, ops.Tensor))
    with self.assertRaisesRegexp(TypeError, "not iterable"):
      for _ in t:
        pass


class SparseTensorTest(test_util.TensorFlowTestCase):

  def testPythonConstruction(self):
    sp = ops.SparseTensor([[1, 2], [2, 0], [3, 4]], ["a", "b", "c"], [4, 5])
    self.assertEqual(sp.indices.dtype, dtypes.int64)
    self.assertEqual(sp.values.dtype, dtypes.string)
    self.assertEqual(sp.shape.dtype, dtypes.int64)


class IndexedSlicesTest(test_util.TensorFlowTestCase):

  def testToTensor(self):
    with self.test_session():
      values = constant_op.constant([2, 3, 5, 7], shape=[2, 2])
      indices = constant_op.constant([0, 2])
      dense_shape = constant_op.constant([3, 2])
      x = ops.IndexedSlices(values, indices, dense_shape)
      tensor = ops.convert_to_tensor(x, name="tensor")
      self.assertAllEqual(tensor.eval(), [[2, 3], [0, 0], [5, 7]])

  def testNegation(self):
    with self.test_session():
      values = constant_op.constant([2, 3, 5, 7], shape=[2, 2])
      indices = constant_op.constant([0, 2])
      x = -ops.IndexedSlices(values, indices)
      self.assertAllEqual(x.values.eval(), [[-2, -3], [-5, -7]])
      self.assertAllEqual(x.indices.eval(), [0, 2])

  def testScalarMul(self):
    with self.test_session():
      values = constant_op.constant([2, 3, 5, 7], shape=[2, 2])
      indices = constant_op.constant([0, 2])
      x = math_ops.scalar_mul(-2, ops.IndexedSlices(values, indices))
      self.assertAllEqual(x.values.eval(), [[-4, -6], [-10, -14]])
      self.assertAllEqual(x.indices.eval(), [0, 2])


class NodeDefConstructorTest(test_util.TensorFlowTestCase):

  def testNoArgs(self):
    nodedef = ops._NodeDef("noop", "bar")
    self.assertProtoEquals("op: 'noop' name: 'bar'", nodedef)

  def testArgs(self):
    nodedef = ops._NodeDef("foo", "bar", device="/device:baz:*")
    self.assertProtoEquals("op:'foo' name:'bar' device:'/device:baz:*'",
                           nodedef)
    nodedef = ops._NodeDef("foo", "bar", device=pydev.Device(job="j"))
    self.assertProtoEquals("op:'foo' name:'bar' device:'/job:j'", nodedef)


# NOTE(mrry): Dummy shape registrations for ops used in the tests.
ops.RegisterShape("a")(None)
ops.RegisterShape("b")(None)
ops.RegisterShape("c")(None)
ops.RegisterShape("add")(None)
ops.RegisterShape("an_op")(None)
ops.RegisterShape("const")(None)
ops.RegisterShape("copy")(None)
ops.RegisterShape("foo")(None)
ops.RegisterShape("identity")(None)
ops.RegisterShape("mul")(None)
ops.RegisterShape("nonrefop")(None)
ops.RegisterShape("noop")(None)
ops.RegisterShape("refop")(None)


def _apply_op(g, *args, **kwargs):
  op = g.create_op(*args, **kwargs)
  if len(op.outputs) == 1:
    return op.outputs[0]
  else:
    return op.outputs


class OperationTest(test_util.TensorFlowTestCase):

  def testNoInputs(self):
    op = ops.Operation(ops._NodeDef("noop", "myop"), ops.Graph(),
                       [],
                       [dtypes.float32, dtypes.string])
    self.assertEqual(2, len(op.values()))
    self.assertEqual(0, len(op.inputs))
    self.assertEqual("myop", op.name)

    float_t, label_str_t = op.values()
    self.assertEqual(dtypes.float32, float_t.dtype)
    self.assertEqual(op, float_t.op)
    self.assertEqual(0, float_t._value_index)
    self.assertEqual(0, len(float_t._consumers))
    self.assertEqual("myop", float_t._as_node_def_input())

    self.assertEqual(dtypes.string, label_str_t.dtype)
    self.assertEqual(op, label_str_t.op)
    self.assertEqual(1, label_str_t._value_index)
    self.assertEqual(0, len(label_str_t._consumers))
    self.assertEqual("myop:1", label_str_t._as_node_def_input())

    self.assertProtoEquals("op:'noop' name:'myop'", op.node_def)

  def testNoOutputs(self):
    g = ops.Graph()
    op1 = ops.Operation(
        ops._NodeDef("noop", "myop1"), g, [], [dtypes.float32])
    float_t, = op1.values()
    op2 = ops.Operation(ops._NodeDef("reop", "myop2"), g, [float_t], [])
    self.assertEqual(0, len(op2.values()))
    self.assertEqual(1, len(op2.inputs))
    self.assertIs(float_t, op2.inputs[0])

    self.assertEqual(1, len(float_t._consumers))
    self.assertEqual(op2, float_t._consumers[0])

    self.assertProtoEquals("op:'noop' name:'myop1'", op1.node_def)
    self.assertProtoEquals("op:'reop' name:'myop2' input:'myop1'",
                           op2.node_def)

  def testInputsAndOutputs(self):
    g = ops.Graph()
    op1 = ops.Operation(
        ops._NodeDef("noop", "myop1"), g, [], [dtypes.float32])
    self.assertEqual(1, len(op1.values()))
    float1_t, = op1.values()

    op2 = ops.Operation(ops._NodeDef("reop", "myop2"), g,
                        [], [dtypes.float32, dtypes.string])
    self.assertEqual(2, len(op2.values()))
    float2_t, label2_str_t = op2.values()

    # Note that we consume label2_str_t twice here.
    op3 = ops.Operation(ops._NodeDef("add", "myop3"), g,
                        [float1_t, label2_str_t, label2_str_t],
                        [dtypes.float32, dtypes.int32])
    self.assertEqual(2, len(op3.values()))

    self.assertEqual(1, len(float1_t._consumers))
    self.assertEqual(op3, float1_t._consumers[0])

    self.assertEqual(0, len(float2_t._consumers))

    self.assertEqual(2, len(label2_str_t._consumers))
    self.assertEqual(op3, label2_str_t._consumers[0])
    self.assertEqual(op3, label2_str_t._consumers[1])

    self.assertProtoEquals("""
    op:'add' name:'myop3'
    input:'myop1' input:'myop2:1' input:'myop2:1'
    """, op3.node_def)

  def testDeviceObject(self):
    op = ops.Operation(ops._NodeDef("noop", "myop"), ops.Graph(), [], [])
    op._set_device("/job:goo/device:GPU:0")
    self.assertProtoEquals(
        "op:'noop' name:'myop' device:'/job:goo/device:GPU:0' ",
        op.node_def)
    op = ops.Operation(ops._NodeDef("noop", "op2"), ops.Graph(), [], [])
    op._set_device(pydev.Device(job="muu", device_type="CPU", device_index=0))
    self.assertProtoEquals(
        "op:'noop' name:'op2' device:'/job:muu/device:CPU:0'",
        op.node_def)

  def testReferenceInput(self):
    g = ops.Graph()
    op1 = ops.Operation(ops._NodeDef("noop", "op1"), g, [],
                        [dtypes.float32_ref, dtypes.float32])
    self.assertProtoEquals("op:'noop' name:'op1'",
                           op1.node_def)
    ref_t, nonref_t = op1.values()
    # NOTE(mrry): Must specify input_types to preserve ref-typed input.
    op2 = ops.Operation(
        ops._NodeDef("refop", "op2"), g, [ref_t, nonref_t], [],
        input_types=[dtypes.float32_ref, dtypes.float32])
    self.assertProtoEquals("op:'refop' name:'op2' input:'op1' input:'op1:1'",
                           op2.node_def)
    op3 = ops.Operation(
        ops._NodeDef("nonrefop", "op3"), g, [ref_t, nonref_t], [])
    self.assertProtoEquals("op:'nonrefop' name:'op3' input:'op1' input:'op1:1'",
                           op3.node_def)

  def testInvalidNames(self):
    g = ops.Graph()
    with self.assertRaises(ValueError):
      ops.Operation(ops._NodeDef("op", ""), g)
    with self.assertRaises(ValueError):
      ops.Operation(ops._NodeDef("op", "_invalid"), g)
    with self.assertRaises(ValueError):
      ops.Operation(ops._NodeDef("op", "-invalid"), g)
    with self.assertRaises(ValueError):
      ops.Operation(ops._NodeDef("op", "/invalid"), g)

  def testShapeFunctionAbsence(self):
    def _test():
      pass
    g = ops.Graph()
    with self.assertRaises(RuntimeError):
      g.create_op("shapeless_op", [], [dtypes.float32])

  def testNoShapeFunction(self):
    g = ops.Graph()
    op = ops.Operation(ops._NodeDef("op", "an_op"), g,
                       output_types = [dtypes.float32])
    self.assertEqual(tensor_shape.unknown_shape(),
                     _apply_op(g, "an_op", [], [dtypes.float32]).get_shape())

class CreateOpTest(test_util.TensorFlowTestCase):

  def testNodeDefArgs(self):
    g = ops.Graph()
    op1 = g.create_op("const", [], [dtypes.float32], None, name="myop1")
    with g.device("/device:GPU:0"):
      op2 = g.create_op("add",
                        [],
                        [dtypes.float32, dtypes.string], None,
                        name="myop2")
    op3 = g.create_op("foo",
                      [list(op1.values())[0], list(op2.values())[1],
                       list(op2.values())[0]],
                      [dtypes.float32, dtypes.int32],
                      None,
                      name="myop3")
    self.assertDeviceEqual(None, op1.device)
    self.assertDeviceEqual("/device:GPU:0", op2.device)
    self.assertDeviceEqual(None, op3.device)
    self.assertProtoEquals("name:'myop1' op:'const'", op1.node_def)
    self.assertProtoEquals("name:'myop2' op:'add' device:'/device:GPU:0'",
                           op2.node_def)
    self.assertProtoEquals(
        "name:'myop3' input:'myop1' input:'myop2:1' input:'myop2' op:'foo'",
        op3.node_def)

  def testReferenceInput(self):
    g = ops.Graph()
    op1 = g.create_op("noop", [],
                      [dtypes.float32_ref, dtypes.float32], name="op1")
    self.assertProtoEquals("op:'noop' name:'op1'", op1.node_def)
    ref_t, nonref_t = op1.values()
    # NOTE(mrry): Must specify input_types to preserve ref-typed input.
    op2 = g.create_op("refop", [ref_t, nonref_t], [],
                      input_types=[dtypes.float32_ref, dtypes.float32],
                      name="op2")
    self.assertProtoEquals("op:'refop' name:'op2' input:'op1' input:'op1:1'",
                           op2.node_def)
    op3 = g.create_op("nonrefop", [ref_t, nonref_t], [], name="op3")
    self.assertProtoEquals("op:'nonrefop' name:'op3' input:'op1' input:'op1:1'",
                           op3.node_def)

  def testFinalized(self):
    g = ops.Graph()
    g.finalize()
    with self.assertRaises(RuntimeError):
      g.create_op("const", [], [dtypes.float32], None, name="myop1")


class ApplyOpTest(test_util.TensorFlowTestCase):

  def testNodeDefArgs(self):
    g = ops.Graph()
    t1 = _apply_op(g, "const", [], [dtypes.float32], name="myop1")
    with g.device("/device:GPU:0"):
      t2 = _apply_op(g, "add",
                     [],
                     [dtypes.float32, dtypes.string],
                     name="myop2")
    t3 = _apply_op(g, "foo", [t1, t2[1], t2[0]],
                   [dtypes.float32, dtypes.int32], name="myop3")
    self.assertTrue(isinstance(t1, ops.Tensor))
    self.assertTrue(isinstance(t2, list))
    self.assertTrue(isinstance(t3, list))
    self.assertTrue(isinstance(t3[0], ops.Tensor))
    self.assertEqual("myop1", t1._as_node_def_input())
    self.assertEqual("myop2", t2[0]._as_node_def_input())
    self.assertEqual("myop2:1", t2[1]._as_node_def_input())
    self.assertEqual("myop3", t3[0]._as_node_def_input())
    # Validate that we got the right ops as well
    self.assertProtoEquals("name:'myop1' op:'const'", t1.op.node_def)
    self.assertProtoEquals("name:'myop2' op:'add' device:'/device:GPU:0'",
                           t2[0].op.node_def)
    self.assertProtoEquals(
        "name:'myop3' input:'myop1' input:'myop2:1' input:'myop2' op:'foo'",
        t3[0].op.node_def)

  def testReferenceInput(self):
    g = ops.Graph()
    ref_t, nonref_t = _apply_op(
        g, "noop", [], [dtypes.float32_ref, dtypes.float32], name="op1")
    self.assertProtoEquals("op:'noop' name:'op1'", ref_t.op.node_def)
    # NOTE(mrry): Must specify input_types to preserve ref-typed input.
    out_2 = _apply_op(g, "refop", [ref_t, nonref_t], [dtypes.int32],
                      input_types=[dtypes.float32_ref, dtypes.float32],
                      name="op2")
    self.assertProtoEquals("op:'refop' name:'op2' input:'op1' input:'op1:1'",
                           out_2.op.node_def)
    out_3 = _apply_op(g, "nonrefop", [ref_t, nonref_t], [dtypes.int32],
                      name="op3")
    self.assertProtoEquals("op:'nonrefop' name:'op3' input:'op1' input:'op1:1'",
                           out_3.op.node_def)


class NameStackTest(test_util.TensorFlowTestCase):

  def testBasics(self):
    g = ops.Graph()
    self.assertEqual("foo", g.unique_name("foo", mark_as_used=False))
    self.assertEqual("foo", g.unique_name("foo", mark_as_used=False))