# 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.
# ============================================================================
"""Interpolation Ops."""

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

import itertools

# Dependency imports
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__ = [
    'interp_regular_1d_grid',
    'batch_interp_regular_1d_grid',
    'batch_interp_regular_nd_grid',
]

def _interp_regular_1d_grid_impl(x,
                                 x_ref_min,
                                 x_ref_max,
                                 y_ref,
                                 axis=-1,
                                 batch_y_ref=False,
                                 fill_value='constant_extension',
                                 fill_value_below=None,
                                 fill_value_above=None,
                                 grid_regularizing_transform=None,
                                 name=None):
  """1-D interpolation that works with/without batching."""
  # To understand the implemention differences between the batch/no-batch
  # versions of this function, you should probably understand the difference
  # between tf.gather and tf.batch_gather.  In particular, we do *not* make the
  # no-batch version a special case of the batch version, because that would
  # an inefficient use of batch_gather with unnecessarily broadcast args.
  with tf.compat.v1.name_scope(
      name,
      values=[
          x, x_ref_min, x_ref_max, y_ref, axis, fill_value, fill_value_below,
          fill_value_above
      ]):

    # Arg checking.
    allowed_fv_st = ('constant_extension', 'extrapolate')
    for fv in (fill_value, fill_value_below, fill_value_above):
      if isinstance(fv, str) and fv not in allowed_fv_st:
        raise ValueError(
            'A fill value ({}) was not an allowed string ({})'.format(
                fv, allowed_fv_st))

    # Separate value fills for below/above incurs extra cost, so keep track of
    # whether this is needed.
    need_separate_fills = (
        fill_value_above is not None or fill_value_below is not None or
        fill_value == 'extrapolate'  # always requries separate below/above
    )
    if need_separate_fills and fill_value_above is None:
      fill_value_above = fill_value
    if need_separate_fills and fill_value_below is None:
      fill_value_below = fill_value

    dtype = dtype_util.common_dtype([x, x_ref_min, x_ref_max, y_ref],
                                    preferred_dtype=tf.float32)
    x = tf.convert_to_tensor(value=x, name='x', dtype=dtype)

    x_ref_min = tf.convert_to_tensor(
        value=x_ref_min, name='x_ref_min', dtype=dtype)
    x_ref_max = tf.convert_to_tensor(
        value=x_ref_max, name='x_ref_max', dtype=dtype)
    if not batch_y_ref:
      _assert_ndims_statically(x_ref_min, expect_ndims=0)
      _assert_ndims_statically(x_ref_max, expect_ndims=0)

    y_ref = tf.convert_to_tensor(value=y_ref, name='y_ref', dtype=dtype)

    if batch_y_ref:
      # If we're batching,
      #   x.shape ~ [A1,...,AN, D],  x_ref_min/max.shape ~ [A1,...,AN]
      # So to add together we'll append a singleton.
      # If not batching, x_ref_min/max are scalar, so this isn't an issue,
      # moreover, if not batching, x can be scalar, and expanding x_ref_min/max
      # would cause a bad expansion of x when added to x (confused yet?).
      x_ref_min = x_ref_min[..., tf.newaxis]
      x_ref_max = x_ref_max[..., tf.newaxis]

    axis = tf.convert_to_tensor(value=axis, name='axis', dtype=tf.int32)
    axis = distribution_util.make_non_negative_axis(axis, tf.rank(y_ref))
    _assert_ndims_statically(axis, expect_ndims=0)

    ny = tf.cast(tf.shape(input=y_ref)[axis], dtype)

    # Map [x_ref_min, x_ref_max] to [0, ny - 1].
    # This is the (fractional) index of x.
    if grid_regularizing_transform is None:
      g = lambda x: x
    else:
      g = grid_regularizing_transform
    fractional_idx = ((g(x) - g(x_ref_min)) / (g(x_ref_max) - g(x_ref_min)))
    x_idx_unclipped = fractional_idx * (ny - 1)

    # Wherever x is NaN, x_idx_unclipped will be NaN as well.
    # Keep track of the nan indices here (so we can impute NaN later).
    # Also eliminate any NaN indices, since there is not NaN in 32bit.
    nan_idx = tf.math.is_nan(x_idx_unclipped)
    x_idx_unclipped = tf.where(nan_idx, tf.zeros_like(x_idx_unclipped),
                               x_idx_unclipped)

    x_idx = tf.clip_by_value(x_idx_unclipped, tf.zeros((), dtype=dtype), ny - 1)

    # Get the index above and below x_idx.
    # Naively we could set idx_below = floor(x_idx), idx_above = ceil(x_idx),
    # however, this results in idx_below == idx_above whenever x is on a grid.
    # This in turn results in y_ref_below == y_ref_above, and then the gradient
    # at this point is zero.  So here we "jitter" one of idx_below, idx_above,
    # so that they are at different values.  This jittering does not affect the
    # interpolated value, but does make the gradient nonzero (unless of course
    # the y_ref values are the same).
    idx_below = tf.floor(x_idx)
    idx_above = tf.minimum(idx_below + 1, ny - 1)
    idx_below = tf.maximum(idx_above - 1, 0)

    # These are the values of y_ref corresponding to above/below indices.
    idx_below_int32 = tf.cast(idx_below, dtype=tf.int32)
    idx_above_int32 = tf.cast(idx_above, dtype=tf.int32)
    if batch_y_ref:
      # If y_ref.shape ~ [A1,...,AN, C, B1,...,BN],
      # and x.shape, x_ref_min/max.shape ~ [A1,...,AN, D]
      # Then y_ref_below.shape ~ [A1,...,AN, D, B1,...,BN]
      y_ref_below = _batch_gather_with_broadcast(y_ref, idx_below_int32, axis)
      y_ref_above = _batch_gather_with_broadcast(y_ref, idx_above_int32, axis)
    else:
      # Here, y_ref_below.shape =
      #   y_ref.shape[:axis] + x.shape + y_ref.shape[axis + 1:]
      y_ref_below = tf.gather(y_ref, idx_below_int32, axis=axis)
      y_ref_above = tf.gather(y_ref, idx_above_int32, axis=axis)

    # Use t to get a convex combination of the below/above values.
    t = x_idx - idx_below

    # x, and tensors shaped like x, need to be added to, and selected with
    # (using tf.where) the output y.  This requires appending singletons.
    # Make functions appropriate for batch/no-batch.
    if batch_y_ref:
      # In the non-batch case, the output shape is going to be
      #   y_ref.shape[:axis] + x.shape + y_ref.shape[axis+1:]
      expand_x_fn = _make_expand_x_fn_for_batch_interpolation(y_ref, axis)
    else:
      # In the batch case, the output shape is going to be
      #   Broadcast(y_ref.shape[:axis], x.shape[:-1]) +
      #   x.shape[-1:] +  y_ref.shape[axis+1:]
      expand_x_fn = _make_expand_x_fn_for_non_batch_interpolation(y_ref, axis)

    t = expand_x_fn(t)
    nan_idx = expand_x_fn(nan_idx, broadcast=True)
    x_idx_unclipped = expand_x_fn(x_idx_unclipped, broadcast=True)

    y = t * y_ref_above + (1 - t) * y_ref_below

    # Now begins a long excursion to fill values outside [x_min, x_max].

    # Re-insert NaN wherever x was NaN.
    y = tf.where(nan_idx,
                 tf.fill(tf.shape(input=y), tf.constant(np.nan, y.dtype)), y)

    if not need_separate_fills:
      if fill_value == 'constant_extension':
        pass  # Already handled by clipping x_idx_unclipped.
      else:
        y = tf.where((x_idx_unclipped < 0) | (x_idx_unclipped > ny - 1),
                     fill_value + tf.zeros_like(y), y)
    else:
      # Fill values below x_ref_min <==> x_idx_unclipped < 0.
      if fill_value_below == 'constant_extension':
        pass  # Already handled by the clipping that created x_idx_unclipped.
      elif fill_value_below == 'extrapolate':
        if batch_y_ref:
          # For every batch member, gather the first two elements of y across
          # `axis`.
          y_0 = tf.gather(y_ref, [0], axis=axis)
          y_1 = tf.gather(y_ref, [1], axis=axis)
        else:
          # If not batching, we want to gather the first two elements, just like
          # above.  However, these results need to be replicated for every
          # member of x.  An easy way to do that is to gather using
          # indices = zeros/ones(x.shape).
          y_0 = tf.gather(
              y_ref, tf.zeros(tf.shape(input=x), dtype=tf.int32), axis=axis)
          y_1 = tf.gather(
              y_ref, tf.ones(tf.shape(input=x), dtype=tf.int32), axis=axis)
        x_delta = (x_ref_max - x_ref_min) / (ny - 1)
        x_factor = expand_x_fn((x - x_ref_min) / x_delta, broadcast=True)
        y = tf.where(x_idx_unclipped < 0, y_0 + x_factor * (y_1 - y_0), y)
      else:
        y = tf.where(x_idx_unclipped < 0, fill_value_below + tf.zeros_like(y),
                     y)
      # Fill values above x_ref_min <==> x_idx_unclipped > ny - 1.
      if fill_value_above == 'constant_extension':
        pass  # Already handled by the clipping that created x_idx_unclipped.
      elif fill_value_above == 'extrapolate':
        ny_int32 = tf.shape(input=y_ref)[axis]
        if batch_y_ref:
          y_n1 = tf.gather(y_ref, [tf.shape(input=y_ref)[axis] - 1], axis=axis)
          y_n2 = tf.gather(y_ref, [tf.shape(input=y_ref)[axis] - 2], axis=axis)
        else:
          y_n1 = tf.gather(
              y_ref, tf.fill(tf.shape(input=x), ny_int32 - 1), axis=axis)
          y_n2 = tf.gather(
              y_ref, tf.fill(tf.shape(input=x), ny_int32 - 2), axis=axis)
        x_delta = (x_ref_max - x_ref_min) / (ny - 1)
        x_factor = expand_x_fn((x - x_ref_max) / x_delta, broadcast=True)
        y = tf.where(x_idx_unclipped > ny - 1,
                     y_n1 + x_factor * (y_n1 - y_n2),
                     y)
      else:
        y = tf.where(x_idx_unclipped > ny - 1,
                     fill_value_above + tf.zeros_like(y), y)

    return y

