# 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.
# ============================================================================
"""Structural Time Series utilities."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# Dependency imports
import tensorflow as tf

from tensorflow_probability.python import distributions as tfd
from tensorflow_probability.python.distributions.mvn_linear_operator import MultivariateNormalLinearOperator
from tensorflow_probability.python.internal import distribution_util as dist_util
from tensorflow_probability.python.internal import prefer_static
from tensorflow_probability.python.sts.internal import missing_values_util

tfl = tf.linalg

def broadcast_batch_shape(distributions):
  """Get broadcast batch shape from distributions, statically if possible."""

  # Static case
  batch_shape = distributions[0].batch_shape
  for distribution in distributions:
    batch_shape = tf.broadcast_static_shape(batch_shape,
                                            distribution.batch_shape)
  if batch_shape.is_fully_defined():
    return batch_shape.as_list()

  # Fallback on dynamic.
  batch_shape = distributions[0].batch_shape_tensor()
  for distribution in distributions:
    batch_shape = tf.broadcast_dynamic_shape(batch_shape,
                                             distribution.batch_shape_tensor())

  return tf.convert_to_tensor(value=batch_shape)

def pad_batch_dimension_for_multiple_chains(
    observed_time_series, model, chain_batch_shape):
  """"Expand the observed time series with extra batch dimension(s)."""
  # Running with multiple chains introduces an extra batch dimension. In
  # general we also need to pad the observed time series with a matching batch
  # dimension.
  #
  # For example, suppose our model has batch shape [3, 4] and
  # the observed time series has shape `concat([[5], [3, 4], [100])`,
  # corresponding to `sample_shape`, `batch_shape`, and `num_timesteps`
  # respectively. The model will produce distributions with batch shape
  # `concat([chain_batch_shape, [3, 4]])`, so we pad `observed_time_series` to
  # have matching shape `[5, 1, 3, 4, 100]`, where the added `1` dimension
  # between the sample and batch shapes will broadcast to `chain_batch_shape`.

  [  # Extract mask and guarantee `event_ndims=2`.
      observed_time_series,
      is_missing
  ] = canonicalize_observed_time_series_with_mask(observed_time_series)

  event_ndims = 2  # event_shape = [num_timesteps, observation_size=1]

  model_batch_ndims = (
      model.batch_shape.ndims if model.batch_shape.ndims is not None else
      tf.shape(input=model.batch_shape_tensor())[0])

  # Compute ndims from chain_batch_shape.
  chain_batch_shape = tf.convert_to_tensor(
      value=chain_batch_shape, name='chain_batch_shape', dtype=tf.int32)
  if not chain_batch_shape.shape.is_fully_defined():
    raise ValueError('Batch shape must have static rank. (given: {})'.format(
        chain_batch_shape))
  if chain_batch_shape.shape.ndims == 0:  # expand int `k` to `[k]`.
    chain_batch_shape = chain_batch_shape[tf.newaxis]
  chain_batch_ndims = tf.compat.dimension_value(chain_batch_shape.shape[0])

  def do_padding(observed_time_series_tensor):
    current_sample_shape = tf.shape(
        input=observed_time_series_tensor)[:-(model_batch_ndims + event_ndims)]
    current_batch_and_event_shape = tf.shape(
        input=observed_time_series_tensor)[-(model_batch_ndims + event_ndims):]
    return tf.reshape(
        tensor=observed_time_series_tensor,
        shape=tf.concat([
            current_sample_shape,
            tf.ones([chain_batch_ndims], dtype=tf.int32),
            current_batch_and_event_shape], axis=0))

  # Padding is only needed if the observed time series has sample shape.
  observed_time_series = prefer_static.cond(
      (dist_util.prefer_static_rank(observed_time_series) >
       model_batch_ndims + event_ndims),
      lambda: do_padding(observed_time_series),
      lambda: observed_time_series)

  if is_missing is not None:
    is_missing = prefer_static.cond(
        (dist_util.prefer_static_rank(is_missing) >
         model_batch_ndims + event_ndims),
        lambda: do_padding(is_missing),
        lambda: is_missing)
    return missing_values_util.MaskedTimeSeries(observed_time_series,
                                                is_missing=is_missing)

  return observed_time_series

def factored_joint_mvn(distributions):
  """Combine MultivariateNormals into a factored joint distribution.

   Given a list of multivariate normal distributions
   `dist[i] = Normal(loc[i], scale[i])`, construct the joint
   distribution given by concatenating independent samples from these
   distributions. This is multivariate normal with mean vector given by the
   concatenation of the component mean vectors, and block-diagonal covariance
   matrix in which the blocks are the component covariances.

   Note that for computational efficiency, multivariate normals are represented
   by a 'scale' (factored covariance) linear operator rather than the full
   covariance matrix.

  Args:
    distributions: Python `iterable` of MultivariateNormal distribution
      instances (e.g., `tfd.MultivariateNormalDiag`,
      `tfd.MultivariateNormalTriL`, etc.). These must be broadcastable to a
      consistent batch shape, but may have different event shapes
      (i.e., defined over spaces of different dimension).

  Returns:
    joint_distribution: An instance of `tfd.MultivariateNormalLinearOperator`
      representing the joint distribution constructed by concatenating
      an independent sample from each input distributions.
  """

  graph_parents = [tensor for distribution in distributions
                   for tensor in distribution._graph_parents]  # pylint: disable=protected-access
  with tf.compat.v1.name_scope('factored_joint_mvn', values=graph_parents):

    # We explicitly broadcast the `locs` so that we can concatenate them.
    # We don't have direct numerical access to the `scales`, which are arbitrary
    # linear operators, but `LinearOperatorBlockDiag` appears to do the right
    # thing without further intervention.
    dtype = tf.debugging.assert_same_float_dtype(distributions)
    broadcast_ones = tf.ones(broadcast_batch_shape(distributions),
                             dtype=dtype)[..., tf.newaxis]
    return MultivariateNormalLinearOperator(
        loc=tf.concat([mvn.mean() * broadcast_ones for mvn in distributions],
                      axis=-1),
        scale=tfl.LinearOperatorBlockDiag([mvn.scale for mvn in distributions],
                                          is_square=True))

def sum_mvns(distributions):
  """Attempt to sum MultivariateNormal distributions.

  The sum of (multivariate) normal random variables is itself (multivariate)
  normal, with mean given by the sum of means and (co)variance given by the
  sum of (co)variances. This method exploits this fact to compute the
  sum of a list of `tfd.MultivariateNormalDiag` objects.

  It may in the future be extended to support summation of other forms of
  (Multivariate)Normal distributions.

  Args:
    distributions: Python `iterable` of `tfd.MultivariateNormalDiag`
      distribution instances. These must all have the same event
      shape, and broadcast to a consistent batch shape.

  Returns:
    sum_distribution: A `tfd.MultivariateNormalDiag` instance with mean
      equal to the sum of input means and covariance equal to the sum of
      input covariances.
  """

  graph_parents = [tensor for distribution in distributions
                   for tensor in distribution._graph_parents]  # pylint: disable=protected-access
  with tf.compat.v1.name_scope('sum_mvns', values=graph_parents):
    if all([isinstance(mvn, tfd.MultivariateNormalDiag)
            for mvn in distributions]):
      return tfd.MultivariateNormalDiag(
          loc=sum([mvn.mean() for mvn in distributions]),
          scale_diag=tf.sqrt(sum([
              mvn.scale.diag**2 for mvn in distributions])))
    else:
      raise NotImplementedError(
          'Sums of distributions other than MultivariateNormalDiag are not '
          'currently implemented. (given: {})'.format(distributions))

