# 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.
# ============================================================================
"""Lewandowski-Kurowicka-Joe distribution on correlation matrices.

The sampler follows the "onion" method from
[1] Daniel Lewandowski, Dorota Kurowicka, and Harry Joe,
"Generating random correlation matrices based on vines and extended
onion method," Journal of Multivariate Analysis 100 (2009), pp
1989-2001.
"""

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 beta
from tensorflow_probability.python.distributions import distribution
from tensorflow_probability.python.distributions import normal
from tensorflow_probability.python.distributions import seed_stream
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import reparameterization
from tensorflow_probability.python.internal import tensorshape_util

__all__ = [
    'LKJ',
]

def _uniform_unit_norm(dimension, shape, dtype, seed):
  """Returns a batch of points chosen uniformly from the unit hypersphere."""
  # This works because the Gaussian distribution is spherically symmetric.
  # raw shape: shape + [dimension]
  raw = normal.Normal(
      loc=dtype_util.as_numpy_dtype(dtype)(0),
      scale=dtype_util.as_numpy_dtype(dtype)(1)).sample(
          tf.concat([shape, [dimension]], axis=0), seed=seed())
  unit_norm = raw / tf.norm(tensor=raw, ord=2, axis=-1)[..., tf.newaxis]
  return unit_norm

def _replicate(n, tensor):
  """Replicate the input tensor n times along a new (major) dimension."""
  # TODO(axch) Does this already exist somewhere?  Should it get contributed?
  multiples = tf.concat([[n], tf.ones_like(tensor.shape)], axis=0)
  return tf.tile(tf.expand_dims(tensor, axis=0), multiples)