def interp_regular_1d_grid(x,
                           x_ref_min,
                           x_ref_max,
                           y_ref,
                           axis=-1,
                           fill_value='constant_extension',
                           fill_value_below=None,
                           fill_value_above=None,
                           grid_regularizing_transform=None,
                           name=None):
  """Linear `1-D` interpolation on a regular (constant spacing) grid.

  Given reference values, this function computes a piecewise linear interpolant
  and evaluates it on a new set of `x` values.

  The interpolant is built from `C` reference values indexed by one dimension
  of `y_ref` (specified by the `axis` kwarg).

  If `y_ref` is a vector, then each value `y_ref[i]` is considered to be equal
  to `f(x_ref[i])`, for `C` (implicitly defined) reference values between
  `x_ref_min` and `x_ref_max`:

  ```none
  x_ref[i] = x_ref_min + i * (x_ref_max - x_ref_min) / (C - 1),
  i = 0, ..., C - 1.
  ```

  If `rank(y_ref) > 1`, then dimension `axis` indexes `C` reference values of
  a shape `y_ref.shape[:axis] + y_ref.shape[axis + 1:]` `Tensor`.

  If `rank(x) > 1`, then the output is obtained by effectively flattening `x`,
  interpolating along `axis`, then expanding the result to shape
  `y_ref.shape[:axis] + x.shape + y_ref.shape[axis + 1:]`.

  These shape semantics are equivalent to `scipy.interpolate.interp1d`.

  Args:
    x: Numeric `Tensor` The x-coordinates of the interpolated output values.
    x_ref_min:  Scalar `Tensor` of same `dtype` as `x`.  The minimum value of
      the (implicitly defined) reference `x_ref`.
    x_ref_max:  Scalar `Tensor` of same `dtype` as `x`.  The maximum value of
      the (implicitly defined) reference `x_ref`.
    y_ref:  `N-D` `Tensor` (`N > 0`) of same `dtype` as `x`. The reference
      output values.
    axis:  Scalar `Tensor` designating the dimension of `y_ref` that indexes
      values of the interpolation table.
      Default value: `-1`, the rightmost axis.
    fill_value:  Determines what values output should take for `x` values that
      are below `x_ref_min` or above `x_ref_max`. `Tensor` or one of the strings
      "constant_extension" ==> Extend as constant function. "extrapolate" ==>
      Extrapolate in a linear fashion.
      Default value: `"constant_extension"`
    fill_value_below:  Optional override of `fill_value` for `x < x_ref_min`.
    fill_value_above:  Optional override of `fill_value` for `x > x_ref_max`.
    grid_regularizing_transform:  Optional transformation `g` which regularizes
      the implied spacing of the x reference points.  In other words, if
      provided, we assume `g(x_ref_i)` is a regular grid between `g(x_ref_min)`
      and `g(x_ref_max)`.
    name:  A name to prepend to created ops.
      Default value: `"interp_regular_1d_grid"`.

  Returns:
    y_interp:  Interpolation between members of `y_ref`, at points `x`.
      `Tensor` of same `dtype` as `x`, and shape
      `y.shape[:axis] + x.shape + y.shape[axis + 1:]`

  Raises:
    ValueError:  If `fill_value` is not an allowed string.
    ValueError:  If `axis` is not a scalar.

  #### Examples

  Interpolate a function of one variable:

  ```python
  y_ref = tf.exp(tf.linspace(start=0., stop=10., num=200))

  interp_regular_1d_grid(
      x=[6.0, 0.5, 3.3], x_ref_min=0., x_ref_max=1., y_ref=y_ref)
  ==> approx [exp(6.0), exp(0.5), exp(3.3)]
  ```

  Interpolate a matrix-valued function of one variable:

  ```python
  mat_0 = [[1., 0.], [0., 1.]]
  mat_1 = [[0., -1], [1, 0]]
  y_ref = [mat_0, mat_1]

  # Get three output matrices at once.
  tfp.math.interp_regular_1d_grid(
      x=[0., 0.5, 1.], x_ref_min=0., x_ref_max=1., y_ref=y_ref, axis=0)
  ==> [mat_0, 0.5 * mat_0 + 0.5 * mat_1, mat_1]
  ```

  Interpolate a scalar valued function, and get a matrix of results:

  ```python
  y_ref = tf.exp(tf.linspace(start=0., stop=10., num=200))
  x = [[1.1, 1.2], [2.1, 2.2]]
  tfp.math.interp_regular_1d_grid(x, x_ref_min=0., x_ref_max=10., y_ref=y_ref)
  ==> tf.exp(x)
  ```

  Interpolate a function of one variable on a log-spaced grid:

  ```python
  x_ref = tf.exp(tf.linspace(tf.log(1.), tf.log(100000.), num_pts))
  y_ref = tf.log(x_ref + x_ref**2)

  interp_regular_1d_grid(x=[1.1, 2.2], x_ref_min=1., x_ref_max=100000., y_ref,
      grid_regularizing_transform=tf.log)
  ==> [tf.log(1.1 + 1.1**2), tf.log(2.2 + 2.2**2)]
  ```

  """

  return _interp_regular_1d_grid_impl(
      x,
      x_ref_min,
      x_ref_max,
      y_ref,
      axis=axis,
      batch_y_ref=False,
      fill_value=fill_value,
      fill_value_below=fill_value_below,
      fill_value_above=fill_value_above,
      grid_regularizing_transform=grid_regularizing_transform,
      name=name or 'interp_regular_1d_grid')

def batch_interp_regular_1d_grid(x,
                                 x_ref_min,
                                 x_ref_max,
                                 y_ref,
                                 axis=-1,
                                 fill_value='constant_extension',
                                 fill_value_below=None,
                                 fill_value_above=None,
                                 grid_regularizing_transform=None,
                                 name=None):
  """Linear `1-D` interpolation on a regular (constant spacing) grid.

  Given [batch of] reference values, this function computes a piecewise linear
  interpolant and evaluates it on a [batch of] of new `x` values.

  The interpolant is built from `C` reference values indexed by one dimension
  of `y_ref` (specified by the `axis` kwarg).

  If `y_ref` is a vector, then each value `y_ref[i]` is considered to be equal
  to `f(x_ref[i])`, for `C` (implicitly defined) reference values between
  `x_ref_min` and `x_ref_max`:

  ```none
  x_ref[i] = x_ref_min + i * (x_ref_max - x_ref_min) / (C - 1),
  i = 0, ..., C - 1.
  ```

  In the general case, dimensions to the left of `axis` in `y_ref` are broadcast
  with leading dimensions in `x`, `x_ref_min`, `x_ref_max`.

  Args:
    x: Numeric `Tensor` The x-coordinates of the interpolated output values for
      each batch.  Shape broadcasts with `[A1, ..., AN, D]`, `N >= 0`.
    x_ref_min:  `Tensor` of same `dtype` as `x`.  The minimum value of the each
      batch of the (implicitly defined) reference `x_ref`. Shape broadcasts with
      `[A1, ..., AN]`, `N >= 0`.
    x_ref_max:  `Tensor` of same `dtype` as `x`.  The maximum value of the each
      batch of the (implicitly defined) reference `x_ref`. Shape broadcasts with
      `[A1, ..., AN]`, `N >= 0`.
    y_ref:  `Tensor` of same `dtype` as `x`.  The reference output values.
      `y_ref.shape[:axis]` broadcasts with the batch shape `[A1, ..., AN]`, and
      `y_ref.shape[axis:]` is `[C, B1, ..., BM]`, so the trailing dimensions
        index `C` reference values of a rank `M` `Tensor` (`M >= 0`).
    axis:  Scalar `Tensor` designating the dimension of `y_ref` that indexes
      values of the interpolation table.
      Default value: `-1`, the rightmost axis.
    fill_value:  Determines what values output should take for `x` values that
      are below `x_ref_min` or above `x_ref_max`. `Tensor` or one of the strings
      "constant_extension" ==> Extend as constant function. "extrapolate" ==>
      Extrapolate in a linear fashion.
      Default value: `"constant_extension"`
    fill_value_below:  Optional override of `fill_value` for `x < x_ref_min`.
    fill_value_above:  Optional override of `fill_value` for `x > x_ref_max`.
    grid_regularizing_transform:  Optional transformation `g` which regularizes
      the implied spacing of the x reference points.  In other words, if
      provided, we assume `g(x_ref_i)` is a regular grid between `g(x_ref_min)`
      and `g(x_ref_max)`.
    name:  A name to prepend to created ops.
      Default value: `"batch_interp_regular_1d_grid"`.

  Returns:
    y_interp:  Interpolation between members of `y_ref`, at points `x`.
      `Tensor` of same `dtype` as `x`, and shape `[A1, ..., AN, D, B1, ..., BM]`

  Raises:
    ValueError:  If `fill_value` is not an allowed string.
    ValueError:  If `axis` is not a scalar.

  #### Examples

  Interpolate a function of one variable:

  ```python
  y_ref = tf.exp(tf.linspace(start=0., stop=10., 20))

  batch_interp_regular_1d_grid(
      x=[6.0, 0.5, 3.3], x_ref_min=0., x_ref_max=1., y_ref=y_ref)
  ==> approx [exp(6.0), exp(0.5), exp(3.3)]
  ```

  Interpolate a batch of functions of one variable.

  ```python
  # First batch member is an exponential function, second is a log.
  implied_x_ref = [tf.linspace(-3., 3.2, 200), tf.linspace(0.5, 3., 200)]
  y_ref = tf.stack(  # Shape [2, 200], 2 batches, 200 reference values per batch
      [tf.exp(implied_x_ref[0]), tf.log(implied_x_ref[1])], axis=0)

  x = [[-1., 1., 0.],  # Shape [2, 3], 2 batches, 3 values per batch.
       [1., 2., 3.]]

  y = tfp.math.batch_interp_regular_1d_grid(  # Shape [2, 3]
      x,
      x_ref_min=[-3., 0.5],
      x_ref_max=[3.2, 3.],
      y_ref=y_ref,
      axis=-1)

  # y[0] approx tf.exp(x[0])
  # y[1] approx tf.log(x[1])
  ```

  Interpolate a function of one variable on a log-spaced grid:

  ```python
  x_ref = tf.exp(tf.linspace(tf.log(1.), tf.log(100000.), num_pts))
  y_ref = tf.log(x_ref + x_ref**2)

  batch_interp_regular_1d_grid(x=[1.1, 2.2], x_ref_min=1., x_ref_max=100000.,
      y_ref, grid_regularizing_transform=tf.log)
  ==> [tf.log(1.1 + 1.1**2), tf.log(2.2 + 2.2**2)]
  ```

  """

  return _interp_regular_1d_grid_impl(
      x,
      x_ref_min,
      x_ref_max,
      y_ref,
      axis=axis,
      batch_y_ref=True,
      fill_value=fill_value,
      fill_value_below=fill_value_below,
      fill_value_above=fill_value_above,
      grid_regularizing_transform=grid_regularizing_transform,
      name=name or 'batch_interp_regular_1d_grid')

