# Copyright 2018 The TensorFlow Probability Authors.
#
# 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.
# ============================================================================
"""The Broyden-Fletcher-Goldfarb-Shanno minimization algorithm.

Quasi Newton methods are a class of popular first order optimization algorithm.
These methods use a positive definite approximation to the exact
Hessian to find the search direction. The Broyden-Fletcher-Goldfarb-Shanno
algorithm (BFGS) is a specific implementation of this general idea.

This module provides an implementation of the BFGS scheme using the Hager Zhang
line search method.
"""

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

import collections

# Dependency imports
import tensorflow as tf

from tensorflow_probability.python.internal import distribution_util
from tensorflow_probability.python.internal import prefer_static
from tensorflow_probability.python.optimizer import bfgs_utils

BfgsOptimizerResults = collections.namedtuple(
    'BfgsOptimizerResults', [
        'converged',  # Scalar boolean tensor indicating whether the minimum
                      # was found within tolerance.
        'failed',  # Scalar boolean tensor indicating whether a line search
                   # step failed to find a suitable step size satisfying Wolfe
                   # conditions. In the absence of any constraints on the
                   # number of objective evaluations permitted, this value will
                   # be the complement of `converged`. However, if there is
                   # a constraint and the search stopped due to available
                   # evaluations being exhausted, both `failed` and `converged`
                   # will be simultaneously False.
        'num_iterations',  # The number of iterations of the BFGS update.
        'num_objective_evaluations',  # The total number of objective
                                      # evaluations performed.
        'position',  # A tensor containing the last argument value found
                     # during the search. If the search converged, then
                     # this value is the argmin of the objective function.
        'objective_value',  # A tensor containing the value of the objective
                            # function at the `position`. If the search
                            # converged, then this is the (local) minimum of
                            # the objective function.
        'objective_gradient',  # A tensor containing the gradient of the
                               # objective function at the
                               # `final_position`. If the search converged
                               # the max-norm of this tensor should be
                               # below the tolerance.
        'inverse_hessian_estimate'  # A tensor containing the inverse of the
                                    # estimated Hessian.
    ])

def minimize(value_and_gradients_function,
             initial_position,
             tolerance=1e-8,
             x_tolerance=0,
             f_relative_tolerance=0,
             initial_inverse_hessian_estimate=None,
             max_iterations=50,
             parallel_iterations=1,
             stopping_condition=None,
             name=None):
  """Applies the BFGS algorithm to minimize a differentiable function.

  Performs unconstrained minimization of a differentiable function using the
  BFGS scheme. For details of the algorithm, see [Nocedal and Wright(2006)][1].

  ### Usage:

  The following example demonstrates the BFGS optimizer attempting to find the
  minimum for a simple two dimensional quadratic objective function.

  ```python
    minimum = np.array([1.0, 1.0])  # The center of the quadratic bowl.
    scales = np.array([2.0, 3.0])  # The scales along the two axes.

    # The objective function and the gradient.
    def quadratic(x):
      value = tf.reduce_sum(scales * (x - minimum) ** 2)
      return value, tf.gradients(value, x)[0]

    start = tf.constant([0.6, 0.8])  # Starting point for the search.
    optim_results = tfp.optimizer.bfgs_minimize(
        quadratic, initial_position=start, tolerance=1e-8)

    with tf.Session() as session:
      results = session.run(optim_results)
      # Check that the search converged
      assert(results.converged)
      # Check that the argmin is close to the actual value.
      np.testing.assert_allclose(results.position, minimum)
      # Print out the total number of function evaluations it took. Should be 6.
      print ("Function evaluations: %d" % results.num_objective_evaluations)
  ```

  ### References:
  [1]: Jorge Nocedal, Stephen Wright. Numerical Optimization. Springer Series in
    Operations Research. pp 136-140. 2006
    http://pages.mtu.edu/~struther/Courses/OLD/Sp2013/5630/Jorge_Nocedal_Numerical_optimization_267490.pdf

  Args:
    value_and_gradients_function:  A Python callable that accepts a point as a
      real `Tensor` and returns a tuple of `Tensor`s of real dtype containing
      the value of the function and its gradient at that point. The function
      to be minimized. The input should be of shape `[..., n]`, where `n` is
      the size of the domain of input points, and all others are batching
      dimensions. The first component of the return value should be a real
      `Tensor` of matching shape `[...]`. The second component (the gradient)
      should also be of shape `[..., n]` like the input value to the function.
    initial_position: real `Tensor` of shape `[..., n]`. The starting point, or
      points when using batching dimensions, of the search procedure. At these
      points the function value and the gradient norm should be finite.
    tolerance: Scalar `Tensor` of real dtype. Specifies the gradient tolerance
      for the procedure. If the supremum norm of the gradient vector is below
      this number, the algorithm is stopped.
    x_tolerance: Scalar `Tensor` of real dtype. If the absolute change in the
      position between one iteration and the next is smaller than this number,
      the algorithm is stopped.
    f_relative_tolerance: Scalar `Tensor` of real dtype. If the relative change
      in the objective value between one iteration and the next is smaller
      than this value, the algorithm is stopped.
    initial_inverse_hessian_estimate: Optional `Tensor` of the same dtype
      as the components of the output of the `value_and_gradients_function`.
      If specified, the shape should broadcastable to shape `[..., n, n]`; e.g.
      if a single `[n, n]` matrix is provided, it will be automatically
      broadcasted to all batches. Alternatively, one can also specify a
      different hessian estimate for each batch member.
      For the correctness of the algorithm, it is required that this parameter
      be symmetric and positive definite. Specifies the starting estimate for
      the inverse of the Hessian at the initial point. If not specified,
      the identity matrix is used as the starting estimate for the
      inverse Hessian.
    max_iterations: Scalar positive int32 `Tensor`. The maximum number of
      iterations for BFGS updates.
    parallel_iterations: Positive integer. The number of iterations allowed to
      run in parallel.
    stopping_condition: (Optional) A Python function that takes as input two
      Boolean tensors of shape `[...]`, and returns a Boolean scalar tensor.
      The input tensors are `converged` and `failed`, indicating the current
      status of each respective batch member; the return value states whether
      the algorithm should stop. The default is tfp.optimizer.converged_all
      which only stops when all batch members have either converged or failed.
      An alternative is tfp.optimizer.converged_any which stops as soon as one
      batch member has converged, or when all have failed.
    name: (Optional) Python str. The name prefixed to the ops created by this
      function. If not supplied, the default name 'minimize' is used.

  Returns:
    optimizer_results: A namedtuple containing the following items:
      converged: boolean tensor of shape `[...]` indicating for each batch
        member whether the minimum was found within tolerance.
      failed:  boolean tensor of shape `[...]` indicating for each batch
        member whether a line search step failed to find a suitable step size
        satisfying Wolfe conditions. In the absence of any constraints on the
        number of objective evaluations permitted, this value will
        be the complement of `converged`. However, if there is
        a constraint and the search stopped due to available
        evaluations being exhausted, both `failed` and `converged`
        will be simultaneously False.
      num_objective_evaluations: The total number of objective
        evaluations performed.
      position: A tensor of shape `[..., n]` containing the last argument value
        found during the search from each starting point. If the search
        converged, then this value is the argmin of the objective function.
      objective_value: A tensor of shape `[...]` with the value of the
        objective function at the `position`. If the search converged, then
        this is the (local) minimum of the objective function.
      objective_gradient: A tensor of shape `[..., n]` containing the gradient
        of the objective function at the `position`. If the search converged
        the max-norm of this tensor should be below the tolerance.
      inverse_hessian_estimate: A tensor of shape `[..., n, n]` containing the
        inverse of the estimated Hessian.
  """
  with tf.compat.v1.name_scope(
      name, 'minimize',
      [initial_position, tolerance, initial_inverse_hessian_estimate]):
    initial_position = tf.convert_to_tensor(
        value=initial_position, name='initial_position')
    dtype = initial_position.dtype.base_dtype
    tolerance = tf.convert_to_tensor(
        value=tolerance, dtype=dtype, name='grad_tolerance')
    f_relative_tolerance = tf.convert_to_tensor(
        value=f_relative_tolerance, dtype=dtype, name='f_relative_tolerance')
    x_tolerance = tf.convert_to_tensor(
        value=x_tolerance, dtype=dtype, name='x_tolerance')
    max_iterations = tf.convert_to_tensor(
        value=max_iterations, name='max_iterations')

    input_shape = distribution_util.prefer_static_shape(initial_position)
    batch_shape, domain_size = input_shape[:-1], input_shape[-1]

    if stopping_condition is None:
      stopping_condition = bfgs_utils.converged_all

    # Control inputs are an optional list of tensors to evaluate before
    # the start of the search procedure. These can be used to assert the
    # validity of inputs to the search procedure.
    control_inputs = None

    if initial_inverse_hessian_estimate is None:
      # Create a default initial inverse Hessian.
      initial_inv_hessian = tf.eye(domain_size,
                                   batch_shape=batch_shape,
                                   dtype=dtype,
                                   name='initial_inv_hessian')
    else:
      # If an initial inverse Hessian is supplied, compute some control inputs
      # to ensure that it is positive definite and symmetric.
      initial_inv_hessian = tf.convert_to_tensor(
          value=initial_inverse_hessian_estimate,
          dtype=dtype,
          name='initial_inv_hessian')
      control_inputs = _inv_hessian_control_inputs(initial_inv_hessian)
      hessian_shape = tf.concat([batch_shape, [domain_size, domain_size]], 0)
      initial_inv_hessian = tf.broadcast_to(initial_inv_hessian, hessian_shape)

    # The `state` here is a `BfgsOptimizerResults` tuple with values for the
    # current state of the algorithm computation.
    def _cond(state):
      """Continue if iterations remain and stopping condition is not met."""
      return ((state.num_iterations < max_iterations) &
              tf.logical_not(stopping_condition(state.converged, state.failed)))

    def _body(state):
      """Main optimization loop."""
      search_direction = _get_search_direction(state.inverse_hessian_estimate,
                                               state.objective_gradient)
      derivative_at_start_pt = tf.reduce_sum(
          input_tensor=state.objective_gradient * search_direction, axis=-1)

      # If the derivative at the start point is not negative, recompute the
      # search direction with the initial inverse Hessian.
      needs_reset = (~state.failed & ~state.converged &
                     (derivative_at_start_pt >= 0))

      search_direction_reset = _get_search_direction(
          initial_inv_hessian, state.objective_gradient)

      actual_serch_direction = tf.where(
          needs_reset, search_direction_reset, search_direction)
      actual_inv_hessian = tf.where(
          needs_reset, initial_inv_hessian, state.inverse_hessian_estimate)

      # Replace the hessian estimate in the state, in case it had to be reset.
      current_state = bfgs_utils.update_fields(
          state, inverse_hessian_estimate=actual_inv_hessian)

      next_state = bfgs_utils.line_search_step(
          current_state,
          value_and_gradients_function, actual_serch_direction,
          tolerance, f_relative_tolerance, x_tolerance, stopping_condition)

      # Update the inverse Hessian if needed and continue.
      return [_update_inv_hessian(current_state, next_state)]

    kwargs = bfgs_utils.get_initial_state_args(
        value_and_gradients_function,
        initial_position,
        tolerance,
        control_inputs)
    kwargs['inverse_hessian_estimate'] = initial_inv_hessian
    initial_state = BfgsOptimizerResults(**kwargs)
    return tf.while_loop(
        cond=_cond,
        body=_body,
        loop_vars=[initial_state],
        parallel_iterations=parallel_iterations)[0]

