# 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.
# ==============================================================================
"""Gradients for operators defined in linalg_ops.py.

Useful reference for derivative formulas is (Mike Giles, 2008).

Ionescu et al. (2015) provide a detailed derivation of formulas for
backpropagating through spectral layers (SVD and Eig).

References:
  An extended collection of matrix derivative results for
  forward and reverse mode automatic differentiation:
    [Mike Giles, 2008]
    (https://ora.ox.ac.uk/objects/uuid:8d0c0a29-c92b-4153-a1d2-38b276e93124)
    ([pdf](http://eprints.maths.ox.ac.uk/1079/1/NA-08-01.pdf))
  Matrix Backpropagation for Deep Networks with Structured Layers
    [Ionescu et al., 2015]
    (https://www.cv-foundation.org/openaccess/content_iccv_2015/html/Ionescu_Matrix_Backpropagation_for_ICCV_2015_paper.html)
    ([pdf](https://www.cv-foundation.org/openaccess/content_iccv_2015/papers/Ionescu_Matrix_Backpropagation_for_ICCV_2015_paper.pdf))
  Training Deep Networks with Structured Layers by Matrix Backpropagation:
    [Ionescu et al., 2015](https://arxiv.org/abs/1509.07838)
    ([pdf](https://arxiv.org/pdf/1509.07838.pdf))
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_linalg_ops
from tensorflow.python.ops import linalg_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.linalg import linalg_impl as _linalg


@ops.RegisterGradient("MatrixInverse")
def _MatrixInverseGrad(op, grad):
  """Gradient for MatrixInverse."""
  ainv = op.outputs[0]
  return -math_ops.matmul(  # pylint: disable=invalid-unary-operand-type
      ainv,
      math_ops.matmul(grad, ainv, adjoint_b=True),
      adjoint_a=True)


@ops.RegisterGradient("Einsum")
def _EinsumGrad(op, grad):
  """Gradient for Einsum."""
  ellipsis = "..."

  def _GetAxisFromLabel(subscripts, label):
    """Returns the axis (possibly negative) corresponding to a label.

    Returns the axis index of the axis label if it is before an ellipsis (or if
    the ellipsis is not present), and the negative index if it occurs after the
    ellipsis. E.g. index of `b` in `ab...cd`, is `1`, but that of `c` is `-2`.

    For multiple occurrences, returns the leftmost one. If not found, returns
    None.

    Args:
      subscripts: A string denoting the einsum subscript (e.g. `ab...cd`)
      label: The single character axis label.
    """
    splits = subscripts.split(ellipsis)
    index = splits[0].find(label)
    if index != -1:
      return index
    if len(splits) < 2:
      return None
    index = splits[1].find(label)
    if index != -1:
      return index - len(splits[1])
    return None

  def _GetBcastSubshape(subscripts):
    """Returns a tuple denoting the slice mapping to ellipsis.

    For a given subscript, returns a tuple (start, end) denoting the start
    axis index and the (negative) end axis index respectively. For any input
    Tensor `x` described by the subscript, `x[start:end]` would be the slice
    represented by the ellipsis. E.g. For `ab...cd` returns `[1, -2]`.

    If ellipsis is not present in `subscripts`, returns `(0, 0)`.

    Args:
      subscripts: A string denoting the einsum subscript.
    """
    start = subscripts.find(ellipsis)
    if start == -1:
      return 0, 0
    remaining = len(subscripts) - (start + len(ellipsis))
    end = -remaining if remaining > 0 else None
    return start, end

  def _GetReducedSubscripts(reduced_label_set, input_shape, subscripts):
    """Returns reduced subscripts and their corresponding dimensions and axes.

    Given a set of axis labels, returns their concatenated subscript, their
    corresponding dimensions from input_shape, and their corresponding axes.
    Note that the concatenated subscript `reduced_subs` may have axis labels
    from `reduced_label_set` in any order. For example, for the reduced label
    set `{b, d}`, subscripts `aabbcd` and input shape `[2,2,5,5,3,4]`, returns
    subscripts `bd`, dimensions `[5,4]` and axes `[2,5]`.

    Args:
      reduced_label_set: Set of axis labels which appear in `subscripts`.
      input_shape: A `Tensor` representing the shape of the einsum operand
        corresponding to `subscripts`.
      subscripts: A string denoting the einsum subscript.

    Returns:
      reduced_subs: Subscripts formed by a concatenation of labels in
        `reduced_label_set`.
      reduced_dims: Dimensions from `input_shape` corresponding to each label
        in `reduced_subs`.
      reduced_axes: Axes described by `subscripts` corresponding to each label
        in `reduced_subs`. If there are multiple occurrences in `subscripts`,
        we consider only the leftmost one.

    """
    # Concatenate the sequence of reduced axis labels.
    reduced_subs = "".join(list(reduced_label_set))
    # Get the axis (may be positive, negative or zero) for each of the reduced
    # labels. If the same label appears multiple times, get the left-most axis.
    reduced_axes = [_GetAxisFromLabel(subscripts, s) for s in reduced_subs]
    # Get the corresponding dimensions for each reduced axis.
    reduced_dims = array_ops.stack([input_shape[ax] for ax in reduced_axes])
    return reduced_subs, reduced_dims, reduced_axes

  def _GetGradReduced(output_grad, output_subs, input_subs, input_shape,
                      reduced_label_set):
    """Returns the gradient wrt input for a unary einsum with reductions.

    Args:
      output_grad: The gradient wrt the output of a unary einsum operation.
      output_subs: The output subscript. (E.g. `ac` for equation `abc->ac`).
      input_subs: The input subscript. (E.g. `abc` for equation `abc->ac`).
      input_shape: A `Tensor` representing the shape of the input operand.
      reduced_label_set: The set of axis labels appearing in `input_subs` but
        not in `output_subs`.
    """
    # Let's say the einsum operation was "aabbcd->ca", where axis labels 'b' and
    # 'd' are reduced with input_shape [2,2,5,5,3,4]. Then obtain the reduced
    # subscripts "bd", corresponding dimensions [5,4] and axes [2,5].
    reduced_subs, reduced_dims, reduced_axes = _GetReducedSubscripts(
        reduced_label_set, input_shape, input_subs)
    # Whether either the input or the output subscripts have a repeated label.
    # This is true for "aabbcd->ca" or "abd->cca" but false for "abcd->ca".
    has_repeated_labels = (
        len(set(input_subs)) + len(set(output_subs)) <
        len(input_subs) + len(output_subs))
    # Compute the input subscripts without the reduced axis labels, e.g. "aac"
    # for the equation "aabbcd->ca".
    input_subs_without_reduced_labels = "".join(
        [s for s in input_subs if s not in reduced_label_set])

    # The gradient wrt the input for the equation "abc->ac" (or, equivalently
    # reduce_sum(..., axis=1)) is just the gradient of the output tiled N times
    # along axis 1, where label 'b' represents a dimension of size N.
    #
    # If we're not dealing with repeated labels, and the non-reduced labels
    # doesn't need to be transposed, then just tiling is enough and there is no
    # need to call another einsum. For example, tiling is sufficient for
    # "abcd->ac". But for equations like "aabbcd->ac" (generalized traces) or
    # "abc->ca" (transpose), we'd need another einsum operation after tiling.
    if (not has_repeated_labels and
        input_subs_without_reduced_labels == output_subs):
      # Obtain the shape of the output, as if keepdims=True on reduce sum. E.g.
      # for the equation "abcd->ac" with input shape [2,5,3,4], we get the
      # reduced shape [2,1,3,1].
      reduced_shape = math_ops.reduced_shape(
          input_shape, ops.convert_to_tensor(reduced_axes))
      # Reshaping the gradient (wrt "ac") to [2,1,3,1] and broadcasting it to
      # the shape [2,5,3,4] results in the gradient wrt "abcd".
      return array_ops.broadcast_to(
          array_ops.reshape(output_grad, reduced_shape), input_shape)

    # If we *do* have traces or transpose operations, then prepend the extra
    # reduced dimensions to the front. E.g. Given the equation "aabbcd->ca" we'd
    # first obtain the VJP for "bdca->ca", and then the VJP for "aabbcd->bdca".
    #
    # Obtain the input shape with reduced dimensions prepended, viz. [5,4,3,2].
    # This is the shape of the intermediate "bdca".
    grad_shape_with_reduced_labels = array_ops.concat(
        [reduced_dims, array_ops.shape(output_grad)], axis=0)
    # Obtain the output shape of the reduction-only equation "bdca->ca" as if
    # keepdims=True; viz. [1,1,3,2]. Since we prepended the reduced labels, we
    # just have to prepend that many 1s to the output shape.
    reduced_shape = (
        array_ops.concat([
            array_ops.ones(len(reduced_label_set), dtype=dtypes.int32),
            array_ops.shape(output_grad)
        ],
                         axis=0))
    # Compute the VJP for the intermediate (viz. "bdca->ca") for which
    # broadcasting is sufficient.
    broadcasted_grad = array_ops.broadcast_to(
        array_ops.reshape(output_grad, reduced_shape),
        grad_shape_with_reduced_labels)
    # Compute the VJP for the final step (viz. "aabbcd->bdca"). We can use
    # einsum with the input and output subscripts reversed (viz. "bdca->aabbcd")
    # since the output axis labels now appear in the input subscripts.
    return gen_linalg_ops.einsum([broadcasted_grad],
                                 "{}->{}".format(reduced_subs + output_subs,
                                                 input_subs))

  def _GetGradWrt(output_grad, other_operand, input_shape, input_subs,
                  other_subs, output_subs):
    """Returns the gradient wrt an input operand for a binary einsum.

    This function does not handle (un)broadcasting. This must be done separately
    on the returned gradient.

    Args:
      output_grad: The gradient wrt the output of a binary einsum operation.
      other_operand: The complementary `Tensor` operand i.e. which is not the