def batch_interp_regular_nd_grid(x,
                                 x_ref_min,
                                 x_ref_max,
                                 y_ref,
                                 axis,
                                 fill_value='constant_extension',
                                 name=None):
  """Multi-linear interpolation on a regular (constant spacing) grid.

  Given [a batch of] reference values, this function computes a multi-linear
  interpolant and evaluates it on [a batch of] of new `x` values.

  The interpolant is built from reference values indexed by `nd` dimensions
  of `y_ref`, starting at `axis`.

  For example, take the case of a `2-D` scalar valued function and no leading
  batch dimensions.  In this case, `y_ref.shape = [C1, C2]` and `y_ref[i, j]`
  is the reference value corresponding to grid point

  ```
  [x_ref_min[0] + i * (x_ref_max[0] - x_ref_min[0]) / (C1 - 1),
   x_ref_min[1] + j * (x_ref_max[1] - x_ref_min[1]) / (C2 - 1)]
  ```

  In the general case, dimensions to the left of `axis` in `y_ref` are broadcast
  with leading dimensions in `x`, `x_ref_min`, `x_ref_max`.

  Args:
    x: Numeric `Tensor` The x-coordinates of the interpolated output values for
      each batch.  Shape `[..., D, nd]`, designating [a batch of] `D`
      coordinates in `nd` space.  `D` must be `>= 1` and is not a batch dim.
    x_ref_min:  `Tensor` of same `dtype` as `x`.  The minimum values of the
      (implicitly defined) reference `x_ref`.  Shape `[..., nd]`.
    x_ref_max:  `Tensor` of same `dtype` as `x`.  The maximum values of the
      (implicitly defined) reference `x_ref`.  Shape `[..., nd]`.
    y_ref:  `Tensor` of same `dtype` as `x`.  The reference output values. Shape
      `[..., C1, ..., Cnd, B1,...,BM]`, designating [a batch of] reference
      values indexed by `nd` dimensions, of a shape `[B1,...,BM]` valued
      function (for `M >= 0`).
    axis:  Scalar integer `Tensor`.  Dimensions `[axis, axis + nd)` of `y_ref`
      index the interpolation table.  E.g. `3-D` interpolation of a scalar
      valued function requires `axis=-3` and a `3-D` matrix valued function
      requires `axis=-5`.
    fill_value:  Determines what values output should take for `x` values that
      are below `x_ref_min` or above `x_ref_max`. Scalar `Tensor` or
      "constant_extension" ==> Extend as constant function.
      Default value: `"constant_extension"`
    name:  A name to prepend to created ops.
      Default value: `"batch_interp_regular_nd_grid"`.

  Returns:
    y_interp:  Interpolation between members of `y_ref`, at points `x`.
      `Tensor` of same `dtype` as `x`, and shape `[..., D, B1, ..., BM].`

  Raises:
    ValueError:  If `rank(x) < 2` is determined statically.
    ValueError:  If `axis` is not a scalar is determined statically.
    ValueError:  If `axis + nd > rank(y_ref)` is determined statically.

  #### Examples

  Interpolate a function of one variable.

  ```python
  y_ref = tf.exp(tf.linspace(start=0., stop=10., 20))

  tfp.math.batch_interp_regular_nd_grid(
      # x.shape = [3, 1], x_ref_min/max.shape = [1].  Trailing `1` for `1-D`.
      x=[[6.0], [0.5], [3.3]], x_ref_min=[0.], x_ref_max=[1.], y_ref=y_ref)
  ==> approx [exp(6.0), exp(0.5), exp(3.3)]
  ```

  Interpolate a scalar function of two variables.

  ```python
  x_ref_min = [0., 2 * np.pi]
  x_ref_max = [0., 2 * np.pi]

  # Build y_ref.
  x0s, x1s = tf.meshgrid(
      tf.linspace(x_ref_min[0], x_ref_max[0], num=100),
      tf.linspace(x_ref_min[1], x_ref_max[1], num=100),
      indexing='ij')

  def func(x0, x1):
    return tf.sin(x0) * tf.cos(x1)

  y_ref = func(x0s, x1s)

  x = np.pi * tf.random_uniform(shape=(10, 2))

  tfp.math.batch_interp_regular_nd_grid(x, x_ref_min, x_ref_max, y_ref, axis=-2)
  ==> tf.sin(x[:, 0]) * tf.cos(x[:, 1])
  ```

  """
  with tf.compat.v1.name_scope(
      name,
      default_name='interp_regular_nd_grid',
      values=[x, x_ref_min, x_ref_max, y_ref, fill_value]):
    dtype = dtype_util.common_dtype([x, x_ref_min, x_ref_max, y_ref],
                                    preferred_dtype=tf.float32)

    # Arg checking.
    if isinstance(fill_value, str):
      if fill_value != 'constant_extension':
        raise ValueError(
            'A fill value ({}) was not an allowed string ({})'.format(
                fill_value, 'constant_extension'))
    else:
      fill_value = tf.convert_to_tensor(
          value=fill_value, name='fill_value', dtype=dtype)
      _assert_ndims_statically(fill_value, expect_ndims=0)

    # x.shape = [..., nd].
    x = tf.convert_to_tensor(value=x, name='x', dtype=dtype)
    _assert_ndims_statically(x, expect_ndims_at_least=2)

    # y_ref.shape = [..., C1,...,Cnd, B1,...,BM]
    y_ref = tf.convert_to_tensor(value=y_ref, name='y_ref', dtype=dtype)

    # x_ref_min.shape = [nd]
    x_ref_min = tf.convert_to_tensor(
        value=x_ref_min, name='x_ref_min', dtype=dtype)
    x_ref_max = tf.convert_to_tensor(
        value=x_ref_max, name='x_ref_max', dtype=dtype)
    _assert_ndims_statically(
        x_ref_min, expect_ndims_at_least=1, expect_static=True)
    _assert_ndims_statically(
        x_ref_max, expect_ndims_at_least=1, expect_static=True)

    # nd is the number of dimensions indexing the interpolation table, it's the
    # "nd" in the function name.
    nd = tf.compat.dimension_value(x_ref_min.shape[-1])
    if nd is None:
      raise ValueError('`x_ref_min.shape[-1]` must be known statically.')
    x_ref_max.shape[-1:].assert_is_compatible_with(x_ref_min.shape[-1:])

    # Convert axis and check it statically.
    axis = tf.convert_to_tensor(value=axis, dtype=tf.int32, name='axis')
    axis = distribution_util.make_non_negative_axis(axis, tf.rank(y_ref))
    axis.shape.assert_has_rank(0)
    axis_ = tf.get_static_value(axis)
    y_ref_rank_ = tf.get_static_value(tf.rank(y_ref))
    if axis_ is not None and y_ref_rank_ is not None:
      if axis_ + nd > y_ref_rank_:
        raise ValueError(
            'Since dims `[axis, axis + nd)` index the interpolation table, we '
            'must have `axis + nd <= rank(y_ref)`.  Found: '
            '`axis`: {},  rank(y_ref): {}, and inferred `nd` from trailing '
            'dimensions of `x_ref_min` to be {}.'.format(
                axis_, y_ref_rank_, nd))

    x_batch_shape = tf.shape(input=x)[:-2]
    x_ref_min_batch_shape = tf.shape(input=x_ref_min)[:-1]
    x_ref_max_batch_shape = tf.shape(input=x_ref_max)[:-1]
    y_ref_batch_shape = tf.shape(input=y_ref)[:axis]

    # Do a brute-force broadcast of batch dims (add zeros).
    batch_shape = y_ref_batch_shape
    for tensor in [x_batch_shape, x_ref_min_batch_shape, x_ref_max_batch_shape]:
      batch_shape = tf.broadcast_dynamic_shape(batch_shape, tensor)

    def _batch_of_zeros_with_rightmost_singletons(n_singletons):
      """Return Tensor of zeros with some singletons on the rightmost dims."""
      ones = tf.ones(shape=[n_singletons], dtype=tf.int32)
      return tf.zeros(shape=tf.concat([batch_shape, ones], axis=0), dtype=dtype)

    x += _batch_of_zeros_with_rightmost_singletons(n_singletons=2)
    x_ref_min += _batch_of_zeros_with_rightmost_singletons(n_singletons=1)
    x_ref_max += _batch_of_zeros_with_rightmost_singletons(n_singletons=1)
    y_ref += _batch_of_zeros_with_rightmost_singletons(
        n_singletons=tf.rank(y_ref) - axis)

    return _batch_interp_with_gather_nd(
        x=x,
        x_ref_min=x_ref_min,
        x_ref_max=x_ref_max,
        y_ref=y_ref,
        nd=nd,
        fill_value=fill_value,
        batch_dims=tf.get_static_value(tf.rank(x)) - 2)