class LKJ(distribution.Distribution):
  """The LKJ distribution on correlation matrices.

  This is a one-parameter family of distributions on correlation matrices.  The
  probability density is proportional to the determinant raised to the power of
  the parameter: `pdf(X; eta) = Z(eta) * det(X) ** (eta - 1)`, where `Z(eta)` is
  a normalization constant.  The uniform distribution on correlation matrices is
  the special case `eta = 1`.

  The distribution is named after Lewandowski, Kurowicka, and Joe, who gave a
  sampler for the distribution in [(Lewandowski, Kurowicka, Joe, 2009)][1].

  #### Examples

  ```python
  # Initialize a single 3x3 LKJ with concentration parameter 1.5
  dist = tfp.distributions.LKJ(dimension=3, concentration=1.5)

  # Evaluate this at a batch of two observations, each in R^{3x3}.
  x = ...  # Shape is [2, 3, 3].
  dist.prob(x)  # Shape is [2].

  # Draw 6 LKJ-distributed 3x3 correlation matrices
  ans = dist.sample(sample_shape=[2, 3], seed=42)
  # shape of ans is [2, 3, 3, 3]
  ```
  """

  def __init__(self,
               dimension,
               concentration,
               input_output_cholesky=False,
               validate_args=False,
               allow_nan_stats=True,
               name='LKJ'):
    """Construct LKJ distributions.

    Args:
      dimension: Python `int`. The dimension of the correlation matrices
        to sample.
      concentration: `float` or `double` `Tensor`. The positive concentration
        parameter of the LKJ distributions. The pdf of a sample matrix `X` is
        proportional to `det(X) ** (concentration - 1)`.
      input_output_cholesky: Python `bool`. If `True`, functions whose input or
        output have the semantics of samples assume inputs are in Cholesky form
        and return outputs in Cholesky form. In particular, if this flag is
        `True`, input to `log_prob` is presumed of Cholesky form and output from
        `sample` is of Cholesky form.  Setting this argument to `True` is purely
        a computational optimization and does not change the underlying
        distribution. Additionally, validation checks which are only defined on
        the multiplied-out form are omitted, even if `validate_args` is `True`.
        Default value: `False` (i.e., input/output does not have Cholesky
        semantics).
      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 `dimension` is negative.
    """
    if dimension < 0:
      raise ValueError(
          'There are no negative-dimension correlation matrices.')
    parameters = dict(locals())
    self._input_output_cholesky = input_output_cholesky
    with tf.name_scope(name):
      concentration = tf.convert_to_tensor(
          value=concentration,
          name='concentration',
          dtype=dtype_util.common_dtype([concentration],
                                        preferred_dtype=tf.float32))
      with tf.control_dependencies([
          # concentration >= 1
          # TODO(b/111451422) Generalize to concentration > 0.
          assert_util.assert_non_negative(concentration - 1.),
      ] if validate_args else []):
        self._dimension = dimension
        self._concentration = tf.identity(concentration, name='concentration')
    super(LKJ, self).__init__(
        dtype=self._concentration.dtype,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
        parameters=parameters,
        graph_parents=[self._concentration],
        name=name)

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

  @property
  def dimension(self):
    """Dimension of returned correlation matrices."""
    return self._dimension

  @property
  def concentration(self):
    """Concentration parameter."""
    return self._concentration

  @property
  def input_output_cholesky(self):
    """Boolean indicating if `Tensor` input/outputs are Cholesky factorized."""
    return self._input_output_cholesky

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

  def _batch_shape(self):
    return self.concentration.shape

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

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

  def _sample_n(self, num_samples, seed=None, name=None):
    """Returns a Tensor of samples from an LKJ distribution.

    Args:
      num_samples: Python `int`. The number of samples to draw.
      seed: Python integer seed for RNG
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      samples: A Tensor of correlation matrices with shape `[n, B, D, D]`,
        where `B` is the shape of the `concentration` parameter, and `D`
        is the `dimension`.

    Raises:
      ValueError: If `dimension` is negative.
    """
    if self.dimension < 0:
      raise ValueError(
          'Cannot sample negative-dimension correlation matrices.')
    # Notation below: B is the batch shape, i.e., tf.shape(concentration)
    seed = seed_stream.SeedStream(seed, 'sample_lkj')
    with tf.name_scope('sample_lkj' or name):
      if not dtype_util.is_floating(self.concentration.dtype):
        raise TypeError(
            'The concentration argument should have floating type, not '
            '{}'.format(dtype_util.name(self.concentration.dtype)))

      concentration = _replicate(num_samples, self.concentration)
      concentration_shape = tf.shape(input=concentration)
      if self.dimension <= 1:
        # For any dimension <= 1, there is only one possible correlation matrix.
        shape = tf.concat([
            concentration_shape, [self.dimension, self.dimension]], axis=0)
        return tf.ones(shape=shape, dtype=self.concentration.dtype)
      beta_conc = concentration + (self.dimension - 2.) / 2.
      beta_dist = beta.Beta(concentration1=beta_conc, concentration0=beta_conc)

      # Note that the sampler below deviates from [1], by doing the sampling in
      # cholesky space. This does not change the fundamental logic of the
      # sampler, but does speed up the sampling.

      # This is the correlation coefficient between the first two dimensions.
      # This is also `r` in reference [1].
      corr12 = 2. * beta_dist.sample(seed=seed()) - 1.

      # Below we construct the Cholesky of the initial 2x2 correlation matrix,
      # which is of the form:
      # [[1, 0], [r, sqrt(1 - r**2)]], where r is the correlation between the
      # first two dimensions.
      # This is the top-left corner of the cholesky of the final sample.
      first_row = tf.concat([
          tf.ones_like(corr12)[..., tf.newaxis],
          tf.zeros_like(corr12)[..., tf.newaxis]], axis=-1)
      second_row = tf.concat([
          corr12[..., tf.newaxis],
          tf.sqrt(1 - corr12**2)[..., tf.newaxis]], axis=-1)

      chol_result = tf.concat([
          first_row[..., tf.newaxis, :],
          second_row[..., tf.newaxis, :]], axis=-2)

      for n in range(2, self.dimension):
        # Loop invariant: on entry, result has shape B + [n, n]
        beta_conc -= 0.5
        # norm is y in reference [1].
        norm = beta.Beta(
            concentration1=n/2.,
            concentration0=beta_conc
        ).sample(seed=seed())
        # distance shape: B + [1] for broadcast
        distance = tf.sqrt(norm)[..., tf.newaxis]
        # direction is u in reference [1].
        # direction shape: B + [n]
        direction = _uniform_unit_norm(
            n, concentration_shape, self.concentration.dtype, seed)
        # raw_correlation is w in reference [1].
        raw_correlation = distance * direction  # shape: B + [n]

        # This is the next row in the cholesky of the result,
        # which differs from the construction in reference [1].
        # In the reference, the new row `z` = chol_result @ raw_correlation^T
        # = C @ raw_correlation^T (where as short hand we use C = chol_result).
        # We prove that the below equation is the right row to add to the
        # cholesky, by showing equality with reference [1].
        # Let S be the sample constructed so far, and let `z` be as in
        # reference [1]. Then at this iteration, the new sample S' will be
        # [[S z^T]
        #  [z 1]]
        # In our case we have the cholesky decomposition factor C, so
        # we want our new row x (same size as z) to satisfy:
        #  [[S z^T]  [[C 0]    [[C^T  x^T]         [[CC^T  Cx^T]
        #   [z 1]] =  [x k]]    [0     k]]  =       [xC^t   xx^T + k**2]]
        # Since C @ raw_correlation^T = z = C @ x^T, and C is invertible,
        # we have that x = raw_correlation. Also 1 = xx^T + k**2, so k
        # = sqrt(1 - xx^T) = sqrt(1 - |raw_correlation|**2) = sqrt(1 -
        # distance**2).
        new_row = tf.concat(
            [raw_correlation, tf.sqrt(1. - norm[..., tf.newaxis])], axis=-1)

        # Finally add this new row, by growing the cholesky of the result.
        chol_result = tf.concat([
            chol_result,
            tf.zeros_like(chol_result[..., 0][..., tf.newaxis])], axis=-1)

        chol_result = tf.concat(
            [chol_result, new_row[..., tf.newaxis, :]], axis=-2)

      if self.input_output_cholesky:
        return chol_result

      result = tf.matmul(chol_result, chol_result, transpose_b=True)
      # The diagonal for a correlation matrix should always be ones. Due to
      # numerical instability the matmul might not achieve that, so manually set
      # these to ones.
      result = tf.linalg.set_diag(
          result,
          tf.ones(shape=tf.shape(input=result)[:-1], dtype=result.dtype))
      # This sampling algorithm can produce near-PSD matrices on which standard
      # algorithms such as `tf.cholesky` or `tf.linalg.self_adjoint_eigvals`
      # fail. Specifically, as documented in b/116828694, around 2% of trials
      # of 900,000 5x5 matrices (distributed according to 9 different
      # concentration parameter values) contained at least one matrix on which
      # the Cholesky decomposition failed.
      return result

  def _validate_dimension(self, x):
    x = tf.convert_to_tensor(value=x, name='x')
    if tensorshape_util.is_fully_defined(x.shape[-2:]):
      if (tensorshape_util.dims(x.shape)[-2] ==
          tensorshape_util.dims(x.shape)[-1] ==
          self.dimension):
        pass
      else:
        raise ValueError(
            'Input dimension mismatch: expected [..., {}, {}], got {}'.format(
                self.dimension, self.dimension, tensorshape_util.dims(x.shape)))
    elif self.validate_args:
      msg = 'Input dimension mismatch: expected [..., {}, {}], got {}'.format(
          self.dimension, self.dimension, tf.shape(input=x))
      with tf.control_dependencies([
          assert_util.assert_equal(
              tf.shape(input=x)[-2], self.dimension, message=msg),
          assert_util.assert_equal(
              tf.shape(input=x)[-1], self.dimension, message=msg)
      ]):
        x = tf.identity(x)
    return x

  def _validate_correlationness(self, x):
    if not self.validate_args or self.input_output_cholesky:
      return x
    checks = [
        assert_util.assert_less_equal(
            dtype_util.as_numpy_dtype(x.dtype)(-1),
            x,
            message='Correlations must be >= -1.'),
        assert_util.assert_less_equal(
            x,
            dtype_util.as_numpy_dtype(x.dtype)(1),
            message='Correlations must be <= 1.'),
        assert_util.assert_near(
            tf.linalg.diag_part(x),
            dtype_util.as_numpy_dtype(x.dtype)(1),
            message='Self-correlations must be = 1.'),
        assert_util.assert_near(
            x,
            tf.linalg.matrix_transpose(x),
            message='Correlation matrices must be symmetric')
    ]
    with tf.control_dependencies(checks):
      return tf.identity(x)

  def _log_prob(self, x):
    # Despite what one might infer from Eq 15 in [1], the formula
    # given for the normalization constant should be read in the sense
    # of division, not multiplication.
    x = self._validate_dimension(x)
    x = self._validate_correlationness(x)
    normalizer = self._log_normalization()
    return self._log_unnorm_prob(x) - normalizer

  def _log_unnorm_prob(self, x, name=None):
    """Returns the unnormalized log density of an LKJ distribution.

    Args:
      x: `float` or `double` `Tensor` of correlation matrices.  The shape of `x`
        must be `B + [D, D]`, where `B` broadcasts with the shape of
        `concentration`.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      log_p: A Tensor of the unnormalized log density of each matrix element of
        `x`, with respect to an LKJ distribution with parameter the
        corresponding element of `concentration`.
    """
    with tf.name_scope(name or 'log_unnorm_prob_lkj'):
      x = tf.convert_to_tensor(value=x, name='x')
      # The density is det(matrix) ** (concentration - 1).
      # Computing the determinant with `logdet` is usually fine, since
      # correlation matrices are Hermitian and PSD. But in some cases, for a
      # PSD matrix whose eigenvalues are close to zero, `logdet` raises an error
      # complaining that it is not PSD. The root cause is the computation of the
      # cholesky decomposition in `logdet`. Hence, we use the less efficient but
      # more robust `slogdet` which does not use `cholesky`.
      #
      # An alternative would have been to check allow_nan_stats and use
      #   eigenvalues = tf.linalg.self_adjoint_eigvals(x)
      #   psd_mask = tf.cast(
      #     tf.reduce_min(eigenvalues, axis=-1) >= 0, dtype=x.dtype)
      #   tf.where(psd_mask, answer, float('-inf'))
      # to emit probability 0 for inputs that are not PSD, without ever raising
      # an error. More care must be taken, as due to numerical stability issues,
      # self_adjoint_eigvals can return slightly negative eigenvalues even for
      # a PSD matrix.
      if self.input_output_cholesky:
        logdet = 2.0 * tf.reduce_sum(
            input_tensor=tf.math.log(tf.linalg.diag_part(x)), axis=[-1])
      else:
        _, logdet = tf.linalg.slogdet(x)
      answer = (self.concentration - 1.) * logdet
      return answer

  def _log_normalization(self, name='log_normalization'):
    """Returns the log normalization of an LKJ distribution.

    Args:
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      log_z: A Tensor of the same shape and dtype as `concentration`, containing
        the corresponding log normalizers.
    """
    # The formula is from D. Lewandowski et al [1], p. 1999, from the
    # proof that eqs 16 and 17 are equivalent.
    with tf.name_scope(name or 'log_normalization_lkj'):
      logpi = np.log(np.pi)
      ans = tf.zeros_like(self.concentration)
      for k in range(1, self.dimension):
        ans += logpi * (k / 2.)
        ans += tf.math.lgamma(self.concentration +
                              (self.dimension - 1 - k) / 2.)
        ans -= tf.math.lgamma(self.concentration + (self.dimension - 1) / 2.)
      return ans

  def _mean(self):
    # The mean of the LKJ distribution (with any concentration parameter) is the
    # identity matrix.  Proof: Imagine a correlation matrix on D variables, and
    # imagine reversing the sense of the kth of those variables.  The
    # off-diagonal entries in row and column k change sign, but LKJ is symmetric
    # with respect to this operation (because the determinant doesn't change).
    # Ergo, the mean must be invariant under it (for any k), and hence all the
    # off-diagonal entries must be 0.
    return self._identity()

  def _identity(self):
    batch = tf.shape(input=self.concentration)
    answer = tf.eye(
        num_rows=self.dimension, batch_shape=batch,
        dtype=self.concentration.dtype)
    # set_shape only necessary because tf.eye doesn't do it itself: b/111413915
    shape = answer.shape[:-2].concatenate([self.dimension, self.dimension])
    tensorshape_util.set_shape(answer, shape)
    return answer


# 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.
# ============================================================================
"""Helper functions for numpy backend."""

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

import importlib
import types

import tensorflow as tf

__all__ = [
    'common_dtype',
    'copy_docstring',
    'numpy_dtype',
    'try_import',
]

# TODO(jvdillon): Get decoration working. Eg,
# # Dependency imports
# import decorator

def copy_docstring(original_fn, new_fn):  # pylint: disable=unused-argument
  return new_fn
  # TODO(jvdillon): Get decoration working. Eg,
  # @decorator.decorator
  # def wrap(wrapped_fn, *args, **kwargs):
  #   del wrapped_fn
  #   return new_fn(*args, **kwargs)
  # return wrap(original_fn)

