# 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.
# ==============================================================================
"""Utilities for testing 
`LinearOperator` and sub-classes."""

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

import abc
import itertools
import numpy as np
import six

from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from ten
sorflow.python.framework import random_seed
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_util
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import linalg_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops.linalg import linalg_impl as linalg
from tensorflow.python.ops.linalg import linear_operator_util
from tensorflow.python.platform import test


class OperatorShapesInfo(object):
  """Object encoding expected shape for a test.

  Encodes the expected shape of a matrix for a test. Also
  allows additional metadata for the test harness.
  """

  def __init__(self, shape, **kwargs):
    self.shape = shape
    self.__dict__.update(kwargs)


class CheckTapeSafeSkipOptions(object):

  # Skip checking this particular method.
  DETERMINANT = "determinant"
  DIAG_PART = "diag_part"
  LOG_ABS_DETERMINANT = "log_abs_determinant"
  TRACE = "trace"


@six.add_metaclass(abc.ABCMeta)  # pylint: disable=no-init
class LinearOperatorDerivedClassTest(test.TestCase):
  """Tests for derived classes.

  Subclasses should implement every abstractmethod, and this will enable all
  test methods to work.
  """

  # Absolute/relative tolerance for tests.
  _atol = {
      dtypes.float16: 1e-3,
      dtypes.float32: 1e-6,
      dtypes.float64: 1e-12,
      dtypes.complex64: 1e-6,
      dtypes.complex128: 1e-12
  }

  _rtol = {
      dtypes.float16: 1e-3,
      dtypes.float32: 1e-6,
      dtypes.float64: 1e-12,
      dtypes.complex64: 1e-6,
      dtypes.complex128: 1e-12
  }

  def assertAC(self, x, y):
    """Derived classes can set _atol, _rtol to get different tolerance."""
    dtype = dtypes.as_dtype(x.dtype)
    atol = self._atol[dtype]
    rtol = self._rtol[dtype]
    self.assertAllClose(x, y, atol=atol, rtol=rtol)

  @staticmethod
  def adjoint_options():
    return [False, True]

  @staticmethod
  def adjoint_arg_options():
    return [False, True]

  @staticmethod
  def dtypes_to_test():
    # TODO(langmore) Test tf.float16 once tf.linalg.solve works in 16bit.
    return [dtypes.float32, dtypes.float64, dtypes.complex64, dtypes.complex128]

  @staticmethod
  def use_placeholder_options():
    return [False, True]

  @staticmethod
  def operator_shapes_infos():
    """Returns list of OperatorShapesInfo, encapsulating the shape to test."""
    raise NotImplementedError("operator_shapes_infos has not been implemented.")

  @abc.abstractmethod
  def operator_and_matrix(
      self, shapes_info, dtype, use_placeholder,
      ensure_self_adjoint_and_pd=False):
    """Build a batch matrix and an Operator that should have similar behavior.

    Every operator acts like a (batch) matrix.  This method returns both
    together, and is used by tests.

    Args:
      shapes_info: `OperatorShapesInfo`, encoding shape information about the
        operator.
      dtype:  Numpy dtype.  Data type of returned array/operator.
      use_placeholder:  Python bool.  If True, initialize the operator with a
        placeholder of undefined shape and correct dtype.
      ensure_self_adjoint_and_pd: If `True`,
        construct this operator to be Hermitian Positive Definite, as well
        as ensuring the hints `is_positive_definite` and `is_self_adjoint`
        are set.
        This is useful for testing methods such as `cholesky`.

    Returns:
      operator:  `LinearOperator` subclass instance.
      mat:  `Tensor` representing operator.
    """
    # Create a matrix as a numpy array with desired shape/dtype.
    # Create a LinearOperator that should have the same behavior as the matrix.
    raise NotImplementedError("Not implemented yet.")

  @abc.abstractmethod
  def make_rhs(self, operator, adjoint, with_batch=True):
    """Make a rhs appropriate for calling operator.solve(rhs).

    Args:
      operator:  A `LinearOperator`
      adjoint:  Python `bool`.  If `True`, we are making a 'rhs' value for the
        adjoint operator.
      with_batch: Python `bool`. If `True`, create `rhs` with the same batch
        shape as operator, and otherwise create a matrix without any batch
        shape.

    Returns:
      A `Tensor`
    """
    raise NotImplementedError("make_rhs is not defined.")

  @abc.abstractmethod
  def make_x(self, operator, adjoint, with_batch=True):
    """Make an 'x' appropriate for calling operator.matmul(x).

    Args:
      operator:  A `LinearOperator`
      adjoint:  Python `bool`.  If `True`, we are making an 'x' value for the
        adjoint operator.
      with_batch: Python `bool`. If `True`, create `x` with the same batch shape
        as operator, and otherwise create a matrix without any batch shape.

    Returns:
      A `Tensor`
    """
    raise NotImplementedError("make_x is not defined.")

  @staticmethod
  def skip_these_tests():
    """List of test names to skip."""
    # Subclasses should over-ride if they want to skip some tests.
    # To skip "test_foo", add "foo" to this list.
    return []

  def assertRaisesError(self, msg):
    """assertRaisesRegexp or OpError, depending on context.executing_eagerly."""
    if context.executing_eagerly():
      return self.assertRaisesRegexp(Exception, msg)
    return self.assertRaisesOpError(msg)

  def check_tape_safe(self, operator, skip_options=None):
    """Check gradients are not None w.r.t. operator.variables.

    Meant to be called from the derived class.

    This ensures grads are not w.r.t every variable in operator.variables.  If
    more fine-grained testing is needed, a custom test should be written.

    Args:
      operator: LinearOperator.  Exact checks done will depend on hints.
      skip_options: Optional list of CheckTapeSafeSkipOptions.
        Makes this test skip particular checks.
    """
    skip_options = skip_options or []

    if not operator.variables:
      raise AssertionError("`operator.variables` was empty")

    def _assert_not_none(iterable):
      for item in iterable:
        self.assertIsNotNone(item)

    # Tape tests that can be run on every operator below.
    with backprop.GradientTape() as tape:
      _assert_not_none(tape.gradient(operator.to_dense(), operator.variables))

    with backprop.GradientTape() as tape:
      _assert_not_none(
          tape.gradient(operator.adjoint().to_dense(), operator.variables))

    x = math_ops.cast(
        array_ops.ones(shape=operator.H.shape_tensor()[:-1]), operator.dtype)

    with backprop.GradientTape() as tape:
      _assert_not_none(tape.gradient(operator.matvec(x), operator.variables))

    # Tests for square, but possibly non-singular operators below.
    if not operator.is_square:
      return

    for option in [
        CheckTapeSafeSkipOptions.DETERMINANT,
        CheckTapeSafeSkipOptions.LOG_ABS_DETERMINANT,
        CheckTapeSafeSkipOptions.DIAG_PART,
        CheckTapeSafeSkipOptions.TRACE,
    ]:
      with backprop.GradientTape() as tape:
        if option not in skip_options:
          _assert_not_none(