def _inv_hessian_control_inputs(inv_hessian):
  """Computes control inputs to validate a provided inverse Hessian.

  These ensure that the provided inverse Hessian is positive definite and
  symmetric.

  Args:
    inv_hessian: The starting estimate for the inverse of the Hessian at the
      initial point.

  Returns:
    A list of tf.Assert ops suitable for use with tf.control_dependencies.
  """
  # The easiest way to validate if the inverse Hessian is positive definite is
  # to compute its Cholesky decomposition.
  is_positive_definite = tf.reduce_all(
      input_tensor=tf.math.is_finite(tf.linalg.cholesky(inv_hessian)),
      axis=[-1, -2])

  # Then check that the supplied inverse Hessian is symmetric.
  is_symmetric = tf.equal(bfgs_utils.norm(
      inv_hessian - _batch_transpose(inv_hessian), dims=2), 0)

  # Simply adding a control dependencies on these results is not enough to
  # trigger them, we need to add asserts on the results.
  return [tf.Assert(is_positive_definite,
                    ['Initial inverse Hessian is not positive definite.',
                     inv_hessian]),
          tf.Assert(is_symmetric,
                    ['Initial inverse Hessian is not symmetric',
                     inv_hessian])]

def _get_search_direction(inv_hessian_approx, gradient):
  """Computes the direction along which to perform line search."""
  return -_mul_right(inv_hessian_approx, gradient)

def _update_inv_hessian(prev_state, next_state):
  """Update the BGFS state by computing the next inverse hessian estimate."""
  # Only update the inverse Hessian if not already failed or converged.
  should_update = ~next_state.converged & ~next_state.failed

  # Compute the normalization term (y^T . s), should not update if is singular.
  gradient_delta = next_state.objective_gradient - prev_state.objective_gradient
  position_delta = next_state.position - prev_state.position
  normalization_factor = tf.reduce_sum(
      input_tensor=gradient_delta * position_delta, axis=-1)
  should_update = should_update & ~tf.equal(normalization_factor, 0)

  def _do_update_inv_hessian():
    next_inv_hessian = _bfgs_inv_hessian_update(
        gradient_delta, position_delta, normalization_factor,
        prev_state.inverse_hessian_estimate)
    return bfgs_utils.update_fields(
        next_state,
        inverse_hessian_estimate=tf.where(should_update,
                                          next_inv_hessian,
                                          prev_state.inverse_hessian_estimate))

  return prefer_static.cond(
      tf.reduce_any(input_tensor=should_update),
      _do_update_inv_hessian,
      lambda: next_state)

def _bfgs_inv_hessian_update(grad_delta, position_delta, normalization_factor,
                             inv_hessian_estimate):
  """Applies the BFGS update to the inverse Hessian estimate.

  The BFGS update rule is (note A^T denotes the transpose of a vector/matrix A).

  ```None
    rho = 1/(grad_delta^T * position_delta)
    U = (I - rho * position_delta * grad_delta^T)
    H_1 =  U * H_0 * U^T + rho * position_delta * position_delta^T
  ```

  Here, `H_0` is the inverse Hessian estimate at the previous iteration and
  `H_1` is the next estimate. Note that `*` should be interpreted as the
  matrix multiplication (with the understanding that matrix multiplication for
  scalars is usual multiplication and for matrix with vector is the action of
  the matrix on the vector.).

  The implementation below utilizes an expanded version of the above formula
  to avoid the matrix multiplications that would be needed otherwise. By
  expansion it is easy to see that one only needs matrix-vector or
  vector-vector operations. The expanded version is:

  ```None
    f = 1 + rho * (grad_delta^T * H_0 * grad_delta)
    H_1 - H_0 = - rho * [position_delta * (H_0 * grad_delta)^T +
                        (H_0 * grad_delta) * position_delta^T] +
                  rho * f * [position_delta * position_delta^T]
  ```

  All the terms in square brackets are matrices and are constructed using
  vector outer products. All the other terms on the right hand side are scalars.
  Also worth noting that the first and second lines are both rank 1 updates
  applied to the current inverse Hessian estimate.

  Args:
    grad_delta: Real `Tensor` of shape `[..., n]`. The difference between the
      gradient at the new position and the old position.
    position_delta: Real `Tensor` of shape `[..., n]`. The change in position
      from the previous iteration to the current one.
    normalization_factor: Real `Tensor` of shape `[...]`. Should be equal to
      `grad_delta^T * position_delta`, i.e. `1/rho` as defined above.
    inv_hessian_estimate: Real `Tensor` of shape `[..., n, n]`. The previous
      estimate of the inverse Hessian. Should be positive definite and
      symmetric.

  Returns:
    A tuple containing the following fields
      is_valid: A Boolean `Tensor` of shape `[...]` indicating batch members
        where the update succeeded. The update can fail if the position change
        becomes orthogonal to the gradient change.
      next_inv_hessian_estimate: A `Tensor` of shape `[..., n, n]`. The next
        Hessian estimate updated using the BFGS update scheme. If the
        `inv_hessian_estimate` is symmetric and positive definite, the
        `next_inv_hessian_estimate` is guaranteed to satisfy the same
        conditions.
  """
  # The quadratic form: y^T.H.y; where H is the inverse Hessian and y is the
  # gradient change.
  conditioned_grad_delta = _mul_right(inv_hessian_estimate, grad_delta)
  conditioned_grad_delta_norm = tf.reduce_sum(
      input_tensor=conditioned_grad_delta * grad_delta, axis=-1)

  # The first rank 1 update term requires the outer product: s.y^T.
  cross_term = _tensor_product(position_delta, conditioned_grad_delta)

  def _expand_scalar(s):
    # Expand dimensions of a batch of scalars to multiply or divide a matrix.
    return s[..., tf.newaxis, tf.newaxis]

  # Symmetrize
  cross_term += _tensor_product(conditioned_grad_delta, position_delta)
  position_term = _tensor_product(position_delta, position_delta)
  with tf.control_dependencies([position_term]):
    position_term *= _expand_scalar(
        1 + conditioned_grad_delta_norm / normalization_factor)

  return (inv_hessian_estimate +
          (position_term - cross_term) / _expand_scalar(normalization_factor))

def _mul_right(mat, vec):
  """Computes the product of a matrix with a vector on the right.

  Note this supports dynamic shapes and batched computation.

  Examples:

    M = tf.reshape(tf.range(6), shape=(3, 2))
    # => [[0, 1],
    #     [2, 3],
    #     [4, 5]]
    v = tf.constant([1, 2])  # Shape: (2,)
    _mul_right(M, v)
    # => [ 2,  8, 14]  # Shape: (3,)

    M = tf.reshape(tf.range(30), shape=(2, 3, 5))
    # => [[[ 0,  1,  2,  3,  4],
    #     [ 5,  6,  7,  8,  9],
    #     [10, 11, 12, 13, 14]],
    #
    #    [[15, 16, 17, 18, 19],
    #     [20, 21, 22, 23, 24],
    #     [25, 26, 27, 28, 29]]]
    v = tf.reshape(tf.range(10), shape=(2, 5))
    # => [[0, 1, 2, 3, 4],
    #     [5, 6, 7, 8, 9]]
    _mul_right(M, v)
    # => [[ 30,  80, 130],
    #     [605, 780, 955]]  # Shape: (2, 3)

  Args:
    mat: A `tf.Tensor` of shape `[..., n, m]`.
    vec: A `tf.Tensor` of shape `[..., m]`.

  Returns:
    A tensor of shape `[..., n]` with matching batch dimensions.
  """
  return tf.squeeze(tf.matmul(mat, tf.expand_dims(vec, axis=-1)), axis=-1)

def _tensor_product(t1, t2):
  """Computes the outer product of two possibly batched vectors.

  Args:
    t1: A `tf.Tensor` of shape `[..., n]`.
    t2: A `tf.Tensor` of shape `[..., m]`.

  Returns:
    A tensor of shape `[..., n, m]` with matching batch dimensions, let's call
    it `r`, whose components are:

    ```None
      r[..., i, j] = t1[..., i] * t2[..., j]
    ```
  """
  return tf.matmul(tf.expand_dims(t1, axis=-1), tf.expand_dims(t2, axis=-2))

def _batch_transpose(mat):
  """Transpose a possibly batched matrix.

  Args:
    mat: A `tf.Tensor` of shape `[..., n, m]`.

  Returns:
    A tensor of shape `[..., m, n]` with matching batch dimensions.
  """
  n = distribution_util.prefer_static_rank(mat)
  perm = tf.range(n)
  perm = tf.concat([perm[:-2], [perm[-1], perm[-2]]], axis=0)
  return tf.transpose(a=mat, perm=perm)