def empirical_statistics(observed_time_series):
  """Compute statistics of a provided time series, as heuristic initialization.

  Args:
    observed_time_series: `Tensor` representing a time series, or batch of time
       series, of shape either `batch_shape + [num_timesteps, 1]` or
       `batch_shape + [num_timesteps]` (allowed if `num_timesteps > 1`).

  Returns:
    observed_mean: `Tensor` of shape `batch_shape`, giving the empirical
      mean of each time series in the batch.
    observed_stddev: `Tensor` of shape `batch_shape`, giving the empirical
      standard deviation of each time series in the batch.
    observed_initial_centered: `Tensor of shape `batch_shape`, giving the
      initial value of each time series in the batch after centering
      (subtracting the mean).
  """

  with tf.compat.v1.name_scope(
      'empirical_statistics', values=[observed_time_series]):

    [
        observed_time_series,
        mask
    ] = canonicalize_observed_time_series_with_mask(observed_time_series)

    squeezed_series = observed_time_series[..., 0]
    if mask is None:
      observed_mean, observed_variance = tf.nn.moments(
          x=squeezed_series, axes=-1)
      observed_initial = squeezed_series[..., 0]
    else:
      broadcast_mask = tf.broadcast_to(tf.cast(mask, tf.bool),
                                       tf.shape(input=squeezed_series))
      observed_mean, observed_variance = (
          missing_values_util.moments_of_masked_time_series(
              squeezed_series, broadcast_mask=broadcast_mask))
      try:
        observed_initial = (
            missing_values_util.initial_value_of_masked_time_series(
                squeezed_series, broadcast_mask=broadcast_mask))
      except NotImplementedError:
        tf.compat.v1.logging.warn(
            'Cannot compute initial values for a masked time series'
            'with dynamic shape; using the mean instead. This will'
            'affect heuristic priors and may change the results of'
            'inference.')
        observed_initial = observed_mean

    observed_stddev = tf.sqrt(observed_variance)
    observed_initial_centered = observed_initial - observed_mean
    return observed_mean, observed_stddev, observed_initial_centered

def _maybe_expand_trailing_dim(observed_time_series_tensor):
  """Ensures `observed_time_series_tensor` has a trailing dimension of size 1.

  The `tfd.LinearGaussianStateSpaceModel` Distribution has event shape of
  `[num_timesteps, observation_size]`, but canonical BSTS models
  are univariate, so their observation_size is always `1`. The extra trailing
  dimension gets annoying, so this method allows arguments with or without the
  extra dimension. There is no ambiguity except in the trivial special case
  where  `num_timesteps = 1`; this can be avoided by specifying any unit-length
  series in the explicit `[num_timesteps, 1]` style.

  Most users should not call this method directly, and instead call
  `canonicalize_observed_time_series_with_mask`, which handles converting
  to `Tensor` and specifying an optional missingness mask.

  Args:
    observed_time_series_tensor: `Tensor` of shape
      `batch_shape + [num_timesteps, 1]` or `batch_shape + [num_timesteps]`,
      where `num_timesteps > 1`.

  Returns:
    expanded_time_series: `Tensor` of shape `batch_shape + [num_timesteps, 1]`.
  """

  with tf.compat.v1.name_scope(
      'maybe_expand_trailing_dim', values=[observed_time_series_tensor]):
    if (observed_time_series_tensor.shape.ndims is not None and
        tf.compat.dimension_value(
            observed_time_series_tensor.shape[-1]) is not None):
      expanded_time_series = (
          observed_time_series_tensor
          if observed_time_series_tensor.shape[-1] == 1
          else observed_time_series_tensor[..., tf.newaxis])
    else:
      expanded_time_series = tf.cond(
          pred=tf.equal(tf.shape(input=observed_time_series_tensor)[-1], 1),
          true_fn=lambda: observed_time_series_tensor,
          false_fn=lambda: observed_time_series_tensor[..., tf.newaxis])
    return expanded_time_series

def canonicalize_observed_time_series_with_mask(
    maybe_masked_observed_time_series):
  """Extract a Tensor with canonical shape and optional mask.

  Args:
    maybe_masked_observed_time_series: a `Tensor`-like object with shape
      `[..., num_timesteps]` or `[..., num_timesteps, 1]`, or a
      `tfp.sts.MaskedTimeSeries` containing such an object.
  Returns:
    masked_time_series: a `tfp.sts.MaskedTimeSeries` namedtuple, in which
      the `observed_time_series` is converted to `Tensor` with canonical shape
      `[..., num_timesteps, 1]`, and `is_missing` is either `None` or a boolean
      `Tensor`.
  """

  with tf.compat.v1.name_scope('canonicalize_observed_time_series_with_mask'):
    if hasattr(maybe_masked_observed_time_series, 'is_missing'):
      observed_time_series = (
          maybe_masked_observed_time_series.time_series)
      is_missing = maybe_masked_observed_time_series.is_missing
    else:
      observed_time_series = maybe_masked_observed_time_series
      is_missing = None

    observed_time_series = tf.convert_to_tensor(value=observed_time_series,
                                                name='observed_time_series')
    observed_time_series = _maybe_expand_trailing_dim(observed_time_series)

    if is_missing is not None:
      is_missing = tf.convert_to_tensor(
          value=is_missing, name='is_missing', dtype_hint=tf.bool)

    return missing_values_util.MaskedTimeSeries(observed_time_series,
                                                is_missing=is_missing)

def mix_over_posterior_draws(means, variances):
  """Construct a predictive normal distribution that mixes over posterior draws.

  Args:
    means: float `Tensor` of shape
      `[num_posterior_draws, ..., num_timesteps]`.
    variances: float `Tensor` of shape
      `[num_posterior_draws, ..., num_timesteps]`.

  Returns:
    mixture_dist: `tfd.MixtureSameFamily(tfd.Independent(tfd.Normal))` instance
      representing a uniform mixture over the posterior samples, with
      `batch_shape = ...` and `event_shape = [num_timesteps]`.

  """
  # The inputs `means`, `variances` have shape
  #   `concat([
  #      [num_posterior_draws],
  #      sample_shape,
  #      batch_shape,
  #      [num_timesteps]])`
  # Because MixtureSameFamily mixes over the rightmost batch dimension,
  # we need to move the `num_posterior_draws` dimension to be rightmost
  # in the batch shape. This requires use of `Independent` (to preserve
  # `num_timesteps` as part of the event shape) and `move_dimension`.
  # TODO(b/120245392): enhance `MixtureSameFamily` to reduce along an
  # arbitrary axis, and eliminate `move_dimension` calls here.

  with tf.compat.v1.name_scope(
      'mix_over_posterior_draws', values=[means, variances]):
    num_posterior_draws = dist_util.prefer_static_value(
        tf.shape(input=means))[0]

    component_observations = tfd.Independent(
        distribution=tfd.Normal(
            loc=dist_util.move_dimension(means, 0, -2),
            scale=tf.sqrt(dist_util.move_dimension(variances, 0, -2))),
        reinterpreted_batch_ndims=1)

    return tfd.MixtureSameFamily(
        mixture_distribution=tfd.Categorical(
            logits=tf.zeros([num_posterior_draws],
                            dtype=component_observations.dtype)),
        components_distribution=component_observations)


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

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