def numpy_dtype(dtype):
  if dtype is None:
    return None
  if hasattr(dtype, 'as_numpy_dtype'):
    return dtype.as_numpy_dtype
  return dtype

def common_dtype(args_list, preferred_dtype=None):
  """Returns explict dtype from `args_list` if exists, else preferred_dtype."""
  dtype = None
  preferred_dtype = (None if preferred_dtype is None
                     else tf.as_dtype(preferred_dtype))
  for a in tf.nest.flatten(args_list):
    if hasattr(a, 'dtype'):
      dt = tf.as_dtype(a.dtype)
    else:
      continue
    if dtype is None:
      dtype = dt
    elif dtype != dt:
      raise TypeError('Found incompatible dtypes, {} and {}.'.format(dtype, dt))
  if dtype is None and preferred_dtype is None:
    return None
  return (preferred_dtype if dtype is None else dtype).as_numpy_dtype

class _FakeModule(types.ModuleType):
  """Dummy module which raises `NotImplementedError` on `getattr` access."""

  def __init__(self, name, doc):
    self._name = name
    self._doc = doc
    types.ModuleType.__init__(self, name, doc)  # pylint: disable=non-parent-init-called

  def __dir__(self):
    return []

  def __getattr__(self, attr):
    raise NotImplementedError(self._doc)

def try_import(name):  # pylint: disable=invalid-name
  try:
    return importlib.import_module(name)
  except ImportError:
    return _FakeModule(name, 'Error loading module "{}".'.format(name))


# 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 Sample 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 as distribution_lib
from tensorflow_probability.python.distributions import kullback_leibler
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import distribution_util
from tensorflow_probability.python.internal import prefer_static
from tensorflow_probability.python.internal import tensorshape_util

def _make_summary_statistic(attr):
  """Factory for implementing summary statistics, eg, mean, stddev, mode."""
  def _fn(self, **kwargs):
    """Implements summary statistic, eg, mean, stddev, mode."""
    x = getattr(self.distribution, attr)(**kwargs)
    shape = prefer_static.concat([
        self.distribution.batch_shape_tensor(),
        prefer_static.ones(prefer_static.rank_from_shape(self.sample_shape),
                           dtype=self.sample_shape.dtype),
        self.distribution.event_shape_tensor(),
    ], axis=0)
    x = tf.reshape(x, shape=shape)
    shape = prefer_static.concat([
        self.distribution.batch_shape_tensor(),
        self.sample_shape,
        self.distribution.event_shape_tensor(),
    ], axis=0)
    return tf.broadcast_to(x, shape)
  return _fn

class Sample(distribution_lib.Distribution):
  """Sample distribution via independent draws.

  This distribution is useful for reducing over a collection of independent,
  identical draws. It is otherwise identical to the input distribution.

  #### Mathematical Details

  The probability function is,

  ```none
  p(x) = prod{ p(x[i]) : i = 0, ..., (n - 1) }
  ```

  #### Examples

  ```python
  tfd = tfp.distributions

  # Example 1: Five scalar draws.

  s = tfd.Sample(
      tfd.Normal(loc=0, scale=1),
      sample_shape=5)
  x = s.sample()
  # ==> x.shape: [5]

  lp = s.log_prob(x)
  # ==> lp.shape: []
  #     Equivalently: tf.reduce_sum(s.distribution.log_prob(x), axis=[0, 1])
  #
  # `Sample.log_prob` computes the per-{sample, batch} `log_prob`s then sums
  # over the `Sample.sample_shape` dimensions. In the above example `log_prob`
  # dims `[0, 1]` are summed out. Conceptually, first dim `1` is summed (this
  # being the intrinsic `event`) then we sum over `Sample.sample_shape` dims, in
  # this case dim `0`.

  # Example 2: `[5, 4]`-draws of a bivariate Normal.

  s = tfd.Sample(
      tfd.Independent(tfd.Normal(loc=tf.zeros([3, 2]), scale=1),
                      reinterpreted_batch_ndims=1),
      sample_shape=[5, 4])
  x = s.sample([6, 1])
  # ==> x.shape: [6, 1, 5, 4, 3, 2]

  lp = s.log_prob(x)
  # ==> lp.shape: [6, 1, 3]
  #     Equivalently: tf.reduce_sum(s.distribution.log_prob(x), axis=[2, 3])
  #
  # `s.log_prob` will reduce over (intrinsic) event dims, i.e., dim `5`, then
  # sums over `s.sample_shape` dims `[2, 3]` corresponding to shape (slice)
  # `[5, 4]`.
  ```

  """

  def __init__(
      self,
      distribution,
      sample_shape=(),
      validate_args=False,
      name=None):
    """Construct the `Sample` distribution.

    Args:
      distribution: The base distribution instance to transform. Typically an
        instance of `Distribution`.
      sample_shape: `int` scalar or vector `Tensor` representing the shape of a
        single sample.
      validate_args: Python `bool`.  Whether to validate input with asserts.
        If `validate_args` is `False`, and the inputs are invalid,
        correct behavior is not guaranteed.
      name: The name for ops managed by the distribution.
        Default value: `None` (i.e., `'Sample' + distribution.name`).
    """
    parameters = dict(locals())
    name = name or 'Sample' + distribution.name
    self._distribution = distribution
    with tf.name_scope(name) as name:
      sample_shape = distribution_util.expand_to_vector(tf.convert_to_tensor(
          value=sample_shape, dtype_hint=tf.int32, name='sample_shape'))
      self._sample_shape = sample_shape
      super(Sample, self).__init__(
          dtype=self._distribution.dtype,
          reparameterization_type=self._distribution.reparameterization_type,
          validate_args=validate_args,
          allow_nan_stats=self._distribution.allow_nan_stats,
          parameters=parameters,
          graph_parents=([sample_shape] + distribution._graph_parents),  # pylint: disable=protected-access
          name=name)

  @property
  def distribution(self):
    return self._distribution

  @property
  def sample_shape(self):
    return self._sample_shape

  def _batch_shape_tensor(self):
    return self.distribution.batch_shape_tensor()

  def _batch_shape(self):
    return self.distribution.batch_shape

  def _event_shape_tensor(self):
    return prefer_static.concat([
        self.sample_shape,
        self.distribution.event_shape_tensor(),
    ], axis=0)

  def _event_shape(self):
    sample_shape = tf.TensorShape(tf.get_static_value(self.sample_shape))
    if (tensorshape_util.rank(sample_shape) is None or
        tensorshape_util.rank(self.distribution.event_shape) is None):
      return tf.TensorShape(None)
    return tensorshape_util.concatenate(sample_shape,
                                        self.distribution.event_shape)

  def _sample_n(self, n, seed, **kwargs):
    fake_sample_ndims = prefer_static.rank_from_shape(self.sample_shape)
    event_ndims = prefer_static.rank_from_shape(
        self.distribution.event_shape_tensor, self.distribution.event_shape)
    batch_ndims = prefer_static.rank_from_shape(
        self.distribution.batch_shape_tensor, self.distribution.batch_shape)
    perm = prefer_static.concat([
        [0],
        prefer_static.range(1 + fake_sample_ndims,
                            1 + fake_sample_ndims + batch_ndims),
        prefer_static.range(1, 1 + fake_sample_ndims),
        prefer_static.range(1 + fake_sample_ndims + batch_ndims,
                            1 + fake_sample_ndims + batch_ndims + event_ndims),
    ], axis=0)
    x = self.distribution.sample(
        prefer_static.concat([[n], self.sample_shape], axis=0),
        seed=seed,
        **kwargs)
    return tf.transpose(a=x, perm=perm)

  def _log_prob(self, x, **kwargs):
    batch_ndims = prefer_static.rank_from_shape(
        self.distribution.batch_shape_tensor,
        self.distribution.batch_shape)
    extra_sample_ndims = prefer_static.rank_from_shape(self.sample_shape)
    event_ndims = prefer_static.rank_from_shape(
        self.distribution.event_shape_tensor,
        self.distribution.event_shape)
    ndims = prefer_static.rank(x)
    # (1) Expand x's dims.
    d = ndims - batch_ndims - extra_sample_ndims - event_ndims
    x = tf.reshape(x, shape=tf.pad(
        tensor=tf.shape(input=x),
        paddings=[[prefer_static.maximum(0, -d), 0]],
        constant_values=1))
    sample_ndims = prefer_static.maximum(0, d)
    # (2) Transpose x's dims.
    sample_dims = prefer_static.range(0, sample_ndims)
    batch_dims = prefer_static.range(sample_ndims, sample_ndims + batch_ndims)
    extra_sample_dims = prefer_static.range(
        sample_ndims + batch_ndims,
        sample_ndims + batch_ndims + extra_sample_ndims)
    event_dims = prefer_static.range(
        sample_ndims + batch_ndims + extra_sample_ndims,
        ndims)
    perm = prefer_static.concat(
        [sample_dims, extra_sample_dims, batch_dims, event_dims], axis=0)
    x = tf.transpose(a=x, perm=perm)
    # (3) Compute x's log_prob.
    lp = self.distribution.log_prob(x, **kwargs)
    # (4) Make the final reduction in x.
    axis = prefer_static.range(sample_ndims, sample_ndims + extra_sample_ndims)
    return tf.reduce_sum(input_tensor=lp, axis=axis)

  def _entropy(self, **kwargs):
    h = self.distribution.entropy(**kwargs)
    n = prefer_static.reduce_prod(self.sample_shape)
    return tf.cast(x=n, dtype=h.dtype) * h

  _mean = _make_summary_statistic('mean')
  _stddev = _make_summary_statistic('stddev')
  _variance = _make_summary_statistic('variance')
  _mode = _make_summary_statistic('mode')