# Copyright 2018 The TensorFlow Probability Authors.
#
# 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.
# ============================================================================
"""Positive-Semidefinite Kernel library utilities."""

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

import numpy as np
import tensorflow as tf

from tensorflow_probability.python.internal import distribution_util
from tensorflow_probability.python.internal import dtype_util

__all__ = [
    'pad_shape_right_with_ones',
    'maybe_get_common_dtype',
    'sum_rightmost_ndims_preserving_shape',
]

def pad_shape_right_with_ones(x, ndims):
  """Maybe add `ndims` ones to `x.shape` on the right.

  If `ndims` is zero, this is a no-op; otherwise, we will create and return a
  new `Tensor` whose shape is that of `x` with `ndims` ones concatenated on the
  right side. If the shape of `x` is known statically, the shape of the return
  value will be as well.

  Args:
    x: The `Tensor` we'll return a reshaping of.
    ndims: Python `integer` number of ones to pad onto `x.shape`.
  Returns:
    If `ndims` is zero, `x`; otherwise, a `Tensor` whose shape is that of `x`
    with `ndims` ones concatenated on the right side. If possible, returns a
    `Tensor` whose shape is known statically.
  Raises:
    ValueError: if `ndims` is not a Python `integer` greater than or equal to
    zero.
  """
  if not (isinstance(ndims, int) and ndims >= 0):
    raise ValueError(
        '`ndims` must be a Python `integer` greater than zero. Got: {}'
        .format(ndims))
  if ndims == 0:
    return x
  x = tf.convert_to_tensor(value=x)
  original_shape = x.shape
  new_shape = distribution_util.pad(
      tf.shape(input=x), axis=0, back=True, value=1, count=ndims)
  x = tf.reshape(x, new_shape)
  x.set_shape(original_shape.concatenate([1]*ndims))
  return x

def sum_rightmost_ndims_preserving_shape(x, ndims):
  """Return `Tensor` with right-most ndims summed.

  Args:
    x: the `Tensor` whose right-most `ndims` dimensions to sum
    ndims: number of right-most dimensions to sum.

  Returns:
    A `Tensor` resulting from calling `reduce_sum` on the `ndims` right-most
    dimensions. If the shape of `x` is statically known, the result will also
    have statically known shape. Otherwise, the resulting shape will only be
    known at runtime.
  """
  x = tf.convert_to_tensor(value=x)
  if x.shape.ndims is not None:
    axes = tf.range(x.shape.ndims - ndims, x.shape.ndims)
  else:
    axes = tf.range(tf.rank(x) - ndims, tf.rank(x))
  return tf.reduce_sum(input_tensor=x, axis=axes)

@tf.custom_gradient
def sqrt_with_finite_grads(x, name=None):
  """A sqrt function whose gradient at zero is very large but finite.

  Args:
    x: a `Tensor` whose sqrt is to be computed.
    name: a Python `str` prefixed to all ops created by this function.
      Default `None` (i.e., "sqrt_with_finite_grads").

  Returns:
    sqrt: the square root of `x`, with an overridden gradient at zero
    grad: a gradient function, which is the same as sqrt's gradient everywhere
      except at zero, where it is given a large finite value, instead of `inf`.

  Raises:
    TypeError: if `tf.convert_to_tensor(x)` is not a `float` type.

  Often in kernel functions, we need to compute the L2 norm of the difference
  between two vectors, `x` and `y`: `sqrt(sum_i((x_i - y_i) ** 2))`. In the
  case where `x` and `y` are identical, e.g., on the diagonal of a kernel
  matrix, we get `NaN`s when we take gradients with respect to the inputs. To
  see, this consider the forward pass:

    ```
    [x_1 ... x_N]  -->  [x_1 ** 2 ... x_N ** 2]  -->
        (x_1 ** 2 + ... + x_N ** 2)  -->  sqrt((x_1 ** 2 + ... + x_N ** 2))
    ```

  When we backprop through this forward pass, the `sqrt` yields an `inf` because
  `grad_z(sqrt(z)) = 1 / (2 * sqrt(z))`. Continuing the backprop to the left, at
  the `x ** 2` term, we pick up a `2 * x`, and when `x` is zero, we get
  `0 * inf`, which is `NaN`.

  We'd like to avoid these `NaN`s, since they infect the rest of the connected
  computation graph. Practically, when two inputs to a kernel function are
  equal, we are in one of two scenarios:
    1. We are actually computing k(x, x), in which case norm(x - x) is
       identically zero, independent of x. In this case, we'd like the
       gradient to reflect this independence: it should be zero.
    2. We are computing k(x, y), and x just *happens* to have the same value
       as y. The gradient at such inputs is in fact ill-defined (there is a
       cusp in the sqrt((x - y) ** 2) surface along the line x = y). There are,
       however, an infinite number of sub-gradients, all of which are valid at
       all such inputs. By symmetry, there is exactly one which is "special":
       zero, and we elect to use that value here. In practice, having two
       identical inputs to a kernel matrix is probably a pathological
       situation to be avoided, but that is better resolved at a higher level
       than this.

  To avoid the infinite gradient at zero, we use tf.custom_gradient to redefine
  the gradient at zero. We assign it to be a very large value, specifically
  the sqrt of the max value of the floating point dtype of the input. We use
  the sqrt (as opposed to just using the max floating point value) to avoid
  potential overflow when combining this value with others downstream.
  """
  with tf.compat.v1.name_scope(name, 'sqrt_with_finite_grads', [x]):
    x = tf.convert_to_tensor(value=x, name='x')
    if not x.dtype.is_floating:
      raise TypeError('Input `x` must be floating type.')
    def grad(grad_ys):
      large_float_like_x = np.sqrt(np.finfo(x.dtype.as_numpy_dtype()).max)
      safe_grads = tf.where(
          tf.equal(x, 0), tf.fill(tf.shape(input=x), large_float_like_x),
          0.5 * tf.math.rsqrt(x))
      return grad_ys * safe_grads
    return tf.sqrt(x), grad

def maybe_get_common_dtype(arg_list):
  """Return common dtype of arg_list, or None.

  Args:
    arg_list: an iterable of items which are either `None` or have a `dtype`
      property.

  Returns:
    dtype: The common dtype of items in `arg_list`, or `None` if the list is
      empty or all items are `None`.
  """
  # Note that `all` defaults to `True` if `arg_list` is empty.
  if all(a is None for a in arg_list):
    return None
  return dtype_util.common_dtype(arg_list, tf.float32)