import math

import numpy as np
import tensorflow.compat.v2 as tf
from tensorflow_probability.python.distributions import distribution
from tensorflow_probability.python.distributions import kullback_leibler
from tensorflow_probability.python.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

class Uniform(distribution.Distribution):
  """Uniform distribution with `low` and `high` parameters.

  #### Mathematical Details

  The probability density function (pdf) is,

  ```none
  pdf(x; a, b) = I[a <= x < b] / Z
  Z = b - a
  ```

  where

  - `low = a`,
  - `high = b`,
  - `Z` is the normalizing constant, and
  - `I[predicate]` is the [indicator function](
    https://en.wikipedia.org/wiki/Indicator_function) for `predicate`.

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

  #### Examples

  ```python
  # Without broadcasting:
  u1 = Uniform(low=3.0, high=4.0)  # a single uniform distribution [3, 4]
  u2 = Uniform(low=[1.0, 2.0],
               high=[3.0, 4.0])  # 2 distributions [1, 3], [2, 4]
  u3 = Uniform(low=[[1.0, 2.0],
                    [3.0, 4.0]],
               high=[[1.5, 2.5],
                     [3.5, 4.5]])  # 4 distributions
  ```

  ```python
  # With broadcasting:
  u1 = Uniform(low=3.0, high=[5.0, 6.0, 7.0])  # 3 distributions
  ```

  """

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

    Args:
      low: Floating point tensor, lower boundary of the output interval. Must
        have `low < high`.
      high: Floating point tensor, upper boundary of the output interval. Must
        have `low < high`.
      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:
      InvalidArgumentError: if `low >= high` and `validate_args=False`.
    """
    parameters = dict(locals())
    with tf.name_scope(name) as name:
      dtype = dtype_util.common_dtype([low, high], tf.float32)
      low = tf.convert_to_tensor(value=low, name="low", dtype=dtype)
      high = tf.convert_to_tensor(value=high, name="high", dtype=dtype)
      with tf.control_dependencies([
          assert_util.assert_less(
              low, high, message="uniform not defined when low >= high.")
      ] if validate_args else []):
        self._low = tf.identity(low)
        self._high = tf.identity(high)
        dtype_util.assert_same_float_dtype([self._low, self._high])
    super(Uniform, 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],
        name=name)

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

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

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

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

  def range(self, name="range"):
    """`high - low`."""
    with self._name_scope(name):
      return self.high - self.low

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

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

  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):
    shape = tf.concat([[n], self.batch_shape_tensor()], 0)
    samples = tf.random.uniform(shape=shape, dtype=self.dtype, seed=seed)
    return self.low + self.range() * samples

  def _prob(self, x):
    broadcasted_x = x * tf.ones(
        self.batch_shape_tensor(), dtype=x.dtype)
    return tf.where(
        tf.math.is_nan(broadcasted_x),
        broadcasted_x,
        tf.where(
            tf.logical_or(
                broadcasted_x < self.low,
                # This > is only sound for continuous uniform
                broadcasted_x > self.high),
            tf.zeros_like(broadcasted_x),
            tf.ones_like(broadcasted_x) / self.range()))

  def _cdf(self, x):
    broadcast_shape = tf.broadcast_dynamic_shape(
        tf.shape(input=x), self.batch_shape_tensor())
    zeros = tf.zeros(broadcast_shape, dtype=self.dtype)
    ones = tf.ones(broadcast_shape, dtype=self.dtype)
    broadcasted_x = x * ones
    result_if_not_big = tf.where(
        x < self.low, zeros, (broadcasted_x - self.low) / self.range())
    return tf.where(x >= self.high, ones, result_if_not_big)

  def _quantile(self, value):
    broadcast_shape = tf.broadcast_dynamic_shape(
        tf.shape(input=value), self.batch_shape_tensor())
    ones = tf.ones(broadcast_shape, dtype=self.dtype)
    broadcasted_value = value * ones
    return (1. - broadcasted_value) * self.low + broadcasted_value * self.high

  def _entropy(self):
    return tf.math.log(self.range())

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

  def _variance(self):
    return tf.square(self.range()) / 12.

  def _stddev(self):
    return self.range() / math.sqrt(12.)

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

  Note that the KL divergence is infinite if the support of `a` is not a subset
  of the support of `b`.

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

  Returns:
    Batchwise KL(a || b)
  """
  with tf.name_scope(name or "kl_uniform_uniform"):
    # Consistent with
    # http://www.mast.queensu.ca/~communications/Papers/gil-msc11.pdf, page 60
    # Watch out for the change in conventions--they use 'a' and 'b' to refer to
    # lower and upper bounds respectively there.
    final_batch_shape = distribution_util.get_broadcast_shape(
        a.low, b.low, a.high, b.high)
    dtype = dtype_util.common_dtype(
        [a.low, a.high, b.low, b.high], tf.float32)
    return tf.where((b.low <= a.low) & (a.high <= b.high),
                    tf.math.log(b.high - b.low) - tf.math.log(a.high - a.low),
                    tf.broadcast_to(
                        dtype_util.as_numpy_dtype(dtype)(np.inf),
                        final_batch_shape))


# 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 `JointDistributionSequential` class."""

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

import collections
import functools

import tensorflow.compat.v2 as tf

from tensorflow_probability.python.distributions import joint_distribution as joint_distribution_lib
from tensorflow_probability.python.distributions import kullback_leibler
from tensorflow_probability.python.distributions import seed_stream
from tensorflow_probability.python.internal import distribution_util
from tensorflow.python.util import tf_inspect  # pylint: disable=g-direct-tensorflow-import

__all__ = [
    'JointDistributionSequential',
]

def _make_summary_statistic(attr):
  """Factory for making summary statistics, eg, mean, mode, stddev."""
  def _fn(self):
    if any(self._dist_fn_args):  # pylint: disable=protected-access
      raise ValueError(
          'Can only compute ' + attr + ' when all distributions are '
          'independent; {}'.format(self.model))
    return self._unflatten(getattr(d(), attr)() for d in self._dist_fn_wrapped)  # pylint: disable=protected-access
  return _fn