@kullback_leibler.RegisterKL(Sample, Sample)
def _kl_sample(a, b, name='kl_sample'):
  """Batched KL divergence `KL(a || b)` for Sample distributions.

  We can leverage the fact that:

  ```
  KL(Sample(a) || Sample(b)) = sum(KL(a || b))
  ```

  where the sum is over the `sample_shape` dims.

  Args:
    a: Instance of `Sample` distribution.
    b: Instance of `Sample` distribution.
    name: (optional) name to use for created ops.
      Default value: `"kl_sample"`'.

  Returns:
    kldiv: Batchwise `KL(a || b)`.

  Raises:
    ValueError: If the `sample_shape` of `a` and `b` don't match.
  """
  assertions = []
  a_ss = tf.get_static_value(a.sample_shape)
  b_ss = tf.get_static_value(b.sample_shape)
  msg = '`a.sample_shape` must be identical to `b.sample_shape`.'
  if a_ss is not None and b_ss is not None:
    if not np.array_equal(a_ss, b_ss):
      raise ValueError(msg)
  elif a.validate_args or b.validate_args:
    assertions.append(assert_util.assert_equal(
        a.sample_shape, b.sample_shape, message=msg))
  with tf.control_dependencies(assertions):
    kl = kullback_leibler.kl_divergence(
        a.distribution, b.distribution, name=name)
    n = prefer_static.reduce_prod(a.sample_shape)
    return tf.cast(x=n, dtype=kl.dtype) * kl


# 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 Triangular 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 seed_stream
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import reparameterization

def _broadcast_to(tensor_to_broadcast, target_tensors):
  """Helper to broadcast a tensor using a list of target tensors."""
  output = tensor_to_broadcast
  for tensor in target_tensors:
    output += tf.zeros_like(tensor)
  return output

class Triangular(distribution.Distribution):
  r"""Triangular distribution with `low`, `high` and `peak` parameters.

  #### Mathematical Details

  The Triangular distribution is specified by two line segments in the plane,
  such that:

    * The first line segment starts at `(a, 0)` and ends at `(c, z)`.
    * The second line segment starts at `(c, z)` and ends at `(b, 0)`.

    ```none
    y

    ^
  z |           o  (c,z)
    |          / \
    |         /   \
    |        /     \
    | (a,0) /       \ (b,0)
  0 +------o---------o-------> x
    0      a    c    b
  ```

  where:
  * a <= c <= b, a < b
  * `low = a`,
  * `high = b`,
  * `peak = c`,
  * `z = 2 / (b - a)`

  The parameters `low`, `high` and `peak` must be shaped in a way that supports
  broadcasting (e.g., `high - low` is a valid operation).

  #### Examples

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

  # Specify a single Triangular distribution.
  u1 = tfd.Triangular(low=3., high=4., peak=3.5)
  u1.mean()
  # ==> 3.5

  # Specify two different Triangular distributions.
  u2 = tfd.Triangular(low=[1., 2.], high=[3., 4.], peak=[2., 3.])
  u2.mean()
  # ==> [2., 3.]

  # Specify three different Triangular distributions by leveraging broadcasting.
  u3 = tfd.Triangular(low=3., high=[5., 6., 7.], peak=3.)
  u3.mean()
  # ==> [3.6666, 4., 4.3333]
  ```

  """

  def __init__(self,
               low=0.,
               high=1.,
               peak=0.5,
               validate_args=False,
               allow_nan_stats=True,
               name="Triangular"):
    """Initialize a batch of Triangular distributions.

    Args:
      low: Floating point tensor, lower boundary of the output interval. Must
        have `low < high`.
        Default value: `0`.
      high: Floating point tensor, upper boundary of the output interval. Must
        have `low < high`.
        Default value: `1`.
      peak: Floating point tensor, mode of the output interval. Must have
        `low <= peak` and `peak <= high`.
        Default value: `0.5`.
      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.
        Default value: `False`.
      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.
        Default value: `True`.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: `'Triangular'`.

    Raises:
      InvalidArgumentError: if `validate_args=True` and one of the following is
        True:
        * `low >= high`.
        * `peak > high`.
        * `low > peak`.
    """
    parameters = dict(locals())
    with tf.name_scope(name) as name:
      dtype = dtype_util.common_dtype([low, high, peak], tf.float32)
      low = tf.convert_to_tensor(value=low, name="low", dtype=dtype)
      high = tf.convert_to_tensor(value=high, name="high", dtype=dtype)
      peak = tf.convert_to_tensor(value=peak, name="peak", dtype=dtype)

      with tf.control_dependencies([
          assert_util.assert_less(
              low, high, message="triangular not defined when low >= high."),
          assert_util.assert_less_equal(
              low, peak, message="triangular not defined when low > peak."),
          assert_util.assert_less_equal(
              peak, high, message="triangular not defined when peak > high."),
      ] if validate_args else []):
        self._low = tf.identity(low, name="low")
        self._high = tf.identity(high, name="high")
        self._peak = tf.identity(peak, name="peak")
        dtype_util.assert_same_float_dtype(
            [self._low, self._high, self._peak])
    super(Triangular, self).__init__(
        dtype=self._low.dtype,
        reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        parameters=parameters,
        graph_parents=[self._low, self._high, self._peak],
        name=name)

  @classmethod
  def _params_event_ndims(cls):
    return dict(low=0, high=0, peak=0)

  @property
  def low(self):
    """Lower boundary of the interval."""
    return self._low

  @property
  def high(self):
    """Upper boundary of the interval."""
    return self._high

  @property
  def peak(self):
    """Peak of the distribution. Lies in the interval."""
    return self._peak

  def _pdf_at_peak(self):
    """Pdf evaluated at the peak."""
    return (self.peak - self.low) / (self.high - self.low)

  def _batch_shape_tensor(self):
    return tf.broadcast_dynamic_shape(
        tf.shape(input=self.peak),
        tf.broadcast_dynamic_shape(
            tf.shape(input=self.low), tf.shape(input=self.high)))

  def _batch_shape(self):
    return tf.broadcast_static_shape(
        self.peak.shape,
        tf.broadcast_static_shape(
            self.low.shape, self.high.shape))

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

  def _sample_n(self, n, seed=None):
    stream = seed_stream.SeedStream(seed, salt="triangular")
    shape = tf.concat([[n], self.batch_shape_tensor()], axis=0)
    samples = tf.random.uniform(shape=shape, dtype=self.dtype, seed=stream())
    # We use Inverse CDF sampling here. Because the CDF is a quadratic function,
    # we must use sqrts here.
    interval_length = self.high - self.low
    return tf.where(
        # Note the CDF on the left side of the peak is
        # (x - low) ** 2 / ((high - low) * (peak - low)).
        # If we plug in peak for x, we get that the CDF at the peak
        # is (peak - low) / (high - low). Because of this we decide
        # which part of the piecewise CDF we should use based on the cdf samples
        # we drew.
        samples < (self.peak - self.low) / interval_length,
        # Inverse of (x - low) ** 2 / ((high - low) * (peak - low)).
        self.low + tf.sqrt(
            samples * interval_length * (self.peak - self.low)),
        # Inverse of 1 - (high - x) ** 2 / ((high - low) * (high - peak))
        self.high - tf.sqrt(
            (1. - samples) * interval_length * (self.high - self.peak)))

  def _prob(self, x):
    if self.validate_args:
      with tf.control_dependencies([
          assert_util.assert_greater_equal(x, self.low),
          assert_util.assert_less_equal(x, self.high)
      ]):
        x = tf.identity(x)

    broadcast_x_to_high = _broadcast_to(x, [self.high])
    left_of_peak = tf.logical_and(
        broadcast_x_to_high > self.low, broadcast_x_to_high <= self.peak)

    interval_length = self.high - self.low
    # This is the pdf function when a low <= high <= x. This looks like
    # a triangle, so we have to treat each line segment separately.
    result_inside_interval = tf.where(
        left_of_peak,
        # Line segment from (self.low, 0) to (self.peak, 2 / (self.high -
        # self.low).
        2. * (x - self.low) / (interval_length * (self.peak - self.low)),
        # Line segment from (self.peak, 2 / (self.high - self.low)) to
        # (self.high, 0).
        2. * (self.high - x) / (interval_length * (self.high - self.peak)))

    broadcast_x_to_peak = _broadcast_to(x, [self.peak])
    outside_interval = tf.logical_or(
        broadcast_x_to_peak < self.low, broadcast_x_to_peak > self.high)

    broadcast_shape = tf.broadcast_dynamic_shape(
        tf.shape(input=x), self.batch_shape_tensor())

    return tf.where(
        outside_interval,
        tf.zeros(broadcast_shape, dtype=self.dtype),
        result_inside_interval)

  def _cdf(self, x):
    broadcast_shape = tf.broadcast_dynamic_shape(
        tf.shape(input=x), self.batch_shape_tensor())

    broadcast_x_to_high = _broadcast_to(x, [self.high])
    left_of_peak = tf.logical_and(
        broadcast_x_to_high > self.low, broadcast_x_to_high <= self.peak)

    interval_length = self.high - self.low
    # Due to the PDF being not smooth at the peak, we have to treat each side
    # somewhat differently. The PDF is two line segments, and thus we get
    # quadratics here for the CDF.
    result_inside_interval = tf.where(
        left_of_peak,
        # (x - low) ** 2 / ((high - low) * (peak - low))
        tf.math.squared_difference(x, self.low) / (interval_length *
                                                   (self.peak - self.low)),
        # 1 - (high - x) ** 2 / ((high - low) * (high - peak))
        1. - tf.math.squared_difference(self.high, x) /
        (interval_length * (self.high - self.peak)))

    broadcast_x_to_high_peak = _broadcast_to(broadcast_x_to_high, [self.peak])
    zeros = tf.zeros(broadcast_shape, dtype=self.dtype)
    # We now add that the left tail is 0 and the right tail is 1.
    result_if_not_big = tf.where(
        broadcast_x_to_high_peak < self.low, zeros, result_inside_interval)

    broadcast_x_to_peak_low = _broadcast_to(x, [self.low, self.peak])
    ones = tf.ones(broadcast_shape, dtype=self.dtype)
    return tf.where(
        broadcast_x_to_peak_low >= self.high, ones, result_if_not_big)

  def _entropy(self):
    return 0.5 - np.log(2.) + tf.math.log(self.high - self.low)

  def _mean(self):
    return (self.low + self.high + self.peak) / 3.

  def _variance(self):
    # ((high - low) ** 2 + (peak - low) ** 2 + (peak - high) ** 2) / 36
    return (tf.math.squared_difference(self.high, self.low) +
            tf.math.squared_difference(self.high, self.peak) +
            tf.math.squared_difference(self.peak, self.low)) / 36.