def _batch_interp_with_gather_nd(x, x_ref_min, x_ref_max, y_ref, nd, fill_value,
                                 batch_dims):
  """N-D interpolation that works with leading batch dims."""
  dtype = x.dtype

  # In this function,
  # x.shape = [A1, ..., An, D, nd], where n = batch_dims
  # and
  # y_ref.shape = [A1, ..., An, C1, C2,..., Cnd, B1,...,BM]
  # y_ref[A1, ..., An, i1,...,ind] is a shape [B1,...,BM] Tensor with the value
  # at index [i1,...,ind] in the interpolation table.
  #  and x_ref_max have shapes [A1, ..., An, nd].

  # ny[k] is number of y reference points in interp dim k.
  ny = tf.cast(tf.shape(input=y_ref)[batch_dims:batch_dims + nd], dtype)

  # Map [x_ref_min, x_ref_max] to [0, ny - 1].
  # This is the (fractional) index of x.
  # x_idx_unclipped[A1, ..., An, d, k] is the fractional index into dim k of
  # interpolation table for the dth x value.
  x_ref_min_expanded = tf.expand_dims(x_ref_min, axis=-2)
  x_ref_max_expanded = tf.expand_dims(x_ref_max, axis=-2)
  x_idx_unclipped = (ny - 1) * (x - x_ref_min_expanded) / (
      x_ref_max_expanded - x_ref_min_expanded)

  # Wherever x is NaN, x_idx_unclipped will be NaN as well.
  # Keep track of the nan indices here (so we can impute NaN later).
  # Also eliminate any NaN indices, since there is not NaN in 32bit.
  nan_idx = tf.math.is_nan(x_idx_unclipped)
  x_idx_unclipped = tf.where(nan_idx, tf.zeros_like(x_idx_unclipped),
                             x_idx_unclipped)

  # x_idx.shape = [A1, ..., An, D, nd]
  x_idx = tf.clip_by_value(x_idx_unclipped, tf.zeros((), dtype=dtype), ny - 1)

  # Get the index above and below x_idx.
  # Naively we could set idx_below = floor(x_idx), idx_above = ceil(x_idx),
  # however, this results in idx_below == idx_above whenever x is on a grid.
  # This in turn results in y_ref_below == y_ref_above, and then the gradient
  # at this point is zero.  So here we "jitter" one of idx_below, idx_above,
  # so that they are at different values.  This jittering does not affect the
  # interpolated value, but does make the gradient nonzero (unless of course
  # the y_ref values are the same).
  idx_below = tf.floor(x_idx)
  idx_above = tf.minimum(idx_below + 1, ny - 1)
  idx_below = tf.maximum(idx_above - 1, 0)

  # These are the values of y_ref corresponding to above/below indices.
  # idx_below_int32.shape = x.shape[:-1] + [nd]
  idx_below_int32 = tf.cast(idx_below, dtype=tf.int32)
  idx_above_int32 = tf.cast(idx_above, dtype=tf.int32)

  # idx_below_list is a length nd list of shape x.shape[:-1] int32 tensors.
  idx_below_list = tf.unstack(idx_below_int32, axis=-1)
  idx_above_list = tf.unstack(idx_above_int32, axis=-1)

  # Use t to get a convex combination of the below/above values.
  # t.shape = [A1, ..., An, D, nd]
  t = x_idx - idx_below

  # x, and tensors shaped like x, need to be added to, and selected with
  # (using tf.where) the output y.  This requires appending singletons.
  def _expand_x_fn(tensor):
    # Reshape tensor to tensor.shape + [1] * M.
    extended_shape = tf.concat([
        tf.shape(input=tensor),
        tf.ones_like(tf.shape(input=y_ref)[batch_dims + nd:])
    ],
                               axis=0)
    return tf.reshape(tensor, extended_shape)

  # Now, t.shape = [A1, ..., An, D, nd] + [1] * (rank(y_ref) - nd - batch_dims)
  t = _expand_x_fn(t)
  s = 1 - t

  # Re-insert NaN wherever x was NaN.
  nan_idx = _expand_x_fn(nan_idx)
  t = tf.where(nan_idx,
               tf.fill(tf.shape(input=t), tf.constant(np.nan, dtype)),
               t)

  terms = []
  # Our work above has located x's fractional index inside a cube of above/below
  # indices. The distance to the below indices is t, and to the above indices
  # is s.
  # Drawing lines from x to the cube walls, we get 2**nd smaller cubes. Each
  # term in the result is a product of a reference point, gathered from y_ref,
  # multiplied by a volume.  The volume is that of the cube opposite to the
  # reference point.  E.g. if the reference point is below x in every axis, the
  # volume is that of the cube with corner above x in every axis, s[0]*...*s[nd]
  # We could probably do this with one massive gather, but that would be very
  # unreadable and un-debuggable.  It also would create a large Tensor.
  for zero_ones_list in _binary_count(nd):
    gather_from_y_ref_idx = []
    opposite_volume_t_idx = []
    opposite_volume_s_idx = []
    for k, zero_or_one in enumerate(zero_ones_list):
      if zero_or_one == 0:
        # If the kth iterate has zero_or_one = 0,
        # Will gather from the "below" reference point along axis k.
        gather_from_y_ref_idx.append(idx_below_list[k])
        # Now append the index to gather for computing opposite_volume.
        # This could be done by initializing opposite_volume to 1, then here:
        #  opposite_volume *= tf.gather(s, indices=k, axis=tf.rank(x) - 1)
        # but that puts a gather in the "inner loop."  Better to append the
        # index and do one larger gather down below.
        opposite_volume_s_idx.append(k)
      else:
        gather_from_y_ref_idx.append(idx_above_list[k])
        # Append an index to gather, having the same effect as
        #   opposite_volume *= tf.gather(t, indices=k, axis=tf.rank(x) - 1)
        opposite_volume_t_idx.append(k)

    # Compute opposite_volume (volume of cube opposite the ref point):
    # Recall t.shape = s.shape = [D, nd] + [1, ..., 1]
    # Gather from t and s along the "nd" axis, which is rank(x) - 1.
    ov_axis = tf.rank(x) - 1
    opposite_volume = (
        tf.reduce_prod(
            input_tensor=tf.gather(
                t, indices=tf.cast(opposite_volume_t_idx, dtype=tf.int32),
                axis=ov_axis),
            axis=ov_axis) *
        tf.reduce_prod(
            input_tensor=tf.gather(
                s, indices=tf.cast(opposite_volume_s_idx, dtype=tf.int32),
                axis=ov_axis),
            axis=ov_axis)
    )  # pyformat: disable

    y_ref_pt = tf.gather_nd(
        y_ref,
        tf.stack(gather_from_y_ref_idx, axis=-1),
        batch_dims=batch_dims)

    terms.append(y_ref_pt * opposite_volume)

  y = tf.math.add_n(terms)

  if tf.debugging.is_numeric_tensor(fill_value):
    # Recall x_idx_unclipped.shape = [D, nd],
    # so here we check if it was out of bounds in any of the nd dims.
    # Thus, oob_idx.shape = [D].
    oob_idx = tf.reduce_any(
        input_tensor=(x_idx_unclipped < 0) | (x_idx_unclipped > ny - 1),
        axis=-1)

    # Now, y.shape = [D, B1,...,BM], so we'll have to broadcast oob_idx.

    oob_idx = _expand_x_fn(oob_idx)  # Shape [D, 1,...,1]
    oob_idx |= tf.fill(tf.shape(input=y), False)
    y = tf.where(oob_idx, tf.fill(tf.shape(input=y), fill_value), y)
  return y

def _assert_ndims_statically(x,
                             expect_ndims=None,
                             expect_ndims_at_least=None,
                             expect_static=False):
  """Assert that Tensor x has expected number of dimensions."""
  ndims = x.shape.ndims
  if ndims is None:
    if expect_static:
      raise ValueError('Expected static ndims. Found: {}'.format(x))
    return
  if expect_ndims is not None and ndims != expect_ndims:
    raise ValueError('ndims must be {}.  Found: {}'.format(expect_ndims, ndims))
  if expect_ndims_at_least is not None and ndims < expect_ndims_at_least:
    raise ValueError('ndims must be at least {}. Found {}'.format(
        expect_ndims_at_least, ndims))

def _make_expand_x_fn_for_non_batch_interpolation(y_ref, axis):
  """Make func to expand left/right (of axis) dims of tensors shaped like x."""
  # This expansion is to help x broadcast with `y`, the output.
  # In the non-batch case, the output shape is going to be
  #   y_ref.shape[:axis] + x.shape + y_ref.shape[axis+1:]

  # Recall we made axis non-negative
  y_ref_shape = tf.shape(input=y_ref)
  y_ref_shape_left = y_ref_shape[:axis]
  y_ref_shape_right = y_ref_shape[axis + 1:]

  def expand_ends(x, broadcast=False):
    """Expand x so it can bcast w/ tensors of output shape."""
    # Assume out_shape = A + x.shape + B, and rank(A) = axis.
    # Expand with singletons with same rank as A, B.
    expanded_shape = tf.pad(
        tensor=tf.shape(input=x),
        paddings=[[axis, tf.size(input=y_ref_shape_right)]],
        constant_values=1)
    x_expanded = tf.reshape(x, expanded_shape)

    if broadcast:
      out_shape = tf.concat((
          y_ref_shape_left,
          tf.shape(input=x),
          y_ref_shape_right,
      ),
                            axis=0)
      if x.dtype.is_bool:
        x_expanded = x_expanded | tf.cast(tf.zeros(out_shape), tf.bool)
      else:
        x_expanded += tf.zeros(out_shape, dtype=x.dtype)
    return x_expanded

  return expand_ends

def _make_expand_x_fn_for_batch_interpolation(y_ref, axis):
  """Make func to expand left/right (of axis) dims of tensors shaped like x."""
  # This expansion is to help x broadcast with `y`, the output.
  # In the batch case, the output shape is going to be
  #   Broadcast(y_ref.shape[:axis], x.shape[:-1]) +
  #   x.shape[-1:] +  y_ref.shape[axis+1:]

  # Recall we made axis non-negative
  y_ref_shape = tf.shape(input=y_ref)
  y_ref_shape_left = y_ref_shape[:axis]
  y_ref_shape_right = y_ref_shape[axis + 1:]

  def expand_right_dims(x, broadcast=False):
    """Expand x so it can bcast w/ tensors of output shape."""
    expanded_shape_left = tf.broadcast_dynamic_shape(
        tf.shape(input=x)[:-1],
        tf.ones([tf.size(input=y_ref_shape_left)], dtype=tf.int32))
    expanded_shape = tf.concat(
        (expanded_shape_left, tf.shape(input=x)[-1:],
         tf.ones([tf.size(input=y_ref_shape_right)], dtype=tf.int32)),
        axis=0)
    x_expanded = tf.reshape(x, expanded_shape)
    if broadcast:
      broadcast_shape_left = tf.broadcast_dynamic_shape(
          tf.shape(input=x)[:-1], y_ref_shape_left)
      broadcast_shape = tf.concat(
          (broadcast_shape_left, tf.shape(input=x)[-1:], y_ref_shape_right),
          axis=0)
      if x.dtype.is_bool:
        x_expanded = x_expanded | tf.cast(tf.zeros(broadcast_shape), tf.bool)
      else:
        x_expanded += tf.zeros(broadcast_shape, dtype=x.dtype)
    return x_expanded

  return expand_right_dims

def _batch_gather_with_broadcast(params, indices, axis):
  """Like batch_gather, but broadcasts to the left of axis."""
  # batch_gather assumes...
  #   params.shape =  [A1,...,AN, B1,...,BM]
  #   indices.shape = [A1,...,AN, C]
  # which gives output of shape
  #                   [A1,...,AN, C, B1,...,BM]
  # Here we broadcast dims of each to the left of `axis` in params, and left of
  # the rightmost dim in indices, e.g. we can
  # have
  #   params.shape =  [A1,...,AN, B1,...,BM]
  #   indices.shape = [a1,...,aN, C],
  # where ai broadcasts with Ai.

  # leading_bcast_shape is the broadcast of [A1,...,AN] and [a1,...,aN].
  leading_bcast_shape = tf.broadcast_dynamic_shape(
      tf.shape(input=params)[:axis],
      tf.shape(input=indices)[:-1])
  params += tf.zeros(
      tf.concat((leading_bcast_shape, tf.shape(input=params)[axis:]), axis=0),
      dtype=params.dtype)
  indices += tf.zeros(
      tf.concat((leading_bcast_shape, tf.shape(input=indices)[-1:]), axis=0),
      dtype=indices.dtype)
  return tf.compat.v1.batch_gather(params, indices)

def _binary_count(n):
  """Count `n` binary digits from [0...0] to [1...1]."""
  return list(itertools.product([0, 1], repeat=n))


# 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 Categorical distribution class."""

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

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.internal import distribution_util
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import reparameterization
from tensorflow_probability.python.internal import tensorshape_util

from tensorflow.python.util import deprecation  # pylint: disable=g-direct-tensorflow-import