class JointDistributionSequential(joint_distribution_lib.JointDistribution):
  """Joint distribution parameterized by distribution-making functions.

  This distribution enables both sampling and joint probability computation from
  a single model specification.

  A joint distribution is a collection of possibly interdependent distributions.
  Like `tf.keras.Sequential`, the `JointDistributionSequential` can be specified
  via a `list` of functions (each responsible for making a
  `tfp.distributions.Distribution`-like instance).  Unlike
  `tf.keras.Sequential`, each function can depend on the output of all previous
  elements rather than only the immediately previous.

  #### Mathematical Details

  The `JointDistributionSequential` implements the chain rule of probability.
  That is, the probability function of a length-`d` vector `x` is,

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

  The `JointDistributionSequential` is parameterized by a `list` comprised of
  either:

  1. `tfp.distributions.Distribution`-like instances or,
  2. `callable`s which return a `tfp.distributions.Distribution`-like instance.

  Each `list` element implements the `i`-th *full conditional distribution*,
  `p(x[i] | x[:i])`. The "conditioned on" elements are represented by the
  `callable`'s required arguments. Directly providing a `Distribution`-like
  instance is a convenience and is semantically identical a zero argument
  `callable`.

  Denote the `i`-th `callable`s non-default arguments as `args[i]`. Since the
  `callable` is the conditional manifest, `0 <= len(args[i]) <= i - 1`. When
  `len(args[i]) < i - 1`, the `callable` only depends on a subset of the
  previous distributions, specifically those at indexes:
  `range(i - 1, i - 1 - num_args[i], -1)`.
  (See "Examples" and "Discussion" for why the order is reversed.)

  #### Examples

  ```python
  tfd = tfp.distributions

  # Consider the following generative model:
  #     e ~ Exponential(rate=[100,120])
  #     g ~ Gamma(concentration=e[0], rate=e[1])
  #     n ~ Normal(loc=0, scale=2.)
  #     m ~ Normal(loc=n, scale=g)
  #     for i = 1, ..., 12:
  #       x[i] ~ Bernoulli(logits=m)

  # In TFP, we can write this as:
  joint = tfd.JointDistributionSequential([
                   tfd.Independent(tfd.Exponential(rate=[100, 120]), 1),  # e
      lambda    e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]),    # g
                   tfd.Normal(loc=0, scale=2.),                           # n
      lambda n, g: tfd.Normal(loc=n, scale=g)                             # m
      lambda    m: tfd.Sample(tfd.Bernoulli(logits=m), 12)                # x
  ])
  # (Notice the 1:1 correspondence between "math" and "code".)

  x = joint.sample()
  # ==> A length-5 list of tfd.Distribution instances
  joint.log_prob(x)
  # ==> A scalar `Tensor` representing the total log prob under all four
  #     distributions.

  joint._resolve_graph()
  # ==> (('e', ()),
  #      ('g', ('e',)),
  #      ('n', ()),
  #      ('m', ('n', 'g')),
  #      ('x', ('m',)))
  ```

  #### Discussion

  `JointDistributionSequential` builds each distribution in `list` order; list
  items must be either a:

  1. `tfd.Distribution`-like instance (e.g., `e` and `n`), or a
  2. Python `callable` (e.g., `g`, `m`, `x`).

  Regarding #1, an object is deemed "`tfd.Distribution`-like" if it has a
  `sample`, `log_prob`, and distribution properties, e.g., `batch_shape`,
  `event_shape`, `dtype`.

  Regarding #2, in addition to using a function (or `lambda`), supplying a TFD
  "`class`" is also permissible, this also being a "Python `callable`." For
  example, instead of writing:
  `lambda loc, scale: tfd.Normal(loc=loc, scale=scale)`
  one could have simply written `tfd.Normal`.

  Notice that directly providing a `tfd.Distribution`-like instance means there
  cannot exist a (dynamic) dependency on other distributions; it is
  "independent" both "computationally" and "statistically." The same is
  self-evidently true of zero-argument `callable`s.

  A distribution instance depends on other distribution instances through the
  distribution making function's *required arguments*. If the distribution maker
  has `k` required arguments then the `JointDistributionSequential` calls the
  maker with samples produced by the previous `k` distributions.

  **Note**: **maker arguments are provided in reverse order** of the previous
  elements in the list. In the example, notice that `m` depends on `n` and `g`
  in this order. The order is reversed for convenience. We reverse arguments
  under the heuristic observation that many graphical models have chain-like
  dependencies which are self-evidently topologically sorted from the human
  cognitive process of postulating the generative process. By feeding the
  previous num required args in reverse order we (often) enable a simpler maker
  function signature. If the maker needs to depend on distribution previous to
  one which is not a dependency, one must use a dummy arg, to "gobble up" the
  unused dependency, e.g.,
  `lambda _ignore, actual_dependency: SomeDistribution(actual_dependency)`.

  **Note**: unlike other non-`JointDistribution` distributions in
  `tfp.distributions`, `JointDistribution.sample` (and subclasses) return a
  structure of  `Tensor`s rather than a `Tensor`.  A structure can be anything
  which is `list`-like, e.g., a `list` or `tuple` of `distribution` makers.
  Accordingly `joint.batch_shape` returns a `list`-like structure of
  `TensorShape`s for each of the distributions' batch shapes and
  `joint.batch_shape_tensor()` returns a `list`-like structure of `Tensor`s for
  each of the distributions' event shapes. (Same with `event_shape` analogues.)
  """

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

    Args:
      model: Python list of either tfd.Distribution instances and/or
        lambda functions which take the `k` previous distributions and returns a
        new tfd.Distribution instance.
      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.
        Default value: `False`.
      name: The name for ops managed by the distribution.
        Default value: `None` (i.e., `"JointDistributionSequential"`).
    """
    parameters = dict(locals())
    with tf.name_scope(name or 'JointDistributionSequential') as name:
      self._model = model
      self._build(model)
      self._most_recently_built_distributions = [
          None if a else d() for d, a
          in zip(self._dist_fn_wrapped, self._dist_fn_args)]
      self._always_use_specified_sample_shape = False
      super(JointDistributionSequential, self).__init__(
          dtype=None,  # Ignored; we'll override.
          reparameterization_type=None,  # Ignored; we'll override.
          validate_args=validate_args,
          allow_nan_stats=False,
          parameters=parameters,
          graph_parents=[],
          name=name)
      # Check valid structure.
      self._unflatten(self._flatten(model))

  def _build(self, model):
    """Creates `dist_fn`, `dist_fn_wrapped`, `dist_fn_args`."""
    if not isinstance(model, collections.Sequence):
      raise TypeError('`model` must be `list`-like (saw: {}).'.format(
          type(model).__name__))
    self._dist_fn = model
    self._dist_fn_wrapped, self._dist_fn_args = zip(*[
        _unify_call_signature(i, dist_fn)
        for i, dist_fn in enumerate(model)])

  def _flat_sample_distributions(self, sample_shape=(), seed=None, value=None):
    # This function additionally depends on:
    #   self._dist_fn_wrapped
    #   self._dist_fn_args
    #   self._always_use_specified_sample_shape
    seed = seed_stream.SeedStream('JointDistributionSequential', seed)
    ds = []
    xs = [None]*len(self._dist_fn_wrapped) if value is None else list(value)
    if len(xs) != len(self._dist_fn_wrapped):
      raise ValueError('Number of `xs`s must match number of '
                       'distributions.')
    for i, (dist_fn, args) in enumerate(zip(self._dist_fn_wrapped,
                                            self._dist_fn_args)):
      ds.append(dist_fn(*xs[:i]))  # Chain rule of probability.
      if xs[i] is None:
        # TODO(b/129364796): We should ignore args prefixed with `_`; this
        # would mean we more often identify when to use `sample_shape=()`
        # rather than `sample_shape=sample_shape`.
        xs[i] = ds[-1].sample(
            () if args and not self._always_use_specified_sample_shape
            else sample_shape, seed=seed())
      else:
        xs[i] = tf.convert_to_tensor(value=xs[i], dtype_hint=ds[-1].dtype)
        seed()  # Ensure reproducibility even when xs are (partially) set.
    # Note: we could also resolve distributions up to the first non-`None` in
    # `self._flatten(value)`, however we omit this feature for simplicity,
    # speed, and because it has not yet been requested.
    return ds, xs

  def _unflatten(self, xs):
    try:
      return type(self.model)(xs)
    except TypeError:
      raise TypeError(
          'Unable to unflatten like `model` with type "{}".'.format(
              type(self.model).__name__))

  def _flatten(self, xs):
    if xs is None:
      return (None,) * len(self._dist_fn_wrapped)
    try:
      xs = tuple(xs)
    except TypeError:
      raise TypeError(
          'Unable to flatten like `model` with type "{}".'.format(
              type(self.model).__name__))
    return xs + (None,)*(len(self._dist_fn_args) - len(xs))

  def _call_attr(self, attr):
    if any(self._dist_fn_args):
      # Const seed for maybe CSE.
      ds, _ = self._flat_sample_distributions(seed=42)
    else:
      ds = tuple(d() for d in self._dist_fn_wrapped)
    return (getattr(d, attr)() for d in ds)

  def _resolve_graph(self, distribution_names=None, leaf_name='x'):
    """Creates a `tuple` of `tuple`s of dependencies.

    This function is **experimental**. That said, we encourage its use
    and ask that you report problems to `tfprobability@tensorflow.org`.

    Args:
      distribution_names: `list` of `str` or `None` names corresponding to each
        of `model` elements. (`None`s are expanding into the
        appropriate `str`.)
      leaf_name: `str` used when no maker depends on a particular
        `model` element.

    Returns:
      graph: `tuple` of `(str tuple)` pairs representing the name of each
        distribution (maker) and the names of its dependencies.

    #### Example

    ```python
    d = tfd.JointDistributionSequential([
                     tfd.Independent(tfd.Exponential(rate=[100, 120]), 1),
        lambda    e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]),
                     tfd.Normal(loc=0, scale=2.),
        lambda n, g: tfd.Normal(loc=n, scale=g),
    ])
    d._resolve_graph()
    # ==> (
    #       ('e', ()),
    #       ('g', ('e',)),
    #       ('n', ()),
    #       ('x', ('n', 'g')),
    #     )
    ```

    """
    # This function additionally depends on:
    #   self._dist_fn_args
    #   self._dist_fn_wrapped
    # TODO(b/129008220): Robustify this procedure. Eg, handle collisions better,
    # ignore args prefixed with `_`.
    if distribution_names is None or any(self._dist_fn_args):
      distribution_names = _resolve_distribution_names(
          self._dist_fn_args, distribution_names, leaf_name)
    if len(set(distribution_names)) != len(distribution_names):
      raise ValueError('Distribution names must be unique: {}'.format(
          distribution_names))
    if len(distribution_names) != len(self._dist_fn_wrapped):
      raise ValueError('Distribution names must be 1:1 with `rvs`.')
    return tuple(zip(distribution_names,
                     tuple(() if a is None else a for a in self._dist_fn_args)))

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

  def _entropy(self):
    """Shannon entropy in nats."""
    if any(self._dist_fn_args):
      raise ValueError(
          'Can only compute entropy when all distributions are independent.')
    return sum(joint_distribution_lib.maybe_check_wont_broadcast(
        (d().entropy() for d in self._dist_fn_wrapped),
        self.validate_args))

  def _cross_entropy(self, other):
    if (not isinstance(other, JointDistributionSequential) or
        len(self.model) != len(other.model)):
      raise ValueError(
          'Can only compute cross entropy between `JointDistributionSequential`s '
          'with the same number of component distributions.')
    if any(self._dist_fn_args) or any(other._dist_fn_args):  # pylint: disable=protected-access
      raise ValueError(
          'Can only compute cross entropy when all component distributions '
          'are independent.')
    return sum(joint_distribution_lib.maybe_check_wont_broadcast(
        (d0().cross_entropy(d1()) for d0, d1
         in zip(self._dist_fn_wrapped, other._dist_fn_wrapped)),  # pylint: disable=protected-access
        self.validate_args))

  def __getitem__(self, slices):
    # There will be some users that who won't like how we've implemented
    # slicing. The currently implemented policy is to slice "upstream" and let
    # post sliced distributions propagate forward. While this is efficient it
    # also means that downstream makers *might* no longer correctly broadcast
    # with the new inputs.
    #
    # For example:
    #
    # ```python
    # d = JointDistributionSequential([
    #   tfd.Normal(0, scale=[0.5, 1, 1.5]),
    #   lambda m: tfd.Normal(loc=m, scale=[1, 2, 3]),
    # ])
    # d[1:]
    # # ==> An exception: the following is not possible:
    # #     Normal(loc=Normal(0, [1, 1.5]).sample(), scale=[1, 2, 3])
    # #     since shape `[2]` cannot broadcast with shape `[3]`.
    # ```
    #
    # While not supporting this is technically not a bug (its an API decision)
    # it certainly won't be universally loved. That this is ok follows from the
    # fact that it will at least fail loudly--unless you slice down so a
    # dimension now has size `1`--this will broadcast. Here again we argue this
    # is "ok" because at this level of TFP users are ultimately responsible for
    # ensuring the correct shapes in their models. Moreover, this problem can
    # be worked around by using a `Deterministic` node in lieu of any constants.
    #
    # Once JointDistributionSequential can be parameterized by a single maker
    # function, we can choose to support more slicing scenarios by building
    # every distribution as we would normally then slicing distributions once
    # built (and passing that as the sole maker to `self.copy`.) Note however
    # that this would not be as performant as we've currently implemented
    # things.
    dfn = []
    def _sliced_maker(d):
      def _fn():
        return d()[slices]
      return _fn
    for d, args in zip(self._dist_fn, self._dist_fn_args):
      if args is None:
        dfn.append(d[slices])
      elif args:
        # Don't slice makers which have inputs; they'll be sliced via upstream.
        dfn.append(d)
      else:
        # Makers which have no deps need slicing, just like distribution
        # instances.
        dfn.append(_sliced_maker(d))
    return self.copy(model=self._unflatten(dfn))

def _unify_call_signature(i, dist_fn):
  """Creates `dist_fn_wrapped` which calls `dist_fn` with all prev nodes.

  Args:
    i: Python `int` corresponding to position in topologically sorted DAG.
    dist_fn: Python `callable` which takes a subset of previously constructed
      distributions (in reverse order) and produces a new distribution instance.

  Returns:
    dist_fn_wrapped: Python `callable` which takes all previous distributions
      (in non reverse order) and produces a  new distribution instance.
    args: `tuple` of `str` representing the arg names of `dist_fn` (and in non
      wrapped, "natural" order). `None` is returned only if the input is not a
      `callable`.
  """
  if distribution_util.is_distribution_instance(dist_fn):
    return (lambda *_: dist_fn), None

  if not callable(dist_fn):
    raise TypeError('{} must be either `tfd.Distribution`-like or '
                    '`callable`.'.format(dist_fn))

  args = _get_required_args(dist_fn)
  if not args:
    return (lambda *_: dist_fn()), ()

  @functools.wraps(dist_fn)
  def dist_fn_wrapped(*xs):
    """Calls `dist_fn` with reversed and truncated args."""
    if i != len(xs):
      raise ValueError(
          'Internal Error: Unexpected number of inputs provided to {}-th '
          'distribution maker (dist_fn: {}, expected: {}, saw: {}).'.format(
              i, dist_fn, i, len(xs)))
    if len(xs) < len(args):
      raise ValueError(
          'Internal Error: Too few inputs provided to {}-th distribution maker '
          '(dist_fn: {}, expected: {}, saw: {}).'.format(
              i, dist_fn, len(args), len(xs)))
    return dist_fn(*reversed(xs[-len(args):]))
  return dist_fn_wrapped, args

def _resolve_distribution_names(dist_fn_args, dist_names, leaf_name):
  """Uses arg names to resolve distribution names."""
  if dist_names is None:
    dist_names = []
  else:
    dist_names = dist_names.copy()
  n = len(dist_fn_args)
  dist_names.extend([None]*(n - len(dist_names)))
  for i_, args in enumerate(reversed(dist_fn_args)):
    if not args:
      continue  # There's no args to analyze.
    i = n - i_ - 1
    for j, arg_name in enumerate(args):
      dist_names[i - j - 1] = arg_name
  j = 0
  for i_ in range(len(dist_names)):
    i = n - i_ - 1
    if dist_names[i] is None:
      dist_names[i] = leaf_name if j == 0 else leaf_name + str(j)
      j += 1
  return tuple(dist_names)

def _get_required_args(fn):
  """Returns the distribution's required args."""
  argspec = tf_inspect.getfullargspec(fn)
  args = argspec.args
  if tf_inspect.isclass(fn):
    args = args[1:]  # Remove the `self` arg.
  if argspec.defaults:
    # Remove the args which have defaults. By convention we only feed
    # *required args*. This means some distributions must always be wrapped
    # with a `lambda`, e.g., `lambda logits: tfd.Bernoulli(logits=logits)`
    # or `lambda probs: tfd.Bernoulli(probs=probs)`.
    args = args[:-len(argspec.defaults)]
  return tuple(args)