# 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.
# ============================================================================
"""Utilities for Markov Chain Monte Carlo (MCMC) sampling.

@@effective_sample_size
@@potential_scale_reduction
"""

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

import tensorflow as tf

from tensorflow_probability.python import stats

__all__ = [
    'effective_sample_size',
    'potential_scale_reduction',
]

def effective_sample_size(states,
                          filter_threshold=0.,
                          filter_beyond_lag=None,
                          name=None):
  """Estimate a lower bound on effective sample size for each independent chain.

  Roughly speaking, "effective sample size" (ESS) is the size of an iid sample
  with the same variance as `state`.

  More precisely, given a stationary sequence of possibly correlated random
  variables `X_1, X_2,...,X_N`, each identically distributed ESS is the number
  such that

  ```Variance{ N**-1 * Sum{X_i} } = ESS**-1 * Variance{ X_1 }.```

  If the sequence is uncorrelated, `ESS = N`.  In general, one should expect
  `ESS <= N`, with more highly correlated sequences having smaller `ESS`.

  Args:
    states:  `Tensor` or list of `Tensor` objects.  Dimension zero should index
      identically distributed states.
    filter_threshold:  `Tensor` or list of `Tensor` objects.
      Must broadcast with `state`.  The auto-correlation sequence is truncated
      after the first appearance of a term less than `filter_threshold`.
      Setting to `None` means we use no threshold filter.  Since `|R_k| <= 1`,
      setting to any number less than `-1` has the same effect.
    filter_beyond_lag:  `Tensor` or list of `Tensor` objects.  Must be
      `int`-like and scalar valued.  The auto-correlation sequence is truncated
      to this length.  Setting to `None` means we do not filter based on number
      of lags.
    name:  `String` name to prepend to created ops.

  Returns:
    ess:  `Tensor` or list of `Tensor` objects.  The effective sample size of
      each component of `states`.  Shape will be `states.shape[1:]`.

  Raises:
    ValueError:  If `states` and `filter_threshold` or `states` and
      `filter_beyond_lag` are both lists with different lengths.

  #### Examples

  We use ESS to estimate standard error.

  ```
  import tensorflow as tf
  import tensorflow_probability as tfp
  tfd = tfp.distributions

  target = tfd.MultivariateNormalDiag(scale_diag=[1., 2.])

  # Get 1000 states from one chain.
  states = tfp.mcmc.sample_chain(
      num_burnin_steps=200,
      num_results=1000,
      current_state=tf.constant([0., 0.]),
      kernel=tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=target.log_prob,
        step_size=0.05,
        num_leapfrog_steps=20))
  states.shape
  ==> (1000, 2)

  ess = effective_sample_size(states)
  ==> Shape (2,) Tensor

  mean, variance = tf.nn.moments(states, axis=0)
  standard_error = tf.sqrt(variance / ess)
  ```

  Some math shows that, with `R_k` the auto-correlation sequence,
  `R_k := Covariance{X_1, X_{1+k}} / Variance{X_1}`, we have

  ```ESS(N) =  N / [ 1 + 2 * ( (N - 1) / N * R_1 + ... + 1 / N * R_{N-1}  ) ]```

  This function estimates the above by first estimating the auto-correlation.
  Since `R_k` must be estimated using only `N - k` samples, it becomes
  progressively noisier for larger `k`.  For this reason, the summation over
  `R_k` should be truncated at some number `filter_beyond_lag < N`.  Since many
  MCMC methods generate chains where `R_k > 0`, a reasonable criteria is to
  truncate at the first index where the estimated auto-correlation becomes
  negative.

  The arguments `filter_beyond_lag`, `filter_threshold` are filters intended to
  remove noisy tail terms from `R_k`.  They combine in an "OR" manner meaning
  terms are removed if they were to be filtered under the `filter_beyond_lag` OR
  `filter_threshold` criteria.
  """
  states_was_list = _is_list_like(states)

  # Convert all args to lists.
  if not states_was_list:
    states = [states]

  filter_beyond_lag = _broadcast_maybelist_arg(states, filter_beyond_lag,
                                               'filter_beyond_lag')
  filter_threshold = _broadcast_maybelist_arg(states, filter_threshold,
                                              'filter_threshold')

  # Process items, one at a time.
  with tf.compat.v1.name_scope(name, 'effective_sample_size'):
    ess_list = [
        _effective_sample_size_single_state(s, ml, mlt)
        for (s, ml, mlt) in zip(states, filter_beyond_lag, filter_threshold)
    ]

  if states_was_list:
    return ess_list
  return ess_list[0]