def _broadcast_cat_event_and_params(event, params, base_dtype):
  """Broadcasts the event or distribution parameters."""
  if dtype_util.is_integer(event.dtype):
    pass
  elif dtype_util.is_floating(event.dtype):
    # When `validate_args=True` we've already ensured int/float casting
    # is closed.
    event = tf.cast(event, dtype=tf.int32)
  else:
    raise TypeError("`value` should have integer `dtype` or "
                    "`self.dtype` ({})".format(base_dtype))
  shape_known_statically = (
      tensorshape_util.rank(params.shape) is not None and
      tensorshape_util.is_fully_defined(params.shape[:-1]) and
      tensorshape_util.is_fully_defined(event.shape))
  if not shape_known_statically or params.shape[:-1] != event.shape:
    params *= tf.ones_like(event[..., tf.newaxis],
                           dtype=params.dtype)
    params_shape = tf.shape(input=params)[:-1]
    event *= tf.ones(params_shape, dtype=event.dtype)
    if tensorshape_util.rank(params.shape) is not None:
      tensorshape_util.set_shape(event, params.shape[:-1])

  return event, params

class Categorical(distribution.Distribution):
  """Categorical distribution over integers.

  The Categorical distribution is parameterized by either probabilities or
  log-probabilities of a set of `K` classes. It is defined over the integers
  `{0, 1, ..., K-1}`.

  The Categorical distribution is closely related to the `OneHotCategorical` and
  `Multinomial` distributions.  The Categorical distribution can be intuited as
  generating samples according to `argmax{ OneHotCategorical(probs) }` itself
  being identical to `argmax{ Multinomial(probs, total_count=1) }`.

  #### Mathematical Details

  The probability mass function (pmf) is,

  ```none
  pmf(k; pi) = prod_j pi_j**[k == j]
  ```

  #### Pitfalls

  The number of classes, `K`, must not exceed:
  - the largest integer representable by `self.dtype`, i.e.,
    `2**(mantissa_bits+1)` (IEEE 754),
  - the maximum `Tensor` index, i.e., `2**31-1`.

  In other words,

  ```python
  K <= min(2**31-1, {
    tf.float16: 2**11,
    tf.float32: 2**24,
    tf.float64: 2**53 }[param.dtype])
  ```

  Note: This condition is validated only when `self.validate_args = True`.

  #### Examples

  Creates a 3-class distribution with the 2nd class being most likely.

  ```python
  dist = Categorical(probs=[0.1, 0.5, 0.4])
  n = 1e4
  empirical_prob = tf.cast(
      tf.histogram_fixed_width(
        dist.sample(int(n)),
        [0., 2],
        nbins=3),
      dtype=tf.float32) / n
  # ==> array([ 0.1005,  0.5037,  0.3958], dtype=float32)
  ```

  Creates a 3-class distribution with the 2nd class being most likely.
  Parameterized by [logits](https://en.wikipedia.org/wiki/Logit) rather than
  probabilities.

  ```python
  dist = Categorical(logits=np.log([0.1, 0.5, 0.4])
  n = 1e4
  empirical_prob = tf.cast(
      tf.histogram_fixed_width(
        dist.sample(int(n)),
        [0., 2],
        nbins=3),
      dtype=tf.float32) / n
  # ==> array([0.1045,  0.5047, 0.3908], dtype=float32)
  ```

  Creates a 3-class distribution with the 3rd class being most likely.
  The distribution functions can be evaluated on counts.

  ```python
  # counts is a scalar.
  p = [0.1, 0.4, 0.5]
  dist = Categorical(probs=p)
  dist.prob(0)  # Shape []

  # p will be broadcast to [[0.1, 0.4, 0.5], [0.1, 0.4, 0.5]] to match counts.
  counts = [1, 0]
  dist.prob(counts)  # Shape [2]

  # p will be broadcast to shape [3, 5, 7, 3] to match counts.
  counts = [[...]] # Shape [5, 7, 3]
  dist.prob(counts)  # Shape [5, 7, 3]
  ```

  """

  def __init__(
      self,
      logits=None,
      probs=None,
      dtype=tf.int32,
      validate_args=False,
      allow_nan_stats=True,
      name="Categorical"):
    """Initialize Categorical distributions using class log-probabilities.

    Args:
      logits: An N-D `Tensor`, `N >= 1`, representing the log probabilities
        of a set of Categorical distributions. The first `N - 1` dimensions
        index into a batch of independent distributions and the last dimension
        represents a vector of logits for each class. Only one of `logits` or
        `probs` should be passed in.
      probs: An N-D `Tensor`, `N >= 1`, representing the probabilities
        of a set of Categorical distributions. The first `N - 1` dimensions
        index into a batch of independent distributions and the last dimension
        represents a vector of probabilities for each class. Only one of
        `logits` or `probs` should be passed in.
      dtype: The type of the event samples (default: int32).
      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.
    """
    parameters = dict(locals())
    with tf.name_scope(name) as name:
      self._logits, self._probs = distribution_util.get_logits_and_probs(
          logits=logits,
          probs=probs,
          validate_args=validate_args,
          multidimensional=True,
          name=name)

      if validate_args:
        self._logits = distribution_util.embed_check_categorical_event_shape(
            self._logits)

      logits_shape_static = tensorshape_util.with_rank_at_least(
          self._logits.shape, 1)
      if tensorshape_util.rank(logits_shape_static) is not None:
        self._batch_rank = tf.convert_to_tensor(
            value=tensorshape_util.rank(logits_shape_static) - 1,
            dtype=tf.int32,
            name="batch_rank")
      else:
        with tf.name_scope("batch_rank"):
          self._batch_rank = tf.rank(self._logits) - 1

      logits_shape = tf.shape(input=self._logits, name="logits_shape")
      num_categories = tf.compat.dimension_value(logits_shape_static[-1])
      if num_categories is not None:
        self._num_categories = tf.convert_to_tensor(
            value=num_categories, dtype=tf.int32, name="num_categories")
      else:
        with tf.name_scope("num_categories"):
          self._num_categories = logits_shape[self._batch_rank]

      if tensorshape_util.is_fully_defined(logits_shape_static[:-1]):
        self._batch_shape_val = tf.constant(
            logits_shape_static[:-1].as_list(),
            dtype=tf.int32,
            name="batch_shape")
      else:
        with tf.name_scope("batch_shape"):
          self._batch_shape_val = logits_shape[:-1]
    super(Categorical, self).__init__(
        dtype=dtype,
        reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        parameters=parameters,
        graph_parents=[self._logits,
                       self._probs],
        name=name)

  @classmethod
  def _params_event_ndims(cls):
    return dict(logits=1, probs=1)

  @property
  @deprecation.deprecated(
      "2019-05-19", "The `event_size` property is deprecated.  Use "
      "`num_categories` instead.  They have the same value, but `event_size` "
      "is misnamed.")
  def event_size(self):
    """Scalar `int32` tensor: the number of categories."""
    return self._num_categories

  @property
  def num_categories(self):
    """Scalar `int32` tensor: the number of categories."""
    return self._num_categories

  @property
  def logits(self):
    """Vector of coordinatewise logits."""
    return self._logits

  @property
  def probs(self):
    """Vector of coordinatewise probabilities."""
    return self._probs

  def _batch_shape_tensor(self):
    return tf.identity(self._batch_shape_val)

  def _batch_shape(self):
    return self.logits.shape[:-1]

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

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

  def _sample_n(self, n, seed=None):
    if tensorshape_util.rank(self.logits.shape) == 2:
      logits_2d = self.logits
    else:
      logits_2d = tf.reshape(self.logits, [-1, self.num_categories])
    sample_dtype = tf.int64 if dtype_util.size(self.dtype) > 4 else tf.int32
    draws = tf.random.categorical(
        logits_2d, n, dtype=sample_dtype, seed=seed)
    draws = tf.reshape(
        tf.transpose(a=draws), tf.concat([[n], self.batch_shape_tensor()], 0))
    return tf.cast(draws, self.dtype)

  def _cdf(self, k):
    k = tf.convert_to_tensor(value=k, name="k")

    k, probs = _broadcast_cat_event_and_params(
        k, self.probs, base_dtype=dtype_util.base_dtype(self.dtype))

    # Since the lowest number in the support is 0, any k < 0 should be zero in
    # the output.
    should_be_zero = k < 0

    # Will use k as an index in the gather below, so clip it to {0,...,K-1}.
    k = tf.clip_by_value(tf.cast(k, tf.int32), 0, self.num_categories - 1)

    batch_shape = tf.shape(input=k)

    # tf.gather(..., batch_dims=batch_dims) requires static batch_dims kwarg, so
    # to handle the case where the batch shape is dynamic, flatten the batch
    # dims (so we know batch_dims=1).
    k_flat_batch = tf.reshape(k, [-1])
    probs_flat_batch = tf.reshape(
        probs, tf.concat(([-1], [self.num_categories]), axis=0))

    cdf_flat = tf.gather(
        tf.cumsum(probs_flat_batch, axis=-1),
        k_flat_batch[..., tf.newaxis],
        batch_dims=1)

    cdf = tf.reshape(cdf_flat, shape=batch_shape)

    return tf.where(should_be_zero, tf.zeros_like(cdf), cdf)

  def _log_prob(self, k):
    k = tf.convert_to_tensor(value=k, name="k")
    if self.validate_args:
      k = distribution_util.embed_check_integer_casting_closed(
          k, target_dtype=tf.int32)
    k, logits = _broadcast_cat_event_and_params(
        k, self.logits, base_dtype=dtype_util.base_dtype(self.dtype))

    return -tf.nn.sparse_softmax_cross_entropy_with_logits(labels=k,
                                                           logits=logits)

  def _entropy(self):
    return -tf.reduce_sum(
        input_tensor=tf.math.multiply_no_nan(
            tf.nn.log_softmax(self.logits), self.probs), axis=-1)

  def _mode(self):
    ret = tf.argmax(input=self.logits, axis=self._batch_rank)
    ret = tf.cast(ret, self.dtype)
    tensorshape_util.set_shape(ret, self.batch_shape)
    return ret

@kullback_leibler.RegisterKL(Categorical, Categorical)
def _kl_categorical_categorical(a, b, name=None):
  """Calculate the batched KL divergence KL(a || b) with a and b Categorical.

  Args:
    a: instance of a Categorical distribution object.
    b: instance of a Categorical distribution object.
    name: (optional) Name to use for created operations.
      default is "kl_categorical_categorical".

  Returns:
    Batchwise KL(a || b)
  """
  with tf.name_scope(name or "kl_categorical_categorical"):
    # sum(probs log(probs / (1 - probs)))
    delta_log_probs1 = (tf.nn.log_softmax(a.logits) -
                        tf.nn.log_softmax(b.logits))
    return tf.reduce_sum(
        input_tensor=tf.nn.softmax(a.logits) * delta_log_probs1, axis=-1)


# 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.
# ============================================================================
"""Monte Carlo integration and helpers."""

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

import tensorflow.compat.v2 as tf

__all__ = [
    'expectation',
    'expectation_importance_sampler',
    'expectation_importance_sampler_logspace',
]