@kullback_leibler.RegisterKL(JointDistributionSequential,
                             JointDistributionSequential)
def _kl_joint_joint(d0, d1, name=None):
  """Calculate the KL divergence between two `JointDistributionSequential`s.

  Args:
    d0: instance of a `JointDistributionSequential` object.
    d1: instance of a `JointDistributionSequential` object.
    name: (optional) Name to use for created operations.
      Default value: `"kl_joint_joint"`.

  Returns:
    kl_joint_joint: `Tensor` The sum of KL divergences between elemental
      distributions of two joint distributions.

  Raises:
    ValueError: when joint distributions have a different number of elemental
      distributions.
    ValueError: when either joint distribution has a distribution with dynamic
      dependency, i.e., when either joint distribution is not a collection of
      independent distributions.
  """
  if len(d0._dist_fn_wrapped) != len(d1._dist_fn_wrapped):  # pylint: disable=protected-access
    raise ValueError(
        'Can only compute KL divergence between when each has the'
        'same number of component distributions.')
  if (not all(a is None for a in d0._dist_fn_args) or  # pylint: disable=protected-access
      not all(a is None for a in d1._dist_fn_args)):  # pylint: disable=protected-access
    raise ValueError(
        'Can only compute KL divergence when all distributions are '
        'independent.')
  with tf.name_scope(name or 'kl_jointseq_jointseq'):
    return sum(kullback_leibler.kl_divergence(d0_(), d1_())
               for d0_, d1_ in zip(d0._dist_fn_wrapped, d1._dist_fn_wrapped))  # pylint: disable=protected-access


