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

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

import math
# 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 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__ = [
    "Wishart",
]

class _WishartLinearOperator(distribution.Distribution):
  """The matrix Wishart distribution on positive definite matrices.

  This distribution is defined by a scalar number of degrees of freedom `df` and
  an instance of `LinearOperator`, which provides matrix-free access to a
  symmetric positive definite operator, which defines the scale matrix.

  #### Mathematical Details

  The probability density function (pdf) is,

  ```none
  pdf(X; df, scale) = det(X)**(0.5 (df-k-1)) exp(-0.5 tr[inv(scale) X]) / Z
  Z = 2**(0.5 df k) |det(scale)|**(0.5 df) Gamma_k(0.5 df)
  ```

  where:

  * `df >= k` denotes the degrees of freedom,
  * `scale` is a symmetric, positive definite, `k x k` matrix,
  * `Z` is the normalizing constant, and,
  * `Gamma_k` is the [multivariate Gamma function](
    https://en.wikipedia.org/wiki/Multivariate_gamma_function).

  #### Examples

  See the `Wishart` class for examples of initializing and using this class.
  """

  def __init__(self,
               df,
               scale_operator,
               input_output_cholesky=False,
               validate_args=False,
               allow_nan_stats=True,
               name=None):
    """Construct Wishart distributions.

    Args:
      df: `float` or `double` tensor, the degrees of freedom of the
        distribution(s). `df` must be greater than or equal to `k`.
      scale_operator: `float` or `double` instance of `LinearOperator`.
      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`, `mean`, and `mode` are of Cholesky form.  Setting this
        argument to `True` is purely a computational optimization and does not
        change the underlying distribution; for instance, `mean` returns the
        Cholesky of the mean, not the mean of Cholesky factors. The `variance`
        and `stddev` methods are unaffected by this flag.
        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:
      TypeError: if scale is not floating-type
      TypeError: if scale.dtype != df.dtype
      ValueError: if df < k, where scale operator event shape is
        `(k, k)`
    """
    parameters = dict(locals())
    self._input_output_cholesky = input_output_cholesky
    with tf.name_scope(name) as name:
      with tf.name_scope("init"):
        if not dtype_util.is_floating(scale_operator.dtype):
          raise TypeError(
              "scale_operator.dtype=%s is not a floating-point type" %
              scale_operator.dtype)
        if not scale_operator.is_square:
          print(scale_operator.to_dense().eval())
          raise ValueError("scale_operator must be square.")

        self._scale_operator = scale_operator
        self._df = tf.convert_to_tensor(
            value=df, dtype=scale_operator.dtype, name="df")
        dtype_util.assert_same_float_dtype([self._df, self._scale_operator])
        if tf.compat.dimension_value(self._scale_operator.shape[-1]) is None:
          self._dimension = tf.cast(
              self._scale_operator.domain_dimension_tensor(),
              dtype=self._scale_operator.dtype,
              name="dimension")
        else:
          self._dimension = tf.convert_to_tensor(
              value=tf.compat.dimension_value(self._scale_operator.shape[-1]),
              dtype=self._scale_operator.dtype,
              name="dimension")
        df_val = tf.get_static_value(self._df)
        dim_val = tf.get_static_value(self._dimension)
        if df_val is not None and dim_val is not None:
          df_val = np.asarray(df_val)
          if not df_val.shape:
            df_val = [df_val]
          if np.any(df_val < dim_val):
            raise ValueError(
                "Degrees of freedom (df = %s) cannot be less than "
                "dimension of scale matrix (scale.dimension = %s)"
                % (df_val, dim_val))
        elif validate_args:
          assertions = assert_util.assert_less_equal(
              self._dimension,
              self._df,
              message=("Degrees of freedom (df = %s) cannot be "
                       "less than dimension of scale matrix "
                       "(scale.dimension = %s)" % (self._dimension, self._df)))
          self._df = distribution_util.with_dependencies(
              [assertions], self._df)
    super(_WishartLinearOperator, self).__init__(
        dtype=self._scale_operator.dtype,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
        parameters=parameters,
        graph_parents=(
            [self._df, self._dimension] + self._scale_operator.graph_parents),
        name=name)

  @property
  def df(self):
    """Wishart distribution degree(s) of freedom."""
    return self._df

  def _square_scale_operator(self):
    return self.scale_operator.matmul(
        self.scale_operator.to_dense(), adjoint_arg=True)

  def scale(self):
    """Wishart distribution scale matrix."""
    if self._input_output_cholesky:
      return self.scale_operator.to_dense()
    else:
      return self._square_scale_operator()

  @property
  def scale_operator(self):
    """Wishart distribution scale matrix as an Linear Operator."""
    return self._scale_operator

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

  @property
  def dimension(self):
    """Dimension of underlying vector space. The `p` in `R^(p*p)`."""
    return self._dimension

  def _event_shape_tensor(self):
    dimension = self.scale_operator.domain_dimension_tensor()
    return tf.stack([dimension, dimension])

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

  def _batch_shape_tensor(self):
    return tf.broadcast_dynamic_shape(
        tf.shape(input=self.df), self.scale_operator.batch_shape_tensor())

  def _batch_shape(self):
    return tf.broadcast_static_shape(
        self.df.shape, self.scale_operator.batch_shape)

  def _sample_n(self, n, seed):
    batch_shape = self.batch_shape_tensor()
    event_shape = self.event_shape_tensor()
    batch_ndims = tf.shape(input=batch_shape)[0]

    ndims = batch_ndims + 3  # sample_ndims=1, event_ndims=2
    shape = tf.concat([[n], batch_shape, event_shape], 0)
    stream = seed_stream.SeedStream(seed, salt="Wishart")

    # Complexity: O(nbk**2)
    x = tf.random.normal(
        shape=shape, mean=0., stddev=1., dtype=self.dtype, seed=stream())

    # Complexity: O(nbk)
    # This parametrization is equivalent to Chi2, i.e.,
    # ChiSquared(k) == Gamma(alpha=k/2, beta=1/2)
    expanded_df = self.df * tf.ones(
        self.scale_operator.batch_shape_tensor(),
        dtype=dtype_util.base_dtype(self.df.dtype))

    g = tf.random.gamma(
        shape=[n],
        alpha=self._multi_gamma_sequence(0.5 * expanded_df, self.dimension),
        beta=0.5,
        dtype=self.dtype,
        seed=stream())

    # Complexity: O(nbk**2)
    x = tf.linalg.band_part(x, -1, 0)  # Tri-lower.

    # Complexity: O(nbk)
    x = tf.linalg.set_diag(x, tf.sqrt(g))

    # Make batch-op ready.
    # Complexity: O(nbk**2)
    perm = tf.concat([tf.range(1, ndims), [0]], 0)
    x = tf.transpose(a=x, perm=perm)
    shape = tf.concat([batch_shape, [event_shape[0]], [event_shape[1] * n]], 0)
    x = tf.reshape(x, shape)

    # Complexity: O(nbM) where M is the complexity of the operator solving a
    # vector system. For LinearOperatorLowerTriangular, each matmul is O(k^3) so
    # this step has complexity O(nbk^3).
    x = self.scale_operator.matmul(x)

    # Undo make batch-op ready.
    # Complexity: O(nbk**2)
    shape = tf.concat([batch_shape, event_shape, [n]], 0)
    x = tf.reshape(x, shape)
    perm = tf.concat([[ndims - 1], tf.range(0, ndims - 1)], 0)
    x = tf.transpose(a=x, perm=perm)

    if not self.input_output_cholesky:
      # Complexity: O(nbk**3)
      x = tf.matmul(x, x, adjoint_b=True)

    return x

  def _log_prob(self, x):
    if self.input_output_cholesky:
      x_sqrt = x
    else:
      # Complexity: O(nbk**3)
      x_sqrt = tf.linalg.cholesky(x)

    batch_shape = self.batch_shape_tensor()
    event_shape = self.event_shape_tensor()
    x_ndims = tf.rank(input=x_sqrt)
    num_singleton_axes_to_prepend = (
        tf.maximum(tf.size(input=batch_shape) + 2, x_ndims) - x_ndims)
    x_with_prepended_singletons_shape = tf.concat([
        tf.ones([num_singleton_axes_to_prepend], dtype=tf.int32),
        tf.shape(input=x_sqrt)], 0)
    x_sqrt = tf.reshape(x_sqrt, x_with_prepended_singletons_shape)
    ndims = tf.rank(x_sqrt)
    # sample_ndims = ndims - batch_ndims - event_ndims
    sample_ndims = ndims - tf.size(input=batch_shape) - 2
    sample_shape = tf.shape(input=x_sqrt)[:sample_ndims]

    # We need to be able to pre-multiply each matrix by its corresponding
    # batch scale matrix. Since a Distribution Tensor supports multiple
    # samples per batch, this means we need to reshape the input matrix `x`
    # so that the first b dimensions are batch dimensions and the last two
    # are of shape [dimension, dimensions*number_of_samples]. Doing these
    # gymnastics allows us to do a batch_solve.
    #
    # After we're done with sqrt_solve (the batch operation) we need to undo
    # this reshaping so what we're left with is a Tensor partitionable by
    # sample, batch, event dimensions.

    # Complexity: O(nbk**2) since transpose must access every element.
    scale_sqrt_inv_x_sqrt = x_sqrt
    perm = tf.concat([tf.range(sample_ndims, ndims),
                      tf.range(0, sample_ndims)], 0)
    scale_sqrt_inv_x_sqrt = tf.transpose(a=scale_sqrt_inv_x_sqrt, perm=perm)
    last_dim_size = (
        tf.cast(self.dimension, dtype=tf.int32) *
        tf.reduce_prod(
            input_tensor=x_with_prepended_singletons_shape[:sample_ndims]))
    shape = tf.concat(
        [x_with_prepended_singletons_shape[sample_ndims:-2],
         [tf.cast(self.dimension, dtype=tf.int32), last_dim_size]],
        axis=0)
    scale_sqrt_inv_x_sqrt = tf.reshape(scale_sqrt_inv_x_sqrt, shape)

    # Complexity: O(nbM*k) where M is the complexity of the operator solving a
    # vector system. For LinearOperatorLowerTriangular, each solve is O(k**2) so
    # this step has complexity O(nbk^3).
    scale_sqrt_inv_x_sqrt = self.scale_operator.solve(scale_sqrt_inv_x_sqrt)

    # Undo make batch-op ready.
    # Complexity: O(nbk**2)
    shape = tf.concat(
        [tf.shape(input=scale_sqrt_inv_x_sqrt)[:-2], event_shape, sample_shape],
        axis=0)
    scale_sqrt_inv_x_sqrt = tf.reshape(scale_sqrt_inv_x_sqrt, shape)
    perm = tf.concat([
        tf.range(ndims - sample_ndims, ndims),
        tf.range(0, ndims - sample_ndims)
    ], 0)
    scale_sqrt_inv_x_sqrt = tf.transpose(a=scale_sqrt_inv_x_sqrt, perm=perm)

    # Write V = SS', X = LL'. Then:
    # tr[inv(V) X] = tr[inv(S)' inv(S) L L']
    #              = tr[inv(S) L L' inv(S)']
    #              = tr[(inv(S) L) (inv(S) L)']
    #              = sum_{ik} (inv(S) L)_{ik}**2
    # The second equality follows from the cyclic permutation property.
    # Complexity: O(nbk**2)
    trace_scale_inv_x = tf.reduce_sum(
        input_tensor=tf.square(scale_sqrt_inv_x_sqrt), axis=[-2, -1])

    # Complexity: O(nbk)
    half_log_det_x = tf.reduce_sum(
        input_tensor=tf.math.log(tf.linalg.diag_part(x_sqrt)), axis=[-1])

    # Complexity: O(nbk**2)
    log_prob = ((self.df - self.dimension - 1.) * half_log_det_x -
                0.5 * trace_scale_inv_x -
                self.log_normalization())

    # Set shape hints.
    # Try to merge what we know from the input x with what we know from the
    # parameters of this distribution.
    if tensorshape_util.rank(x.shape) is not None and tensorshape_util.rank(
        self.batch_shape) is not None:
      tensorshape_util.set_shape(
          log_prob,
          tf.broadcast_static_shape(x.shape[:-2], self.batch_shape))

    return log_prob

  def _prob(self, x):
    return tf.exp(self._log_prob(x))

  def _entropy(self):
    half_dp1 = 0.5 * self.dimension + 0.5
    half_df = 0.5 * self.df
    return (self.dimension * (half_df + half_dp1 * math.log(2.)) +
            2 * half_dp1 * self.scale_operator.log_abs_determinant() +
            self._multi_lgamma(half_df, self.dimension) +
            (half_dp1 - half_df) * self._multi_digamma(half_df, self.dimension))

  def _mean(self):
    # Because df is a scalar, we need to expand dimensions to match
    # scale_operator. We use ellipses notation (...) to select all dimensions
    # and add two dimensions to the end.
    df = self.df[..., tf.newaxis, tf.newaxis]
    if self.input_output_cholesky:
      return tf.sqrt(df) * self.scale_operator.to_dense()
    return df * self._square_scale_operator()

  def _variance(self):
    # Because df is a scalar, we need to expand dimensions to match
    # scale_operator. We use ellipses notation (...) to select all dimensions
    # and add two dimensions to the end.
    df = self.df[..., tf.newaxis, tf.newaxis]
    x = tf.sqrt(df) * self._square_scale_operator()
    d = tf.expand_dims(tf.linalg.diag_part(x), -1)
    v = tf.square(x) + tf.matmul(d, d, adjoint_b=True)
    return v

  def _mode(self):
    s = self.df - self.dimension - 1.
    s = tf.where(
        tf.less(s, 0.), tf.constant(float("NaN"), dtype=self.dtype, name="nan"),
        s)
    if self.input_output_cholesky:
      return tf.sqrt(s) * self.scale_operator.to_dense()
    return s * self._square_scale_operator()

  def mean_log_det(self, name="mean_log_det"):
    """Computes E[log(det(X))] under this Wishart distribution."""
    with self._name_scope(name):
      return (self._multi_digamma(0.5 * self.df, self.dimension) +
              self.dimension * math.log(2.) +
              2 * self.scale_operator.log_abs_determinant())

  def log_normalization(self, name="log_normalization"):
    """Computes the log normalizing constant, log(Z)."""
    with self._name_scope(name):
      return (self.df * self.scale_operator.log_abs_determinant() +
              0.5 * self.df * self.dimension * math.log(2.) +
              self._multi_lgamma(0.5 * self.df, self.dimension))

  def _multi_gamma_sequence(self, a, p, name="multi_gamma_sequence"):
    """Creates sequence used in multivariate (di)gamma; shape = shape(a)+[p]."""
    with self._name_scope(name):
      # Linspace only takes scalars, so we'll add in the offset afterwards.
      seq = tf.linspace(
          tf.constant(0., dtype=self.dtype), 0.5 - 0.5 * p, tf.cast(
              p, tf.int32))
      return seq + tf.expand_dims(a, [-1])

  def _multi_lgamma(self, a, p, name="multi_lgamma"):
    """Computes the log multivariate gamma function; log(Gamma_p(a))."""
    with self._name_scope(name):
      seq = self._multi_gamma_sequence(a, p)
      return (0.25 * p * (p - 1.) * math.log(math.pi) +
              tf.reduce_sum(input_tensor=tf.math.lgamma(seq), axis=[-1]))

  def _multi_digamma(self, a, p, name="multi_digamma"):
    """Computes the multivariate digamma function; Psi_p(a)."""
    with self._name_scope(name):
      seq = self._multi_gamma_sequence(a, p)
      return tf.reduce_sum(input_tensor=tf.math.digamma(seq), axis=[-1])

