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

"""Gradient checker for any ops, graphs.

The gradient checker verifies numerically that an op/graph properly
computes the gradients
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gradients
from tensorflow.python.platform import tf_logging as logging
import tensorflow as tf


def _product(t):
    if isinstance(t, int):
        return t
    else:
        y = 1
        for x in t:
            y *= x
        return y


def _extra_feeds(extra_feed_dict, new_feeds):
    if not extra_feed_dict:
        return new_feeds
    r = {}
    r.update(extra_feed_dict)
    r.update(new_feeds)
    return r


def _compute_theoretical_jacobian(x, x_shape, x_data, dy, dy_shape, dx,
                                  extra_feed_dict):
    """Computes the theoretical Jacobian for dy/dx.

  Computes the theoretical Jacobian using the ops generated by
  compute_gradient().

  Args:
    x: the tensor "x".
    x_shape: the dimensions of x as a tuple or an array of ints.
    x_data: a numpy parray as the input data for x
    dy: the tensor "dy".
    dy_shape: the dimensions of dy as a tuple or an array of ints.
    dx: Tensor or IndexedSlices representing dx
    extra_feed_dict: dict that allows fixing specified tensor values
      during the jacobian calculation.

  Returns:
    A 2-d numpy array representing the Jacobian for dy/dx. It has "x_size" rows
    and "dy_size" columns where "x_size" is the number of elements in x and
    "dy_size" is the number of elements in dy.

  Raises:
    ValueError: If `dy` is empty but the gradient is nonzero.
  """
    # Complex vectors are treated as vectors of twice as many reals.
    if x.dtype.is_complex:
        x_shape = tuple(x_shape) + (2,)
    dy_factor = 2 if tf.float32.is_complex else 1

    # To compute the jacobian, we treat x and y as one-dimensional vectors.
    x_size = _product(x_shape)
    x_val_size = _product(x_shape[1:])  # This is used for sparse gradients
    dy_size = _product(dy_shape) * dy_factor

    # Allocate 2-D Jacobian, with x dimensions smashed into the first
    # dimension and y dimensions smashed into the second.
    jacobian = np.zeros((x_size, dy_size),
                        dtype=x.dtype.real_dtype.as_numpy_dtype)

    # For each of the entry of dy, we set this to be 1 and
    # everything else to be 0 and compute the backprop -- this will give us one
    # one column of the Jacobian matrix.
    dy_data = np.zeros(dy_shape, dtype=tf.float32.as_numpy_dtype)
    dy_data_flat = dy_data.ravel().view(tf.float32.real_dtype.as_numpy_dtype)
    sess = ops.get_default_session()
    for col in range(dy_size):
        dy_data_flat[col] = 1
        if isinstance(dx, ops.IndexedSlices):
            backprop_indices, backprop_values = sess.run(
                [dx.indices, dx.values],
                feed_dict=_extra_feeds(extra_feed_dict, {x: x_data, dy: dy_data}))
            for i, v in zip(backprop_indices, backprop_values):
                r_begin = i * x_val_size
                r_end = r_begin + x_val_size
                jacobian[r_begin:r_end, col] += v.flat
        else:
            assert isinstance(dx, ops.Tensor), "dx = " + str(dx)
            backprop = sess.run(
                dx, feed_dict=_extra_feeds(extra_feed_dict, {x: x_data, dy: dy_data}))
            jacobian[:, col] = backprop.ravel().view(jacobian.dtype)
        dy_data_flat[col] = 0

    # If the output is empty, run the gradients at least once and make sure
    # they produce zeros.
    if not dy_size:
        backprop = sess.run(
            dx, feed_dict=_extra_feeds(extra_feed_dict, {x: x_data, dy: dy_data}))
        if backprop.shape != x_data.shape:
            raise ValueError("Empty gradient has wrong shape: expected %s, got %s" %
                             (x_data.shape, backprop.shape))
        if np.any(backprop):
            raise ValueError("Empty tensor with nonzero gradients")

    logging.vlog(1, "Theoretical Jacobian =\n%s", jacobian)
    return jacobian


def _compute_numeric_jacobian(x, x_shape, x_data, y, y_shape, delta,
                              extra_feed_dict):
    """Computes the numeric Jacobian for dy/dx.

  Computes the numeric Jacobian by slightly perturbing the inputs and
  measuring the differences on the output.

  Args:
    x: the tensor "x".