# 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.
# ============================================================================
"""Half-Cauchy Distribution Class."""

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

import functools
# Dependency imports
import numpy as np
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 dtype_util
from tensorflow_probability.python.internal import reparameterization

__all__ = [
    "HalfCauchy",
]

def check_arg_in_support(f):
  """Decorator function for argument bounds checking.

  This decorator is meant to be used with methods that require the first
  argument to be in the support of the distribution. If `validate_args` is
  `True`, the method is wrapped with an assertion that the first argument is
  greater than or equal to `loc`, since the support of the half-Cauchy
  distribution is given by `[loc, infinity)`.

  Args:
    f: method to be decorated.

  Returns:
    Returns a decorated method that, when `validate_args` attribute of the class
    is `True`, will assert that all elements in the first argument are within
    the support of the distribution before executing the original method.
  """
  @functools.wraps(f)
  def _check_arg_and_apply_f(*args, **kwargs):
    dist = args[0]
    x = args[1]
    with tf.control_dependencies([
        assert_util.assert_greater_equal(
            x, dist.loc, message="x is not in the support of the distribution")
    ] if dist.validate_args else []):
      return f(*args, **kwargs)
  return _check_arg_and_apply_f

class HalfCauchy(distribution.Distribution):
  """Half-Cauchy distribution.

  The half-Cauchy distribution is parameterized by a `loc` and a
  `scale` parameter. It represents the right half of the two symmetric halves in
  a [Cauchy distribution](https://en.wikipedia.org/wiki/Cauchy_distribution).

  #### Mathematical Details
  The probability density function (pdf) for the half-Cauchy distribution
  is given by

  ```none
  pdf(x; loc, scale) = 2 / (pi scale (1 + z**2))
  z = (x - loc) / scale
  ```

  where `loc` is a scalar in `R` and `scale` is a positive scalar in `R`.

  The support of the distribution is given by the interval `[loc, infinity)`.

  """

  def __init__(self,
               loc,
               scale,
               validate_args=False,
               allow_nan_stats=True,
               name="HalfCauchy"):
    """Construct a half-Cauchy distribution with `loc` and `scale`.

    Args:
      loc: Floating-point `Tensor`; the location(s) of the distribution(s).
      scale: Floating-point `Tensor`; the scale(s) of the distribution(s).
        Must contain only positive values.
      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` (i.e. do not validate args).
      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: 'HalfCauchy'.

    Raises:
      TypeError: if `loc` and `scale` have different `dtype`.
    """
    parameters = dict(locals())
    with tf.name_scope(name) as name:
      dtype = dtype_util.common_dtype([loc, scale], preferred_dtype=tf.float32)
      loc = tf.convert_to_tensor(value=loc, name="loc", dtype=dtype)
      scale = tf.convert_to_tensor(value=scale, name="scale", dtype=dtype)
      with tf.control_dependencies(
          [assert_util.assert_positive(scale)] if validate_args else []):
        self._loc = tf.identity(loc, name="loc")
        self._scale = tf.identity(scale, name="loc")
      dtype_util.assert_same_float_dtype([self._loc, self._scale])
    super(HalfCauchy, self).__init__(
        dtype=self._scale.dtype,
        reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        parameters=parameters,
        graph_parents=[self._loc, self._scale],
        name=name)

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

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

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

  @property
  def scale(self):
    """Distribution parameter for the scale."""
    return self._scale

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

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

  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):
    shape = tf.concat([[n], self._batch_shape_tensor()], 0)
    probs = tf.random.uniform(
        shape, minval=0., maxval=1., dtype=self.dtype, seed=seed)
    return self._quantile(probs)

  @check_arg_in_support
  def _log_prob(self, x):
    def log_prob_on_support(x):
      return (np.log(2 / np.pi) - tf.math.log(self.scale) -
              tf.math.log1p(self._z(x)**2))

    return self._extend_support_with_default_value(
        x, log_prob_on_support, default_value=-np.inf)

  @check_arg_in_support
  def _log_cdf(self, x):
    return self._extend_support_with_default_value(
        x,
        lambda x: np.log(2 / np.pi) + tf.math.log(tf.atan(self._z(x))),
        default_value=-np.inf)

  def _z(self, x):
    """Standardize input `x`."""
    with tf.name_scope("standardize"):
      return (x - self.loc) / self.scale

  def _inv_z(self, z):
    """Reconstruct input `x` from a its normalized version."""
    with tf.name_scope("reconstruct"):
      return z * self.scale + self.loc

  def _entropy(self):
    h = np.log(2 * np.pi) + tf.math.log(self.scale)
    return h * tf.ones_like(self.loc)

  def _quantile(self, p):
    return self.loc + self.scale * tf.tan((np.pi / 2) * p)

  def _mode(self):
    return self.loc * tf.ones_like(self.scale)

  def _mean(self):
    if self.allow_nan_stats:
      return tf.fill(self.batch_shape_tensor(),
                     dtype_util.as_numpy_dtype(self.dtype)(np.nan))
    raise ValueError("`mean` is undefined for the half-Cauchy distribution.")

  def _stddev(self):
    if self.allow_nan_stats:
      return tf.fill(self.batch_shape_tensor(),
                     dtype_util.as_numpy_dtype(self.dtype)(np.nan))
    raise ValueError("`stddev` is undefined for the half-Cauchy distribution.")

  def _variance(self):
    if self.allow_nan_stats:
      return tf.fill(self.batch_shape_tensor(),
                     dtype_util.as_numpy_dtype(self.dtype)(np.nan))
    raise ValueError(
        "`variance` is undefined for the half-Cauchy distribution.")

  def _extend_support_with_default_value(self, x, f, default_value):
    """Returns `f(x)` if x is in the support, and `default_value` otherwise.

    Given `f` which is defined on the support of this distribution
    (`x >= loc`), extend the function definition to the real line
    by defining `f(x) = default_value` for `x < loc`.

    Args:
      x: Floating-point `Tensor` to evaluate `f` at.
      f: Callable that takes in a `Tensor` and returns a `Tensor`. This
        represents the function whose domain of definition we want to extend.
      default_value: Python or numpy literal representing the value to use for
        extending the domain.
    Returns:
      `Tensor` representing an extension of `f(x)`.
    """
    with tf.name_scope("extend_support_with_default_value"):
      x = tf.convert_to_tensor(value=x, dtype=self.dtype, name="x")
      loc = self.loc + tf.zeros_like(self.scale) + tf.zeros_like(x)
      x = x + tf.zeros_like(loc)
      # Substitute out-of-support values in x with values that are in the
      # support of the distribution before applying f.
      y = f(tf.where(x < loc, self._inv_z(0.5) + tf.zeros_like(x), x))
      if default_value == 0.:
        default_value = tf.zeros_like(y)
      elif default_value == 1.:
        default_value = tf.ones_like(y)
      else:
        default_value = tf.fill(
            dims=tf.shape(input=y),
            value=dtype_util.as_numpy_dtype(self.dtype)(default_value))
      return tf.where(x < loc, default_value, y)


