# Copyright 2019 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 tridiagonal solve ops."""

import numpy as np

from tensorflow.compiler.tests import xla_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gradients as gradient_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.linalg import linalg_impl
from tensorflow.python.platform import test

_sample_diags = np.array([[2, 1, 4, 0], [1, 3, 2, 2], [0, 1, -1, 1]],
                         dtype=np.float32)
_sample_rhs = np.array([1, 2, 3, 4], dtype=np.float32)
_sample_result = np.array([-9, 5, -4, 4], dtype=np.float32)


def _tfconst(array):
  return constant_op.constant(array, dtype=dtypes.float32)


def _tf_ones(shape):
  return array_ops.ones(shape, dtype=dtypes.float64)


class TridiagonalSolveOpsTest(xla_test.XLATestCase):
  """Test for tri-diagonal matrix related ops."""

  def testTridiagonalSolverSolves1Rhs(self):
    np.random.seed(19)

    batch_size = 8
    num_dims = 11

    diagonals_np = np.random.normal(size=(batch_size, 3,
                                          num_dims)).astype(np.float32)
    rhs_np = np.random.normal(size=(batch_size, num_dims, 1)).astype(np.float32)

    with self.session() as sess, self.test_scope():
      diags = array_ops.placeholder(
          shape=(batch_size, 3, num_dims), dtype=dtypes.float32)
      rhs = array_ops.placeholder(
          shape=(batch_size, num_dims, 1), dtype=dtypes.float32)
      x_np = sess.run(
          linalg_impl.tridiagonal_solve(diags, rhs, partial_pivoting=False),
          feed_dict={
              diags: diagonals_np,
              rhs: rhs_np
          })[:, :, 0]

    superdiag_np = diagonals_np[:, 0]
    diag_np = diagonals_np[:, 1]
    subdiag_np = diagonals_np[:, 2]

    y = np.zeros((batch_size, num_dims), dtype=np.float32)

    for i in range(num_dims):
      if i == 0:
        y[:, i] = (
            diag_np[:, i] * x_np[:, i] + superdiag_np[:, i] * x_np[:, i + 1])
      elif i == num_dims - 1:
        y[:, i] = (
            subdiag_np[:, i] * x_np[:, i - 1] + diag_np[:, i] * x_np[:, i])
      else:
        y[:, i] = (
            subdiag_np[:, i] * x_np[:, i - 1] + diag_np[:, i] * x_np[:, i] +
            superdiag_np[:, i] * x_np[:, i + 1])

    self.assertAllClose(y, rhs_np[:, :, 0], rtol=1e-4, atol=1e-4)

  def testTridiagonalSolverSolvesKRhs(self):
    np.random.seed(19)

    batch_size = 8
    num_dims = 11
    num_rhs = 5

    diagonals_np = np.random.normal(size=(batch_size, 3,
                                          num_dims)).astype(np.float32)
    rhs_np = np.random.normal(size=(batch_size, num_dims,
                                    num_rhs)).astype(np.float32)

    with self.session() as sess, self.test_scope():
      diags = array_ops.placeholder(
          shape=(batch_size, 3, num_dims), dtype=dtypes.float32)
      rhs = array_ops.placeholder(
          shape=(batch_size, num_dims, num_rhs), dtype=dtypes.float32)
      x_np = sess.run(
          linalg_impl.tridiagonal_solve(diags, rhs, partial_pivoting=False),
          feed_dict={
              diags: diagonals_np,
              rhs: rhs_np
          })

    superdiag_np = diagonals_np[:, 0]
    diag_np = diagonals_np[:, 1]
    subdiag_np = diagonals_np[:, 2]

    for eq in range(num_rhs):
      y = np.zeros((batch_size, num_dims), dtype=np.float32)
      for i in range(num_dims):
        if i == 0:
          y[:, i] = (
              diag_np[:, i] * x_np[:, i, eq] +
              superdiag_np[:, i] * x_np[:, i + 1, eq])
        elif i == num_dims - 1:
          y[:, i] = (
              subdiag_np[:, i] * x_np[:, i - 1, eq] +
              diag_np[:, i] * x_np[:, i, eq])
        else:
          y[:, i] = (
              subdiag_np[:, i] * x_np[:, i - 1, eq] +
              diag_np[:, i] * x_np[:, i, eq] +
              superdiag_np[:, i] * x_np[:, i + 1, eq])

      self.assertAllClose(y, rhs_np[:, :, eq], rtol=1e-4, atol=1e-4)

  # All the following is adapted from tridiagonal_solve_op_test.py
  def _test(self,
            diags,
            rhs,
            expected,
            diags_format="compact",
            transpose_rhs=False):
    with self.session() as sess, self.test_scope():
      self.assertAllClose(
          sess.run(
              linalg_impl.tridiagonal_solve(
                  _tfconst(diags),
                  _tfconst(rhs),
                  diags_format,
                  transpose_rhs,
                  conjugate_rhs=False,
                  partial_pivoting=False)),
          np.asarray(expected, dtype=np.float32))

  def _testWithDiagonalLists(self,
                             diags,
                             rhs,
                             expected,
                             diags_format="compact",
                             transpose_rhs=False):
    with self.session() as sess, self.test_scope():
      self.assertAllClose(
          sess.run(
              linalg_impl.tridiagonal_solve([_tfconst(x) for x in diags],
                                            _tfconst(rhs),
                                            diags_format,
                                            transpose_rhs,
                                            conjugate_rhs=False,
                                            partial_pivoting=False)),
          sess.run(_tfconst(expected)))

  def testReal(self):
    self._test(diags=_sample_diags, rhs=_sample_rhs, expected=_sample_result)

  # testComplex is skipped as complex type is not yet supported.

  def test3x3(self):
    self._test(
        diags=[[2.0, -1.0, 0.0], [1.0, 3.0, 1.0], [0.0, -1.0, -2.0]],
        rhs=[1.0, 2.0, 3.0],
        expected=[-3.0, 2.0, 7.0])

  def test2x2(self):
    self._test(
        diags=[[2.0, 0.0], [1.0, 3.0], [0.0, 1.0]],
        rhs=[1.0, 4.0],
        expected=[-5.0, 3.0])

  def test1x1(self):
    self._test(diags=[[0], [3], [0]], rhs=[6], expected=[2])

  def test0x0(self):
    self._test(
        diags=np.zeros(shape=(3, 0), dtype=np.float32),
        rhs=np.zeros(shape=(0, 1), dtype=np.float32),
        expected=np.zeros(shape=(0, 1), dtype=np.float32))

  def test2x2WithMultipleRhs(self):
    self._test(
        diags=[[2, 0], [1, 3], [0, 1]],
        rhs=[[1, 2, 3], [4, 8, 12]],
        expected=[[-5, -10, -15], [3, 6, 9]])

  def test1x1WithMultipleRhs(self):
    self._test(diags=[[0], [3], [0]], rhs=[[6, 9, 12]], expected=[[2, 3, 4]])

  # test1x1NotInvertible is skipped as runtime error not raised for now.

  # test2x2NotInvertible is skipped as runtime error not raised for now.

  @test_util.disable_mlir_bridge("Error messages differ")
  def testPartialPivotingRaises(self):
    np.random.seed(0)
    batch_size = 8
    num_dims = 11
    num_rhs = 5

    diagonals_np = np.random.normal(size=(batch_size, 3,
                                          num_dims)).astype(np.float32)
    rhs_np = np.random.normal(size=(batch_size, num_dims,
                                    num_rhs)).astype(np.float32)

    with self.session() as sess, self.test_scope():
      with self.assertRaisesRegex(
          errors_impl.UnimplementedError,
          "Current implementation does not yet support pivoting."):
        diags = array_ops.placeholder(
            shape=(batch_size, 3, num_dims), dtype=dtypes.float32)
        rhs = array_ops.placeholder(
            shape=(batch_size, num_dims, num_rhs), dtype=dtypes.float32)
        sess.run(
            linalg_impl.tridiagonal_solve(diags, rhs, partial_pivoting=True),
            feed_dict={
                diags: diagonals_np,
                rhs: rhs_np
            })

  # testCaseRequiringPivotingLastRows is skipped as pivoting is not supported
  # for now.

  # testNotInvertible is skipped as runtime error not raised for now.

  def testDiagonal(self):
    self._test(
        diags=[[0, 0, 0, 0], [1, 2, -1, -2], [0, 0, 0, 0]],
        rhs=[1, 2, 3, 4],
        expected=[1, 1, -3, -2])

  def testUpperTriangular(self):
    self._test(
        diags=[[2, 4, -1, 0], [1, 3, 1, 2], [0, 0, 0, 0]],
        rhs=[1, 6, 4, 4],
        expected=[13, -6, 6, 2])

  def testLowerTriangular(self):
    self._test(
        diags=[[0, 0, 0, 0], [2, -1, 3, 1], [0, 1, 4, 2]],
        rhs=[4, 5, 6, 1],
        expected=[2, -3, 6, -11])

  def testWithTwoRightHandSides(self):
    self._test(
        diags=_sample_diags,
        rhs=np.transpose([_sample_rhs, 2 * _sample_rhs]),
        expected=np.transpose([_sample_result, 2 * _sample_result]))

  def testBatching(self):
    self._test(
        diags=np.array([_sample_diags, -_sample_diags]),
        rhs=np.array([_sample_rhs, 2 * _sample_rhs]),
        expected=np.array([_sample_result, -2 * _sample_result]))

  def testWithTwoBatchingDimensions(self):
    self._test(
        diags=np.array([[_sample_diags, -_sample_diags, _sample_diags],
                        [-_sample_diags, _sample_diags, -_sample_diags]]),
        rhs=np.array([[_sample_rhs, 2 * _sample_rhs, 3 * _sample_rhs],
                      [4 * _sample_rhs, 5 * _sample_rhs, 6 * _sample_rhs]]),
        expected=np.array(
            [[_sample_result, -2 * _sample_result, 3 * _sample_result],
             [-4 * _sample_result, 5 * _sample_result, -6 * _sample_result]]))

  def testBatchingAndTwoRightHandSides(self):
    rhs = np.transpose([_sample_rhs, 2 * _sample_rhs])
    expected_result = np.transpose([_sample_result, 2 * _sample_result])
    self._test(
        diags=np.array([_sample_diags, -_sample_diags]),
        rhs=np.array([rhs, 2 * rhs]),
        expected=np.array([expected_result, -2 * expected_result]))

  def testSequenceFormat(self):
    self._testWithDiagonalLists(
        diags=[[2, 1, 4], [1, 3, 2, 2], [1, -1, 1]],
        rhs=[1, 2, 3, 4],
        expected=[-9, 5, -4, 4],
        diags_format="sequence")

  def testSequenceFormatWithDummyElements(self):
    dummy = 20  # Should be ignored by the solver.
    self._testWithDiagonalLists(
        diags=[
            [2, 1, 4, dummy],
            [1, 3, 2, 2],
            [dummy, 1, -1, 1],
        ],
        rhs=[1, 2, 3, 4],
        expected=[-9, 5, -4, 4],
        diags_format="sequence")

  def testSequenceFormatWithBatching(self):
    self._testWithDiagonalLists(
        diags=[[[2, 1, 4], [-2, -1, -4]], [[1, 3, 2, 2], [-1, -3, -2, -2]],
               [[1, -1, 1], [-1, 1, -1]]],
        rhs=[[1, 2, 3, 4], [1, 2, 3, 4]],
        expected=[[-9, 5, -4, 4], [9, -5, 4, -4]],
        diags_format="sequence")

  def testMatrixFormat(self):
    self._test(
        diags=[[1, 2, 0, 0], [1, 3, 1, 0], [0, -1, 2, 4], [0, 0, 1, 2]],
        rhs=[1, 2, 3, 4],
        expected=[-9, 5, -4, 4],
        diags_format="matrix")

  def testMatrixFormatWithMultipleRightHandSides(self):
    self._test(
        diags=[[1, 2, 0, 0], [1, 3, 1, 0], [0, -1, 2, 4], [0, 0, 1, 2]],
        rhs=[[1, -1], [2, -2], [3, -3], [4, -4]],
        expected=[[-9, 9], [5, -5], [-4, 4], [4, -4]],
        diags_format="matrix")

  def testMatrixFormatWithBatching(self):
    self._test(
        diags=[[[1, 2, 0, 0], [1, 3, 1, 0], [0, -1, 2, 4], [0, 0, 1, 2]],
               [[-1, -2, 0, 0], [-1, -3, -1, 0], [0, 1, -2, -4], [0, 0, -1,
                                                                  -2]]],
        rhs=[[1, 2, 3, 4], [1, 2, 3, 4]],
        expected=[[-9, 5, -4, 4], [9, -5, 4, -4]],
        diags_format="matrix")

  def testRightHandSideAsColumn(self):
    self._test(
        diags=_sample_diags,
        rhs=np.transpose([_sample_rhs]),
        expected=np.transpose([_sample_result]),
        diags_format="compact")

  def testTransposeRhs(self):
    self._test(
        diags=_sample_diags,
        rhs=np.array([_sample_rhs, 2 * _sample_rhs]),
        expected=np.array([_sample_result, 2 * _sample_result]).T,
        transpose_rhs=True)

  # testConjugateRhs is skipped as complex type is not yet supported.

  # testAjointRhs is skipped as complex type is not yet supported

  def testTransposeRhsWithRhsAsVector(self):
    self._test(
        diags=_sample_diags,
        rhs=_sample_rhs,
        expected=_sample_result,
        transpose_rhs=True)

  # testConjugateRhsWithRhsAsVector is skipped as complex type is not yet
  # supported.

  def testTransposeRhsWithRhsAsVectorAndBatching(self):
    self._test(
        diags=np.array([_sample_diags, -_sample_diags]),
        rhs=np.array([_sample_rhs, 2 * _sample_rhs]),
        expected=np.array([_sample_result, -2 * _sample_result]),
        transpose_rhs=True)

  # Gradient tests

  def _gradientTest(
      self,
      diags,
      rhs,
      y,  # output = reduce_sum(y * tridiag_solve(diags, rhs))
      expected_grad_diags,  # expected gradient of output w.r.t. diags
      expected_grad_rhs,  # expected gradient of output w.r.t. rhs
      diags_format="compact",
      transpose_rhs=False,
      feed_dict=None):
    expected_grad_diags = np.array(expected_grad_diags).astype(np.float32)
    expected_grad_rhs = np.array(expected_grad_rhs).astype(np.float32)
    with self.session() as sess, self.test_scope():
      diags = _tfconst(diags)
      rhs = _tfconst(rhs)
      y = _tfconst(y)

      x = linalg_impl.tridiagonal_solve(
          diags,
          rhs,
          diagonals_format=diags_format,
          transpose_rhs=transpose_rhs,
          conjugate_rhs=False,
          partial_pivoting=False)

      res = math_ops.reduce_sum(x * y)
      actual_grad_diags = sess.run(
          gradient_ops.gradients(res, diags)[0], feed_dict=feed_dict)
      actual_rhs_diags = sess.run(
          gradient_ops.gradients(res, rhs)[0], feed_dict=feed_dict)
    self.assertAllClose(expected_grad_diags, actual_grad_diags)
    self.assertAllClose(expected_grad_rhs, actual_rhs_diags)

  def testGradientSimple(self):
    self._gradientTest(
        diags=_sample_diags,
        rhs=_sample_rhs,
        y=[1, 3, 2, 4],
        expected_grad_diags=[[-5, 0, 4, 0], [9, 0, -4, -16], [0, 0, 5, 16]],
        expected_grad_rhs=[1, 0, -1, 4])

  def testGradientWithMultipleRhs(self):
    self._gradientTest(
        diags=_sample_diags,
        rhs=[[1, 2], [2, 4], [3, 6], [4, 8]],
        y=[[1, 5], [2, 6], [3, 7], [4, 8]],
        expected_grad_diags=[[-20, 28, -60, 0], [36, -35, 60, 80],
                             [0, 63, -75, -80]],
        expected_grad_rhs=[[0, 2], [1, 3], [1, 7], [0, -10]])

  def _makeDataForGradientWithBatching(self):
    y = np.array([1, 3, 2, 4]).astype(np.float32)
    grad_diags = np.array([[-5, 0, 4, 0], [9, 0, -4, -16],
                           [0, 0, 5, 16]]).astype(np.float32)
    grad_rhs = np.array([1, 0, -1, 4]).astype(np.float32)

    diags_batched = np.array(
        [[_sample_diags, 2 * _sample_diags, 3 * _sample_diags],
         [4 * _sample_diags, 5 * _sample_diags,
          6 * _sample_diags]]).astype(np.float32)
    rhs_batched = np.array([[_sample_rhs, -_sample_rhs, _sample_rhs],
                            [-_sample_rhs, _sample_rhs,
                             -_sample_rhs]]).astype(np.float32)
    y_batched = np.array([[y, y, y], [y, y, y]]).astype(np.float32)
    expected_grad_diags_batched = np.array(
        [[grad_diags, -grad_diags / 4, grad_diags / 9],
         [-grad_diags / 16, grad_diags / 25,
          -grad_diags / 36]]).astype(np.float32)
    expected_grad_rhs_batched = np.array(
        [[grad_rhs, grad_rhs / 2, grad_rhs / 3],
         [grad_rhs / 4, grad_rhs / 5, grad_rhs / 6]]).astype(np.float32)

    return (y_batched, diags_batched, rhs_batched, expected_grad_diags_batched,
            expected_grad_rhs_batched)

  def testGradientWithBatchDims(self):
    y, diags, rhs, expected_grad_diags, expected_grad_rhs = (
        self._makeDataForGradientWithBatching())

    self._gradientTest(
        diags=diags,
        rhs=rhs,
        y=y,
        expected_grad_diags=expected_grad_diags,
        expected_grad_rhs=expected_grad_rhs)

  # testGradientWithUnknownShapes is skipped as shapes should be fully known.

  def _assertRaises(self, diags, rhs, diags_format="compact"):
    with self.assertRaises(ValueError):
      linalg_impl.tridiagonal_solve(diags, rhs, diags_format)

  # Invalid input shapes
  def testInvalidShapesCompactFormat(self):

    def test_raises(diags_shape, rhs_shape):
      self._assertRaises(_tf_ones(diags_shape), _tf_ones(rhs_shape), "compact")

    test_raises((5, 4, 4), (5, 4))
    test_raises((5, 3, 4), (4, 5))
    test_raises((5, 3, 4), (5))
    test_raises((5), (5, 4))

  def testInvalidShapesSequenceFormat(self):

    def test_raises(diags_tuple_shapes, rhs_shape):
      diagonals = tuple(_tf_ones(shape) for shape in diags_tuple_shapes)
      self._assertRaises(diagonals, _tf_ones(rhs_shape), "sequence")

    test_raises(((5, 4), (5, 4)), (5, 4))
    test_raises(((5, 4), (5, 4), (5, 6)), (5, 4))
    test_raises(((5, 3), (5, 4), (5, 6)), (5, 4))
    test_raises(((5, 6), (5, 4), (5, 3)), (5, 4))
    test_raises(((5, 4), (7, 4), (5, 4)), (5, 4))
    test_raises(((5, 4), (7, 4), (5, 4)), (3, 4))

  def testInvalidShapesMatrixFormat(self):

    def test_raises(diags_shape, rhs_shape):
      self._assertRaises(_tf_ones(diags_shape), _tf_ones(rhs_shape), "matrix")

    test_raises((5, 4, 7), (5, 4))
    test_raises((5, 4, 4), (3, 4))
    test_raises((5, 4, 4), (5, 3))

  # Tests involving placeholder with an unknown dimension are all skipped.
  # Dimensions have to be all known statically.


if __name__ == "__main__":
  test.main()