def expectation_importance_sampler(f,
                                   log_p,
                                   sampling_dist_q,
                                   z=None,
                                   n=None,
                                   seed=None,
                                   name='expectation_importance_sampler'):
  r"""Monte Carlo estimate of \\(E_p[f(Z)] = E_q[f(Z) p(Z) / q(Z)]\\).

  With \\(p(z) := exp^{log_p(z)}\\), this `Op` returns

  \\(n^{-1} sum_{i=1}^n [ f(z_i) p(z_i) / q(z_i) ],  z_i ~ q,\\)
  \\(\approx E_q[ f(Z) p(Z) / q(Z) ]\\)
  \\(=       E_p[f(Z)]\\)

  This integral is done in log-space with max-subtraction to better handle the
  often extreme values that `f(z) p(z) / q(z)` can take on.

  If `f >= 0`, it is up to 2x more efficient to exponentiate the result of
  `expectation_importance_sampler_logspace` applied to `Log[f]`.

  User supplies either `Tensor` of samples `z`, or number of samples to draw `n`

  Args:
    f: Callable mapping samples from `sampling_dist_q` to `Tensors` with shape
      broadcastable to `q.batch_shape`.
      For example, `f` works "just like" `q.log_prob`.
    log_p:  Callable mapping samples from `sampling_dist_q` to `Tensors` with
      shape broadcastable to `q.batch_shape`.
      For example, `log_p` works "just like" `sampling_dist_q.log_prob`.
    sampling_dist_q:  The sampling distribution.
      `tfp.distributions.Distribution`.
      `float64` `dtype` recommended.
      `log_p` and `q` should be supported on the same set.
    z:  `Tensor` of samples from `q`, produced by `q.sample` for some `n`.
    n:  Integer `Tensor`.  Number of samples to generate if `z` is not provided.
    seed:  Python integer to seed the random number generator.
    name:  A name to give this `Op`.

  Returns:
    The importance sampling estimate.  `Tensor` with `shape` equal
      to batch shape of `q`, and `dtype` = `q.dtype`.
  """
  q = sampling_dist_q
  with tf.name_scope(name):
    z = _get_samples(q, z, n, seed)

    log_p_z = log_p(z)
    q_log_prob_z = q.log_prob(z)

    def _importance_sampler_positive_f(log_f_z):
      # Same as expectation_importance_sampler_logspace, but using Tensors
      # rather than samples and functions.  Allows us to sample once.
      log_values = log_f_z + log_p_z - q_log_prob_z
      return _logspace_mean(log_values)

    # With \\(f_{plus}(z) = max(0, f(z)), f_{minus}(z) = max(0, -f(z))\\),
    # \\(E_p[f(Z)] = E_p[f_{plus}(Z)] - E_p[f_{minus}(Z)]\\)
    # \\(          = E_p[f_{plus}(Z) + 1] - E_p[f_{minus}(Z) + 1]\\)
    # Without incurring bias, 1 is added to each to prevent zeros in logspace.
    # The logarithm is approximately linear around 1 + epsilon, so this is good
    # for small values of 'z' as well.
    f_z = f(z)
    log_f_plus_z = tf.math.log1p(tf.nn.relu(f_z))
    log_f_minus_z = tf.math.log1p(tf.nn.relu(-1. * f_z))

    log_f_plus_integral = _importance_sampler_positive_f(log_f_plus_z)
    log_f_minus_integral = _importance_sampler_positive_f(log_f_minus_z)

  return tf.math.exp(log_f_plus_integral) - tf.math.exp(log_f_minus_integral)

def expectation_importance_sampler_logspace(
    log_f,
    log_p,
    sampling_dist_q,
    z=None,
    n=None,
    seed=None,
    name='expectation_importance_sampler_logspace'):
  r"""Importance sampling with a positive function, in log-space.

  With \\(p(z) := exp^{log_p(z)}\\), and \\(f(z) = exp{log_f(z)}\\),
  this `Op` returns

  \\(Log[ n^{-1} sum_{i=1}^n [ f(z_i) p(z_i) / q(z_i) ] ],  z_i ~ q,\\)
  \\(\approx Log[ E_q[ f(Z) p(Z) / q(Z) ] ]\\)
  \\(=       Log[E_p[f(Z)]]\\)

  This integral is done in log-space with max-subtraction to better handle the
  often extreme values that `f(z) p(z) / q(z)` can take on.

  In contrast to `expectation_importance_sampler`, this `Op` returns values in
  log-space.

  User supplies either `Tensor` of samples `z`, or number of samples to draw `n`

  Args:
    log_f: Callable mapping samples from `sampling_dist_q` to `Tensors` with
      shape broadcastable to `q.batch_shape`.
      For example, `log_f` works "just like" `sampling_dist_q.log_prob`.
    log_p:  Callable mapping samples from `sampling_dist_q` to `Tensors` with
      shape broadcastable to `q.batch_shape`.
      For example, `log_p` works "just like" `q.log_prob`.
    sampling_dist_q:  The sampling distribution.
      `tfp.distributions.Distribution`.
      `float64` `dtype` recommended.
      `log_p` and `q` should be supported on the same set.
    z:  `Tensor` of samples from `q`, produced by `q.sample` for some `n`.
    n:  Integer `Tensor`.  Number of samples to generate if `z` is not provided.
    seed:  Python integer to seed the random number generator.
    name:  A name to give this `Op`.

  Returns:
    Logarithm of the importance sampling estimate.  `Tensor` with `shape` equal
      to batch shape of `q`, and `dtype` = `q.dtype`.
  """
  q = sampling_dist_q
  with tf.name_scope(name):
    z = _get_samples(q, z, n, seed)
    log_values = log_f(z) + log_p(z) - q.log_prob(z)
    return _logspace_mean(log_values)

def _logspace_mean(log_values):
  """Evaluate `Log[E[values]]` in a stable manner.

  Args:
    log_values:  `Tensor` holding `Log[values]`.

  Returns:
    `Tensor` of same `dtype` as `log_values`, reduced across dim 0.
      `Log[Mean[values]]`.
  """
  # center = Max[Log[values]],  with stop-gradient
  # The center hopefully keep the exponentiated term small.  It is canceled
  # from the final result, so putting stop gradient on it will not change the
  # final result.  We put stop gradient on to eliminate unnecessary computation.
  center = tf.stop_gradient(_sample_max(log_values))

  # centered_values = exp{Log[values] - E[Log[values]]}
  centered_values = tf.math.exp(log_values - center)

  # log_mean_of_values = Log[ E[centered_values] ] + center
  #                    = Log[ E[exp{log_values - E[log_values]}] ] + center
  #                    = Log[E[values]] - E[log_values] + center
  #                    = Log[E[values]]
  log_mean_of_values = tf.math.log(_sample_mean(centered_values)) + center

  return log_mean_of_values

def expectation(f, samples, log_prob=None, use_reparametrization=True,
                axis=0, keep_dims=False, name=None):
  r"""Computes the Monte-Carlo approximation of \\(E_p[f(X)]\\).

  This function computes the Monte-Carlo approximation of an expectation, i.e.,

  \\(E_p[f(X)] \approx= m^{-1} sum_i^m f(x_j),  x_j\  ~iid\ p(X)\\)

  where:

  - `x_j = samples[j, ...]`,
  - `log(p(samples)) = log_prob(samples)` and
  - `m = prod(shape(samples)[axis])`.

  Tricks: Reparameterization and Score-Gradient

  When p is "reparameterized", i.e., a diffeomorphic transformation of a
  parameterless distribution (e.g.,
  `Normal(Y; m, s) <=> Y = sX + m, X ~ Normal(0,1)`), we can swap gradient and
  expectation, i.e.,
  grad[ Avg{ \\(s_i : i=1...n\\) } ] = Avg{ grad[\\(s_i\\)] : i=1...n } where
  S_n = Avg{\\(s_i\\)}` and `\\(s_i = f(x_i), x_i ~ p\\).

  However, if p is not reparameterized, TensorFlow's gradient will be incorrect
  since the chain-rule stops at samples of non-reparameterized distributions.
  (The non-differentiated result, `approx_expectation`, is the same regardless
  of `use_reparametrization`.) In this circumstance using the Score-Gradient
  trick results in an unbiased gradient, i.e.,

  ```none
  grad[ E_p[f(X)] ]
  = grad[ int dx p(x) f(x) ]
  = int dx grad[ p(x) f(x) ]
  = int dx [ p'(x) f(x) + p(x) f'(x) ]
  = int dx p(x) [p'(x) / p(x) f(x) + f'(x) ]
  = int dx p(x) grad[ f(x) p(x) / stop_grad[p(x)] ]
  = E_p[ grad[ f(x) p(x) / stop_grad[p(x)] ] ]
  ```

  Unless p is not reparametrized, it is usually preferable to
  `use_reparametrization = True`.

  Warning: users are responsible for verifying `p` is a "reparameterized"
  distribution.

  Example Use:

  ```python
  import tensorflow_probability as tfp
  tfd = tfp.distributions

  # Monte-Carlo approximation of a reparameterized distribution, e.g., Normal.

  num_draws = int(1e5)
  p = tfd.Normal(loc=0., scale=1.)
  q = tfd.Normal(loc=1., scale=2.)
  exact_kl_normal_normal = tfd.kl_divergence(p, q)
  # ==> 0.44314718
  approx_kl_normal_normal = tfp.monte_carlo.expectation(
      f=lambda x: p.log_prob(x) - q.log_prob(x),
      samples=p.sample(num_draws, seed=42),
      log_prob=p.log_prob,
      use_reparametrization=(p.reparameterization_type
                             == distribution.FULLY_REPARAMETERIZED))
  # ==> 0.44632751
  # Relative Error: <1%

  # Monte-Carlo approximation of non-reparameterized distribution, e.g., Gamma.

  num_draws = int(1e5)
  p = ds.Gamma(concentration=1., rate=1.)
  q = ds.Gamma(concentration=2., rate=3.)
  exact_kl_gamma_gamma = tfd.kl_divergence(p, q)
  # ==> 0.37999129
  approx_kl_gamma_gamma = tfp.monte_carlo.expectation(
      f=lambda x: p.log_prob(x) - q.log_prob(x),
      samples=p.sample(num_draws, seed=42),
      log_prob=p.log_prob,
      use_reparametrization=(p.reparameterization_type
                             == distribution.FULLY_REPARAMETERIZED))
  # ==> 0.37696719
  # Relative Error: <1%

  # For comparing the gradients, see `monte_carlo_test.py`.
  ```

  Note: The above example is for illustration only. To compute approximate
  KL-divergence, the following is preferred:

  ```python
  approx_kl_p_q = tfp.vi.monte_carlo_csiszar_f_divergence(
      f=bf.kl_reverse,
      p_log_prob=q.log_prob,
      q=p,
      num_draws=num_draws)
  ```

  Args:
    f: Python callable which can return `f(samples)`.
    samples: `Tensor` of samples used to form the Monte-Carlo approximation of
      \\(E_p[f(X)]\\).  A batch of samples should be indexed by `axis`
      dimensions.
    log_prob: Python callable which can return `log_prob(samples)`. Must
      correspond to the natural-logarithm of the pdf/pmf of each sample. Only
      required/used if `use_reparametrization=False`.
      Default value: `None`.
    use_reparametrization: Python `bool` indicating that the approximation
      should use the fact that the gradient of samples is unbiased. Whether
      `True` or `False`, this arg only affects the gradient of the resulting
      `approx_expectation`.
      Default value: `True`.
    axis: The dimensions to average. If `None`, averages all
      dimensions.
      Default value: `0` (the left-most dimension).
    keep_dims: If True, retains averaged dimensions using size `1`.
      Default value: `False`.
    name: A `name_scope` for operations created by this function.
      Default value: `None` (which implies "expectation").

  Returns:
    approx_expectation: `Tensor` corresponding to the Monte-Carlo approximation
      of \\(E_p[f(X)]\\).

  Raises:
    ValueError: if `f` is not a Python `callable`.
    ValueError: if `use_reparametrization=False` and `log_prob` is not a Python
      `callable`.
  """

  with tf.name_scope(name or 'expectation'):
    if not callable(f):
      raise ValueError('`f` must be a callable function.')
    if use_reparametrization:
      return tf.math.reduce_mean(
          input_tensor=f(samples), axis=axis, keepdims=keep_dims)
    else:
      if not callable(log_prob):
        raise ValueError('`log_prob` must be a callable function.')
      stop = tf.stop_gradient  # For readability.
      x = stop(samples)
      logpx = log_prob(x)
      fx = f(x)  # Call `f` once in case it has side-effects.
      # We now rewrite f(x) so that:
      #   `grad[f(x)] := grad[f(x)] + f(x) * grad[logqx]`.
      # To achieve this, we use a trick that
      #   `h(x) - stop(h(x)) == zeros_like(h(x))`
      # but its gradient is grad[h(x)].
      # Note that IEEE754 specifies that `x - x == 0.` and `x + 0. == x`, hence
      # this trick loses no precision. For more discussion regarding the
      # relevant portions of the IEEE754 standard, see the StackOverflow
      # question,
      # "Is there a floating point value of x, for which x-x == 0 is false?"
      # http://stackoverflow.com/q/2686644
      fx += stop(fx) * (logpx - stop(logpx))  # Add zeros_like(logpx).
      return tf.math.reduce_mean(input_tensor=fx, axis=axis, keepdims=keep_dims)