# Copyright 2018 The TensorFlow Probability Authors.
#
# 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.
# ============================================================================
"""The Limited-Memory BFGS minimization algorithm.

Limited-memory quasi-Newton methods are useful for solving large problems
whose Hessian matrices cannot be computed at a reasonable cost or are not
sparse. Instead of storing fully dense n x n approximations of Hessian
matrices, they only save a few vectors of length n that represent the
approximations implicitly.

This module implements the algorithm know as L-BFGS, which, as its name
suggests, is a limited-memory version of the BFGS algorithm.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections

# Dependency imports
import tensorflow as tf

from tensorflow_probability.python.internal import distribution_util
from tensorflow_probability.python.internal import prefer_static
from tensorflow_probability.python.optimizer import bfgs_utils

LBfgsOptimizerResults = collections.namedtuple(
    'LBfgsOptimizerResults', [
        'converged',  # Scalar boolean tensor indicating whether the minimum
                      # was found within tolerance.
        'failed',  # Scalar boolean tensor indicating whether a line search
                   # step failed to find a suitable step size satisfying Wolfe
                   # conditions. In the absence of any constraints on the
                   # number of objective evaluations permitted, this value will
                   # be the complement of `converged`. However, if there is
                   # a constraint and the search stopped due to available
                   # evaluations being exhausted, both `failed` and `converged`
                   # will be simultaneously False.
        'num_iterations',  # The number of iterations of the BFGS update.
        'num_objective_evaluations',  # The total number of objective
                                      # evaluations performed.
        'position',  # A tensor containing the last argument value found
                     # during the search. If the search converged, then
                     # this value is the argmin of the objective function.
        'objective_value',  # A tensor containing the value of the objective
                            # function at the `position`. If the search
                            # converged, then this is the (local) minimum of
                            # the objective function.
        'objective_gradient',  # A tensor containing the gradient of the
                               # objective function at the
                               # `final_position`. If the search converged
                               # the max-norm of this tensor should be
                               # below the tolerance.
        'position_deltas',  # A tensor encoding information about the latest
                            # changes in `position` during the algorithm
                            # execution. Its shape is of the form
                            # `(num_correction_pairs,) + position.shape` where
                            # `num_correction_pairs` is given as an argument to
                            # the minimize function.
        'gradient_deltas',  # A tensor encoding information about the latest
                            # changes in `objective_gradient` during the
                            # algorithm execution. Has the same shape as
                            # position_deltas.
    ])

def minimize(value_and_gradients_function,
             initial_position,
             num_correction_pairs=10,
             tolerance=1e-8,
             x_tolerance=0,
             f_relative_tolerance=0,
             initial_inverse_hessian_estimate=None,
             max_iterations=50,
             parallel_iterations=1,
             stopping_condition=None,
             name=None):
  """Applies the L-BFGS algorithm to minimize a differentiable function.

  Performs unconstrained minimization of a differentiable function using the
  L-BFGS scheme. See [Nocedal and Wright(2006)][1] for details of the algorithm.

  ### Usage:

  The following example demonstrates the L-BFGS optimizer attempting to find the
  minimum for a simple high-dimensional quadratic objective function.

  ```python
    # A high-dimensional quadratic bowl.
    ndims = 60
    minimum = np.ones([ndims], dtype='float64')
    scales = np.arange(ndims, dtype='float64') + 1.0

    # The objective function and the gradient.
    def quadratic(x):
      value = tf.reduce_sum(scales * (x - minimum) ** 2)
      return value, tf.gradients(value, x)[0]

    start = np.arange(ndims, 0, -1, dtype='float64')
    optim_results = tfp.optimizer.lbfgs_minimize(
        quadratic, initial_position=start, num_correction_pairs=10,
        tolerance=1e-8)

    with tf.Session() as session:
      results = session.run(optim_results)
      # Check that the search converged
      assert(results.converged)
      # Check that the argmin is close to the actual value.
      np.testing.assert_allclose(results.position, minimum)
  ```

  ### References:

  [1] Jorge Nocedal, Stephen Wright. Numerical Optimization. Springer Series
      in Operations Research. pp 176-180. 2006

  http://pages.mtu.edu/~struther/Courses/OLD/Sp2013/5630/Jorge_Nocedal_Numerical_optimization_267490.pdf

  Args:
    value_and_gradients_function:  A Python callable that accepts a point as a
      real `Tensor` and returns a tuple of `Tensor`s of real dtype containing
      the value of the function and its gradient at that point. The function
      to be minimized. The input is of shape `[..., n]`, where `n` is the size
      of the domain of input points, and all others are batching dimensions.
      The first component of the return value is a real `Tensor` of matching
      shape `[...]`. The second component (the gradient) is also of shape
      `[..., n]` like the input value to the function.
    initial_position: Real `Tensor` of shape `[..., n]`. The starting point, or
      points when using batching dimensions, of the search procedure. At these
      points the function value and the gradient norm should be finite.
    num_correction_pairs: Positive integer. Specifies the maximum number of
      (position_delta, gradient_delta) correction pairs to keep as implicit
      approximation of the Hessian matrix.
    tolerance: Scalar `Tensor` of real dtype. Specifies the gradient tolerance
      for the procedure. If the supremum norm of the gradient vector is below
      this number, the algorithm is stopped.
    x_tolerance: Scalar `Tensor` of real dtype. If the absolute change in the
      position between one iteration and the next is smaller than this number,
      the algorithm is stopped.
    f_relative_tolerance: Scalar `Tensor` of real dtype. If the relative change
      in the objective value between one iteration and the next is smaller
      than this value, the algorithm is stopped.
    initial_inverse_hessian_estimate: None. Option currently not supported.
    max_iterations: Scalar positive int32 `Tensor`. The maximum number of
      iterations for L-BFGS updates.
    parallel_iterations: Positive integer. The number of iterations allowed to
      run in parallel.
    stopping_condition: (Optional) A Python function that takes as input two
      Boolean tensors of shape `[...]`, and returns a Boolean scalar tensor.
      The input tensors are `converged` and `failed`, indicating the current
      status of each respective batch member; the return value states whether
      the algorithm should stop. The default is tfp.optimizer.converged_all
      which only stops when all batch members have either converged or failed.
      An alternative is tfp.optimizer.converged_any which stops as soon as one
      batch member has converged, or when all have failed.
    name: (Optional) Python str. The name prefixed to the ops created by this
      function. If not supplied, the default name 'minimize' is used.

  Returns:
    optimizer_results: A namedtuple containing the following items:
      converged: Scalar boolean tensor indicating whether the minimum was
        found within tolerance.
      failed:  Scalar boolean tensor indicating whether a line search
        step failed to find a suitable step size satisfying Wolfe
        conditions. In the absence of any constraints on the
        number of objective evaluations permitted, this value will
        be the complement of `converged`. However, if there is
        a constraint and the search stopped due to available
        evaluations being exhausted, both `failed` and `converged`
        will be simultaneously False.
      num_objective_evaluations: The total number of objective
        evaluations performed.
      position: A tensor containing the last argument value found
        during the search. If the search converged, then
        this value is the argmin of the objective function.
      objective_value: A tensor containing the value of the objective
        function at the `position`. If the search converged, then this is
        the (local) minimum of the objective function.
      objective_gradient: A tensor containing the gradient of the objective
        function at the `position`. If the search converged the
        max-norm of this tensor should be below the tolerance.
      position_deltas: A tensor encoding information about the latest
        changes in `position` during the algorithm execution.
      gradient_deltas: A tensor encoding information about the latest
        changes in `objective_gradient` during the algorithm execution.
  """
  if initial_inverse_hessian_estimate is not None:
    raise NotImplementedError(
        'Support of initial_inverse_hessian_estimate arg not yet implemented')

  if stopping_condition is None:
    stopping_condition = bfgs_utils.converged_all

  with tf.compat.v1.name_scope(name, 'minimize', [initial_position, tolerance]):
    initial_position = tf.convert_to_tensor(
        value=initial_position, name='initial_position')
    dtype = initial_position.dtype.base_dtype
    tolerance = tf.convert_to_tensor(
        value=tolerance, dtype=dtype, name='grad_tolerance')
    f_relative_tolerance = tf.convert_to_tensor(
        value=f_relative_tolerance, dtype=dtype, name='f_relative_tolerance')
    x_tolerance = tf.convert_to_tensor(
        value=x_tolerance, dtype=dtype, name='x_tolerance')
    max_iterations = tf.convert_to_tensor(
        value=max_iterations, name='max_iterations')

    # The `state` here is a `LBfgsOptimizerResults` tuple with values for the
    # current state of the algorithm computation.
    def _cond(state):
      """Continue if iterations remain and stopping condition is not met."""
      return ((state.num_iterations < max_iterations) &
              tf.logical_not(stopping_condition(state.converged, state.failed)))

    def _body(current_state):
      """Main optimization loop."""
      search_direction = _get_search_direction(current_state)

      # TODO(b/120134934): Check if the derivative at the start point is not
      # negative, if so then reset position/gradient deltas and recompute
      # search direction.

      next_state = bfgs_utils.line_search_step(
          current_state,
          value_and_gradients_function, search_direction,
          tolerance, f_relative_tolerance, x_tolerance, stopping_condition)

      # If not failed or converged, update the Hessian estimate.
      should_update = ~(next_state.converged | next_state.failed)
      state_after_inv_hessian_update = bfgs_utils.update_fields(
          next_state,
          position_deltas=_queue_push(
              current_state.position_deltas, should_update,
              next_state.position - current_state.position),
          gradient_deltas=_queue_push(
              current_state.gradient_deltas, should_update,
              next_state.objective_gradient - current_state.objective_gradient))
      return [state_after_inv_hessian_update]

    initial_state = _get_initial_state(value_and_gradients_function,
                                       initial_position,
                                       num_correction_pairs,
                                       tolerance)
    return tf.while_loop(
        cond=_cond,
        body=_body,
        loop_vars=[initial_state],
        parallel_iterations=parallel_iterations)[0]

def _get_initial_state(value_and_gradients_function,
                       initial_position,
                       num_correction_pairs,
                       tolerance):
  """Create LBfgsOptimizerResults with initial state of search procedure."""
  init_args = bfgs_utils.get_initial_state_args(
      value_and_gradients_function,
      initial_position,
      tolerance)
  empty_queue = _make_empty_queue_for(num_correction_pairs, initial_position)
  init_args.update(position_deltas=empty_queue, gradient_deltas=empty_queue)
  return LBfgsOptimizerResults(**init_args)

def _get_search_direction(state):
  """Computes the search direction to follow at the current state.

  On the `k`-th iteration of the main L-BFGS algorithm, the state has collected
  the most recent `m` correction pairs in position_deltas and gradient_deltas,
  where `k = state.num_iterations` and `m = min(k, num_correction_pairs)`.

  Assuming these, the code below is an implementation of the L-BFGS two-loop
  recursion algorithm given by [Nocedal and Wright(2006)][1]:

  ```None
    q_direction = objective_gradient
    for i in reversed(range(m)):  # First loop.
      inv_rho[i] = gradient_deltas[i]^T * position_deltas[i]
      alpha[i] = position_deltas[i]^T * q_direction / inv_rho[i]
      q_direction = q_direction - alpha[i] * gradient_deltas[i]

    kth_inv_hessian_factor = (gradient_deltas[-1]^T * position_deltas[-1] /
                              gradient_deltas[-1]^T * gradient_deltas[-1])
    r_direction = kth_inv_hessian_factor * I * q_direction

    for i in range(m):  # Second loop.
      beta = gradient_deltas[i]^T * r_direction / inv_rho[i]
      r_direction = r_direction + position_deltas[i] * (alpha[i] - beta)

    return -r_direction  # Approximates - H_k * objective_gradient.
  ```

  Args:
    state: A `LBfgsOptimizerResults` tuple with the current state of the
      search procedure.

  Returns:
    A real `Tensor` of the same shape as the `state.position`. The direction
    along which to perform line search.
  """
  # The number of correction pairs that have been collected so far.
  num_elements = tf.minimum(
      state.num_iterations,
      distribution_util.prefer_static_shape(state.position_deltas)[0])

  def _two_loop_algorithm():
    """L-BFGS two-loop algorithm."""
    # Correction pairs are always appended to the end, so only the latest
    # `num_elements` vectors have valid position/gradient deltas.
    position_deltas = state.position_deltas[-num_elements:]
    gradient_deltas = state.gradient_deltas[-num_elements:]

    # Pre-compute all `inv_rho[i]`s.
    inv_rhos = tf.reduce_sum(
        input_tensor=gradient_deltas * position_deltas, axis=-1)

    def first_loop(acc, args):
      _, q_direction = acc
      position_delta, gradient_delta, inv_rho = args
      alpha = tf.reduce_sum(
          input_tensor=position_delta * q_direction, axis=-1) / inv_rho
      direction_delta = tf.expand_dims(alpha, axis=-1) * gradient_delta
      return (alpha, q_direction - direction_delta)

    # Run first loop body computing and collecting `alpha[i]`s, while also
    # computing the updated `q_direction` at each step.
    zero = tf.zeros_like(inv_rhos[0])
    alphas, q_directions = tf.scan(
        first_loop, [position_deltas, gradient_deltas, inv_rhos],
        initializer=(zero, state.objective_gradient), reverse=True)

    # We use `H^0_k = gamma_k * I` as an estimate for the initial inverse
    # hessian for the k-th iteration; then `r_direction = H^0_k * q_direction`.
    gamma_k = inv_rhos[-1] / tf.reduce_sum(
        input_tensor=gradient_deltas[-1] * gradient_deltas[-1], axis=-1)
    r_direction = tf.expand_dims(gamma_k, axis=-1) * q_directions[0]

    def second_loop(r_direction, args):
      alpha, position_delta, gradient_delta, inv_rho = args
      beta = tf.reduce_sum(
          input_tensor=gradient_delta * r_direction, axis=-1) / inv_rho
      direction_delta = tf.expand_dims(alpha - beta, axis=-1) * position_delta
      return r_direction + direction_delta

    # Finally, run second loop body computing the updated `r_direction` at each
    # step.
    r_directions = tf.scan(
        second_loop, [alphas, position_deltas, gradient_deltas, inv_rhos],
        initializer=r_direction)
    return -r_directions[-1]

  return prefer_static.cond(tf.equal(num_elements, 0),
                            (lambda: -state.objective_gradient),
                            _two_loop_algorithm)

def _make_empty_queue_for(k, element):
  """Creates a `tf.Tensor` suitable to hold `k` element-shaped tensors.

  For example:

  ```python
    element = tf.constant([[0., 1., 2., 3., 4.],
                           [5., 6., 7., 8., 9.]])

    # A queue capable of holding 3 elements.
    _make_empty_queue_for(3, element)
    # => [[[ 0.,  0.,  0.,  0.,  0.],
    #      [ 0.,  0.,  0.,  0.,  0.]],
    #
    #     [[ 0.,  0.,  0.,  0.,  0.],
    #      [ 0.,  0.,  0.,  0.,  0.]],
    #
    #     [[ 0.,  0.,  0.,  0.,  0.],
    #      [ 0.,  0.,  0.,  0.,  0.]]]
  ```

  Args:
    k: A positive scalar integer, number of elements that each queue will hold.
    element: A `tf.Tensor`, only its shape and dtype information are relevant.

  Returns:
    A zero-filed `tf.Tensor` of shape `(k,) + tf.shape(element)` and same dtype
    as `element`.
  """
  queue_shape = tf.concat(
      [[k], distribution_util.prefer_static_shape(element)], axis=0)
  return tf.zeros(queue_shape, dtype=element.dtype.base_dtype)

def _queue_push(queue, should_update, new_vecs):
  """Conditionally push new vectors into a batch of first-in-first-out queues.

  The `queue` of shape `[k, ..., n]` can be thought of as a batch of queues,
  each holding `k` n-D vectors; while `new_vecs` of shape `[..., n]` is a
  fresh new batch of n-D vectors. The `should_update` batch of Boolean scalars,
  i.e. shape `[...]`, indicates batch members whose corresponding n-D vector in
  `new_vecs` should be added at the back of its queue, pushing out the
  corresponding n-D vector from the front. Batch members in `new_vecs` for
  which `should_update` is False are ignored.

  Note: the choice of placing `k` at the dimension 0 of the queue is
  constrained by the L-BFGS two-loop algorithm above. The algorithm uses
  tf.scan to iterate over the `k` correction pairs simulatneously across all
  batches, and tf.scan itself can only iterate over dimension 0.

  For example:

  ```python
    k, b, n = (3, 2, 5)
    queue = tf.reshape(tf.range(30), (k, b, n))
    # => [[[ 0,  1,  2,  3,  4],
    #      [ 5,  6,  7,  8,  9]],
    #
    #     [[10, 11, 12, 13, 14],
    #      [15, 16, 17, 18, 19]],
    #
    #     [[20, 21, 22, 23, 24],
    #      [25, 26, 27, 28, 29]]]

    element = tf.reshape(tf.range(30, 40), (b, n))
    # => [[30, 31, 32, 33, 34],
          [35, 36, 37, 38, 39]]

    should_update = tf.constant([True, False])  # Shape: (b,)

    _queue_add(should_update, queue, element)
    # => [[[10, 11, 12, 13, 14],
    #      [ 5,  6,  7,  8,  9]],
    #
    #     [[20, 21, 22, 23, 24],
    #      [15, 16, 17, 18, 19]],
    #
    #     [[30, 31, 32, 33, 34],
    #      [25, 26, 27, 28, 29]]]
  ```

  Args:
    queue: A `tf.Tensor` of shape `[k, ..., n]`; a batch of queues each with
      `k` n-D vectors.
    should_update: A Boolean `tf.Tensor` of shape `[...]` indicating batch
      members where new vectors should be added to their queues.
    new_vecs: A `tf.Tensor` of shape `[..., n]`; a batch of n-D vectors to add
      at the end of their respective queues, pushing out the first element from
      each.

  Returns:
    A new `tf.Tensor` of shape `[k, ..., n]`.
  """
  new_queue = tf.concat([queue[1:], [new_vecs]], axis=0)
  update_pattern = tf.broadcast_to(
      should_update[tf.newaxis, ..., tf.newaxis],
      distribution_util.prefer_static_shape(queue))
  return tf.where(update_pattern, new_queue, queue)


# Copyright 2018 The TensorFlow Probability Authors.
#
# 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.
# ============================================================================
# Copyright 2018 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.
# ==============================================================================
"""Estimating the volume of the correlation matrices with bounded determinant.