def _effective_sample_size_single_state(states, filter_beyond_lag,
                                        filter_threshold):
  """ESS computation for one single Tensor argument."""

  with tf.compat.v1.name_scope(
      'effective_sample_size_single_state',
      values=[states, filter_beyond_lag, filter_threshold]):

    states = tf.convert_to_tensor(value=states, name='states')
    dt = states.dtype

    # filter_beyond_lag == None ==> auto_corr is the full sequence.
    auto_corr = stats.auto_correlation(
        states, axis=0, max_lags=filter_beyond_lag)
    if filter_threshold is not None:
      filter_threshold = tf.convert_to_tensor(
          value=filter_threshold, dtype=dt, name='filter_threshold')
      # Get a binary mask to zero out values of auto_corr below the threshold.
      #   mask[i, ...] = 1 if auto_corr[j, ...] > threshold for all j <= i,
      #   mask[i, ...] = 0, otherwise.
      # So, along dimension zero, the mask will look like [1, 1, ..., 0, 0,...]
      # Building step by step,
      #   Assume auto_corr = [1, 0.5, 0.0, 0.3], and filter_threshold = 0.2.
      # Step 1:  mask = [False, False, True, False]
      mask = auto_corr < filter_threshold
      # Step 2:  mask = [0, 0, 1, 1]
      mask = tf.cast(mask, dtype=dt)
      # Step 3:  mask = [0, 0, 1, 2]
      mask = tf.cumsum(mask, axis=0)
      # Step 4:  mask = [1, 1, 0, 0]
      mask = tf.maximum(1. - mask, 0.)
      auto_corr *= mask

    # With R[k] := auto_corr[k, ...],
    # ESS = N / {1 + 2 * Sum_{k=1}^N (N - k) / N * R[k]}
    #     = N / {-1 + 2 * Sum_{k=0}^N (N - k) / N * R[k]} (since R[0] = 1)
    #     approx N / {-1 + 2 * Sum_{k=0}^M (N - k) / N * R[k]}
    # where M is the filter_beyond_lag truncation point chosen above.

    # Get the factor (N - k) / N, and give it shape [M, 1,...,1], having total
    # ndims the same as auto_corr
    n = _axis_size(states, axis=0)
    k = tf.range(0., _axis_size(auto_corr, axis=0))
    nk_factor = (n - k) / n
    if auto_corr.shape.ndims is not None:
      new_shape = [-1] + [1] * (auto_corr.shape.ndims - 1)
    else:
      new_shape = tf.concat(
          ([-1],
           tf.ones([tf.rank(auto_corr) - 1], dtype=tf.int32)),
          axis=0)
    nk_factor = tf.reshape(nk_factor, new_shape)

    return n / (-1 +
                2 * tf.reduce_sum(input_tensor=nk_factor * auto_corr, axis=0))

def potential_scale_reduction(chains_states,
                              independent_chain_ndims=1,
                              name=None):
  """Gelman and Rubin (1992)'s potential scale reduction for chain convergence.

  Given `N > 1` states from each of `C > 1` independent chains, the potential
  scale reduction factor, commonly referred to as R-hat, measures convergence of
  the chains (to the same target) by testing for equality of means.
  Specifically, R-hat measures the degree to which variance (of the means)
  between chains exceeds what one would expect if the chains were identically
  distributed. See [Gelman and Rubin (1992)][1]; [Brooks and Gelman (1998)][2].

  Some guidelines:

  * The initial state of the chains should be drawn from a distribution
    overdispersed with respect to the target.
  * If all chains converge to the target, then as `N --> infinity`, R-hat --> 1.
    Before that, R-hat > 1 (except in pathological cases, e.g. if the chain
    paths were identical).
  * The above holds for any number of chains `C > 1`.  Increasing `C` does
    improves effectiveness of the diagnostic.
  * Sometimes, R-hat < 1.2 is used to indicate approximate convergence, but of
    course this is problem dependent. See [Brooks and Gelman (1998)][2].
  * R-hat only measures non-convergence of the mean. If higher moments, or
    other statistics are desired, a different diagnostic should be used. See
    [Brooks and Gelman (1998)][2].

  Args:
    chains_states:  `Tensor` or Python `list` of `Tensor`s representing the
      state(s) of a Markov Chain at each result step.  The `ith` state is
      assumed to have shape `[Ni, Ci1, Ci2,...,CiD] + A`.
      Dimension `0` indexes the `Ni > 1` result steps of the Markov Chain.
      Dimensions `1` through `D` index the `Ci1 x ... x CiD` independent
      chains to be tested for convergence to the same target.
      The remaining dimensions, `A`, can have any shape (even empty).
    independent_chain_ndims: Integer type `Tensor` with value `>= 1` giving the
      number of giving the number of dimensions, from `dim = 1` to `dim = D`,
      holding independent chain results to be tested for convergence.
    name: `String` name to prepend to created tf.  Default:
      `potential_scale_reduction`.

  Returns:
    `Tensor` or Python `list` of `Tensor`s representing the R-hat statistic for
    the state(s).  Same `dtype` as `state`, and shape equal to
    `state.shape[1 + independent_chain_ndims:]`.

  Raises:
    ValueError:  If `independent_chain_ndims < 1`.

  #### Examples

  Diagnosing convergence by monitoring 10 chains that each attempt to
  sample from a 2-variate normal.

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

  target = tfd.MultivariateNormalDiag(scale_diag=[1., 2.])

  # Get 10 (2x) overdispersed initial states.
  initial_state = target.sample(10) * 2.
  ==> (10, 2)

  # Get 1000 samples from the 10 independent chains.
  chains_states, _ = tfp.mcmc.sample_chain(
      num_burnin_steps=200,
      num_results=1000,
      current_state=initial_state,
      kernel=tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=target.log_prob,
          step_size=0.05,
          num_leapfrog_steps=20))
  chains_states.shape
  ==> (1000, 10, 2)

  rhat = tfp.mcmc.diagnostic.potential_scale_reduction(
      chains_states, independent_chain_ndims=1)

  # The second dimension needed a longer burn-in.
  rhat.eval()
  ==> [1.05, 1.3]
  ```

  To see why R-hat is reasonable, let `X` be a random variable drawn uniformly
  from the combined states (combined over all chains).  Then, in the limit
  `N, C --> infinity`, with `E`, `Var` denoting expectation and variance,

  ```R-hat = ( E[Var[X | chain]] + Var[E[X | chain]] ) / E[Var[X | chain]].```

  Using the law of total variance, the numerator is the variance of the combined
  states, and the denominator is the total variance minus the variance of the
  the individual chain means.  If the chains are all drawing from the same
  distribution, they will have the same mean, and thus the ratio should be one.

  #### References

  [1]: Stephen P. Brooks and Andrew Gelman. General Methods for Monitoring
       Convergence of Iterative Simulations. _Journal of Computational and
       Graphical Statistics_, 7(4), 1998.

  [2]: Andrew Gelman and Donald B. Rubin. Inference from Iterative Simulation
       Using Multiple Sequences. _Statistical Science_, 7(4):457-472, 1992.
  """
  chains_states_was_list = _is_list_like(chains_states)
  if not chains_states_was_list:
    chains_states = [chains_states]

  # tf.get_static_value returns None iff a constant value (as a numpy
  # array) is not efficiently computable.  Therefore, we try constant_value then
  # check for None.
  icn_const_ = tf.get_static_value(
      tf.convert_to_tensor(value=independent_chain_ndims))
  if icn_const_ is not None:
    independent_chain_ndims = icn_const_
    if icn_const_ < 1:
      raise ValueError(
          'Argument `independent_chain_ndims` must be `>= 1`, found: {}'.format(
              independent_chain_ndims))

  with tf.compat.v1.name_scope(name, 'potential_scale_reduction'):
    rhat_list = [
        _potential_scale_reduction_single_state(s, independent_chain_ndims)
        for s in chains_states
    ]

  if chains_states_was_list:
    return rhat_list
  return rhat_list[0]

def _potential_scale_reduction_single_state(state, independent_chain_ndims):
  """potential_scale_reduction for one single state `Tensor`."""
  with tf.compat.v1.name_scope(
      'potential_scale_reduction_single_state',
      values=[state, independent_chain_ndims]):
    # We assume exactly one leading dimension indexes e.g. correlated samples
    # from each Markov chain.
    state = tf.convert_to_tensor(value=state, name='state')
    sample_ndims = 1

    sample_axis = tf.range(0, sample_ndims)
    chain_axis = tf.range(sample_ndims,
                          sample_ndims + independent_chain_ndims)
    sample_and_chain_axis = tf.range(
        0, sample_ndims + independent_chain_ndims)

    n = _axis_size(state, sample_axis)
    m = _axis_size(state, chain_axis)

    # In the language of Brooks and Gelman (1998),
    # B / n is the between chain variance, the variance of the chain means.
    # W is the within sequence variance, the mean of the chain variances.
    b_div_n = _reduce_variance(
        tf.reduce_mean(input_tensor=state, axis=sample_axis, keepdims=True),
        sample_and_chain_axis,
        biased=False)
    w = tf.reduce_mean(
        input_tensor=_reduce_variance(
            state, sample_axis, keepdims=True, biased=True),
        axis=sample_and_chain_axis)

    # sigma^2_+ is an estimate of the true variance, which would be unbiased if
    # each chain was drawn from the target.  c.f. "law of total variance."
    sigma_2_plus = w + b_div_n

    return ((m + 1.) / m) * sigma_2_plus / w - (n - 1.) / (m * n)