class Wishart(_WishartLinearOperator):
  """The matrix Wishart distribution on positive definite matrices.

  This distribution is defined by a scalar degrees of freedom `df` and a scale
  matrix, which can either be a full symmetric matrix or a lower triangular
  Cholesky factor.

  #### Mathematical Details

  The probability density function (pdf) is,

  ```none
  pdf(X; df, scale) = det(X)**(0.5 (df-k-1)) exp(-0.5 tr[inv(scale) X]) / Z
  Z = 2**(0.5 df k) |det(scale)|**(0.5 df) Gamma_k(0.5 df)
  ```

  where:
  * `df >= k` denotes the degrees of freedom,
  * `scale` is a symmetric, positive definite, `k x k` matrix,
  * `Z` is the normalizing constant, and,
  * `Gamma_k` is the [multivariate Gamma function](
    https://en.wikipedia.org/wiki/Multivariate_gamma_function).

  #### Examples

  ```python
  # Initialize a single 3x3 Wishart with Cholesky factored scale matrix and 5
  # degrees-of-freedom.(*)
  df = 5
  chol_scale = tf.cholesky(...)  # Shape is [3, 3].
  dist = tfd.Wishart(df=df, scale_tril=chol_scale)

  # Evaluate this on an observation in R^3, returning a scalar.
  x = ...  # A 3x3 positive definite matrix.
  dist.prob(x)  # Shape is [], a scalar.

  # Evaluate this on a two observations, each in R^{3x3}, returning a length two
  # Tensor.
  x = [x0, x1]  # Shape is [2, 3, 3].
  dist.prob(x)  # Shape is [2].

  # Initialize two 3x3 Wisharts with full scale matrices.
  df = [5, 4]
  scale = ...  # Shape is [2, 3, 3].
  dist = tfd.Wishart(df=df, scale=scale)

  # Evaluate this on four observations.
  x = [[x0, x1], [x2, x3]]  # Shape is [2, 2, 3, 3].
  dist.prob(x)  # Shape is [2, 2].

  # (*) - To efficiently create a trainable covariance matrix, see the example
  #   in tfp.distributions.matrix_diag_transform.
  ```
  """

  def __init__(self,
               df,
               scale=None,
               scale_tril=None,
               input_output_cholesky=False,
               validate_args=False,
               allow_nan_stats=True,
               name="Wishart"):
    """Construct Wishart distributions.

    Args:
      df: `float` or `double` `Tensor`. Degrees of freedom, must be greater than
        or equal to dimension of the scale matrix.
      scale: `float` or `double` `Tensor`. The symmetric positive definite
        scale matrix of the distribution. Exactly one of `scale` and
        'scale_tril` must be passed.
      scale_tril: `float` or `double` `Tensor`. The Cholesky factorization
        of the symmetric positive definite scale matrix of the distribution.
        Exactly one of `scale` and 'scale_tril` must be passed.
      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`, `mean`, and `mode` are of Cholesky form.  Setting this
        argument to `True` is purely a computational optimization and does not
        change the underlying distribution; for instance, `mean` returns the
        Cholesky of the mean, not the mean of Cholesky factors. The `variance`
        and `stddev` methods are unaffected by this flag.
        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 zero or both of 'scale' and 'scale_tril' are passed in.
    """
    parameters = dict(locals())

    with tf.name_scope(name) as name:
      with tf.name_scope("init"):
        if (scale is None) == (scale_tril is None):
          raise ValueError("Must pass scale or scale_tril, but not both.")

        dtype = dtype_util.common_dtype([df, scale, scale_tril], tf.float32)
        df = tf.convert_to_tensor(value=df, name="df", dtype=dtype)
        if scale is not None:
          scale = tf.convert_to_tensor(value=scale, name="scale", dtype=dtype)
          if validate_args:
            scale = distribution_util.assert_symmetric(scale)
          scale_tril = tf.linalg.cholesky(scale)
        else:  # scale_tril is not None
          scale_tril = tf.convert_to_tensor(
              value=scale_tril, name="scale_tril", dtype=dtype)
          if validate_args:
            scale_tril = distribution_util.with_dependencies([
                assert_util.assert_positive(
                    tf.linalg.diag_part(scale_tril),
                    message="scale_tril must be positive definite"),
                assert_util.assert_equal(
                    tf.shape(input=scale_tril)[-1],
                    tf.shape(input=scale_tril)[-2],
                    message="scale_tril must be square")
            ], scale_tril)

      super(Wishart, self).__init__(
          df=df,
          scale_operator=tf.linalg.LinearOperatorLowerTriangular(
              tril=scale_tril,
              is_non_singular=True,
              is_positive_definite=True,
              is_square=True),
          input_output_cholesky=input_output_cholesky,
          validate_args=validate_args,
          allow_nan_stats=allow_nan_stats,
          name=name)
    self._parameters = parameters

  @classmethod
  def _params_event_ndims(cls):
    return dict(df=0, scale=2, scale_tril=2)


# 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 same-family Mixture distribution class."""

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

import collections

import tensorflow.compat.v2 as tf

from tensorflow_probability.python.distributions import distribution
from tensorflow_probability.python.distributions import independent
from tensorflow_probability.python.distributions import seed_stream
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import distribution_util as distribution_utils
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import reparameterization
from tensorflow_probability.python.internal import tensorshape_util

from tensorflow.python.ops import array_ops  # pylint: disable=g-direct-tensorflow-import
from tensorflow.python.ops.parallel_for import gradients  # pylint: disable=g-direct-tensorflow-import