Why?  Because lkj_test.py tests the sampler for the LKJ distribution
by estimating the same volume another way.

How?  Rejection sampling.  Or, more precisely, importance sampling,
proposing from the uniform distribution on symmetric matrices with
diagonal 1s and entries in [-1, 1].  Such a matrix is a correlation
matrix if and only if it is also positive semi-definite.

The samples can then be converted into a confidence interval on the
volume in question by the [Clopper-Pearson
method](https://en.wikipedia.org/wiki/Binomial_proportion_confidence_interval),
also implemented here.
"""

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

import importlib
import sys

# Dependency imports
import numpy as np
import tensorflow as tf

from tensorflow_probability.python.distributions import uniform
from tensorflow_probability.python.internal import distribution_util as util

__all__ = [
    "correlation_matrix_volume_rejection_samples",
    "compute_true_volumes",
]

def try_import(name):  # pylint: disable=invalid-name
  module = None
  try:
    module = importlib.import_module(name)
  except ImportError as e:
    tf.compat.v1.logging.warning("Could not import %s: %s" % (name, str(e)))
  return module

optimize = try_import("scipy.optimize")
stats = try_import("scipy.stats")

def _psd_mask(x):
  """Computes whether each square matrix in the input is positive semi-definite.

  Args:
    x: A floating-point `Tensor` of shape `[B1, ..., Bn, M, M]`.

  Returns:
    mask: A floating-point `Tensor` of shape `[B1, ... Bn]`.  Each
      scalar is 1 if the corresponding matrix was PSD, otherwise 0.
  """
  # Allegedly
  # https://scicomp.stackexchange.com/questions/12979/testing-if-a-matrix-is-positive-semi-definite
  # it is more efficient to test for positive semi-definiteness by
  # trying to compute the Cholesky decomposition -- the matrix is PSD
  # if you succeed and not PSD if you fail.  However, TensorFlow's
  # Cholesky raises an exception if _any_ of the input matrices are
  # not PSD, from which I don't know how to extract _which ones_, so I
  # proceed by explicitly computing all the eigenvalues and checking
  # whether they are all positive or not.
  #
  # Also, as was discussed in the answer, it is somewhat dangerous to
  # treat SPD-ness as binary in floating-point arithmetic. Cholesky
  # factorization can complete and 'look' like everything is fine
  # (e.g., O(1) entries and a diagonal of all ones) but the matrix can
  # have an exponential condition number.
  eigenvalues, _ = tf.linalg.eigh(x)
  return tf.cast(
      tf.reduce_min(input_tensor=eigenvalues, axis=-1) >= 0, dtype=x.dtype)

def _det_large_enough_mask(x, det_bounds):
  """Returns whether the input matches the given determinant limit.

  Args:
    x: A floating-point `Tensor` of shape `[B1, ..., Bn, M, M]`.
    det_bounds: A floating-point `Tensor` that must broadcast to shape
      `[B1, ..., Bn]`, giving the desired lower bound on the
      determinants in `x`.

  Returns:
    mask: A floating-point `Tensor` of shape [B1, ..., Bn].  Each
      scalar is 1 if the corresponding matrix had determinant above
      the corresponding bound, otherwise 0.
  """
  # For the curious: I wonder whether it is possible and desirable to
  # use a Cholesky decomposition-based algorithm for this, since the
  # only matrices whose determinant this code cares about will be PSD.
  # Didn't figure out how to code that in TensorFlow.
  #
  # Expert opinion is that it would be about twice as fast since
  # Cholesky is roughly half the cost of Gaussian Elimination with
  # Partial Pivoting. But this is less of an impact than the switch in
  # _psd_mask.
  return tf.cast(tf.linalg.det(x) > det_bounds, dtype=x.dtype)

def _uniform_correlation_like_matrix(num_rows, batch_shape, dtype, seed):
  """Returns a uniformly random `Tensor` of "correlation-like" matrices.

  A "correlation-like" matrix is a symmetric square matrix with all entries
  between -1 and 1 (inclusive) and 1s on the main diagonal.  Of these,
  the ones that are positive semi-definite are exactly the correlation
  matrices.

  Args:
    num_rows: Python `int` dimension of the correlation-like matrices.
    batch_shape: `Tensor` or Python `tuple` of `int` shape of the
      batch to return.
    dtype: `dtype` of the `Tensor` to return.
    seed: Random seed.

  Returns:
    matrices: A `Tensor` of shape `batch_shape + [num_rows, num_rows]`
      and dtype `dtype`.  Each entry is in [-1, 1], and each matrix
      along the bottom two dimensions is symmetric and has 1s on the
      main diagonal.
  """
  num_entries = num_rows * (num_rows + 1) / 2
  ones = tf.ones(shape=[num_entries], dtype=dtype)
  # It seems wasteful to generate random values for the diagonal since
  # I am going to throw them away, but `fill_triangular` fills the
  # diagonal, so I probably need them.
  # It's not impossible that it would be more efficient to just fill
  # the whole matrix with random values instead of messing with
  # `fill_triangular`.  Then would need to filter almost half out with
  # `matrix_band_part`.
  unifs = uniform.Uniform(-ones, ones).sample(batch_shape, seed=seed)
  tril = util.fill_triangular(unifs)
  symmetric = tril + tf.linalg.matrix_transpose(tril)
  diagonal_ones = tf.ones(
      shape=util.pad(batch_shape, axis=0, back=True, value=num_rows),
      dtype=dtype)
  return tf.linalg.set_diag(symmetric, diagonal_ones)

def correlation_matrix_volume_rejection_samples(
    det_bounds, dim, sample_shape, dtype, seed):
  """Returns rejection samples from trying to get good correlation matrices.

  The proposal being rejected from is the uniform distribution on
  "correlation-like" matrices.  We say a matrix is "correlation-like"
  if it is a symmetric square matrix with all entries between -1 and 1
  (inclusive) and 1s on the main diagonal.  Of these, the ones that
  are positive semi-definite are exactly the correlation matrices.

  The rejection algorithm, then, is to sample a `Tensor` of
  `sample_shape` correlation-like matrices of dimensions `dim` by
  `dim`, and check each one for (i) being a correlation matrix (i.e.,
  PSD), and (ii) having determinant at least the corresponding entry
  of `det_bounds`.

  Args:
    det_bounds: A `Tensor` of lower bounds on the determinants of
      acceptable matrices.  The shape must broadcast with `sample_shape`.
    dim: A Python `int` dimension of correlation matrices to sample.
    sample_shape: Python `tuple` of `int` shape of the samples to
      compute, excluding the two matrix dimensions.
    dtype: The `dtype` in which to do the computation.
    seed: Random seed.

  Returns:
    weights: A `Tensor` of shape `sample_shape`.  Each entry is 0 if the
      corresponding matrix was not a correlation matrix, or had too
      small of a determinant.  Otherwise, the entry is the
      multiplicative inverse of the density of proposing that matrix
      uniformly, i.e., the volume of the set of `dim` by `dim`
      correlation-like matrices.
    volume: The volume of the set of `dim` by `dim` correlation-like
      matrices.
  """
  with tf.compat.v1.name_scope("rejection_sampler"):
    rej_proposals = _uniform_correlation_like_matrix(
        dim, sample_shape, dtype, seed=seed)
    rej_proposal_volume = 2. ** (dim * (dim - 1) / 2.)
    # The density of proposing any given point is 1 / rej_proposal_volume;
    # The weight of that point should be scaled by
    # 1 / density = rej_proposal_volume.
    rej_weights = rej_proposal_volume * _psd_mask(
        rej_proposals) * _det_large_enough_mask(rej_proposals, det_bounds)
    return rej_weights, rej_proposal_volume

def _clopper_pearson_confidence_interval(samples, error_rate):
  """Computes a confidence interval for the mean of the given 1-D distribution.

  Assumes (and checks) that the given distribution is Bernoulli, i.e.,
  takes only two values.  This licenses using the CDF of the binomial
  distribution for the confidence, which is tighter (for extreme
  probabilities) than the DKWM inequality.  The method is known as the
  [Clopper-Pearson method]
  (https://en.wikipedia.org/wiki/Binomial_proportion_confidence_interval).

  Assumes:

  - The given samples were drawn iid from the distribution of interest.

  - The given distribution is a Bernoulli, i.e., supported only on
    low and high.

  Guarantees:

  - The probability (over the randomness of drawing the given sample)
    that the true mean is outside the returned interval is no more
    than the given error_rate.

  Args:
    samples: `np.ndarray` of samples drawn iid from the distribution
      of interest.
    error_rate: Python `float` admissible rate of mistakes.

  Returns:
    low: Lower bound of confidence interval.
    high: Upper bound of confidence interval.

  Raises:
    ValueError: If `samples` has rank other than 1 (batch semantics
      are not implemented), or if `samples` contains values other than
      `low` or `high` (as that makes the distribution not Bernoulli).
  """
  # TODO(b/78025336) Migrate this confidence interval function
  # to statistical_testing.py.  In order to do that
  # - Get the binomial CDF from the Binomial distribution
  # - Implement scalar root finding in TF.  Batch bisection search
  #   shouldn't be too hard, and is definitely good enough for this
  #   problem.  Batching the Brent algorithm (from scipy) that is used
  #   here may be more involved, but may also not be necessary---it's
  #   only used here because scipy made it convenient.  In particular,
  #   robustness is more important than speed here, which may make
  #   bisection search actively better.
  # - The rest is just a matter of rewriting in the appropriate style.
  if optimize is None or stats is None:
    raise ValueError(
        "Scipy is required for computing Clopper-Pearson confidence intervals")
  if len(samples.shape) != 1:
    raise ValueError("Batch semantics not implemented")
  n = len(samples)
  low = np.amin(samples)
  high = np.amax(samples)
  successes = np.count_nonzero(samples - low)
  failures = np.count_nonzero(samples - high)
  if successes + failures != n:
    uniques = np.unique(samples)
    msg = ("Purportedly Bernoulli distribution had distinct samples"
           " {}, {}, and {}".format(uniques[0], uniques[1], uniques[2]))
    raise ValueError(msg)
  def p_small_enough(p):
    prob = stats.binom.logcdf(successes, n, p)
    return prob - np.log(error_rate / 2.)
  def p_big_enough(p):
    prob = stats.binom.logsf(successes, n, p)
    return prob - np.log(error_rate / 2.)
  high_p = optimize.brentq(
      p_small_enough, float(successes) / n, 1., rtol=1e-9)
  low_p = optimize.brentq(
      p_big_enough, 0., float(successes) / n, rtol=1e-9)
  low_interval = low + (high - low) * low_p
  high_interval = low + (high - low) * high_p
  return (low_interval, high_interval)

def compute_true_volumes(
    det_bounds, dim, num_samples, error_rate=1e-6, seed=42):
  """Returns confidence intervals for the desired correlation matrix volumes.

  The confidence intervals are computed by the [Clopper-Pearson method]
  (https://en.wikipedia.org/wiki/Binomial_proportion_confidence_interval).

  Args:
    det_bounds: A rank-1 numpy array of lower bounds on the
      determinants of acceptable matrices.  Entries must be unique.
    dim: A Python `int` dimension of correlation matrices to sample.
    num_samples: The number of samples to draw.
    error_rate: The statistical significance of the returned
      confidence intervals.  The significance is broadcast: Each
      returned interval separately may be incorrect with probability
      (under the sample of correlation-like matrices drawn internally)
      at most `error_rate`.
    seed: Random seed.

  Returns:
    bounds: A Python `dict` mapping each determinant bound to the low, high
      tuple giving the confidence interval.
  """
  bounds = {}
  with tf.compat.v1.Session() as sess:
    rej_weights, _ = correlation_matrix_volume_rejection_samples(
        det_bounds, dim, [num_samples, len(det_bounds)], np.float32, seed=seed)
    rej_weights = sess.run(rej_weights)
    for rw, det in zip(np.rollaxis(rej_weights, 1), det_bounds):
      template = ("Estimating volume of {}x{} correlation "
                  "matrices with determinant >= {}.")
      print(template.format(dim, dim, det))
      sys.stdout.flush()
      bounds[det] = _clopper_pearson_confidence_interval(
          rw, error_rate=error_rate)
    return bounds


# Copyright 2018 The TensorFlow Probability Authors.
#
# 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.
# ============================================================================
"""The von Mises distribution class."""

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

# Dependency imports
import numpy as np
import tensorflow.compat.v2 as tf
from tensorflow_probability.python.distributions import distribution
from tensorflow_probability.python.distributions import kullback_leibler
from tensorflow_probability.python.distributions import normal
from tensorflow_probability.python.distributions.seed_stream import SeedStream
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import distribution_util
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import reparameterization
from tensorflow_probability.python.math.gradient import value_and_gradient

class VonMises(distribution.Distribution):
  """The von Mises distribution over angles.

  The von Mises distribution is a univariate directional distribution.
  Similarly to Normal distribution, it is a maximum entropy distribution.
  The samples of this distribution are angles, measured in radians.
  They are 2 pi-periodic: x = 0 and x = 2pi are equivalent.
  This means that the density is also 2 pi-periodic.
  The generated samples, however, are guaranteed to be in [-pi, pi) range.

  When `concentration = 0`, this distribution becomes a Uniform distribuion on
  the [-pi, pi) domain.

  The von Mises distribution is a special case of von Mises-Fisher distribution
  for n=2. However, the TFP's VonMisesFisher implementation represents the
  samples and location as (x, y) points on a circle, while VonMises represents
  them as scalar angles.

  #### Mathematical details

  The probability density function (pdf) of this distribution is,

  ```none
  pdf(x; loc, concentration) = exp(concentration cos(x - loc)) / Z
  Z = 2 * pi * I_0 (concentration)
  ```

  where:
  * `I_0 (concentration)` is the modified Bessel function of order zero;
  * `loc` the circular mean of the distribution, a scalar. It can take arbitrary
    values, but it is 2pi-periodic: loc and loc + 2pi result in the same
    distribution.
  * `concentration >= 0` parameter is the concentration parameter. When
  `concentration = 0`,
    this distribution becomes a Uniform distribution on [-pi, pi).

  #### Examples

  Examples of initialization of one or a batch of distributions.

  ```python
  # Create a batch of three von Mises distributions.
  loc = [1, 2, 3]
  concentration = [1, 2, 3]
  dist = tfp.distributions.VonMises(loc=[1.0, 2.0], concentration=[0.5, 2.0])

  dist.sample([3])  # Shape [3, 2]
  ```

  Arguments are broadcast when possible.

  ```python
  dist = tfp.distributions.VonMises(loc=1.0, concentration=[0.5, 2.0])

  # Evaluate the pdf of both distributions on the same point, 3.0,
  # returning a length 2 tensor.
  dist.prob(3.0)
  ```

  """

  def __init__(self,
               loc,
               concentration,
               validate_args=False,
               allow_nan_stats=True,
               name="VonMises"):
    """Construct von Mises distributions with given location and concentration.

    The parameters `loc` and `concentration` must be shaped in a way that
    supports broadcasting (e.g. `loc + concentration` is a valid operation).

    Args:
      loc: Floating point tensor, the circular means of the distribution(s).
      concentration: Floating point tensor, the level of concentration of the
        distribution(s) around `loc`. Must take non-negative values.
        `concentration = 0` defines a Uniform distribution, while
        `concentration = +inf` indicates a Deterministic distribution at `loc`.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
        result is undefined. When `False`, an exception is raised if one or more
        of the statistic's batch members are undefined.
      name: Python `str` name prefixed to Ops created by this class.

    Raises:
      TypeError: if loc and concentration are different dtypes.
    """
    parameters = dict(locals())
    with tf.name_scope(name) as name:
      dtype = dtype_util.common_dtype([loc, concentration],
                                      preferred_dtype=tf.float32)
      loc = tf.convert_to_tensor(value=loc, name="loc", dtype=dtype)
      concentration = tf.convert_to_tensor(
          value=concentration, name="concentration", dtype=dtype)
      with tf.control_dependencies([
          assert_util.assert_non_negative(concentration)
      ] if validate_args else []):
        self._loc = tf.identity(loc, name="loc")
        self._concentration = tf.identity(concentration, name="concentration")
        dtype_util.assert_same_float_dtype([self._loc, self._concentration])
    super(VonMises, self).__init__(
        dtype=self._concentration.dtype,
        reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        parameters=parameters,
        graph_parents=[self._loc, self._concentration],
        name=name)

  @staticmethod
  def _param_shapes(sample_shape):
    return dict(
        zip(("loc", "concentration"),
            ([tf.convert_to_tensor(value=sample_shape, dtype=tf.int32)] * 2)))

  @classmethod
  def _params_event_ndims(cls):
    return dict(loc=0, concentration=0)

  @property
  def loc(self):
    """Distribution parameter for the circular mean (loc)."""
    return self._loc

  @property
  def concentration(self):
    """Distribution parameter for the concentration."""
    return self._concentration

  def _batch_shape_tensor(self):
    return tf.broadcast_dynamic_shape(
        tf.shape(input=self.loc), tf.shape(input=self.concentration))

  def _batch_shape(self):
    return tf.broadcast_static_shape(self.loc.shape, self.concentration.shape)

  def _event_shape_tensor(self):
    return tf.constant([], dtype=tf.int32)

  def _event_shape(self):
    return tf.TensorShape([])

  def _log_prob(self, x):
    return self._log_unnormalized_prob(x) - self._log_normalization()

  def _log_unnormalized_prob(self, x):
    z = self._z(x)
    return self.concentration * (tf.cos(z) - 1)

  def _log_normalization(self):
    return np.log(2. * np.pi) + tf.math.log(
        tf.math.bessel_i0e(self.concentration))

  def _prob(self, x):
    unnormalized_prob = tf.exp(self._log_unnormalized_prob(x))
    normalization = (2. * np.pi) * tf.math.bessel_i0e(self.concentration)
    return unnormalized_prob / normalization

  def _cdf(self, x):
    z = self._z(x) + tf.zeros_like(self.concentration)
    concentration = self.concentration + tf.zeros_like(z)
    return von_mises_cdf(z, concentration)

  def _entropy(self):
    i0e = tf.math.bessel_i0e(self.concentration)
    i1e = tf.math.bessel_i1e(self.concentration)
    entropy = (
        self.concentration * (1 - i1e / i0e) + tf.math.log(i0e) +
        np.log(2 * np.pi))
    return entropy + tf.zeros_like(self.loc)

  @distribution_util.AppendDocstring(
      """Note: This function computes the circular mean, defined as
      atan2(E sin(x), E cos(x)). The values are in the [-pi, pi] range.""")
  def _mean(self):
    return self.loc + tf.zeros_like(self.concentration)

  def _mode(self):
    return self.mean()

  @distribution_util.AppendDocstring(
      """Note: This function computes the circular variance defined as
      1 - E cos(x - loc). The values are in the [0, 1] range.""")
  def _variance(self):
    concentration = self.concentration + tf.zeros_like(self.loc)
    return 1. - tf.math.bessel_i1e(concentration) / tf.math.bessel_i0e(
        concentration)

  def _z(self, x):
    """Standardize input `x` to a zero-loc von Mises."""
    with tf.name_scope("standardize"):
      return x - self.loc

  def _sample_n(self, n, seed=None):
    # random_von_mises does not work for zero concentration, so round it up to
    # something very small.
    tiny = np.finfo(dtype_util.as_numpy_dtype(self.dtype)).tiny
    concentration = tf.maximum(self.concentration, tiny)

    sample_batch_shape = tf.concat([[n], self._batch_shape_tensor()], axis=0)
    samples = random_von_mises(
        sample_batch_shape, concentration, dtype=self.dtype, seed=seed)

    # vonMises(0, concentration) -> vonMises(loc, concentration)
    samples += self.loc
    # Map the samples to [-pi, pi].
    samples -= 2. * np.pi * tf.round(samples / (2. * np.pi))
    return samples

@kullback_leibler.RegisterKL(VonMises, VonMises)
def _kl_von_mises_von_mises(d1, d2, name=None):
  """Batchwise KL divergence KL(d1 || d2) with d1 and d2 von Mises.

  Args:
    d1: instance of a von Mises distribution object.
    d2: instance of a a von Mises distribution object.
    name: (optional) Name to use for created operations.
      default is "kl_von_mises_von_mises".

  Returns:
    Batchwise KL(d1 || d2)
  """
  with tf.name_scope(name or "kl_von_mises_von_mises"):
    # The density of von Mises is (abbreviating the concentration for conc):
    #   vonMises(x; loc, conc) = exp(conc cos(x - loc)) / (2 pi I_0 (conc) )
    # We need two properties:
    # 1. Standardization: if z ~ vonMises(0, conc), then
    #    z + loc ~ vonMises(loc, conc).
    # 2. Expectation of cosine:
    #    E_q(z | 0, conc) cos z = I_1 (conc) / I_0 (conc)
    # Now,
    # KL(d1 || d2)
    #   = E_vonMises(x; loc1, conc1) log vonMises(x; loc1, conc1)
    #                                     / vonMises(x; loc2, conc2)
    # Plugging the densities and rearranging, we have
    #   log I_0(conc2) / I_0(conc1)
    #     + E_vonMises(x; loc1, conc1) [ conc1 cos (z - loc1)
    #                                     - conc2 cos (z - loc2) ]
    # Let's transform the second term using the standardization property:
    #   E_vonMises(x; 0, conc1) [conc1 cos z - conc2 cos (z - (loc2 - loc1))]
    # Applying the cos (x - y) = cos x cos y + sin x sin y expansion, we get
    #   E_vonMises(x; 0, conc1) [conc1 cos z - conc2 cos (loc2 - loc1) cos z
    #     - conc2 sin(loc2 - loc1) sin z]
    # Because the distribution is symmetric around zero, the last term vanishes
    # in expectation. The remaining two terms are computed using the
    # "expectation of cosine" property:
    #   (conc1 - conc2 cos (loc2 - loc1) E_vonMises(x; 0, conc1) cos z
    #     = (conc1 - conc2 cos (loc2 - loc1)) I_1(conc1) / I_0(conc1)
    # In total, we have
    #   KL(d1 || d2) = log I_0(conc2) / I_0(conc1)
    #     + (conc1 - conc2 cos (loc2 - loc1)) I_1(conc1) / I_0(conc1)
    # To improve the numerical stability, we can replace I_j(k) functions with
    # the exponentially scaled versions using the equality
    # I_j(k) = I_j^E(k) exp(k) (which holds for k >= 0):
    #   KL(d1 || d2) = (conc2 - conc1) + log I_0^E(conc2) / I_0^E(conc1)
    #     + (conc1 - conc2 cos (loc2 - loc1)) I_1^E(conc1) / I_0^E(conc1)
    # Note that this formula is numerically stable for conc1 = 0 and/or
    # conc2 = 0 because I_0 (0) = I_0^E (0) = 1.
    i0e_concentration1 = tf.math.bessel_i0e(d1.concentration)
    i1e_concentration1 = tf.math.bessel_i1e(d1.concentration)
    i0e_concentration2 = tf.math.bessel_i0e(d2.concentration)
    return ((d2.concentration - d1.concentration) +
            tf.math.log(i0e_concentration2 / i0e_concentration1) +
            (d1.concentration - d2.concentration * tf.cos(d1.loc - d2.loc)) *
            (i1e_concentration1 / i0e_concentration1))

@tf.custom_gradient
def von_mises_cdf(x, concentration):
  """Computes the cumulative density function (CDF) of von Mises distribution.

  Denote the density of vonMises(loc=0, concentration=concentration) by p(t).
  Note that p(t) is periodic, p(t) = p(t + 2 pi).
  The CDF at the point x is defined as int_{-pi}^x p(t) dt.
  Thus, when x in [-pi, pi], the CDF is in [0, 1]; when x is in [pi, 3pi], the
  CDF is in [1, 2], etc.

  The CDF is not available in closed form. Instead, we use the method [1]
  which uses either a series expansion or a Normal approximation, depending on
  the value of concentration.

  We also compute the derivative of the CDF w.r.t. both x and concentration.
  The derivative w.r.t. x is p(x), while the derivative w.r.t. concentration is
  computed
  using automatic differentiation. We use forward mode for the series case
  (which allows to save memory) and backward mode for the Normal approximation.

  Arguments:
    x: The point at which to evaluate the CDF.
    concentration: The concentration parameter of the von Mises distribution.

  Returns:
    The value of the CDF computed elementwise.

  References:
    [1] G. Hill "Algorithm 518: Incomplete Bessel Function I_0. The Von Mises
    Distribution." ACM Transactions on Mathematical Software, 1977
  """
  x = tf.convert_to_tensor(value=x)
  concentration = tf.convert_to_tensor(value=concentration)
  dtype = x.dtype

  # Map x to [-pi, pi].
  num_periods = tf.round(x / (2. * np.pi))
  x -= (2. * np.pi) * num_periods

  # We take the hyperparameters from Table I of [1], the row for D=8
  # decimal digits of accuracy. ck is the cut-off for concentration:
  # if concentration < ck,  the series expansion is used;
  # otherwise, the Normal approximation is used.
  ck = 10.5
  # The number of terms in the series expansion. [1] chooses it as a function
  # of concentration, n(concentration). This is hard to implement in TF.
  # Instead, we upper bound it over concentrations:
  #   num_terms = ceil ( max_{concentration <= ck} n(concentration) ).
  # The maximum is achieved for concentration = ck.
  num_terms = 20

  cdf_series, dcdf_dconcentration_series = _von_mises_cdf_series(
      x, concentration, num_terms, dtype)
  cdf_normal, dcdf_dconcentration_normal = _von_mises_cdf_normal(
      x, concentration, dtype)

  use_series = concentration < ck
  cdf = tf.where(use_series, cdf_series, cdf_normal)
  cdf += num_periods
  dcdf_dconcentration = tf.where(use_series, dcdf_dconcentration_series,
                                 dcdf_dconcentration_normal)

  def grad(dy):
    prob = tf.exp(concentration * (tf.cos(x) - 1.)) / (
        (2. * np.pi) * tf.math.bessel_i0e(concentration))
    return dy * prob, dy * dcdf_dconcentration

  return cdf, grad

def _von_mises_cdf_series(x, concentration, num_terms, dtype):
  """Computes the von Mises CDF and its derivative via series expansion."""
  # Keep the number of terms as a float. It should be a small integer, so
  # exactly representable as a float.
  num_terms = tf.cast(num_terms, dtype=dtype)

  def loop_body(n, rn, drn_dconcentration, vn, dvn_dconcentration):
    """One iteration of the series loop."""

    denominator = 2. * n / concentration + rn
    ddenominator_dk = -2. * n / concentration ** 2 + drn_dconcentration
    rn = 1. / denominator
    drn_dconcentration = -ddenominator_dk / denominator ** 2

    multiplier = tf.sin(n * x) / n + vn
    vn = rn * multiplier
    dvn_dconcentration = (drn_dconcentration * multiplier +
                          rn * dvn_dconcentration)
    n -= 1.

    return n, rn, drn_dconcentration, vn, dvn_dconcentration

  (_, _, _, vn, dvn_dconcentration) = tf.while_loop(
      cond=lambda n, *_: n > 0.,
      body=loop_body,
      loop_vars=(
          num_terms,  # n
          tf.zeros_like(x, name="rn"),
          tf.zeros_like(x, name="drn_dconcentration"),
          tf.zeros_like(x, name="vn"),
          tf.zeros_like(x, name="dvn_dconcentration"),
      ),
  )

  cdf = .5 + x / (2. * np.pi) + vn / np.pi
  dcdf_dconcentration = dvn_dconcentration / np.pi

  # Clip the result to [0, 1].
  cdf_clipped = tf.clip_by_value(cdf, 0., 1.)
  # The clipped values do not depend on concentration anymore, so set their
  # derivative to zero.
  dcdf_dconcentration *= tf.cast((cdf >= 0.) & (cdf <= 1.), dtype)

  return cdf_clipped, dcdf_dconcentration

def _von_mises_cdf_normal(x, concentration, dtype):
  """Computes the von Mises CDF and its derivative via Normal approximation."""

  def cdf_func(concentration):
    """A helper function that is passed to value_and_gradient."""
    # z is an "almost Normally distributed" random variable.
    z = ((np.sqrt(2. / np.pi) / tf.math.bessel_i0e(concentration)) *
         tf.sin(.5 * x))

    # This is the correction described in [1] which reduces the error
    # of the Normal approximation.
    z2 = z ** 2
    z3 = z2 * z
    z4 = z2 ** 2
    c = 24. * concentration
    c1 = 56.

    xi = z - z3 / ((c - 2. * z2 - 16.) / 3. -
                   (z4 + (7. / 4.) * z2 + 167. / 2.) / (c - c1 - z2 + 3.)) ** 2

    distrib = normal.Normal(tf.cast(0., dtype), tf.cast(1., dtype))

    return distrib.cdf(xi)

  return value_and_gradient(cdf_func, concentration)

def random_von_mises(shape, concentration, dtype=tf.float32, seed=None):
  """Samples from the standardized von Mises distribution.

  The distribution is vonMises(loc=0, concentration=concentration), so the mean
  is zero.
  The location can then be changed by adding it to the samples.

  The sampling algorithm is rejection sampling with wrapped Cauchy proposal [1].
  The samples are pathwise differentiable using the approach of [2].

  Arguments:
    shape: The output sample shape.
    concentration: The concentration parameter of the von Mises distribution.
    dtype: The data type of concentration and the outputs.
    seed: (optional) The random seed.

  Returns:
    Differentiable samples of standardized von Mises.

  References:
    [1] Luc Devroye "Non-Uniform Random Variate Generation", Springer-Verlag,
    1986; Chapter 9, p. 473-476.
    http://www.nrbook.com/devroye/Devroye_files/chapter_nine.pdf
    + corrections http://www.nrbook.com/devroye/Devroye_files/errors.pdf
    [2] Michael Figurnov, Shakir Mohamed, Andriy Mnih. "Implicit
    Reparameterization Gradients", 2018.
  """
  seed = SeedStream(seed, salt="von_mises")
  concentration = tf.convert_to_tensor(
      value=concentration, dtype=dtype, name="concentration")

  @tf.custom_gradient
  def rejection_sample_with_gradient(concentration):
    """Performs rejection sampling for standardized von Mises.

    A nested function is required because @tf.custom_gradient does not handle
    non-tensor inputs such as dtype. Instead, they are captured by the outer
    scope.

    Arguments:
      concentration: The concentration parameter of the distribution.

    Returns:
      Differentiable samples of standardized von Mises.
    """
    r = 1. + tf.sqrt(1. + 4. * concentration ** 2)
    rho = (r - tf.sqrt(2. * r)) / (2. * concentration)

    s_exact = (1. + rho ** 2) / (2. * rho)

    # For low concentration, s becomes numerically unstable.
    # To fix that, we use an approximation. Here is the derivation.
    # First-order Taylor expansion at conc = 0 gives
    #   sqrt(1 + 4 concentration^2) ~= 1 + (2 concentration)^2 / 2.
    # Therefore, r ~= 2 + 2 concentration. By plugging this into rho, we have
    #   rho ~= conc + 1 / conc - sqrt(1 + 1 / concentration^2).
    # Let's expand the last term at concentration=0 up to the linear term:
    #   sqrt(1 + 1 / concentration^2) ~= 1 / concentration + concentration / 2
    # Thus, rho ~= concentration / 2. Finally,
    #   s = 1 / (2 rho) + rho / 2 ~= 1 / concentration + concentration / 4.
    # Since concentration is small, we drop the second term and simply use
    #   s ~= 1 / concentration.
    s_approximate = 1. / concentration

    # To compute the cutoff, we compute s_exact using mpmath with 30 decimal
    # digits precision and compare that to the s_exact and s_approximate
    # computed with dtype. Then, the cutoff is the largest concentration for
    # which abs(s_exact - s_exact_mpmath) > abs(s_approximate - s_exact_mpmath).
    s_concentration_cutoff_dict = {
        tf.float16: 1.8e-1,
        tf.float32: 2e-2,
        tf.float64: 1.2e-4,
    }
    s_concentration_cutoff = s_concentration_cutoff_dict[dtype]

    s = tf.where(concentration > s_concentration_cutoff, s_exact, s_approximate)

    def loop_body(done, u, w):
      """Resample the non-accepted points."""
      # We resample u each time completely. Only its sign is used outside the
      # loop, which is random.
      u = tf.random.uniform(
          shape, minval=-1., maxval=1., dtype=dtype, seed=seed())
      z = tf.cos(np.pi * u)
      # Update the non-accepted points.
      w = tf.where(done, w, (1. + s * z) / (s + z))
      y = concentration * (s - w)

      v = tf.random.uniform(
          shape, minval=0., maxval=1., dtype=dtype, seed=seed())
      accept = (y * (2. - y) >= v) | (tf.math.log(y / v) + 1. >= y)

      return done | accept, u, w

    _, u, w = tf.while_loop(
        cond=lambda done, *_: ~tf.reduce_all(input_tensor=done),
        body=loop_body,
        loop_vars=(
            tf.zeros(shape, dtype=tf.bool, name="done"),
            tf.zeros(shape, dtype=dtype, name="u"),
            tf.zeros(shape, dtype=dtype, name="w"),
        ),
        # The expected number of iterations depends on concentration.
        # It monotonically increases from one iteration for concentration = 0 to
        # sqrt(2 pi / e) ~= 1.52 iterations for concentration = +inf [1].
        # We use a limit of 100 iterations to avoid infinite loops
        # for very large / nan concentration.
        maximum_iterations=100,
        parallel_iterations=1 if seed.original_seed is None else 10,
    )

    x = tf.sign(u) * tf.math.acos(w)

    def grad(dy):
      """The gradient of the von Mises samples w.r.t. concentration."""
      broadcast_concentration = concentration + tf.zeros_like(x)
      _, dcdf_dconcentration = value_and_gradient(
          lambda conc: von_mises_cdf(x, conc), broadcast_concentration)
      inv_prob = tf.exp(-broadcast_concentration * (tf.cos(x) - 1.)) * (
          (2. * np.pi) * tf.math.bessel_i0e(broadcast_concentration))
      # Compute the implicit reparameterization gradient [2],
      # dz/dconc = -(dF(z; conc) / dconc) / p(z; conc)
      ret = dy * (-inv_prob * dcdf_dconcentration)
      # Sum over the sample dimensions. Assume that they are always the first
      # ones.
      num_sample_dimensions = (tf.rank(broadcast_concentration) -
                               tf.rank(concentration))
      return tf.reduce_sum(
          input_tensor=ret, axis=tf.range(num_sample_dimensions))

    return x, grad

  return rejection_sample_with_gradient(concentration)