# TODO(b/72873233) Move some variant of this to tfd.sample_stats.
def _reduce_variance(x, axis=None, biased=True, keepdims=False):
  with tf.compat.v1.name_scope('reduce_variance'):
    x = tf.convert_to_tensor(value=x, name='x')
    mean = tf.reduce_mean(input_tensor=x, axis=axis, keepdims=True)
    biased_var = tf.reduce_mean(
        input_tensor=tf.math.squared_difference(x, mean),
        axis=axis,
        keepdims=keepdims)
    if biased:
      return biased_var
    n = _axis_size(x, axis)
    return (n / (n - 1.)) * biased_var

def _axis_size(x, axis=None):
  """Get number of elements of `x` in `axis`, as type `x.dtype`."""
  if axis is None:
    return tf.cast(tf.size(input=x), x.dtype)
  return tf.cast(
      tf.reduce_prod(input_tensor=tf.gather(tf.shape(input=x), axis)), x.dtype)

def _is_list_like(x):
  """Helper which returns `True` if input is `list`-like."""
  return isinstance(x, (tuple, list))

def _broadcast_maybelist_arg(states, secondary_arg, name):
  """Broadcast a listable secondary_arg to that of states."""
  if _is_list_like(secondary_arg):
    if len(secondary_arg) != len(states):
      raise ValueError('Argument `%s` was a list of different length ({}) than '
                       '`states` ({})'.format(name, len(states)))
  else:
    secondary_arg = [secondary_arg] * len(states)

  return secondary_arg


# 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 PoissonLogNormalQuadratureCompound 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.bijectors import exp as exp_bijector
from tensorflow_probability.python.distributions import categorical
from tensorflow_probability.python.distributions import distribution
from tensorflow_probability.python.distributions import normal
from tensorflow_probability.python.distributions import poisson
from tensorflow_probability.python.distributions import seed_stream
from tensorflow_probability.python.distributions import transformed_distribution
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

__all__ = [
    "PoissonLogNormalQuadratureCompound",
    "quadrature_scheme_lognormal_gauss_hermite",
    "quadrature_scheme_lognormal_quantiles",
]

def quadrature_scheme_lognormal_gauss_hermite(
    loc, scale, quadrature_size,
    validate_args=False, name=None):  # pylint: disable=unused-argument
  """Use Gauss-Hermite quadrature to form quadrature on positive-reals.

  Note: for a given `quadrature_size`, this method is generally less accurate
  than `quadrature_scheme_lognormal_quantiles`.

  Args:
    loc: `float`-like (batch of) scalar `Tensor`; the location parameter of
      the LogNormal prior.
    scale: `float`-like (batch of) scalar `Tensor`; the scale parameter of
      the LogNormal prior.
    quadrature_size: Python `int` scalar representing the number of quadrature
      points.
    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.
    name: Python `str` name prefixed to Ops created by this class.

  Returns:
    grid: (Batch of) length-`quadrature_size` vectors representing the
      `log_rate` parameters of a `Poisson`.
    probs: (Batch of) length-`quadrature_size` vectors representing the
      weight associate with each `grid` value.
  """
  with tf.name_scope(
      name or "vector_diffeomixture_quadrature_gauss_hermite"):
    grid, probs = np.polynomial.hermite.hermgauss(deg=quadrature_size)
    npdt = dtype_util.as_numpy_dtype(loc.dtype)
    grid = grid.astype(npdt)
    probs = probs.astype(npdt)
    probs /= np.linalg.norm(probs, ord=1, keepdims=True)
    probs = tf.convert_to_tensor(value=probs, name="probs", dtype=loc.dtype)
    # The following maps the broadcast of `loc` and `scale` to each grid
    # point, i.e., we are creating several log-rates that correspond to the
    # different Gauss-Hermite quadrature points and (possible) batches of
    # `loc` and `scale`.
    grid = (loc[..., tf.newaxis] + np.sqrt(2.) * scale[..., tf.newaxis] * grid)
    return grid, probs

def quadrature_scheme_lognormal_quantiles(
    loc, scale, quadrature_size,
    validate_args=False, name=None):
  """Use LogNormal quantiles to form quadrature on positive-reals.

  Args:
    loc: `float`-like (batch of) scalar `Tensor`; the location parameter of
      the LogNormal prior.
    scale: `float`-like (batch of) scalar `Tensor`; the scale parameter of
      the LogNormal prior.
    quadrature_size: Python `int` scalar representing the number of quadrature
      points.
    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.
    name: Python `str` name prefixed to Ops created by this class.

  Returns:
    grid: (Batch of) length-`quadrature_size` vectors representing the
      `log_rate` parameters of a `Poisson`.
    probs: (Batch of) length-`quadrature_size` vectors representing the
      weight associate with each `grid` value.
  """
  with tf.name_scope(name or "quadrature_scheme_lognormal_quantiles"):
    # Create a LogNormal distribution.
    dist = transformed_distribution.TransformedDistribution(
        distribution=normal.Normal(loc=loc, scale=scale),
        bijector=exp_bijector.Exp(),
        validate_args=validate_args)
    batch_ndims = tensorshape_util.rank(dist.batch_shape)
    if batch_ndims is None:
      batch_ndims = tf.shape(input=dist.batch_shape_tensor())[0]

    def _compute_quantiles():
      """Helper to build quantiles."""
      # Omit {0, 1} since they might lead to Inf/NaN.
      zero = tf.zeros([], dtype=dist.dtype)
      edges = tf.linspace(zero, 1., quadrature_size + 3)[1:-1]
      # Expand edges so its broadcast across batch dims.
      edges = tf.reshape(
          edges,
          shape=tf.concat(
              [[-1], tf.ones([batch_ndims], dtype=tf.int32)], axis=0))
      quantiles = dist.quantile(edges)
      # Cyclically permute left by one.
      perm = tf.concat([tf.range(1, 1 + batch_ndims), [0]], axis=0)
      quantiles = tf.transpose(a=quantiles, perm=perm)
      return quantiles
    quantiles = _compute_quantiles()

    # Compute grid as quantile midpoints.
    grid = (quantiles[..., :-1] + quantiles[..., 1:]) / 2.
    # Set shape hints.
    new_shape = tensorshape_util.concatenate(dist.batch_shape,
                                             [quadrature_size])
    tensorshape_util.set_shape(grid, new_shape)

    # By construction probs is constant, i.e., `1 / quadrature_size`. This is
    # important, because non-constant probs leads to non-reparameterizable
    # samples.
    probs = tf.fill(
        dims=[quadrature_size], value=1. / tf.cast(quadrature_size, dist.dtype))

    return grid, probs