class MixtureSameFamily(distribution.Distribution):
  """Mixture (same-family) distribution.

  The `MixtureSameFamily` distribution implements a (batch of) mixture
  distribution where all components are from different parameterizations of the
  same distribution type. It is parameterized by a `Categorical` "selecting
  distribution" (over `k` components) and a components distribution, i.e., a
  `Distribution` with a rightmost batch shape (equal to `[k]`) which indexes
  each (batch of) component.

  #### Examples

  ```python
  tfd = tfp.distributions

  ### Create a mixture of two scalar Gaussians:

  gm = tfd.MixtureSameFamily(
      mixture_distribution=tfd.Categorical(
          probs=[0.3, 0.7]),
      components_distribution=tfd.Normal(
        loc=[-1., 1],       # One for each component.
        scale=[0.1, 0.5]))  # And same here.

  gm.mean()
  # ==> 0.4

  gm.variance()
  # ==> 1.018

  # Plot PDF.
  x = np.linspace(-2., 3., int(1e4), dtype=np.float32)
  import matplotlib.pyplot as plt
  plt.plot(x, gm.prob(x).eval());

  ### Create a mixture of two Bivariate Gaussians:

  gm = tfd.MixtureSameFamily(
      mixture_distribution=tfd.Categorical(
          probs=[0.3, 0.7]),
      components_distribution=tfd.MultivariateNormalDiag(
          loc=[[-1., 1],  # component 1
               [1, -1]],  # component 2
          scale_identity_multiplier=[.3, .6]))

  gm.mean()
  # ==> array([ 0.4, -0.4], dtype=float32)

  gm.covariance()
  # ==> array([[ 1.119, -0.84],
  #            [-0.84,  1.119]], dtype=float32)

  # Plot PDF contours.
  def meshgrid(x, y=x):
    [gx, gy] = np.meshgrid(x, y, indexing='ij')
    gx, gy = np.float32(gx), np.float32(gy)
    grid = np.concatenate([gx.ravel()[None, :], gy.ravel()[None, :]], axis=0)
    return grid.T.reshape(x.size, y.size, 2)
  grid = meshgrid(np.linspace(-2, 2, 100, dtype=np.float32))
  plt.contour(grid[..., 0], grid[..., 1], gm.prob(grid).eval());

  ```

  """

  def __init__(self,
               mixture_distribution,
               components_distribution,
               reparameterize=False,
               validate_args=False,
               allow_nan_stats=True,
               name="MixtureSameFamily"):
    """Construct a `MixtureSameFamily` distribution.

    Args:
      mixture_distribution: `tfp.distributions.Categorical`-like instance.
        Manages the probability of selecting components. The number of
        categories must match the rightmost batch dimension of the
        `components_distribution`. Must have either scalar `batch_shape` or
        `batch_shape` matching `components_distribution.batch_shape[:-1]`.
      components_distribution: `tfp.distributions.Distribution`-like instance.
        Right-most batch dimension indexes components.
      reparameterize: Python `bool`, default `False`. Whether to reparameterize
        samples of the distribution using implicit reparameterization gradients
        [(Figurnov et al., 2018)][1]. The gradients for the mixture logits are
        equivalent to the ones described by [(Graves, 2016)][2]. The gradients
        for the components parameters are also computed using implicit
        reparameterization (as opposed to ancestral sampling), meaning that
        all components are updated every step.
        Only works when:
          (1) components_distribution is fully reparameterized;
          (2) components_distribution is either a scalar distribution or
          fully factorized (tfd.Independent applied to a scalar distribution);
          (3) batch shape has a known rank.
        Experimental, may be slow and produce infs/NaNs.
      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 not dtype_util.is_integer(mixture_distribution.dtype)`.
      ValueError: if mixture_distribution does not have scalar `event_shape`.
      ValueError: if `mixture_distribution.batch_shape` and
        `components_distribution.batch_shape[:-1]` are both fully defined and
        the former is neither scalar nor equal to the latter.
      ValueError: if `mixture_distribution` categories does not equal
        `components_distribution` rightmost batch shape.

    #### References

    [1]: Michael Figurnov, Shakir Mohamed and Andriy Mnih. Implicit
         reparameterization gradients. In _Neural Information Processing
         Systems_, 2018. https://arxiv.org/abs/1805.08498

    [2]: Alex Graves. Stochastic Backpropagation through Mixture Density
         Distributions. _arXiv_, 2016. https://arxiv.org/abs/1607.05690
    """
    parameters = dict(locals())
    with tf.name_scope(name) as name:
      self._mixture_distribution = mixture_distribution
      self._components_distribution = components_distribution
      self._runtime_assertions = []

      s = components_distribution.event_shape_tensor()
      self._event_ndims = tf.compat.dimension_value(s.shape[0])
      if self._event_ndims is None:
        self._event_ndims = tf.size(input=s)
      self._event_size = tf.reduce_prod(input_tensor=s)

      if not dtype_util.is_integer(mixture_distribution.dtype):
        raise ValueError(
            "`mixture_distribution.dtype` ({}) is not over integers".format(
                dtype_util.name(mixture_distribution.dtype)))

      if (tensorshape_util.rank(mixture_distribution.event_shape) is not None
          and tensorshape_util.rank(mixture_distribution.event_shape) != 0):
        raise ValueError("`mixture_distribution` must have scalar `event_dim`s")
      elif validate_args:
        self._runtime_assertions += [
            assert_util.assert_equal(
                tf.size(input=mixture_distribution.event_shape_tensor()),
                0,
                message="`mixture_distribution` must have scalar `event_dim`s"),
        ]

      mdbs = mixture_distribution.batch_shape
      cdbs = tensorshape_util.with_rank_at_least(
          components_distribution.batch_shape, 1)[:-1]
      if tensorshape_util.is_fully_defined(
          mdbs) and tensorshape_util.is_fully_defined(cdbs):
        if tensorshape_util.rank(mdbs) != 0 and mdbs != cdbs:
          raise ValueError(
              "`mixture_distribution.batch_shape` (`{}`) is not "
              "compatible with `components_distribution.batch_shape` "
              "(`{}`)".format(
                  tensorshape_util.as_list(mdbs),
                  tensorshape_util.as_list(cdbs)))
      elif validate_args:
        mdbs = mixture_distribution.batch_shape_tensor()
        cdbs = components_distribution.batch_shape_tensor()[:-1]
        self._runtime_assertions += [
            assert_util.assert_equal(
                distribution_utils.pick_vector(
                    mixture_distribution.is_scalar_batch(), cdbs, mdbs),
                cdbs,
                message=(
                    "`mixture_distribution.batch_shape` is not "
                    "compatible with `components_distribution.batch_shape`"))
        ]

      km = tf.compat.dimension_value(
          tensorshape_util.with_rank_at_least(mixture_distribution.logits.shape,
                                              1)[-1])
      kc = tf.compat.dimension_value(
          tensorshape_util.with_rank_at_least(
              components_distribution.batch_shape, 1)[-1])
      if km is not None and kc is not None and km != kc:
        raise ValueError("`mixture_distribution components` ({}) does not "
                         "equal `components_distribution.batch_shape[-1]` "
                         "({})".format(km, kc))
      elif validate_args:
        km = tf.shape(input=mixture_distribution.logits)[-1]
        kc = components_distribution.batch_shape_tensor()[-1]
        self._runtime_assertions += [
            assert_util.assert_equal(
                km,
                kc,
                message=("`mixture_distribution components` does not equal "
                         "`components_distribution.batch_shape[-1:]`")),
        ]
      elif km is None:
        km = tf.shape(input=mixture_distribution.logits)[-1]

      self._num_components = km

      self._reparameterize = reparameterize
      if reparameterize:
        # Note: tfd.Independent passes through the reparameterization type hence
        # we do not need separate logic for Independent.
        if (self._components_distribution.reparameterization_type !=
            reparameterization.FULLY_REPARAMETERIZED):
          raise ValueError("Cannot reparameterize a mixture of "
                           "non-reparameterized components.")
        reparameterization_type = reparameterization.FULLY_REPARAMETERIZED
      else:
        reparameterization_type = reparameterization.NOT_REPARAMETERIZED

      super(MixtureSameFamily, self).__init__(
          dtype=self._components_distribution.dtype,
          reparameterization_type=reparameterization_type,
          validate_args=validate_args,
          allow_nan_stats=allow_nan_stats,
          parameters=parameters,
          graph_parents=(
              self._mixture_distribution._graph_parents  # pylint: disable=protected-access
              + self._components_distribution._graph_parents),  # pylint: disable=protected-access
          name=name)

  @property
  def mixture_distribution(self):
    return self._mixture_distribution

  @property
  def components_distribution(self):
    return self._components_distribution

  def __getitem__(self, slices):
    # Because slicing is parameterization-dependent, we only implement slicing
    # for instances of MSF, not subclasses thereof.
    if type(self) is not MixtureSameFamily:  # pylint: disable=unidiomatic-typecheck
      return super(MixtureSameFamily, self).__getitem__(slices)

    slices = (
        list(slices) if isinstance(slices, collections.Sequence) else [slices])
    mixture_rank = tensorshape_util.rank(self.mixture_distribution.batch_shape)
    if mixture_rank is None:
      raise NotImplementedError("Cannot slice MixtureSameFamily with unknown "
                                "mixture_distribution rank")
    elif mixture_rank > 0:
      sliced_mixture_dist = self.mixture_distribution.__getitem__(slices)
    else:  # must be scalar
      sliced_mixture_dist = self.mixture_distribution

    # The components distribution has the component axis as the last batch dim,
    # and this must be preserved.
    if Ellipsis not in slices:
      slices.append(Ellipsis)
    slices.append(slice(None))
    sliced_components_dist = self.components_distribution.__getitem__(slices)
    return self.copy(
        mixture_distribution=sliced_mixture_dist,
        components_distribution=sliced_components_dist)

  def _batch_shape_tensor(self):
    with tf.control_dependencies(self._runtime_assertions):
      return self.components_distribution.batch_shape_tensor()[:-1]

  def _batch_shape(self):
    return tensorshape_util.with_rank_at_least(
        self.components_distribution.batch_shape, 1)[:-1]

  def _event_shape_tensor(self):
    with tf.control_dependencies(self._runtime_assertions):
      return self.components_distribution.event_shape_tensor()

  def _event_shape(self):
    return self.components_distribution.event_shape

  def _sample_n(self, n, seed):
    with tf.control_dependencies(self._runtime_assertions):
      seed = seed_stream.SeedStream(seed, salt="MixtureSameFamily")
      x = self.components_distribution.sample(n, seed=seed())  # [n, B, k, E]
      # TODO(jvdillon): Consider using tf.gather (by way of index unrolling).
      npdt = dtype_util.as_numpy_dtype(x.dtype)
      mask = tf.one_hot(
          indices=self.mixture_distribution.sample(n, seed=seed()),  # [n, B]
          depth=self._num_components,  # == k
          on_value=npdt(1),
          off_value=npdt(0))  # [n, B, k]
      mask = distribution_utils.pad_mixture_dimensions(
          mask, self, self.mixture_distribution,
          self._event_ndims)                         # [n, B, k, [1]*e]
      x = tf.reduce_sum(
          input_tensor=x * mask, axis=-1 - self._event_ndims)  # [n, B, E]
      if self._reparameterize:
        x = self._reparameterize_sample(x)
      return x

  def _log_prob(self, x):
    with tf.control_dependencies(self._runtime_assertions):
      x = self._pad_sample_dims(x)
      log_prob_x = self.components_distribution.log_prob(x)  # [S, B, k]
      log_mix_prob = tf.nn.log_softmax(
          self.mixture_distribution.logits, axis=-1)  # [B, k]
      return tf.reduce_logsumexp(
          input_tensor=log_prob_x + log_mix_prob, axis=-1)  # [S, B]

  def _mean(self):
    with tf.control_dependencies(self._runtime_assertions):
      probs = distribution_utils.pad_mixture_dimensions(
          self.mixture_distribution.probs, self, self.mixture_distribution,
          self._event_ndims)                         # [B, k, [1]*e]
      return tf.reduce_sum(
          input_tensor=probs * self.components_distribution.mean(),
          axis=-1 - self._event_ndims)  # [B, E]

  def _log_cdf(self, x):
    x = self._pad_sample_dims(x)
    log_cdf_x = self.components_distribution.log_cdf(x)      # [S, B, k]
    log_mix_prob = tf.nn.log_softmax(
        self.mixture_distribution.logits, axis=-1)  # [B, k]
    return tf.reduce_logsumexp(
        input_tensor=log_cdf_x + log_mix_prob, axis=-1)  # [S, B]

  def _variance(self):
    with tf.control_dependencies(self._runtime_assertions):
      # Law of total variance: Var(Y) = E[Var(Y|X)] + Var(E[Y|X])
      probs = distribution_utils.pad_mixture_dimensions(
          self.mixture_distribution.probs, self, self.mixture_distribution,
          self._event_ndims)                         # [B, k, [1]*e]
      mean_cond_var = tf.reduce_sum(
          input_tensor=probs * self.components_distribution.variance(),
          axis=-1 - self._event_ndims)  # [B, E]
      var_cond_mean = tf.reduce_sum(
          input_tensor=probs *
          tf.math.squared_difference(self.components_distribution.mean(),
                                     self._pad_sample_dims(self._mean())),
          axis=-1 - self._event_ndims)  # [B, E]
      return mean_cond_var + var_cond_mean                   # [B, E]

  def _covariance(self):
    static_event_ndims = tensorshape_util.rank(self.event_shape)
    if static_event_ndims is not None and static_event_ndims != 1:
      # Covariance is defined only for vector distributions.
      raise NotImplementedError("covariance is not implemented")

    with tf.control_dependencies(self._runtime_assertions):
      # Law of total variance: Var(Y) = E[Var(Y|X)] + Var(E[Y|X])
      probs = distribution_utils.pad_mixture_dimensions(
          distribution_utils.pad_mixture_dimensions(
              self.mixture_distribution.probs, self, self.mixture_distribution,
              self._event_ndims),
          self, self.mixture_distribution,
          self._event_ndims)                         # [B, k, 1, 1]
      mean_cond_var = tf.reduce_sum(
          input_tensor=probs * self.components_distribution.covariance(),
          axis=-3)  # [B, e, e]
      var_cond_mean = tf.reduce_sum(
          input_tensor=probs *
          _outer_squared_difference(self.components_distribution.mean(),
                                    self._pad_sample_dims(self._mean())),
          axis=-3)  # [B, e, e]
      return mean_cond_var + var_cond_mean                   # [B, e, e]

  def _pad_sample_dims(self, x):
    with tf.name_scope("pad_sample_dims"):
      ndims = tensorshape_util.rank(
          x.shape) if tensorshape_util.rank(x.shape) is not None else tf.rank(x)
      shape = tf.shape(input=x)
      d = ndims - self._event_ndims
      x = tf.reshape(x, shape=tf.concat([shape[:d], [1], shape[d:]], axis=0))
      return x

  def _reparameterize_sample(self, x):
    """Adds reparameterization (pathwise) gradients to samples of the mixture.

    Implicit reparameterization gradients are
       dx/dphi = -(d transform(x, phi) / dx)^-1 * d transform(x, phi) / dphi,
    where transform(x, phi) is distributional transform that removes all
    parameters from samples x.

    We implement them by replacing x with
      -stop_gradient(d transform(x, phi) / dx)^-1 * transform(x, phi)]
    for the backward pass (gradient computation).
    The derivative of this quantity w.r.t. phi is then the implicit
    reparameterization gradient.
    Note that this replaces the gradients w.r.t. both the mixture
    distribution parameters and components distributions parameters.

    Limitations:
      1. Fundamental: components must be fully reparameterized.
      2. Distributional transform is currently only implemented for
        factorized components.
      3. Distributional transform currently only works for known rank of the
        batch tensor.

    Arguments:
      x: Sample of mixture distribution

    Returns:
      Tensor with same value as x, but with reparameterization gradients
    """
    # Remove the existing gradients of x wrt parameters of the components.
    x = tf.stop_gradient(x)

    x_2d_shape = [-1, self._event_size]  # [S*prod(B), prod(E)]

    # Perform distributional transform of x in [S, B, E] shape,
    # but have Jacobian of size [S*prod(B), prod(E), prod(E)].
    def reshaped_distributional_transform(x_2d):
      return tf.reshape(
          self._distributional_transform(tf.reshape(x_2d, tf.shape(input=x))),
          x_2d_shape)

    # transform_2d: [S*prod(B), prod(E)]
    # jacobian: [S*prod(B), prod(E), prod(E)]
    transform_2d, jacobian = _value_and_batch_jacobian(
        reshaped_distributional_transform, tf.reshape(x, x_2d_shape))

    # We only provide the first derivative; the second derivative computed by
    # autodiff would be incorrect, so we raise an error if it is requested.
    transform_2d = _prevent_2nd_derivative(transform_2d)

    # Compute [- stop_gradient(jacobian)^-1 * transform] by solving a linear
    # system. The Jacobian is lower triangular because the distributional
    # transform for i-th event dimension does not depend on the next
    # dimensions.
    surrogate_x_2d = -tf.linalg.triangular_solve(
        tf.stop_gradient(jacobian), tf.expand_dims(transform_2d, axis=-1),
        lower=True)  # [S*prod(B), prod(E), 1]
    surrogate_x = tf.reshape(surrogate_x_2d, tf.shape(input=x))

    # Replace gradients of x with gradients of surrogate_x, but keep the value.
    return x + (surrogate_x - tf.stop_gradient(surrogate_x))

  def _distributional_transform(self, x):
    """Performs distributional transform of the mixture samples.

    Distributional transform removes the parameters from samples of a
    multivariate distribution by applying conditional CDFs:
      (F(x_1), F(x_2 | x1_), ..., F(x_d | x_1, ..., x_d-1))
    (the indexing is over the "flattened" event dimensions).
    The result is a sample of product of Uniform[0, 1] distributions.

    We assume that the components are factorized, so the conditional CDFs become
      F(x_i | x_1, ..., x_i-1) = sum_k w_i^k F_k (x_i),
    where w_i^k is the posterior mixture weight: for i > 0
      w_i^k = w_k prob_k(x_1, ..., x_i-1) / sum_k' w_k' prob_k'(x_1, ..., x_i-1)
    and w_0^k = w_k is the mixture probability of the k-th component.

    Arguments:
      x: Sample of mixture distribution

    Returns:
      Result of the distributional transform
    """

    if tensorshape_util.rank(x.shape) is None:
      # tf.nn.softmax raises an error when applied to inputs of undefined rank.
      raise ValueError("Distributional transform does not support inputs of "
                       "undefined rank.")

    # Obtain factorized components distribution and assert that it's
    # a scalar distribution.
    if isinstance(self._components_distribution, independent.Independent):
      univariate_components = self._components_distribution.distribution
    else:
      univariate_components = self._components_distribution

    with tf.control_dependencies([
        assert_util.assert_equal(
            univariate_components.is_scalar_event(),
            True,
            message="`univariate_components` must have scalar event")
    ]):
      x_padded = self._pad_sample_dims(x)  # [S, B, 1, E]
      log_prob_x = univariate_components.log_prob(x_padded)  # [S, B, k, E]
      cdf_x = univariate_components.cdf(x_padded)  # [S, B, k, E]

      # log prob_k (x_1, ..., x_i-1)
      cumsum_log_prob_x = tf.reshape(
          tf.math.cumsum(
              # [S*prod(B)*k, prod(E)]
              tf.reshape(log_prob_x, [-1, self._event_size]),
              exclusive=True,
              axis=-1),
          tf.shape(input=log_prob_x))  # [S, B, k, E]

      logits_mix_prob = distribution_utils.pad_mixture_dimensions(
          self.mixture_distribution.logits, self, self.mixture_distribution,
          self._event_ndims)  # [B, k, 1]

      # Logits of the posterior weights: log w_k + log prob_k (x_1, ..., x_i-1)
      log_posterior_weights_x = logits_mix_prob + cumsum_log_prob_x

      component_axis = tensorshape_util.rank(x.shape) - self._event_ndims
      posterior_weights_x = tf.nn.softmax(log_posterior_weights_x,
                                          axis=component_axis)
      return tf.reduce_sum(
          input_tensor=posterior_weights_x * cdf_x, axis=component_axis)