# 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.
# ============================================================================
"""Elliptic Slice sampler transition kernel."""

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

import collections
# Dependency imports

import numpy as np

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__ = [
    'EllipticalSliceSampler',
]

def _rotate_on_ellipse(state_parts, vectors, angle):
  new_state_parts = []
  padded_angle = _right_pad_with_ones(angle, tf.rank(state_parts[0]))
  for state, vector in zip(state_parts, vectors):
    new_state_parts.append(
        state * tf.cos(padded_angle) + vector * tf.sin(padded_angle))
  return new_state_parts

def  _right_pad_with_ones(x, target_rank):
  static_target_rank = tf.get_static_value(target_rank)
  if x.shape.is_fully_defined() and static_target_rank is not None:
    return tf.reshape(x, x.shape.concatenate([1] * (
        int(static_target_rank) - x.shape.ndims)))
  return tf.reshape(
      x, tf.concat(
          [tf.shape(x),
           tf.ones(
               [target_rank - tf.rank(x)], dtype=target_rank.dtype)], axis=0))

EllipticalSliceSamplerKernelResults = collections.namedtuple(
    'EllipticalSliceSamplerKernelResults',
    [
        'log_likelihood',  # For "next_state".
        'angle',  # Angle previous state was rotated by.
        'normal_samples',  # Normal samples used to generate the next state.
    ])