class PoissonLogNormalQuadratureCompound(distribution.Distribution):
  """`PoissonLogNormalQuadratureCompound` distribution.

  The `PoissonLogNormalQuadratureCompound` is an approximation to a
  Poisson-LogNormal [compound distribution](
  https://en.wikipedia.org/wiki/Compound_probability_distribution), i.e.,

  ```none
  p(k|loc, scale)
  = int_{R_+} dl LogNormal(l | loc, scale) Poisson(k | l)
  approx= sum{ prob[d] Poisson(k | lambda(grid[d])) : d=0, ..., deg-1 }
  ```

  By default, the `grid` is chosen as quantiles of the `LogNormal` distribution
  parameterized by `loc`, `scale` and the `prob` vector is
  `[1. / quadrature_size]*quadrature_size`.

  In the non-approximation case, a draw from the LogNormal prior represents the
  Poisson rate parameter. Unfortunately, the non-approximate distribution lacks
  an analytical probability density function (pdf). Therefore the
  `PoissonLogNormalQuadratureCompound` class implements an approximation based
  on [quadrature](https://en.wikipedia.org/wiki/Numerical_integration).

  Note: although the `PoissonLogNormalQuadratureCompound` is approximately the
  Poisson-LogNormal compound distribution, it is itself a valid distribution.
  Viz., it possesses a `sample`, `log_prob`, `mean`, `variance`, etc. which are
  all mutually consistent.

  #### Mathematical Details

  The `PoissonLogNormalQuadratureCompound` approximates a Poisson-LogNormal
  [compound distribution](
  https://en.wikipedia.org/wiki/Compound_probability_distribution). Using
  variable-substitution and [numerical quadrature](
  https://en.wikipedia.org/wiki/Numerical_integration) (default:
  based on `LogNormal` quantiles) we can redefine the distribution to be a
  parameter-less convex combination of `deg` different Poisson samples.

  That is, defined over positive integers, this distribution is parameterized
  by a (batch of) `loc` and `scale` scalars.

  The probability density function (pdf) is,

  ```none
  pdf(k | loc, scale, deg)
    = sum{ prob[d] Poisson(k | lambda=exp(grid[d]))
          : d=0, ..., deg-1 }
  ```

  #### Examples

  ```python
  tfd = tfp.distributions

  # Create two batches of PoissonLogNormalQuadratureCompounds, one with
  # prior `loc = 0.` and another with `loc = 1.` In both cases `scale = 1.`
  pln = tfd.PoissonLogNormalQuadratureCompound(
      loc=[0., -0.5],
      scale=1.,
      quadrature_size=10,
      validate_args=True)
  """

  def __init__(self,
               loc,
               scale,
               quadrature_size=8,
               quadrature_fn=quadrature_scheme_lognormal_quantiles,
               validate_args=False,
               allow_nan_stats=True,
               name="PoissonLogNormalQuadratureCompound"):
    """Constructs the PoissonLogNormalQuadratureCompound`.

    Note: `probs` returned by (optional) `quadrature_fn` are presumed to be
    either a length-`quadrature_size` vector or a batch of vectors in 1-to-1
    correspondence with the returned `grid`. (I.e., broadcasting is only
    partially supported.)

    Args:
      loc: `float`-like (batch of) scalar `Tensor`; the location parameter of
        the LogNormal prior.
      scale: `float`-like (batch of) scalar `Tensor`; the scale parameter of
        the LogNormal prior.
      quadrature_size: Python `int` scalar representing the number of quadrature
        points.
      quadrature_fn: Python callable taking `loc`, `scale`,
        `quadrature_size`, `validate_args` and returning `tuple(grid, probs)`
        representing the LogNormal grid and corresponding normalized weight.
        normalized) weight.
        Default value: `quadrature_scheme_lognormal_quantiles`.
      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 `quadrature_grid` and `quadrature_probs` have different base
        `dtype`.
    """
    parameters = dict(locals())
    with tf.name_scope(name) as name:
      dtype = dtype_util.common_dtype([loc, scale], tf.float32)
      if loc is not None:
        loc = tf.convert_to_tensor(value=loc, name="loc", dtype=dtype)
      if scale is not None:
        scale = tf.convert_to_tensor(value=scale, dtype=dtype, name="scale")
      self._quadrature_grid, self._quadrature_probs = tuple(quadrature_fn(
          loc, scale, quadrature_size, validate_args))

      dt = self._quadrature_grid.dtype
      if not dtype_util.base_equal(dt, self._quadrature_probs.dtype):
        raise TypeError("Quadrature grid dtype ({}) does not match quadrature "
                        "probs dtype ({}).".format(
                            dtype_util.name(dt),
                            dtype_util.name(self._quadrature_probs.dtype)))

      self._distribution = poisson.Poisson(
          log_rate=self._quadrature_grid,
          validate_args=validate_args,
          allow_nan_stats=allow_nan_stats)

      self._mixture_distribution = categorical.Categorical(
          logits=tf.math.log(self._quadrature_probs),
          validate_args=validate_args,
          allow_nan_stats=allow_nan_stats)

      self._loc = loc
      self._scale = scale
      self._quadrature_size = quadrature_size

      super(PoissonLogNormalQuadratureCompound, self).__init__(
          dtype=dt,
          reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
          validate_args=validate_args,
          allow_nan_stats=allow_nan_stats,
          parameters=parameters,
          graph_parents=[loc, scale],
          name=name)

  @property
  def mixture_distribution(self):
    """Distribution which randomly selects a Poisson with quadrature param."""
    return self._mixture_distribution

  @property
  def distribution(self):
    """Base Poisson parameterized by a quadrature grid."""
    return self._distribution

  @property
  def loc(self):
    """Location parameter of the LogNormal prior."""
    return self._loc

  @property
  def scale(self):
    """Scale parameter of the LogNormal prior."""
    return self._scale

  @property
  def quadrature_size(self):
    return self._quadrature_size

  def _batch_shape_tensor(self):
    return tf.broadcast_dynamic_shape(
        self.distribution.batch_shape_tensor(),
        tf.shape(input=self.mixture_distribution.logits))[:-1]

  def _batch_shape(self):
    return tf.broadcast_static_shape(
        self.distribution.batch_shape,
        self.mixture_distribution.logits.shape)[:-1]

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

  def _sample_n(self, n, seed=None):
    # Get ids as a [n, batch_size]-shaped matrix, unless batch_shape=[] then get
    # ids as a [n]-shaped vector.
    batch_size = tensorshape_util.num_elements(self.batch_shape)
    if batch_size is None:
      batch_size = tf.reduce_prod(input_tensor=self.batch_shape_tensor())
    # We need to "sample extra" from the mixture distribution if it doesn't
    # already specify a probs vector for each batch coordinate.
    # We only support this kind of reduced broadcasting, i.e., there is exactly
    # one probs vector for all batch dims or one for each.
    stream = seed_stream.SeedStream(
        seed, salt="PoissonLogNormalQuadratureCompound")
    ids = self._mixture_distribution.sample(
        sample_shape=concat_vectors(
            [n],
            distribution_util.pick_vector(
                self.mixture_distribution.is_scalar_batch(),
                [batch_size],
                np.int32([]))),
        seed=stream())
    # We need to flatten batch dims in case mixture_distribution has its own
    # batch dims.
    ids = tf.reshape(
        ids,
        shape=concat_vectors([n],
                             distribution_util.pick_vector(
                                 self.is_scalar_batch(), np.int32([]),
                                 np.int32([-1]))))

    # Stride `quadrature_size` for `batch_size` number of times.
    offset = tf.range(
        start=0,
        limit=batch_size * self._quadrature_size,
        delta=self._quadrature_size,
        dtype=ids.dtype)
    ids += offset
    rate = tf.gather(tf.reshape(self.distribution.rate, shape=[-1]), ids)
    rate = tf.reshape(
        rate, shape=concat_vectors([n], self.batch_shape_tensor()))
    return tf.random.poisson(lam=rate, shape=[], dtype=self.dtype, seed=seed)

  def _log_prob(self, x):
    return tf.reduce_logsumexp(
        input_tensor=(self.mixture_distribution.logits +
                      self.distribution.log_prob(x[..., tf.newaxis])),
        axis=-1)

  def _mean(self):
    return tf.exp(
        tf.reduce_logsumexp(
            input_tensor=self.mixture_distribution.logits +
            self.distribution.log_rate,
            axis=-1))

  def _variance(self):
    return tf.exp(self._log_variance())

  def _stddev(self):
    return tf.exp(0.5 * self._log_variance())

  def _log_variance(self):
    # Following calculation is based on law of total variance:
    #
    # Var[Z] = E[Var[Z | V]] + Var[E[Z | V]]
    #
    # where,
    #
    # Z|v ~ interpolate_affine[v](distribution)
    # V ~ mixture_distribution
    #
    # thus,
    #
    # E[Var[Z | V]] = sum{ prob[d] Var[d] : d=0, ..., deg-1 }
    # Var[E[Z | V]] = sum{ prob[d] (Mean[d] - Mean)**2 : d=0, ..., deg-1 }
    v = tf.stack(
        [
            # log(self.distribution.variance()) = log(Var[d]) = log(rate[d])
            self.distribution.log_rate,
            # log((Mean[d] - Mean)**2)
            2. * tf.math.log(
                tf.abs(self.distribution.mean() -
                       self._mean()[..., tf.newaxis])),
        ],
        axis=-1)
    return tf.reduce_logsumexp(
        input_tensor=self.mixture_distribution.logits[..., tf.newaxis] + v,
        axis=[-2, -1])

def concat_vectors(*args):
  """Concatenates input vectors, statically if possible."""
  args_ = [tf.get_static_value(x) for x in args]
  if any(vec is None for vec in args_):
    return tf.concat(args, axis=0)
  return [val for vec in args_ for val in vec]