def _outer_squared_difference(x, y):
  """Convenience function analogous to tf.squared_difference."""
  z = x - y
  return z[..., tf.newaxis, :] * z[..., tf.newaxis]

def _value_and_batch_jacobian(f, x):
  """Enables uniform interface to value and batch jacobian calculation.

  Works in both eager and graph modes.

  Arguments:
    f: The scalar function to evaluate.
    x: The value at which to compute the value and the batch jacobian.

  Returns:
    A tuple (f(x), J(x)), where J(x) is the batch jacobian.
  """
  if tf.executing_eagerly():
    with tf.GradientTape() as tape:
      tape.watch(x)
      value = f(x)
    batch_jacobian = tape.batch_jacobian(value, x)
  else:
    value = f(x)
    batch_jacobian = gradients.batch_jacobian(value, x)
  return value, batch_jacobian

@tf.custom_gradient
def _prevent_2nd_derivative(x):
  """Disables computation of the second derivatives for a tensor.

  NB: you need to apply a non-identity function to the output tensor for the
  exception to be raised.

  Arguments:
    x: A tensor.

  Returns:
    A tensor with the same value and the same derivative as x, but that raises
    LookupError when trying to compute the second derivatives.
  """
  def grad(dy):
    return array_ops.prevent_gradient(
        dy, message="Second derivative is not implemented.")

  return tf.identity(x), grad