class EllipticalSliceSampler(kernel_base.TransitionKernel):
  """Runs one step of the elliptic slice sampler.

  Elliptical Slice Sampling is a Markov Chain Monte Carlo (MCMC) algorithm
  based, as stated in [Murray, 2010][1].

  Given `log_likelihood_fn` and `normal_sampler_fn`, the goal of Elliptical
  Slice Sampling is to sample from:

  ```none
  p(f) = N(f; 0, Sigma)L(f) / Z
  ```

  where:
    * `L = log_likelihood_fn`
    * `Sigma` is a covariance matrix.
    * Samples from `normal_sampler_fn` are distributed as `N(f; 0, Sigma)`.
    * `Z` is a normalizing constant.

  In other words, sampling from a posterior distribution that is proportional
  to a multivariate gaussian prior multiplied by some likelihood function.

  The `one_step` function can update multiple chains in parallel. It assumes
  that all leftmost dimensions of `current_state` index independent chain states
  (and are therefore updated independently). The output of
  `log_likelihood_fn(*current_state)` should sum log-probabilities across all
  event dimensions. Slices along the rightmost dimensions may have different
  target distributions; for example, `current_state[0, :]` could have a
  different target distribution from `current_state[1, :]`.
  These semantics are governed both by `log_likelihood_fn(*current_state)` and
  `normal_sampler_fn`.

  Note that the sampler only supports states where all components have a common
  dtype.

  ### Examples:

  #### Simple chain with warm-up.

  In this example we have the following model.

  ```none
    p(loc | loc0, scale0) ~ N(loc0, scale0)
    p(x | loc, sigma) ~ N(mu, sigma)
  ```

  What we would like to do is sample from `p(loc | x, loc0, scale0)`. In other
  words, given some data, we would like to infer the posterior distribution
  of the mean that generated that data point.

  We can use elliptical slice sampling here.

  ```python
    import tensorflow as tf
    import tensorflow_probability as tfp
    import numpy as np

    tfd = tfp.distributions

    dtype = np.float64

    # loc0 = 0, scale0 = 1
    normal_sampler_fn = lambda seed: return tfd.Normal(
        loc=dtype(0), scale=dtype(1)).sample(seed=seed)

    # We saw the following data.
    data_points = np.random.randn(20)

    # scale = 2.
    log_likelihood_fn = lambda state: return tf.reduce_sum(
        tfd.Normal(state, dtype(2.)).log_prob(data_points))

    kernel = tfp.mcmc.EllipticalSliceSampler(
        normal_sampler_fn=normal_sampler_fn,
        log_likelihood_fn=log_likelihood_fn,
        seed=1234)

    samples, _ = tfp.mcmc.sample_chain(
        num_results=int(3e5),
        current_state=dtype(1),
        kernel=kernel,
        num_burnin_steps=1000,
        parallel_iterations=1)  # For determinism.

    sample_mean = tf.reduce_mean(samples, axis=0)
    sample_std = tf.sqrt(
      tf.reduce_mean(tf.squared_difference(samples, sample_mean),
                     axis=0))

    with tf.Session() as sess:
      [sample_mean, sample_std] = sess.run([sample_mean, sample_std])

    print("Sample mean: ", sample_mean)
    print("Sample Std: ", sample_std)
  ```

  ### References

  [1]: Ian Murray, Ryan P. Adams, David J.C. MacKay. Elliptical slice sampling.
       proceedings.mlr.press/v9/murray10a/murray10a.pdf
  """

  def __init__(self,
               normal_sampler_fn,
               log_likelihood_fn,
               seed=None,
               name=None):
    """Initializes this transition kernel.

    Args:
      normal_sampler_fn: Python callable that takes in a seed and returns a
        sample from a multivariate normal distribution. Note that the shape of
        the samples must agree with `log_likelihood_fn`.
      log_likelihood_fn: Python callable which takes an argument like
        `current_state` (or `*current_state` if it is a list) and returns its
        (possibly unnormalized) log-likelihood.
      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., 'slice_sampler_kernel').

    Returns:
      next_state: Tensor or Python list of `Tensor`s representing the state(s)
        of the Markov chain(s) at each result step. Has same shape as
        `current_state`.
      kernel_results: `collections.namedtuple` of internal calculations used to
        advance the chain.
    """
    self._seed_stream = distributions.SeedStream(
        seed, salt='elliptical_slice_sampler')
    self._parameters = dict(
        normal_sampler_fn=normal_sampler_fn,
        log_likelihood_fn=log_likelihood_fn,
        seed=seed,
        name=name)

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

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

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

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

  @property
  def parameters(self):
    """Returns `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):
    """Runs one iteration of the Elliptical Slice Sampler.

    Args:
      current_state: `Tensor` or Python `list` of `Tensor`s representing the
        current state(s) of the Markov chain(s). The first `r` dimensions
        index independent chains,
        `r = tf.rank(log_likelihood_fn(*normal_sampler_fn()))`.
      previous_kernel_results: `collections.namedtuple` containing `Tensor`s
        representing values from previous calls to this function (or from the
        `bootstrap_results` function.)

    Returns:
      next_state: Tensor or Python list of `Tensor`s representing the state(s)
        of the Markov chain(s) after taking exactly one step. Has same type and
        shape as `current_state`.
      kernel_results: `collections.namedtuple` of internal calculations used to
        advance the chain.

    Raises:
      TypeError: if `not log_likelihood.dtype.is_floating`.
    """
    with tf.compat.v1.name_scope(
        name=mcmc_util.make_name(self.name, 'elliptical_slice', 'one_step'),
        values=[self._seed_stream,
                current_state,
                previous_kernel_results.log_likelihood]):
      with tf.compat.v1.name_scope('initialize'):
        [
            init_state_parts,
            init_log_likelihood
        ] = _prepare_args(
            self.log_likelihood_fn,
            current_state,
            previous_kernel_results.log_likelihood)

      normal_samples = self.normal_sampler_fn(self._seed_stream())  # pylint: disable=not-callable
      normal_samples = list(normal_samples) if mcmc_util.is_list_like(
          normal_samples) else [normal_samples]
      u = tf.random.uniform(
          shape=tf.shape(init_log_likelihood),
          seed=self._seed_stream(),
          dtype=init_log_likelihood.dtype.base_dtype,
      )
      threshold = init_log_likelihood + tf.math.log(u)

      starting_angle = tf.random.uniform(
          shape=tf.shape(init_log_likelihood),
          minval=0.,
          maxval=2 * np.pi,
          name='angle',
          seed=self._seed_stream(),
          dtype=init_log_likelihood.dtype.base_dtype,
      )
      starting_angle_min = starting_angle - 2 * np.pi
      starting_angle_max = starting_angle

      starting_state_parts = _rotate_on_ellipse(
          init_state_parts, normal_samples, starting_angle)
      starting_log_likelihood = self.log_likelihood_fn(*starting_state_parts)  # pylint: disable=not-callable

      def chain_not_done(
          angle,
          angle_min,
          angle_max,
          current_state_parts,
          current_log_likelihood):
        del angle, angle_min, angle_max, current_state_parts
        return tf.reduce_any(current_log_likelihood < threshold)

      def sample_next_angle(
          angle,
          angle_min,
          angle_max,
          current_state_parts,
          current_log_likelihood):
        """Slice sample a new angle, and rotate init_state by that amount."""
        chain_not_done = current_log_likelihood < threshold
        # Box in on angle. Only update angles for which we haven't generated a
        # point that beats the threshold.
        angle_min = tf.where(
            tf.math.logical_and(angle < 0, chain_not_done),
            angle,
            angle_min)
        angle_max = tf.where(
            tf.math.logical_and(angle >= 0, chain_not_done),
            angle,
            angle_max)
        new_angle = tf.random.uniform(
            shape=tf.shape(current_log_likelihood),
            minval=angle_min,
            maxval=angle_max,
            seed=self._seed_stream(),
            dtype=angle.dtype.base_dtype
        )
        angle = tf.where(chain_not_done, new_angle, angle)
        next_state_parts = _rotate_on_ellipse(
            init_state_parts, normal_samples, angle)

        new_state_parts = []
        broadcasted_chain_not_done = _right_pad_with_ones(
            chain_not_done, tf.rank(next_state_parts[0]))
        for n_state, c_state in zip(next_state_parts, current_state_parts):
          new_state_part = tf.where(
              tf.broadcast_to(
                  broadcasted_chain_not_done,
                  tf.shape(n_state)),
              n_state,
              c_state)
          new_state_parts.append(new_state_part)

        return (
            angle,
            angle_min,
            angle_max,
            new_state_parts,
            self.log_likelihood_fn(*new_state_parts)  # pylint: disable=not-callable
        )

      [
          next_angle,
          _,
          _,
          next_state_parts,
          next_log_likelihood,
      ] = tf.while_loop(
          cond=chain_not_done,
          body=sample_next_angle,
          loop_vars=[
              starting_angle,
              starting_angle_min,
              starting_angle_max,
              starting_state_parts,
              starting_log_likelihood
          ])

      return [
          next_state_parts if mcmc_util.is_list_like(
              current_state) else next_state_parts[0],
          EllipticalSliceSamplerKernelResults(
              log_likelihood=next_log_likelihood,
              angle=next_angle,
              normal_samples=normal_samples,
          ),
      ]

  def bootstrap_results(self, init_state):
    with tf.compat.v1.name_scope(
        name=mcmc_util.make_name(
            self.name, 'elliptical_slice', 'bootstrap_results'),
        values=[init_state]):
      if not mcmc_util.is_list_like(init_state):
        init_state = [init_state]
      init_state = [tf.convert_to_tensor(x) for x in init_state]
      init_log_likelihood = self.log_likelihood_fn(*init_state)  # pylint:disable=not-callable
      return EllipticalSliceSamplerKernelResults(
          log_likelihood=init_log_likelihood,
          angle=tf.zeros_like(init_log_likelihood),
          normal_samples=[tf.zeros_like(x) for x in init_state]
      )

def _maybe_call_fn(fn,
                   fn_arg_list,
                   fn_result=None,
                   description='log_likelihood'):
  """Helper which computes `fn_result` if needed."""
  fn_arg_list = (list(fn_arg_list) if mcmc_util.is_list_like(fn_arg_list)
                 else [fn_arg_list])
  if fn_result is None:
    fn_result = fn(*fn_arg_list)
  if not fn_result.dtype.is_floating:
    raise TypeError('`{}` must be a `Tensor` with `float` `dtype`.'.format(
        description))
  return fn_result

def _prepare_args(log_likelihood_fn, state,
                  log_likelihood=None, description='log_likelihood'):
  """Processes input args to meet list-like assumptions."""
  state_parts = list(state) if mcmc_util.is_list_like(state) else [state]
  state_parts = [tf.convert_to_tensor(s, name='current_state')
                 for s in state_parts]

  log_likelihood = _maybe_call_fn(
      log_likelihood_fn,
      state_parts,
      log_likelihood,
      description)
  return [state_parts, log_likelihood]