def _sample_mean(values):
  """Mean over sample indices.  In this module this is always [0]."""
  return tf.math.reduce_mean(input_tensor=values, axis=[0])

def _sample_max(values):
  """Max over sample indices.  In this module this is always [0]."""
  return tf.math.reduce_max(input_tensor=values, axis=[0])

def _get_samples(dist, z, n, seed):
  """Check args and return samples."""
  with tf.name_scope('get_samples'):
    if (n is None) == (z is None):
      raise ValueError(
          'Must specify exactly one of arguments "n" and "z".  Found: '
          'n = %s, z = %s' % (n, z))
    if n is not None:
      return dist.sample(n, seed=seed)
    else:
      return tf.convert_to_tensor(value=z, name='z')


# 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 Empirical distribution class."""

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

import tensorflow.compat.v2 as tf
from tensorflow_probability.python.distributions import distribution
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.internal import tensorshape_util
from tensorflow.python.ops import gen_array_ops  # pylint: disable=g-direct-tensorflow-import

__all__ = [
    'Empirical'
]

def _broadcast_event_and_samples(event, samples, event_ndims):
  """Broadcasts the event or samples."""
  # This is the shape of self.samples, without the samples axis, i.e. the shape
  # of the result of a call to dist.sample(). This way we can broadcast it with
  # event to get a properly-sized event, then add the singleton dim back at
  # -event_ndims - 1.
  samples_shape = tf.concat(
      [tf.shape(input=samples)[:-event_ndims - 1],
       tf.shape(input=samples)[tf.rank(samples) - event_ndims:]],
      axis=0)
  event *= tf.ones(samples_shape, dtype=event.dtype)
  event = tf.expand_dims(event, axis=-event_ndims - 1)
  samples *= tf.ones_like(event, dtype=samples.dtype)

  return event, samples

class Empirical(distribution.Distribution):
  """Empirical distribution.

  The Empirical distribution is parameterized by a [batch] multiset of samples.
  It describes the empirical measure (observations) of a variable.

  Note: some methods (log_prob, prob, cdf, mode, entropy) are not differentiable
  with regard to samples.

  #### Mathematical Details

  The probability mass function (pmf) and cumulative distribution function (cdf)
  are

  ```none
  pmf(k; s1, ..., sn) = sum_i I(k)^{k == si} / n
  I(k)^{k == si} == 1, if k == si, else 0.
  cdf(k; s1, ..., sn) = sum_i I(k)^{k >= si} / n
  I(k)^{k >= si} == 1, if k >= si, else 0.
  ```

  #### Examples

  ```python

  # Initialize a empirical distribution with 4 scalar samples.
  dist = Empirical(samples=[0., 1., 1., 2.])
  dist.cdf(1.)
  ==> 0.75
  dist.prob([0., 1.])
  ==> [0.25, 0.5] # samples will be broadcast to
                    [[0., 1., 1., 2.], [0., 1., 1., 2.]] to match event.

  # Initialize a empirical distribution with a [2] batch of scalar samples.
  dist = Empirical(samples=[[0., 1.], [1., 2.]])
  dist.cdf([0., 2.])
  ==> [0.5, 1.]
  dist.prob(0.)
  ==> [0.5, 0] # event will be broadcast to [0., 0.] to match samples.

  # Initialize a empirical distribution with 4 vector-like samples.
  dist = Empirical(samples=[[0., 0.], [0., 1.], [0., 1.], [1., 2.]],
                   event_ndims=1)
  dist.cdf([0., 1.])
  ==> 0.75
  dist.prob([[0., 1.], [1., 2.]])
  ==> [0.5, 0.25] # samples will be broadcast to shape [2, 4, 2] to match event.

  # Initialize a empirical distribution with a [2] batch of vector samples.
  dist = Empirical(samples=[[[0., 0.], [0., 1.]], [[0., 1.], [1., 2.]]],
                   event_ndims=1)
  dist.cdf([[0., 0.], [0., 1.]])
  ==> [0.5, 0.5]
  dist.prob([0., 1.])
  ==> [0.5, 1.] # event will be broadcast to shape [[0., 1.], [0., 1.]]
                  to match samples.
  ```

  """

  def __init__(self,
               samples,
               event_ndims=0,
               validate_args=False,
               allow_nan_stats=True,
               name='Empirical'):
    """Initialize `Empirical` distributions.

    Args:
      samples: Numeric `Tensor` of shape [B1, ..., Bk, S, E1, ..., En]`,
        `k, n >= 0`. Samples or batches of samples on which the distribution
        is based. The first `k` dimensions index into a batch of independent
        distributions. Length of `S` dimension determines number of samples
        in each multiset. The last `n` dimension represents samples for each
        distribution. n is specified by argument event_ndims.
      event_ndims: Python `int32`, default `0`. number of dimensions for each
        event. When `0` this distribution has scalar samples. When `1` this
        distribution has vector-like samples.
      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:
      ValueError: if the rank of `samples` < event_ndims + 1.
    """

    parameters = dict(locals())
    with tf.name_scope(name):
      self._samples = tf.convert_to_tensor(value=samples, name='samples')
      self._event_ndims = event_ndims
      self._samples_axis = (
          (tensorshape_util.rank(self.samples.shape) or tf.rank(self.samples)) -
          self._event_ndims - 1)
      with tf.control_dependencies(
          [assert_util.assert_rank_at_least(self._samples, event_ndims + 1)]):
        samples_shape = distribution_util.prefer_static_shape(self._samples)
        self._num_samples = samples_shape[self._samples_axis]

    super(Empirical, self).__init__(
        dtype=self._samples.dtype,
        reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        parameters=parameters,
        graph_parents=[self._samples],
        name=name)

  @staticmethod
  def _param_shapes(sample_shape):
    return {'samples': tf.convert_to_tensor(value=sample_shape, dtype=tf.int32)}

  def _params_event_ndims(self):
    return dict(samples=self._event_ndims + 1)

  @property
  def samples(self):
    """Distribution parameter."""
    return self._samples

  @property
  def num_samples(self):
    """Number of samples."""
    return self._num_samples

  def _batch_shape_tensor(self):
    return tf.shape(input=self.samples)[:self._samples_axis]

  def _batch_shape(self):
    if tensorshape_util.rank(self.samples.shape) is None:
      return tf.TensorShape(None)
    return self.samples.shape[:self._samples_axis]

  def _event_shape_tensor(self):
    return tf.shape(input=self.samples)[self._samples_axis + 1:]

  def _event_shape(self):
    if tensorshape_util.rank(self.samples.shape) is None:
      return tf.TensorShape(None)
    return self.samples.shape[self._samples_axis + 1:]

  def _mean(self):
    return tf.reduce_mean(input_tensor=self.samples, axis=self._samples_axis)

  def _stddev(self):
    r = self.samples - tf.expand_dims(self.mean(), axis=self._samples_axis)
    var = tf.reduce_mean(input_tensor=tf.square(r), axis=self._samples_axis)
    return tf.sqrt(var)

  def _sample_n(self, n, seed=None):
    indices = tf.random.uniform([n], maxval=self.num_samples,
                                dtype=tf.int32, seed=seed)
    draws = tf.gather(self.samples, indices, axis=self._samples_axis)
    axes = tf.concat(
        [[self._samples_axis],
         tf.range(self._samples_axis),
         tf.range(self._event_ndims) + self._samples_axis + 1],
        axis=0)
    draws = tf.transpose(a=draws, perm=axes)
    return draws

  def _mode(self):
    # Samples count can vary by batch member. Use map_fn to compute mode for
    # each batch separately.
    def _get_mode(samples):
      # TODO(b/123985779): Swith to tf.unique_with_counts_v2 when exposed
      count = gen_array_ops.unique_with_counts_v2(samples, axis=[0]).count
      return tf.argmax(input=count)

    # Flatten samples for each batch.
    if self._event_ndims == 0:
      samples = tf.reshape(self.samples, [-1, self.num_samples])
      mode_shape = self.batch_shape_tensor()
    else:
      event_size = tf.reduce_prod(input_tensor=self.event_shape_tensor())
      samples = tf.reshape(self.samples, [-1, self.num_samples, event_size])
      mode_shape = tf.concat(
          [self.batch_shape_tensor(), self.event_shape_tensor()],
          axis=0)

    indices = tf.map_fn(_get_mode, samples, dtype=tf.int64)
    full_indices = tf.stack(
        [tf.range(tf.shape(input=indices)[0]),
         tf.cast(indices, tf.int32)],
        axis=1)

    mode = tf.gather_nd(samples, full_indices)
    return tf.reshape(mode, mode_shape)

  def _entropy(self):
    # Use map_fn to compute entropy for each batch separately.
    def _get_entropy(samples):
      # TODO(b/123985779): Swith to tf.unique_with_counts_v2 when exposed
      count = gen_array_ops.unique_with_counts_v2(samples, axis=[0]).count
      prob = count / self.num_samples
      entropy = tf.reduce_sum(input_tensor=-prob * tf.math.log(prob))
      return entropy

    # Flatten samples for each batch.
    if self._event_ndims == 0:
      samples = tf.reshape(self.samples, [-1, self.num_samples])
    else:
      event_size = tf.reduce_prod(input_tensor=self.event_shape_tensor())
      samples = tf.reshape(self.samples, [-1, self.num_samples, event_size])

    entropy = tf.map_fn(_get_entropy, samples)
    entropy_shape = self.batch_shape_tensor()
    if dtype_util.is_floating(self.dtype):
      entropy = tf.cast(entropy, self.dtype)
    return tf.reshape(entropy, entropy_shape)

  def _cdf(self, event):
    event = tf.convert_to_tensor(value=event, name='event', dtype=self.dtype)
    event, samples = _broadcast_event_and_samples(event, self.samples,
                                                  event_ndims=self._event_ndims)
    cdf = tf.reduce_sum(
        input_tensor=tf.cast(
            tf.reduce_all(
                input_tensor=samples <= event,
                axis=tf.range(-self._event_ndims, 0)),
            dtype=tf.int32),
        axis=-1) / self.num_samples
    if dtype_util.is_floating(self.dtype):
      cdf = tf.cast(cdf, self.dtype)
    return cdf

  def _prob(self, event):
    event = tf.convert_to_tensor(value=event, name='event', dtype=self.dtype)
    event, samples = _broadcast_event_and_samples(event, self.samples,
                                                  event_ndims=self._event_ndims)
    prob = tf.reduce_sum(
        input_tensor=tf.cast(
            tf.reduce_all(
                input_tensor=tf.equal(samples, event),
                axis=tf.range(-self._event_ndims, 0)),
            dtype=tf.int32),
        axis=-1) / self.num_samples
    if dtype_util.is_floating(self.dtype):
      prob = tf.cast(prob, self.dtype)
    return prob


