# Copyright 2016 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.
# ==============================================================================

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

import numpy as np
import tensorflow as tf

from tensorflow.contrib.solvers.python.ops import util


class UtilTest(tf.test.TestCase):

  def _testCreateOperator(self, use_static_shape_):
    a_np = np.array([[1., 2.], [3., 4.], [5., 6.]])
    x = np.array([[2.], [-3.]])
    y = np.array([[2], [-3.], [5.]])
    with self.test_session() as sess:
      if use_static_shape_:
        a = tf.constant(a_np, dtype=tf.float32)
      else:
        a = tf.placeholder(tf.float32)
      op = util.create_operator(a)
      ax = op.apply(x)
      aty = op.apply_adjoint(y)
      op_shape = tf.convert_to_tensor(op.shape)
      if use_static_shape_:
        op_shape_val, ax_val, aty_val = sess.run([op_shape, ax, aty])
      else:
        op_shape_val, ax_val, aty_val = sess.run([op_shape, ax, aty],
                                                 feed_dict={a: a_np})
      self.assertAllEqual(op_shape_val, [3, 2])
      self.assertAllClose(ax_val, [[-4], [-6], [-8]])
      self.assertAllClose(aty_val, [[18], [22]])

  def testCreateOperator(self):
    self._testCreateOperator(True)

  def testCreateOperatorUnknownShape(self):
    self._testCreateOperator(False)

  def testL2Norm(self):
    with self.test_session():
      x_np = np.array([[2], [-3.], [5.]])
      x_norm_np = np.linalg.norm(x_np)
      x_normalized_np = x_np / x_norm_np
      x = tf.constant(x_np)
      l2norm = util.l2norm(x)
      l2norm_squared = util.l2norm_squared(x)
      x_normalized, x_norm = util.l2normalize(x)
      self.assertAllClose(l2norm.eval(), x_norm_np)
      self.assertAllClose(l2norm_squared.eval(), np.square(x_norm_np))
      self.assertAllClose(x_norm.eval(), x_norm_np)
      self.assertAllClose(x_normalized.eval(), x_normalized_np)


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