# 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.
# ============================================================================
"""Methods for decomposing StructuralTimeSeries models."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections

# Dependency imports
import tensorflow as tf

from tensorflow_probability.python.internal import distribution_util as dist_util
from tensorflow_probability.python.sts.internal import util as sts_util

def _split_covariance_into_marginals(covariance, block_sizes):
  """Split a covariance matrix into block-diagonal marginals of given sizes."""
  start_dim = 0
  marginals = []
  for size in block_sizes:
    end_dim = start_dim + size
    marginals.append(covariance[..., start_dim:end_dim, start_dim:end_dim])
    start_dim = end_dim
  return marginals

def _decompose_from_posterior_marginals(
    model, posterior_means, posterior_covs, parameter_samples):
  """Utility method to decompose a joint posterior into components.

  Args:
    model: `tfp.sts.Sum` instance defining an additive STS model.
    posterior_means: float `Tensor` of shape `concat(
      [[num_posterior_draws], batch_shape, num_timesteps, latent_size])`
      representing the posterior mean over latents in an
      `AdditiveStateSpaceModel`.
    posterior_covs: float `Tensor` of shape `concat(
      [[num_posterior_draws], batch_shape, num_timesteps,
      latent_size, latent_size])`
      representing the posterior marginal covariances over latents in an
      `AdditiveStateSpaceModel`.
    parameter_samples: Python `list` of `Tensors` representing posterior
      samples of model parameters, with shapes `[concat([
      [num_posterior_draws], param.prior.batch_shape,
      param.prior.event_shape]) for param in model.parameters]`. This may
      optionally also be a map (Python `dict`) of parameter names to
      `Tensor` values.

  Returns:
    component_dists: A `collections.OrderedDict` instance mapping
      component StructuralTimeSeries instances (elements of `model.components`)
      to `tfd.Distribution` instances representing the posterior marginal
      distributions on the process modeled by each component. Each distribution
      has batch shape matching that of `posterior_means`/`posterior_covs`, and
      event shape of `[num_timesteps]`.
  """

  try:
    model.components
  except AttributeError:
    raise ValueError('Model decomposed into components must be an instance of'
                     '`tfp.sts.Sum` (passed model {})'.format(model))

  with tf.compat.v1.name_scope('decompose_from_posterior_marginals'):

    # Extract the component means/covs from the joint latent posterior.
    latent_sizes = [component.latent_size for component in model.components]
    component_means = tf.split(posterior_means, latent_sizes, axis=-1)
    component_covs = _split_covariance_into_marginals(
        posterior_covs, latent_sizes)

    # Instantiate per-component state space models, and use them to push the
    # posterior means/covs through the observation model for each component.
    num_timesteps = dist_util.prefer_static_value(
        tf.shape(input=posterior_means))[-2]
    component_ssms = model.make_component_state_space_models(
        num_timesteps=num_timesteps,
        param_vals=parameter_samples)
    component_predictive_dists = collections.OrderedDict()
    for (component, component_ssm,
         component_mean, component_cov) in zip(model.components, component_ssms,
                                               component_means, component_covs):
      component_obs_mean, component_obs_cov = (
          component_ssm.latents_to_observations(
              latent_means=component_mean,
              latent_covs=component_cov))

      # Using the observation means and covs, build a mixture distribution
      # that integrates over the posterior draws.
      component_predictive_dists[component] = sts_util.mix_over_posterior_draws(
          means=component_obs_mean[..., 0],
          variances=component_obs_cov[..., 0, 0])
  return component_predictive_dists

def decompose_by_component(model, observed_time_series, parameter_samples):
  """Decompose an observed time series into contributions from each component.

  This method decomposes a time series according to the posterior represention
  of a structural time series model. In particular, it:
    - Computes the posterior marginal mean and covariances over the additive
      model's latent space.
    - Decomposes the latent posterior into the marginal blocks for each
      model component.
    - Maps the per-component latent posteriors back through each component's
      observation model, to generate the time series modeled by that component.

  Args:
    model: An instance of `tfp.sts.Sum` representing a structural time series
      model.
    observed_time_series: `float` `Tensor` of shape
      `batch_shape + [num_timesteps, 1]` (omitting the trailing unit dimension
      is also supported when `num_timesteps > 1`), specifying an observed time
      series. May optionally be an instance of `tfp.sts.MaskedTimeSeries`, which
      includes a mask `Tensor` to specify timesteps with missing observations.
    parameter_samples: Python `list` of `Tensors` representing posterior
      samples of model parameters, with shapes `[concat([
      [num_posterior_draws], param.prior.batch_shape,
      param.prior.event_shape]) for param in model.parameters]`. This may
      optionally also be a map (Python `dict`) of parameter names to
      `Tensor` values.
  Returns:
    component_dists: A `collections.OrderedDict` instance mapping
      component StructuralTimeSeries instances (elements of `model.components`)
      to `tfd.Distribution` instances representing the posterior marginal
      distributions on the process modeled by each component. Each distribution
      has batch shape matching that of `posterior_means`/`posterior_covs`, and
      event shape of `[num_timesteps]`.

  #### Examples

  Suppose we've built a model and fit it to data:

  ```python
    day_of_week = tfp.sts.Seasonal(
        num_seasons=7,
        observed_time_series=observed_time_series,
        name='day_of_week')
    local_linear_trend = tfp.sts.LocalLinearTrend(
        observed_time_series=observed_time_series,
        name='local_linear_trend')
    model = tfp.sts.Sum(components=[day_of_week, local_linear_trend],
                        observed_time_series=observed_time_series)

    num_steps_forecast = 50
    samples, kernel_results = tfp.sts.fit_with_hmc(model, observed_time_series)
  ```

  To extract the contributions of individual components, pass the time series
  and sampled parameters into `decompose_by_component`:

  ```python
    component_dists = decompose_by_component(
      model,
      observed_time_series=observed_time_series,
      parameter_samples=samples)

    # Component mean and stddev have shape `[len(observed_time_series)]`.
    day_of_week_effect_mean = component_dists[day_of_week].mean()
    day_of_week_effect_stddev = component_dists[day_of_week].stddev()
  ```

  Using the component distributions, we can visualize the uncertainty for
  each component:

  ```
  from matplotlib import pylab as plt
  num_components = len(component_dists)
  xs = np.arange(len(observed_time_series))
  fig = plt.figure(figsize=(12, 3 * num_components))
  for i, (component, component_dist) in enumerate(component_dists.items()):

    # If in graph mode, replace `.numpy()` with `.eval()` or `sess.run()`.
    component_mean = component_dist.mean().numpy()
    component_stddev = component_dist.stddev().numpy()

    ax = fig.add_subplot(num_components, 1, 1 + i)
    ax.plot(xs, component_mean, lw=2)
    ax.fill_between(xs,
                    component_mean - 2 * component_stddev,
                    component_mean + 2 * component_stddev,
                    alpha=0.5)
    ax.set_title(component.name)
  ```

  """

  with tf.compat.v1.name_scope('decompose_by_component',
                               values=[observed_time_series]):
    [
        observed_time_series,
        is_missing
    ] = sts_util.canonicalize_observed_time_series_with_mask(
        observed_time_series)

    # Run smoothing over the training timesteps to extract the
    # posterior on latents.
    num_timesteps = dist_util.prefer_static_value(
        tf.shape(input=observed_time_series))[-2]
    ssm = model.make_state_space_model(num_timesteps=num_timesteps,
                                       param_vals=parameter_samples)
    posterior_means, posterior_covs = ssm.posterior_marginals(
        observed_time_series, mask=is_missing)

    return _decompose_from_posterior_marginals(
        model, posterior_means, posterior_covs, parameter_samples)

def decompose_forecast_by_component(model, forecast_dist, parameter_samples):
  """Decompose a forecast distribution into contributions from each component.

  Args:
    model: An instance of `tfp.sts.Sum` representing a structural time series
      model.
    forecast_dist: A `Distribution` instance returned by `tfp.sts.forecast()`.
      (specifically, must be a `tfd.MixtureSameFamily` over a
      `tfd.LinearGaussianStateSpaceModel` parameterized by posterior samples).
    parameter_samples: Python `list` of `Tensors` representing posterior samples
      of model parameters, with shapes `[concat([[num_posterior_draws],
      param.prior.batch_shape, param.prior.event_shape]) for param in
      model.parameters]`. This may optionally also be a map (Python `dict`) of
      parameter names to `Tensor` values.
  Returns:
    component_forecasts: A `collections.OrderedDict` instance mapping
      component StructuralTimeSeries instances (elements of `model.components`)
      to `tfd.Distribution` instances representing the marginal forecast for
      each component. Each distribution has batch and event shape matching
      `forecast_dist` (specifically, the event shape is
      `[num_steps_forecast]`).

  #### Examples

  Suppose we've built a model, fit it to data, and constructed a forecast
  distribution:

  ```python
    day_of_week = tfp.sts.Seasonal(
        num_seasons=7,
        observed_time_series=observed_time_series,
        name='day_of_week')
    local_linear_trend = tfp.sts.LocalLinearTrend(
        observed_time_series=observed_time_series,
        name='local_linear_trend')
    model = tfp.sts.Sum(components=[day_of_week, local_linear_trend],
                        observed_time_series=observed_time_series)

    num_steps_forecast = 50
    samples, kernel_results = tfp.sts.fit_with_hmc(model, observed_time_series)
    forecast_dist = tfp.sts.forecast(model, observed_time_series,
                                 parameter_samples=samples,
                                 num_steps_forecast=num_steps_forecast)
  ```

  To extract the forecast for individual components, pass the forecast
  distribution into `decompose_forecast_by_components`:

  ```python
    component_forecasts = decompose_forecast_by_component(
      model, forecast_dist, samples)

    # Component mean and stddev have shape `[num_steps_forecast]`.
    day_of_week_effect_mean = forecast_components[day_of_week].mean()
    day_of_week_effect_stddev = forecast_components[day_of_week].stddev()
  ```

  Using the component forecasts, we can visualize the uncertainty for each
  component:

  ```
  from matplotlib import pylab as plt
  num_components = len(component_forecasts)
  xs = np.arange(num_steps_forecast)
  fig = plt.figure(figsize=(12, 3 * num_components))
  for i, (component, component_dist) in enumerate(component_forecasts.items()):

    # If in graph mode, replace `.numpy()` with `.eval()` or `sess.run()`.
    component_mean = component_dist.mean().numpy()
    component_stddev = component_dist.stddev().numpy()

    ax = fig.add_subplot(num_components, 1, 1 + i)
    ax.plot(xs, component_mean, lw=2)
    ax.fill_between(xs,
                    component_mean - 2 * component_stddev,
                    component_mean + 2 * component_stddev,
                    alpha=0.5)
    ax.set_title(component.name)
  ```

  """

  with tf.compat.v1.name_scope('decompose_forecast_by_component'):
    try:
      forecast_lgssm = forecast_dist.components_distribution
      forecast_latent_mean, _ = forecast_lgssm._joint_mean()  # pylint: disable=protected-access
      forecast_latent_covs, _ = forecast_lgssm._joint_covariances()  # pylint: disable=protected-access
    except AttributeError as e:
      raise ValueError(
          'Forecast distribution must be a MixtureSameFamily of'
          'LinearGaussianStateSpaceModel distributions, such as returned by'
          '`tfp.sts.forecast()`. (saw exception: {})'.format(e))

    # Since `parameter_samples` will have sample shape `[num_posterior_draws]`,
    # we need to move the `num_posterior_draws` dimension of the forecast
    # moments from the trailing batch dimension, where it's currently put by
    # `sts.forecast`, back to the leading (sample shape) dimension.
    forecast_latent_mean = dist_util.move_dimension(
        forecast_latent_mean, source_idx=-3, dest_idx=0)
    forecast_latent_covs = dist_util.move_dimension(
        forecast_latent_covs, source_idx=-4, dest_idx=0)

    return _decompose_from_posterior_marginals(
        model, forecast_latent_mean, forecast_latent_covs, parameter_samples)


# 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.
# ============================================================================
"""Functions for working with `tf.SparseTensor`."""

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

import tensorflow as tf

__all__ = [
    'dense_to_sparse',
]

def dense_to_sparse(x, ignore_value=None, name=None):
  """Converts dense `Tensor` to `SparseTensor`, dropping `ignore_value` cells.

  Args:
    x: A `Tensor`.
    ignore_value: Entries in `x` equal to this value will be
      absent from the return `SparseTensor`. If `None`, default value of
      `x` dtype will be used (e.g. '' for `str`, 0 for `int`).
    name: Python `str` prefix for ops created by this function.

  Returns:
    sparse_x: A `tf.SparseTensor` with the same shape as `x`.

  Raises:
    ValueError: when `x`'s rank is `None`.
  """
  # Copied (with modifications) from:
  # tensorflow/contrib/layers/python/ops/sparse_ops.py.
  with tf.compat.v1.name_scope(name, 'dense_to_sparse', [x, ignore_value]):
    x = tf.convert_to_tensor(value=x, name='x')
    if ignore_value is None:
      if x.dtype.base_dtype == tf.string:
        # Exception due to TF strings are converted to numpy objects by default.
        ignore_value = ''
      else:
        ignore_value = x.dtype.as_numpy_dtype(0)
      ignore_value = tf.cast(ignore_value, x.dtype, name='ignore_value')
    indices = tf.where(tf.not_equal(x, ignore_value), name='indices')
    return tf.SparseTensor(
        indices=indices,
        values=tf.gather_nd(x, indices, name='values'),
        dense_shape=tf.shape(input=x, out_type=tf.int64, name='dense_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.
# ============================================================================
"""Random variable."""

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

import functools
import tensorflow as tf

from tensorflow.python.client import session as tf_session
from tensorflow.python.framework import ops

__all__ = [
    "RandomVariable",
]

def _operator(attr):
  """Defers an operator overload to `attr`.

  Args:
    attr: Operator attribute to use.

  Returns:
    Function calling operator attribute.
  """
  @functools.wraps(attr)
  def func(a, *args):
    return attr(a.value, *args)
  return func

class RandomVariable(object):
  """Class for random variables.

  `RandomVariable` encapsulates properties of a random variable, namely, its
  distribution, sample shape, and (optionally overridden) value. Its `value`
  property is a `tf.Tensor`, which embeds the `RandomVariable` object into the
  TensorFlow graph. `RandomVariable` also features operator overloading and
  registration to TensorFlow sessions, enabling idiomatic usage as if one were
  operating on `tf.Tensor`s.

  The random variable's shape is given by

  `sample_shape + distribution.batch_shape + distribution.event_shape`,

  where `sample_shape` is an optional argument describing the shape of
  independent, identical draws from the distribution (default is `()`, meaning
  a single draw); `distribution.batch_shape` describes the shape of
  independent-but-not-identical draws (determined by the shape of the
  distribution's parameters); and `distribution.event_shape` describes the
  shape of dependent dimensions (e.g., `Normal` has scalar `event_shape`;
  `Dirichlet` has vector `event_shape`).

  #### Examples

  ```python
  import tensorflow_probability as tfp
  from tensorflow_probability import edward2 as ed
  tfd = tfp.distributions

  z1 = tf.constant([[1.0, -0.8], [0.3, -1.0]])
  z2 = tf.constant([[0.9, 0.2], [2.0, -0.1]])
  x = ed.RandomVariable(tfd.Bernoulli(logits=tf.matmul(z1, z2)))

  loc = ed.RandomVariable(tfd.Normal(0., 1.))
  x = ed.RandomVariable(tfd.Normal(loc, 1.), sample_shape=50)
  assert x.shape.as_list() == [50]
  assert x.sample_shape.as_list() == [50]
  assert x.distribution.batch_shape.as_list() == []
  assert x.distribution.event_shape.as_list() == []
  ```
  """

  def __init__(self,
               distribution,
               sample_shape=(),
               value=None):
    """Create a new random variable.

    Args:
      distribution: tfd.Distribution governing the distribution of the random
        variable, such as sampling and log-probabilities.
      sample_shape: tf.TensorShape of samples to draw from the random variable.
        Default is `()` corresponding to a single sample.
      value: Fixed tf.Tensor to associate with random variable. Must have shape
        `sample_shape + distribution.batch_shape + distribution.event_shape`.
        Default is to sample from random variable according to `sample_shape`.

    Raises:
      ValueError: `value` has incompatible shape with
        `sample_shape + distribution.batch_shape + distribution.event_shape`.
    """
    self._distribution = distribution
    self._sample_shape = sample_shape
    if value is not None:
      value = tf.cast(value, self.distribution.dtype)
      value_shape = value.shape
      expected_value_shape = self.sample_shape.concatenate(
          self.distribution.batch_shape).concatenate(
              self.distribution.event_shape)
      if not value_shape.is_compatible_with(expected_value_shape):
        raise ValueError(
            "Incompatible shape for initialization argument 'value'. "
            "Expected %s, got %s." % (expected_value_shape, value_shape))
    self._value = value

  @property
  def distribution(self):
    """Distribution of random variable."""
    return self._distribution

  @property
  def dtype(self):
    """`Dtype` of elements in this random variable."""
    return self.value.dtype

  @property
  def sample_shape(self):
    """Sample shape of random variable as a `TensorShape`."""
    if isinstance(self._sample_shape, tf.Tensor):
      return tf.TensorShape(tf.get_static_value(self._sample_shape))
    return tf.TensorShape(self._sample_shape)

  def sample_shape_tensor(self, name="sample_shape_tensor"):
    """Sample shape of random variable as a 1-D `Tensor`.

    Args:
      name: name to give to the op

    Returns:
      sample_shape: `Tensor`.
    """
    with tf.compat.v1.name_scope(name):
      if isinstance(self._sample_shape, tf.Tensor):
        return self._sample_shape
      return tf.convert_to_tensor(
          value=self.sample_shape.as_list(), dtype=tf.int32)

  @property
  def shape(self):
    """Shape of random variable."""
    return self.value.shape

  @property
  def value(self):
    """Get tensor that the random variable corresponds to."""
    if self._value is None:
      try:
        self._value = self.distribution.sample(self.sample_shape_tensor())
      except NotImplementedError:
        raise NotImplementedError(
            "sample is not implemented for {0}. You must either pass in the "
            "value argument or implement sample for {0}."
            .format(self.distribution.__class__.__name__))
    return self._value

  def __str__(self):
    if not isinstance(self.value, ops.EagerTensor):
      name = self.distribution.name
    else:
      name = _numpy_text(self.value)
    return "RandomVariable(\"%s\"%s%s%s)" % (
        name,
        ", shape=%s" % self.shape if self.shape.ndims is not None else "",
        ", dtype=%s" % self.dtype.name if self.dtype else "",
        ", device=%s" % self.value.device if self.value.device else "")

  def __repr__(self):
    string = "ed.RandomVariable '%s' shape=%s dtype=%s" % (
        self.distribution.name, self.shape, self.dtype.name)
    if hasattr(self.value, "numpy"):
      string += " numpy=%s" % _numpy_text(self.value, is_repr=True)
    return "<%s>" % string

  # Overload operators following tf.Tensor.
  __add__ = _operator(tf.Tensor.__add__)
  __radd__ = _operator(tf.Tensor.__radd__)
  __sub__ = _operator(tf.Tensor.__sub__)
  __rsub__ = _operator(tf.Tensor.__rsub__)
  __mul__ = _operator(tf.Tensor.__mul__)
  __rmul__ = _operator(tf.Tensor.__rmul__)
  __div__ = _operator(tf.Tensor.__div__)
  __rdiv__ = _operator(tf.Tensor.__rdiv__)
  __truediv__ = _operator(tf.Tensor.__truediv__)
  __rtruediv__ = _operator(tf.Tensor.__rtruediv__)
  __floordiv__ = _operator(tf.Tensor.__floordiv__)
  __rfloordiv__ = _operator(tf.Tensor.__rfloordiv__)
  __mod__ = _operator(tf.Tensor.__mod__)
  __rmod__ = _operator(tf.Tensor.__rmod__)
  __lt__ = _operator(tf.Tensor.__lt__)
  __le__ = _operator(tf.Tensor.__le__)
  __gt__ = _operator(tf.Tensor.__gt__)
  __ge__ = _operator(tf.Tensor.__ge__)
  __and__ = _operator(tf.Tensor.__and__)
  __rand__ = _operator(tf.Tensor.__rand__)
  __or__ = _operator(tf.Tensor.__or__)
  __ror__ = _operator(tf.Tensor.__ror__)
  __xor__ = _operator(tf.Tensor.__xor__)
  __rxor__ = _operator(tf.Tensor.__rxor__)
  __getitem__ = _operator(tf.Tensor.__getitem__)
  __pow__ = _operator(tf.Tensor.__pow__)
  __rpow__ = _operator(tf.Tensor.__rpow__)
  __invert__ = _operator(tf.Tensor.__invert__)
  __neg__ = _operator(tf.Tensor.__neg__)
  __abs__ = _operator(tf.Tensor.__abs__)
  __matmul__ = _operator(tf.Tensor.__matmul__)
  __rmatmul__ = _operator(tf.Tensor.__rmatmul__)
  __iter__ = _operator(tf.Tensor.__iter__)
  __bool__ = _operator(tf.Tensor.__bool__)
  __nonzero__ = _operator(tf.Tensor.__nonzero__)

  def __hash__(self):
    return id(self)

  def __eq__(self, other):
    return id(self) == id(other)

  def __ne__(self, other):
    return not self == other

  def eval(self, session=None, feed_dict=None):
    """In a session, computes and returns the value of this random variable.

    This is not a graph construction method, it does not add ops to the graph.

    This convenience method requires a session where the graph
    containing this variable has been launched. If no session is
    passed, the default session is used.

    Args:
      session: tf.BaseSession.
        The `tf.Session` to use to evaluate this random variable. If
        none, the default session is used.
      feed_dict: dict.
        A dictionary that maps `tf.Tensor` objects to feed values. See
        `tf.Session.run()` for a description of the valid feed values.

    Returns:
      Value of the random variable.

    #### Examples

    ```python
    x = Normal(0.0, 1.0)
    with tf.Session() as sess:
      # Usage passing the session explicitly.
      print(x.eval(sess))
      # Usage with the default session.  The 'with' block
      # above makes 'sess' the default session.
      print(x.eval())
    ```
    """
    return self.value.eval(session=session, feed_dict=feed_dict)

  def numpy(self):
    """Value as NumPy array, only available for TF Eager."""
    if not isinstance(self.value, ops.EagerTensor):
      raise NotImplementedError("value argument must be a EagerTensor.")

    return self.value.numpy()

  def get_shape(self):
    """Get shape of random variable."""
    return self.shape

  # This enables the RandomVariable's overloaded "right" binary operators to
  # run when the left operand is an ndarray, because it accords the
  # RandomVariable class higher priority than an ndarray, or a numpy matrix.
  __array_priority__ = 100

def _numpy_text(tensor, is_repr=False):
  """Human-readable representation of a tensor's numpy value."""
  if tensor.dtype.is_numpy_compatible:
    text = repr(tensor.numpy()) if is_repr else str(tensor.numpy())
  else:
    text = "<unprintable>"
  if "\n" in text:
    text = "\n" + text
  return text

def _session_run_conversion_fetch_function(tensor):
  return ([tensor.value], lambda val: val[0])

def _session_run_conversion_feed_function(feed, feed_val):
  return [(feed.value, feed_val)]

def _session_run_conversion_feed_function_for_partial_run(feed):
  return [feed.value]

def _tensor_conversion_function(v, dtype=None, name=None, as_ref=False):
  del name, as_ref  # unused
  if dtype and not dtype.is_compatible_with(v.dtype):
    raise ValueError(
        "Incompatible type conversion requested to type '%s' for variable "
        "of type '%s'" % (dtype.name, v.dtype.name))
  return v.value

tf_session.register_session_run_conversion_functions(  # enable sess.run, eval
    RandomVariable,
    _session_run_conversion_fetch_function,
    _session_run_conversion_feed_function,
    _session_run_conversion_feed_function_for_partial_run)

tf.register_tensor_conversion_function(  # enable tf.convert_to_tensor
    RandomVariable, _tensor_conversion_function)


# 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 Normal distribution: conjugate posterior closed form calculations."""

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

import tensorflow.compat.v2 as tf
from tensorflow_probability.python.distributions import normal

def normal_conjugates_known_scale_posterior(prior, scale, s, n):
  """Posterior Normal distribution with conjugate prior on the mean.

  This model assumes that `n` observations (with sum `s`) come from a
  Normal with unknown mean `loc` (described by the Normal `prior`)
  and known variance `scale**2`. The "known scale posterior" is
  the distribution of the unknown `loc`.

  Accepts a prior Normal distribution object, having parameters
  `loc0` and `scale0`, as well as known `scale` values of the predictive
  distribution(s) (also assumed Normal),
  and statistical estimates `s` (the sum(s) of the observations) and
  `n` (the number(s) of observations).

  Returns a posterior (also Normal) distribution object, with parameters
  `(loc', scale'**2)`, where:

  ```
  mu ~ N(mu', sigma'**2)
  sigma'**2 = 1/(1/sigma0**2 + n/sigma**2),
  mu' = (mu0/sigma0**2 + s/sigma**2) * sigma'**2.
  ```

  Distribution parameters from `prior`, as well as `scale`, `s`, and `n`.
  will broadcast in the case of multidimensional sets of parameters.

  Args:
    prior: `Normal` object of type `dtype`:
      the prior distribution having parameters `(loc0, scale0)`.
    scale: tensor of type `dtype`, taking values `scale > 0`.
      The known stddev parameter(s).
    s: Tensor of type `dtype`. The sum(s) of observations.
    n: Tensor of type `int`. The number(s) of observations.

  Returns:
    A new Normal posterior distribution object for the unknown observation
    mean `loc`.

  Raises:
    TypeError: if dtype of `s` does not match `dtype`, or `prior` is not a
      Normal object.
  """
  if not isinstance(prior, normal.Normal):
    raise TypeError("Expected prior to be an instance of type Normal")

  if s.dtype != prior.dtype:
    raise TypeError(
        "Observation sum s.dtype does not match prior dtype: %s vs. %s"
        % (s.dtype, prior.dtype))

  n = tf.cast(n, prior.dtype)
  scale0_2 = tf.square(prior.scale)
  scale_2 = tf.square(scale)
  scalep_2 = 1.0/(1/scale0_2 + n/scale_2)
  return normal.Normal(
      loc=(prior.loc / scale0_2 + s / scale_2) * scalep_2,
      scale=tf.sqrt(scalep_2))

def normal_conjugates_known_scale_predictive(prior, scale, s, n):
  """Posterior predictive Normal distribution w. conjugate prior on the mean.

  This model assumes that `n` observations (with sum `s`) come from a
  Normal with unknown mean `loc` (described by the Normal `prior`)
  and known variance `scale**2`. The "known scale predictive"
  is the distribution of new observations, conditioned on the existing
  observations and our prior.

  Accepts a prior Normal distribution object, having parameters
  `loc0` and `scale0`, as well as known `scale` values of the predictive
  distribution(s) (also assumed Normal),
  and statistical estimates `s` (the sum(s) of the observations) and
  `n` (the number(s) of observations).

  Calculates the Normal distribution(s) `p(x | sigma**2)`:

  ```
  p(x | sigma**2) = int N(x | mu, sigma**2)N(mu | prior.loc, prior.scale**2) dmu
                  = N(x | prior.loc, 1 / (sigma**2 + prior.scale**2))
  ```

  Returns the predictive posterior distribution object, with parameters
  `(loc', scale'**2)`, where:

  ```
  sigma_n**2 = 1/(1/sigma0**2 + n/sigma**2),
  mu' = (mu0/sigma0**2 + s/sigma**2) * sigma_n**2.
  sigma'**2 = sigma_n**2 + sigma**2,
  ```

  Distribution parameters from `prior`, as well as `scale`, `s`, and `n`.
  will broadcast in the case of multidimensional sets of parameters.

  Args:
    prior: `Normal` object of type `dtype`:
      the prior distribution having parameters `(loc0, scale0)`.
    scale: tensor of type `dtype`, taking values `scale > 0`.
      The known stddev parameter(s).
    s: Tensor of type `dtype`. The sum(s) of observations.
    n: Tensor of type `int`. The number(s) of observations.

  Returns:
    A new Normal predictive distribution object.

  Raises:
    TypeError: if dtype of `s` does not match `dtype`, or `prior` is not a
      Normal object.
  """
  if not isinstance(prior, normal.Normal):
    raise TypeError("Expected prior to be an instance of type Normal")

  if s.dtype != prior.dtype:
    raise TypeError(
        "Observation sum s.dtype does not match prior dtype: %s vs. %s"
        % (s.dtype, prior.dtype))

  n = tf.cast(n, prior.dtype)
  scale0_2 = tf.square(prior.scale)
  scale_2 = tf.square(scale)
  scalep_2 = 1.0/(1/scale0_2 + n/scale_2)
  return normal.Normal(
      loc=(prior.loc / scale0_2 + s / scale_2) * scalep_2,
      scale=tf.sqrt(scalep_2 + scale_2))


# 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.
# ============================================================================
"""Real NVP bijector."""

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

import tensorflow as tf

from tensorflow_probability.python.bijectors import bijector
from tensorflow_probability.python.internal import tensorshape_util

__all__ = [
    "RealNVP",
    "real_nvp_default_template"
]

class RealNVP(bijector.Bijector):
  """RealNVP "affine coupling layer" for vector-valued events.

  Real NVP models a normalizing flow on a `D`-dimensional distribution via a
  single `D-d`-dimensional conditional distribution [(Dinh et al., 2017)][1]:

  `y[d:D] = x[d:D] * tf.exp(log_scale_fn(x[0:d])) + shift_fn(x[0:d])`
  `y[0:d] = x[0:d]`

  The last `D-d` units are scaled and shifted based on the first `d` units only,
  while the first `d` units are 'masked' and left unchanged. Real NVP's
  `shift_and_log_scale_fn` computes vector-valued quantities. For
  scale-and-shift transforms that do not depend on any masked units, i.e.
  `d=0`, use the `tfb.Affine` bijector with learned parameters instead.

  Masking is currently only supported for base distributions with
  `event_ndims=1`. For more sophisticated masking schemes like checkerboard or
  channel-wise masking [(Papamakarios et al., 2016)[4], use the `tfb.Permute`
  bijector to re-order desired masked units into the first `d` units. For base
  distributions with `event_ndims > 1`, use the `tfb.Reshape` bijector to
  flatten the event shape.

  Recall that the MAF bijector [(Papamakarios et al., 2016)][4] implements a
  normalizing flow via an autoregressive transformation. MAF and IAF have
  opposite computational tradeoffs - MAF can train all units in parallel but
  must sample units sequentially, while IAF must train units sequentially but
  can sample in parallel. In contrast, Real NVP can compute both forward and
  inverse computations in parallel. However, the lack of an autoregressive
  transformations makes it less expressive on a per-bijector basis.

  A "valid" `shift_and_log_scale_fn` must compute each `shift` (aka `loc` or
  "mu" in [Papamakarios et al. (2016)][4]) and `log(scale)` (aka "alpha" in
  [Papamakarios et al. (2016)][4]) such that each are broadcastable with the
  arguments to `forward` and `inverse`, i.e., such that the calculations in
  `forward`, `inverse` [below] are possible. For convenience,
  `real_nvp_default_nvp` is offered as a possible `shift_and_log_scale_fn`
  function.

  NICE [(Dinh et al., 2014)][2] is a special case of the Real NVP bijector
  which discards the scale transformation, resulting in a constant-time
  inverse-log-determinant-Jacobian. To use a NICE bijector instead of Real
  NVP, `shift_and_log_scale_fn` should return `(shift, None)`, and
  `is_constant_jacobian` should be set to `True` in the `RealNVP` constructor.
  Calling `real_nvp_default_template` with `shift_only=True` returns one such
  NICE-compatible `shift_and_log_scale_fn`.

  Caching: the scalar input depth `D` of the base distribution is not known at
  construction time. The first call to any of `forward(x)`, `inverse(x)`,
  `inverse_log_det_jacobian(x)`, or `forward_log_det_jacobian(x)` memoizes
  `D`, which is re-used in subsequent calls. This shape must be known prior to
  graph execution (which is the case if using tf.layers).

  #### Example Use

  ```python
  tfd = tfp.distributions
  tfb = tfp.bijectors

  # A common choice for a normalizing flow is to use a Gaussian for the base
  # distribution. (However, any continuous distribution would work.) E.g.,
  nvp = tfd.TransformedDistribution(
      distribution=tfd.MultivariateNormalDiag(loc=[0., 0., 0.]),
      bijector=tfb.RealNVP(
          num_masked=2,
          shift_and_log_scale_fn=tfb.real_nvp_default_template(
              hidden_layers=[512, 512])))

  x = nvp.sample()
  nvp.log_prob(x)
  nvp.log_prob(0.)
  ```

  For more examples, see [Jang (2018)][3].

  #### References

  [1]: Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density Estimation
       using Real NVP. In _International Conference on Learning
       Representations_, 2017. https://arxiv.org/abs/1605.08803

  [2]: Laurent Dinh, David Krueger, and Yoshua Bengio. NICE: Non-linear
       Independent Components Estimation. _arXiv preprint arXiv:1410.8516_,
       2014. https://arxiv.org/abs/1410.8516

  [3]: Eric Jang. Normalizing Flows Tutorial, Part 2: Modern Normalizing Flows.
       _Technical Report_, 2018. http://blog.evjang.com/2018/01/nf2.html

  [4]: George Papamakarios, Theo Pavlakou, and Iain Murray. Masked
       Autoregressive Flow for Density Estimation. In _Neural Information
       Processing Systems_, 2017. https://arxiv.org/abs/1705.07057
  """

  def __init__(self,
               num_masked,
               shift_and_log_scale_fn,
               is_constant_jacobian=False,
               validate_args=False,
               name=None):
    """Creates the Real NVP or NICE bijector.

    Args:
      num_masked: Python `int` indicating that the first `d` units of the event
        should be masked. Must be in the closed interval `[1, D-1]`, where `D`
        is the event size of the base distribution.
      shift_and_log_scale_fn: Python `callable` which computes `shift` and
        `log_scale` from both the forward domain (`x`) and the inverse domain
        (`y`). Calculation must respect the "autoregressive property" (see class
        docstring). Suggested default
        `masked_autoregressive_default_template(hidden_layers=...)`.
        Typically the function contains `tf.Variables` and is wrapped using
        `tf.make_template`. Returning `None` for either (both) `shift`,
        `log_scale` is equivalent to (but more efficient than) returning zero.
      is_constant_jacobian: Python `bool`. Default: `False`. When `True` the
        implementation assumes `log_scale` does not depend on the forward domain
        (`x`) or inverse domain (`y`) values. (No validation is made;
        `is_constant_jacobian=False` is always safe but possibly computationally
        inefficient.)
      validate_args: Python `bool` indicating whether arguments should be
        checked for correctness.
      name: Python `str`, name given to ops managed by this object.

    Raises:
      ValueError: If num_masked < 1.
    """
    name = name or "real_nvp"
    if num_masked <= 0:
      raise ValueError("num_masked must be a positive integer.")
    self._num_masked = num_masked
    # At construction time, we don't know input_depth.
    self._input_depth = None
    self._shift_and_log_scale_fn = shift_and_log_scale_fn
    super(RealNVP, self).__init__(
        forward_min_event_ndims=1,
        is_constant_jacobian=is_constant_jacobian,
        validate_args=validate_args,
        name=name)

  def _cache_input_depth(self, x):
    if self._input_depth is None:
      self._input_depth = tf.compat.dimension_value(
          tensorshape_util.with_rank_at_least(x.shape, 1)[-1])
      if self._input_depth is None:
        raise NotImplementedError(
            "Rightmost dimension must be known prior to graph execution.")
      if self._num_masked >= self._input_depth:
        raise ValueError(
            "Number of masked units must be smaller than the event size.")

  def _forward(self, x, **condition_kwargs):
    self._cache_input_depth(x)
    # Performs scale and shift.
    x0, x1 = x[..., :self._num_masked], x[..., self._num_masked:]
    shift, log_scale = self._shift_and_log_scale_fn(
        x0, self._input_depth - self._num_masked, **condition_kwargs)
    y1 = x1
    if log_scale is not None:
      y1 *= tf.exp(log_scale)
    if shift is not None:
      y1 += shift
    y = tf.concat([x0, y1], axis=-1)
    return y

  def _inverse(self, y, **condition_kwargs):
    self._cache_input_depth(y)
    # Performs un-shift and un-scale.
    y0, y1 = y[..., :self._num_masked], y[..., self._num_masked:]
    shift, log_scale = self._shift_and_log_scale_fn(
        y0, self._input_depth - self._num_masked, **condition_kwargs)
    x1 = y1
    if shift is not None:
      x1 -= shift
    if log_scale is not None:
      x1 *= tf.exp(-log_scale)
    x = tf.concat([y0, x1], axis=-1)
    return x

  def _inverse_log_det_jacobian(self, y, **condition_kwargs):
    self._cache_input_depth(y)
    y0 = y[..., :self._num_masked]
    _, log_scale = self._shift_and_log_scale_fn(
        y0, self._input_depth - self._num_masked, **condition_kwargs)
    if log_scale is None:
      return tf.constant(0., dtype=y.dtype, name="ildj")
    return -tf.reduce_sum(input_tensor=log_scale, axis=-1)

  def _forward_log_det_jacobian(self, x, **condition_kwargs):
    self._cache_input_depth(x)
    x0 = x[..., :self._num_masked]
    _, log_scale = self._shift_and_log_scale_fn(
        x0, self._input_depth - self._num_masked, **condition_kwargs)
    if log_scale is None:
      return tf.constant(0., dtype=x.dtype, name="fldj")
    return tf.reduce_sum(input_tensor=log_scale, axis=-1)

def real_nvp_default_template(hidden_layers,
                              shift_only=False,
                              activation=tf.nn.relu,
                              name=None,
                              *args,  # pylint: disable=keyword-arg-before-vararg
                              **kwargs):
  """Build a scale-and-shift function using a multi-layer neural network.

  This will be wrapped in a make_template to ensure the variables are only
  created once. It takes the `d`-dimensional input x[0:d] and returns the `D-d`
  dimensional outputs `loc` ("mu") and `log_scale` ("alpha").

  The default template does not support conditioning and will raise an
  exception if `condition_kwargs` are passed to it. To use conditioning in
  real nvp bijector, implement a conditioned shift/scale template that
  handles the `condition_kwargs`.

  Arguments:
    hidden_layers: Python `list`-like of non-negative integer, scalars
      indicating the number of units in each hidden layer. Default: `[512, 512].
    shift_only: Python `bool` indicating if only the `shift` term shall be
      computed (i.e. NICE bijector). Default: `False`.
    activation: Activation function (callable). Explicitly setting to `None`
      implies a linear activation.
    name: A name for ops managed by this function. Default:
      "real_nvp_default_template".
    *args: `tf.layers.dense` arguments.
    **kwargs: `tf.layers.dense` keyword arguments.

  Returns:
    shift: `Float`-like `Tensor` of shift terms ("mu" in
      [Papamakarios et al.  (2016)][1]).
    log_scale: `Float`-like `Tensor` of log(scale) terms ("alpha" in
      [Papamakarios et al. (2016)][1]).

  Raises:
    NotImplementedError: if rightmost dimension of `inputs` is unknown prior to
      graph execution, or if `condition_kwargs` is not empty.

  #### References

  [1]: George Papamakarios, Theo Pavlakou, and Iain Murray. Masked
       Autoregressive Flow for Density Estimation. In _Neural Information
       Processing Systems_, 2017. https://arxiv.org/abs/1705.07057
  """

  with tf.compat.v2.name_scope(name or "real_nvp_default_template"):

    def _fn(x, output_units, **condition_kwargs):
      """Fully connected MLP parameterized via `real_nvp_template`."""
      if condition_kwargs:
        raise NotImplementedError(
            "Conditioning not implemented in the default template.")

      if tensorshape_util.rank(x.shape) == 1:
        x = x[tf.newaxis, ...]
        reshape_output = lambda x: x[0]
      else:
        reshape_output = lambda x: x
      for units in hidden_layers:
        x = tf.compat.v1.layers.dense(
            inputs=x,
            units=units,
            activation=activation,
            *args,  # pylint: disable=keyword-arg-before-vararg
            **kwargs)
      x = tf.compat.v1.layers.dense(
          inputs=x,
          units=(1 if shift_only else 2) * output_units,
          activation=None,
          *args,  # pylint: disable=keyword-arg-before-vararg
          **kwargs)
      if shift_only:
        return reshape_output(x), None
      shift, log_scale = tf.split(x, 2, axis=-1)
      return reshape_output(shift), reshape_output(log_scale)

    return tf.compat.v1.make_template("real_nvp_default_template", _fn)