# 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.
# ============================================================================
"""Metropolis-Hastings Transition Kernel."""

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

import collections
import warnings

import tensorflow as tf

from tensorflow_probability.python import distributions
from tensorflow_probability.python.mcmc import kernel as kernel_base
from tensorflow_probability.python.mcmc.internal import util as mcmc_util

__all__ = [
    'MetropolisHastings',
]

# Cause all warnings to always be triggered.
# Not having this means subsequent calls wont trigger the warning.
warnings.filterwarnings('always',
                        module='tensorflow_probability.*metropolis_hastings',
                        append=True)  # Don't override user-set filters.

MetropolisHastingsKernelResults = collections.namedtuple(
    'MetropolisHastingsKernelResults',
    [
        'accepted_results',
        'is_accepted',
        'log_accept_ratio',
        'proposed_state',
        'proposed_results',
        'extra',
    ])

class MetropolisHastings(kernel_base.TransitionKernel):
  """Runs one step of the Metropolis-Hastings algorithm.

  The [Metropolis-Hastings algorithm](
  https://en.wikipedia.org/wiki/Metropolis%E2%80%93Hastings_algorithm) is a
  Markov chain Monte Carlo (MCMC) technique which uses a proposal distribution
  to eventually sample from a target distribution.

  Note: `inner_kernel.one_step` must return `kernel_results` as a
  `collections.namedtuple` which must:

  - have a `target_log_prob` field,
  - optionally have a `log_acceptance_correction` field, and,
  - have only fields which are `Tensor`-valued.

  The Metropolis-Hastings log acceptance-probability is computed as:

  ```python
  log_accept_ratio = (current_kernel_results.target_log_prob
                      - previous_kernel_results.target_log_prob
                      + current_kernel_results.log_acceptance_correction)
  ```

  If `current_kernel_results.log_acceptance_correction` does not exist, it is
  presumed `0.` (i.e., that the proposal distribution is symmetric).

  The most common use-case for `log_acceptance_correction` is in the
  Metropolis-Hastings algorithm, i.e.,

  ```none
  accept_prob(x' | x) = p(x') / p(x) (g(x|x') / g(x'|x))

  where,
    p  represents the target distribution,
    g  represents the proposal (conditional) distribution,
    x' is the proposed state, and,
    x  is current state
  ```

  The log of the parenthetical term is the `log_acceptance_correction`.

  The `log_acceptance_correction` may not necessarily correspond to the ratio of
  proposal distributions, e.g, `log_acceptance_correction` has a different
  interpretation in Hamiltonian Monte Carlo.

  #### Examples

  ```python
  import tensorflow_probability as tfp
  hmc = tfp.mcmc.MetropolisHastings(
      tfp.mcmc.UncalibratedHamiltonianMonteCarlo(
          target_log_prob_fn=lambda x: -x - x**2,
          step_size=0.1,
          num_leapfrog_steps=3))
  # ==> functionally equivalent to:
  # hmc = tfp.mcmc.HamiltonianMonteCarlo(
  #     target_log_prob_fn=lambda x: -x - x**2,
  #     step_size=0.1,
  #     num_leapfrog_steps=3)
  ```

  """

  def __init__(self, inner_kernel, seed=None, name=None):
    """Instantiates this object.

    Args:
      inner_kernel: `TransitionKernel`-like object which has
        `collections.namedtuple` `kernel_results` and which contains a
        `target_log_prob` member and optionally a `log_acceptance_correction`
        member.
      seed: Python integer to seed the random number generator.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., "mh_kernel").

    Returns:
      metropolis_hastings_kernel: Instance of `TransitionKernel` which wraps the
        input transition kernel with the Metropolis-Hastings algorithm.
    """
    if inner_kernel.is_calibrated:
      warnings.warn('Supplied `TransitionKernel` is already calibrated. '
                    'Composing `MetropolisHastings` `TransitionKernel` '
                    'may not be required.')
    self._seed_stream = distributions.SeedStream(
        seed, 'metropolis_hastings_one_step')
    self._parameters = dict(
        inner_kernel=inner_kernel,
        seed=seed,
        name=name)

  @property
  def inner_kernel(self):
    return self._parameters['inner_kernel']

  @property
  def seed(self):
    return self._parameters['seed']

  @property
  def name(self):
    return self._parameters['name']

  @property
  def parameters(self):
    """Return `dict` of ``__init__`` arguments and their values."""
    return self._parameters

  @property
  def is_calibrated(self):
    return True

  def one_step(self, current_state, previous_kernel_results):
    """Takes one step of the TransitionKernel.

    Args:
      current_state: `Tensor` or Python `list` of `Tensor`s representing the
        current state(s) of the Markov chain(s).
      previous_kernel_results: A (possibly nested) `tuple`, `namedtuple` or
        `list` of `Tensor`s representing internal calculations made within the
        previous call to this function (or as returned by `bootstrap_results`).

    Returns:
      next_state: `Tensor` or Python `list` of `Tensor`s representing the
        next state(s) of the Markov chain(s).
      kernel_results: A (possibly nested) `tuple`, `namedtuple` or `list` of
        `Tensor`s representing internal calculations made within this function.

    Raises:
      ValueError: if `inner_kernel` results doesn't contain the member
        "target_log_prob".
    """
    with tf.compat.v1.name_scope(
        name=mcmc_util.make_name(self.name, 'mh', 'one_step'),
        values=[current_state, previous_kernel_results]):
      # Take one inner step.
      [
          proposed_state,
          proposed_results,
      ] = self.inner_kernel.one_step(
          current_state,
          previous_kernel_results.accepted_results)

      if (not has_target_log_prob(proposed_results) or
          not has_target_log_prob(previous_kernel_results.accepted_results)):
        raise ValueError('"target_log_prob" must be a member of '
                         '`inner_kernel` results.')

      # Compute log(acceptance_ratio).
      to_sum = [proposed_results.target_log_prob,
                -previous_kernel_results.accepted_results.target_log_prob]
      try:
        if (not mcmc_util.is_list_like(
            proposed_results.log_acceptance_correction)
            or proposed_results.log_acceptance_correction):
          to_sum.append(proposed_results.log_acceptance_correction)
      except AttributeError:
        warnings.warn('Supplied inner `TransitionKernel` does not have a '
                      '`log_acceptance_correction`. Assuming its value is `0.`')
      log_accept_ratio = mcmc_util.safe_sum(
          to_sum, name='compute_log_accept_ratio')

      # If proposed state reduces likelihood: randomly accept.
      # If proposed state increases likelihood: always accept.
      # I.e., u < min(1, accept_ratio),  where u ~ Uniform[0,1)
      #       ==> log(u) < log_accept_ratio
      log_uniform = tf.math.log(
          tf.random.uniform(
              shape=tf.shape(input=proposed_results.target_log_prob),
              dtype=proposed_results.target_log_prob.dtype.base_dtype,
              seed=self._seed_stream()))
      is_accepted = log_uniform < log_accept_ratio

      next_state = mcmc_util.choose(
          is_accepted,
          proposed_state,
          current_state,
          name='choose_next_state')

      kernel_results = MetropolisHastingsKernelResults(
          accepted_results=mcmc_util.choose(
              is_accepted,
              proposed_results,
              previous_kernel_results.accepted_results,
              name='choose_inner_results'),
          is_accepted=is_accepted,
          log_accept_ratio=log_accept_ratio,
          proposed_state=proposed_state,
          proposed_results=proposed_results,
          extra=[],
      )

      return next_state, kernel_results

  def bootstrap_results(self, init_state):
    """Returns an object with the same type as returned by `one_step`.

    Args:
      init_state: `Tensor` or Python `list` of `Tensor`s representing the
        initial state(s) of the Markov chain(s).

    Returns:
      kernel_results: A (possibly nested) `tuple`, `namedtuple` or `list` of
        `Tensor`s representing internal calculations made within this function.

    Raises:
      ValueError: if `inner_kernel` results doesn't contain the member
        "target_log_prob".
    """
    with tf.compat.v1.name_scope(
        name=mcmc_util.make_name(self.name, 'mh', 'bootstrap_results'),
        values=[init_state]):
      pkr = self.inner_kernel.bootstrap_results(init_state)
      if not has_target_log_prob(pkr):
        raise ValueError(
            '"target_log_prob" must be a member of `inner_kernel` results.')
      x = pkr.target_log_prob
      return MetropolisHastingsKernelResults(
          accepted_results=pkr,
          is_accepted=tf.ones_like(x, dtype=tf.bool),
          log_accept_ratio=tf.zeros_like(x),
          proposed_state=init_state,
          proposed_results=pkr,
          extra=[],
      )

def has_target_log_prob(kernel_results):
  """Returns `True` if `target_log_prob` is a member of input."""
  return getattr(kernel_results, 'target_log_prob', None) is not None