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

from __future__ import absolute_import
from __future__ import division

import tensorflow.compat.v2 as tf

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

__all__ = [
    "HiddenMarkovModel",
]

class HiddenMarkovModel(distribution.Distribution):
  """Hidden Markov model distribution.

  The `HiddenMarkovModel` distribution implements a (batch of) hidden
  Markov models where the initial states, transition probabilities
  and observed states are all given by user-provided distributions.
  This model assumes that the transition matrices are fixed over time.

  In this model, there is a sequence of integer-valued hidden states:
  `z[0], z[1], ..., z[num_steps - 1]` and a sequence of observed states:
  `x[0], ..., x[num_steps - 1]`.
  The distribution of `z[0]` is given by `initial_distribution`.
  The conditional probability of `z[i  +  1]` given `z[i]` is described by
  the batch of distributions in `transition_distribution`.
  For a batch of hidden Markov models, the coordinates before the rightmost one
  of the `transition_distribution` batch correspond to indices into the hidden
  Markov model batch. The rightmost coordinate of the batch is used to select
  which distribution `z[i + 1]` is drawn from.  The distributions corresponding
  to the probability of `z[i + 1]` conditional on `z[i] == k` is given by the
  elements of the batch whose rightmost coordinate is `k`.
  Similarly, the conditional distribution of `z[i]` given `x[i]` is given by
  the batch of `observation_distribution`.
  When the rightmost coordinate of `observation_distribution` is `k` it
  gives the conditional probabilities of `x[i]` given `z[i] == k`.
  The probability distribution associated with the `HiddenMarkovModel`
  distribution is the marginal distribution of `x[0],...,x[num_steps - 1]`.

  #### Examples

  ```python
  tfd = tfp.distributions

  # A simple weather model.

  # Represent a cold day with 0 and a hot day with 1.
  # Suppose the first day of a sequence has a 0.8 chance of being cold.
  # We can model this using the categorical distribution:

  initial_distribution = tfd.Categorical(probs=[0.8, 0.2])

  # Suppose a cold day has a 30% chance of being followed by a hot day
  # and a hot day has a 20% chance of being followed by a cold day.
  # We can model this as:

  transition_distribution = tfd.Categorical(probs=[[0.7, 0.3],
                                                   [0.2, 0.8]])

  # Suppose additionally that on each day the temperature is
  # normally distributed with mean and standard deviation 0 and 5 on
  # a cold day and mean and standard deviation 15 and 10 on a hot day.
  # We can model this with:

  observation_distribution = tfd.Normal(loc=[0., 15.], scale=[5., 10.])

  # We can combine these distributions into a single week long
  # hidden Markov model with:

  model = tfd.HiddenMarkovModel(
      initial_distribution=initial_distribution,
      transition_distribution=transition_distribution,
      observation_distribution=observation_distribution,
      num_steps=7)

  # The expected temperatures for each day are given by:

  model.mean()  # shape [7], elements approach 9.0

  # The log pdf of a week of temperature 0 is:

  model.log_prob(tfp.zeros(shape=[7]))
  ```

  #### References
  [1] https://en.wikipedia.org/wiki/Hidden_Markov_model
  """

  def __init__(self,
               initial_distribution,
               transition_distribution,
               observation_distribution,
               num_steps,
               validate_args=False,
               allow_nan_stats=True,
               name="HiddenMarkovModel"):
    """Initialize hidden Markov model.

    Args:
      initial_distribution: A `Categorical`-like instance.
        Determines probability of first hidden state in Markov chain.
        The number of categories must match the number of categories of
        `transition_distribution` as well as both the rightmost batch
        dimension of `transition_distribution` and the rightmost batch
        dimension of `observation_distribution`.
      transition_distribution: A `Categorical`-like instance.
        The rightmost batch dimension indexes the probability distribution
        of each hidden state conditioned on the previous hidden state.
      observation_distribution: A `tfp.distributions.Distribution`-like
        instance.  The rightmost batch dimension indexes the distribution
        of each observation conditioned on the corresponding hidden state.
      num_steps: The number of steps taken in Markov chain. A python `int`.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
        Default value: `False`.
      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
        result is undefined. When `False`, an exception is raised if one or
        more of the statistic's batch members are undefined.
        Default value: `True`.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: "HiddenMarkovModel".

    Raises:
      ValueError: if `num_steps` is not at least 1.
      ValueError: if `initial_distribution` does not have scalar `event_shape`.
      ValueError: if `transition_distribution` does not have scalar
        `event_shape.`
      ValueError: if `transition_distribution` and `observation_distribution`
        are fully defined but don't have matching rightmost dimension.
    """

    parameters = dict(locals())

    # pylint: disable=protected-access
    with tf.name_scope(name) as name:
      self._runtime_assertions = []  # pylint: enable=protected-access

      if num_steps < 1:
        raise ValueError("num_steps ({}) must be at least 1.".format(num_steps))

      self._initial_distribution = initial_distribution
      self._observation_distribution = observation_distribution
      self._transition_distribution = transition_distribution

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

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

      if (transition_distribution.batch_shape is not None and
          tensorshape_util.rank(transition_distribution.batch_shape) == 0):
        raise ValueError(
            "`transition_distribution` can't have scalar batches")
      elif validate_args:
        self._runtime_assertions += [
            assert_util.assert_greater(
                tf.size(input=transition_distribution.batch_shape_tensor()),
                0,
                message="`transition_distribution` can't have scalar "
                "batches")
        ]

      if (observation_distribution.batch_shape is not None and
          tensorshape_util.rank(observation_distribution.batch_shape) == 0):
        raise ValueError(
            "`observation_distribution` can't have scalar batches")
      elif validate_args:
        self._runtime_assertions += [
            assert_util.assert_greater(
                tf.size(input=observation_distribution.batch_shape_tensor()),
                0,
                message="`observation_distribution` can't have scalar "
                "batches")
        ]

      # Infer number of hidden states and check consistency
      # between transitions and observations
      with tf.control_dependencies(self._runtime_assertions):
        self._num_states = ((transition_distribution.batch_shape and
                             transition_distribution.batch_shape[-1]) or
                            transition_distribution.batch_shape_tensor()[-1])

        observation_states = ((observation_distribution.batch_shape and
                               observation_distribution.batch_shape[-1]) or
                              observation_distribution.batch_shape_tensor()[-1])

      if (tf.is_tensor(self._num_states) or tf.is_tensor(observation_states)):
        if validate_args:
          self._runtime_assertions += [
              assert_util.assert_equal(
                  self._num_states,
                  observation_states,
                  message="`transition_distribution` and "
                  "`observation_distribution` must agree on "
                  "last dimension of batch size")
          ]
      elif self._num_states != observation_states:
        raise ValueError("`transition_distribution` and "
                         "`observation_distribution` must agree on "
                         "last dimension of batch size")

      self._log_init = _extract_log_probs(self._num_states,
                                          initial_distribution)
      self._log_trans = _extract_log_probs(self._num_states,
                                           transition_distribution)

      self._num_steps = num_steps
      self._num_states = tf.shape(input=self._log_init)[-1]

      self._underlying_event_rank = tf.size(
          input=self._observation_distribution.event_shape_tensor())

      self.static_event_shape = tf.TensorShape(
          [num_steps]).concatenate(self._observation_distribution.event_shape)

      with tf.control_dependencies(self._runtime_assertions):
        self.static_batch_shape = tf.broadcast_static_shape(
            self._initial_distribution.batch_shape,
            tf.broadcast_static_shape(
                self._transition_distribution.batch_shape[:-1],
                self._observation_distribution.batch_shape[:-1]))

      # pylint: disable=protected-access
      super(HiddenMarkovModel, self).__init__(
          dtype=self._observation_distribution.dtype,
          reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
          validate_args=validate_args,
          allow_nan_stats=allow_nan_stats,
          parameters=parameters,
          graph_parents=(self._initial_distribution._graph_parents +
                         self._transition_distribution._graph_parents +
                         self._observation_distribution._graph_parents),
          name=name)
      # pylint: enable=protected-access

      self._parameters = parameters

  def _batch_shape_tensor(self):
    with tf.control_dependencies(self._runtime_assertions):
      return tf.broadcast_dynamic_shape(
          self._initial_distribution.batch_shape_tensor(),
          tf.broadcast_dynamic_shape(
              self._transition_distribution.batch_shape_tensor()[:-1],
              self._observation_distribution.batch_shape_tensor()[:-1]))

  def _batch_shape(self):
    return self.static_batch_shape

  def _event_shape_tensor(self):
    with tf.control_dependencies(self._runtime_assertions):
      return tf.concat([[self._num_steps],
                        self.observation_distribution.event_shape_tensor()],
                       axis=0)

  def _event_shape(self):
    return self.static_event_shape

  @property
  def initial_distribution(self):
    return self._initial_distribution

  @property
  def transition_distribution(self):
    return self._transition_distribution

  @property
  def observation_distribution(self):
    return self._observation_distribution

  @property
  def num_steps(self):
    return self._num_steps

  @property
  def num_states(self):
    return self._num_states

  def _sample_n(self, n, seed=None):
    with tf.control_dependencies(self._runtime_assertions):
      seed = seed_stream.SeedStream(seed, salt="HiddenMarkovModel")

      num_states = self._num_states

      batch_shape = self.batch_shape_tensor()
      batch_size = tf.reduce_prod(input_tensor=batch_shape)

      # The batch sizes of the underlying initial distributions and
      # transition distributions might not match the batch size of
      # the HMM distribution.
      # As a result we need to ask for more samples from the
      # underlying distributions and then reshape the results into
      # the correct batch size for the HMM.
      init_repeat = (
          tf.reduce_prod(input_tensor=self.batch_shape_tensor()) //
          tf.reduce_prod(
              input_tensor=self._initial_distribution.batch_shape_tensor()))
      init_state = self._initial_distribution.sample(n * init_repeat,
                                                     seed=seed())
      init_state = tf.reshape(init_state, [n, batch_size])
      # init_state :: n batch_size

      transition_repeat = (
          tf.reduce_prod(input_tensor=self.batch_shape_tensor()) //
          tf.reduce_prod(input_tensor=self._transition_distribution
                         .batch_shape_tensor()[:-1]))

      def generate_step(state, _):
        """Take a single step in Markov chain."""

        gen = self._transition_distribution.sample(n * transition_repeat,
                                                   seed=seed())
        # gen :: (n * transition_repeat) transition_batch

        new_states = tf.reshape(gen,
                                [n, batch_size, num_states])

        # new_states :: n batch_size num_states

        old_states_one_hot = tf.one_hot(state, num_states, dtype=tf.int32)

        # old_states :: n batch_size num_states

        return tf.reduce_sum(
            input_tensor=old_states_one_hot * new_states, axis=-1)

      if self._num_steps > 1:
        dummy_index = tf.zeros(self._num_steps - 1, dtype=tf.float32)
        hidden_states = tf.scan(generate_step, dummy_index,
                                initializer=init_state)

        # TODO(b/115618503): add/use prepend_initializer to tf.scan
        hidden_states = tf.concat([[init_state],
                                   hidden_states], axis=0)
      else:
        hidden_states = init_state[tf.newaxis, ...]

      # hidden_states :: num_steps n batch_size num_states

      hidden_one_hot = tf.one_hot(hidden_states, num_states,
                                  dtype=self._observation_distribution.dtype)
      # hidden_one_hot :: num_steps n batch_size num_states

      # The observation distribution batch size might not match
      # the required batch size so as with the initial and
      # transition distributions we generate more samples and
      # reshape.
      observation_repeat = (
          batch_size //
          tf.reduce_prod(input_tensor=self._observation_distribution
                         .batch_shape_tensor()[:-1]))

      possible_observations = self._observation_distribution.sample(
          [self._num_steps, observation_repeat * n])

      inner_shape = self._observation_distribution.event_shape

      # possible_observations :: num_steps (observation_repeat * n)
      #                          observation_batch[:-1] num_states inner_shape

      possible_observations = tf.reshape(
          possible_observations,
          tf.concat([[self._num_steps, n],
                     batch_shape,
                     [num_states],
                     inner_shape], axis=0))

      # possible_observations :: steps n batch_size num_states inner_shape

      hidden_one_hot = tf.reshape(hidden_one_hot,
                                  tf.concat([[self._num_steps, n],
                                             batch_shape,
                                             [num_states],
                                             tf.ones_like(inner_shape)],
                                            axis=0))

      # hidden_one_hot :: steps n batch_size num_states "inner_shape"

      observations = tf.reduce_sum(
          input_tensor=hidden_one_hot * possible_observations,
          axis=-1 - tf.size(input=inner_shape))

      # observations :: steps n batch_size inner_shape

      observations = distribution_util.move_dimension(
          observations, 0, 1 + tf.size(input=batch_shape))

      # returned :: n batch_shape steps inner_shape

      return observations

  def _log_prob(self, value):
    with tf.control_dependencies(self._runtime_assertions):
      # The argument `value` is a tensor of sequences of observations.
      # `observation_batch_shape` is the shape of that tensor with the
      # sequence part removed.
      # `observation_batch_shape` is then broadcast to the full batch shape
      # to give the `batch_shape` that defines the shape of the result.

      observation_tensor_shape = tf.shape(input=value)
      observation_batch_shape = observation_tensor_shape[
          :-1 - self._underlying_event_rank]
      # value :: observation_batch_shape num_steps observation_event_shape
      batch_shape = tf.broadcast_dynamic_shape(observation_batch_shape,
                                               self.batch_shape_tensor())
      log_init = tf.broadcast_to(self._log_init,
                                 tf.concat([batch_shape,
                                            [self._num_states]], axis=0))
      # log_init :: batch_shape num_states
      log_transition = self._log_trans

      # `observation_event_shape` is the shape of each sequence of observations
      # emitted by the model.
      observation_event_shape = observation_tensor_shape[
          -1 - self._underlying_event_rank:]
      working_obs = tf.broadcast_to(value,
                                    tf.concat([batch_shape,
                                               observation_event_shape],
                                              axis=0))
      # working_obs :: batch_shape observation_event_shape
      r = self._underlying_event_rank

      # Move index into sequence of observations to front so we can apply
      # tf.foldl
      working_obs = distribution_util.move_dimension(working_obs, -1 - r,
                                                     0)[..., tf.newaxis]
      # working_obs :: num_steps batch_shape underlying_event_shape
      observation_probs = (
          self._observation_distribution.log_prob(working_obs))

      def forward_step(log_prev_step, log_prob_observation):
        return _log_vector_matrix(log_prev_step,
                                  log_transition) + log_prob_observation

      fwd_prob = tf.foldl(forward_step, observation_probs, initializer=log_init)
      # fwd_prob :: batch_shape num_states

      log_prob = tf.reduce_logsumexp(input_tensor=fwd_prob, axis=-1)
      # log_prob :: batch_shape

      return log_prob

  def _marginal_hidden_probs(self):
    """Compute marginal pdf for each individual observable."""

    initial_log_probs = tf.broadcast_to(self._log_init,
                                        tf.concat([self.batch_shape_tensor(),
                                                   [self._num_states]],
                                                  axis=0))
    # initial_log_probs :: batch_shape num_states

    if self._num_steps > 1:
      transition_log_probs = self._log_trans

      def forward_step(log_probs, _):
        return _log_vector_matrix(log_probs, transition_log_probs)

      dummy_index = tf.zeros(self._num_steps - 1, dtype=tf.float32)

      forward_log_probs = tf.scan(forward_step, dummy_index,
                                  initializer=initial_log_probs,
                                  name="forward_log_probs")

      forward_log_probs = tf.concat([[initial_log_probs], forward_log_probs],
                                    axis=0)
    else:
      forward_log_probs = initial_log_probs[tf.newaxis, ...]

    # returns :: num_steps batch_shape num_states

    return tf.exp(forward_log_probs)

  def _mean(self):
    with tf.control_dependencies(self._runtime_assertions):
      probs = self._marginal_hidden_probs()
      # probs :: num_steps batch_shape num_states
      means = self._observation_distribution.mean()
      # means :: observation_batch_shape[:-1] num_states
      #          observation_event_shape
      means_shape = tf.concat(
          [self.batch_shape_tensor(),
           [self._num_states],
           self._observation_distribution.event_shape_tensor()],
          axis=0)
      means = tf.broadcast_to(means, means_shape)
      # means :: batch_shape num_states observation_event_shape

      observation_event_shape = (
          self._observation_distribution.event_shape_tensor())
      batch_size = tf.reduce_prod(input_tensor=self.batch_shape_tensor())
      flat_probs_shape = [self._num_steps, batch_size, self._num_states]
      flat_means_shape = [
          batch_size, self._num_states,
          tf.reduce_prod(input_tensor=observation_event_shape)
      ]

      flat_probs = tf.reshape(probs, flat_probs_shape)
      # flat_probs :: num_steps batch_size num_states
      flat_means = tf.reshape(means, flat_means_shape)
      # flat_means :: batch_size num_states observation_event_size
      flat_mean = tf.einsum("ijk,jkl->jil", flat_probs, flat_means)
      # flat_mean :: batch_size num_steps observation_event_size
      unflat_mean_shape = tf.concat(
          [self.batch_shape_tensor(),
           [self._num_steps],
           observation_event_shape],
          axis=0)
      # returns :: batch_shape num_steps observation_event_shape
      return tf.reshape(flat_mean, unflat_mean_shape)

  def _variance(self):
    with tf.control_dependencies(self._runtime_assertions):
      probs = self._marginal_hidden_probs()
      # probs :: num_steps batch_shape num_states
      means = self._observation_distribution.mean()
      # means :: observation_batch_shape[:-1] num_states
      #          observation_event_shape
      means_shape = tf.concat(
          [self.batch_shape_tensor(),
           [self._num_states],
           self._observation_distribution.event_shape_tensor()],
          axis=0)
      means = tf.broadcast_to(means, means_shape)
      # means :: batch_shape num_states observation_event_shape

      observation_event_shape = (
          self._observation_distribution.event_shape_tensor())
      batch_size = tf.reduce_prod(input_tensor=self.batch_shape_tensor())
      flat_probs_shape = [self._num_steps, batch_size, self._num_states]
      flat_means_shape = [
          batch_size, 1, self._num_states,
          tf.reduce_prod(input_tensor=observation_event_shape)
      ]

      flat_probs = tf.reshape(probs, flat_probs_shape)
      # flat_probs :: num_steps batch_size num_states
      flat_means = tf.reshape(means, flat_means_shape)
      # flat_means :: batch_size 1 num_states observation_event_size
      flat_mean = tf.einsum("ijk,jmkl->jiml", flat_probs, flat_means)
      # flat_mean :: batch_size num_steps 1 observation_event_size

      variances = self._observation_distribution.variance()
      variances = tf.broadcast_to(variances, means_shape)
      # variances :: batch_shape num_states observation_event_shape
      flat_variances = tf.reshape(variances, flat_means_shape)
      # flat_variances :: batch_size 1 num_states observation_event_size

      # For a mixture of n distributions with mixture probabilities
      # p[i], and where the individual distributions have means and
      # variances given by mean[i] and var[i], the variance of
      # the mixture is given by:
      #
      # var = sum i=1..n p[i] * ((mean[i] - mean)**2 + var[i]**2)

      flat_variance = tf.einsum("ijk,jikl->jil",
                                flat_probs,
                                (flat_means - flat_mean)**2 + flat_variances)
      # flat_variance :: batch_size num_steps observation_event_size

      unflat_mean_shape = tf.concat(
          [self.batch_shape_tensor(),
           [self._num_steps],
           observation_event_shape],
          axis=0)

      # returns :: batch_shape num_steps observation_event_shape
      return tf.reshape(flat_variance, unflat_mean_shape)

  def _observation_shape_preconditions(self, observation_tensor_shape):
    return tf.control_dependencies([assert_util.assert_equal(
        observation_tensor_shape[-1 - self._underlying_event_rank],
        self._num_steps,
        message="The tensor `observations` must consist of sequences"
                "of observations from `HiddenMarkovModel` of length"
                "`num_steps`.")])

  def posterior_marginals(self, observations, name=None):
    """Compute marginal posterior distribution for each state.

    This function computes, for each time step, the marginal
    conditional probability that the hidden Markov model was in
    each possible state given the observations that were made
    at each time step.
    So if the hidden states are `z[0],...,z[num_steps - 1]` and
    the observations are `x[0], ..., x[num_steps - 1]`, then
    this function computes `P(z[i] | x[0], ..., x[num_steps - 1])`
    for all `i` from `0` to `num_steps - 1`.

    This operation is sometimes called smoothing. It uses a form
    of the forward-backward algorithm.

    Note: the behavior of this function is undefined if the
    `observations` argument represents impossible observations
    from the model.

    Args:
      observations: A tensor representing a batch of observations
        made on the hidden Markov model.  The rightmost dimension of this tensor
        gives the steps in a sequence of observations from a single sample from
        the hidden Markov model. The size of this dimension should match the
        `num_steps` parameter of the hidden Markov model object. The other
        dimensions are the dimensions of the batch and these are broadcast with
        the hidden Markov model's parameters.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: "HiddenMarkovModel".

    Returns:
      posterior_marginal: A `Categorical` distribution object representing the
        marginal probability of the hidden Markov model being in each state at
        each step. The rightmost dimension of the `Categorical` distributions
        batch will equal the `num_steps` parameter providing one marginal
        distribution for each step. The other dimensions are the dimensions
        corresponding to the batch of observations.

    Raises:
      ValueError: if rightmost dimension of `observations` does not
      have size `num_steps`.
    """

    with tf.name_scope(name or "posterior_marginals"):
      with tf.control_dependencies(self._runtime_assertions):
        observation_tensor_shape = tf.shape(input=observations)

        with self._observation_shape_preconditions(observation_tensor_shape):
          observation_batch_shape = observation_tensor_shape[
              :-1 - self._underlying_event_rank]
          observation_event_shape = observation_tensor_shape[
              -1 - self._underlying_event_rank:]

          batch_shape = tf.broadcast_dynamic_shape(observation_batch_shape,
                                                   self.batch_shape_tensor())
          log_init = tf.broadcast_to(self._log_init,
                                     tf.concat([batch_shape,
                                                [self._num_states]],
                                               axis=0))
          log_transition = self._log_trans

          observations = tf.broadcast_to(observations,
                                         tf.concat([batch_shape,
                                                    observation_event_shape],
                                                   axis=0))
          observation_rank = tf.rank(observations)
          underlying_event_rank = self._underlying_event_rank
          observations = distribution_util.move_dimension(
              observations, observation_rank - underlying_event_rank - 1, 0)
          observations = tf.expand_dims(
              observations,
              observation_rank - underlying_event_rank)
          observation_log_probs = self._observation_distribution.log_prob(
              observations)

          log_adjoint_prob = tf.zeros_like(log_init)

          def forward_step(log_previous_step, log_prob_observation):
            return _log_vector_matrix(log_previous_step,
                                      log_transition) + log_prob_observation

          log_prob = log_init + observation_log_probs[0]

          forward_log_probs = tf.scan(forward_step, observation_log_probs[1:],
                                      initializer=log_prob,
                                      name="forward_log_probs")

          forward_log_probs = tf.concat([[log_prob], forward_log_probs], axis=0)

          def backward_step(log_previous_step, log_prob_observation):
            return _log_matrix_vector(log_transition,
                                      log_prob_observation + log_previous_step)

          backward_log_adjoint_probs = tf.scan(
              backward_step,
              observation_log_probs[1:],
              initializer=log_adjoint_prob,
              reverse=True,
              name="backward_log_adjoint_probs")

          total_log_prob = tf.reduce_logsumexp(
              input_tensor=forward_log_probs[-1], axis=-1)

          backward_log_adjoint_probs = tf.concat([backward_log_adjoint_probs,
                                                  [log_adjoint_prob]], axis=0)

          log_likelihoods = forward_log_probs + backward_log_adjoint_probs

          marginal_log_probs = distribution_util.move_dimension(
              log_likelihoods - total_log_prob[..., tf.newaxis], 0, -2)

          return categorical.Categorical(logits=marginal_log_probs)

  def posterior_mode(self, observations, name=None):
    """Compute maximum likelihood sequence of hidden states.

    When this function is provided with a sequence of observations
    `x[0], ..., x[num_steps - 1]`, it returns the sequence of hidden
    states `z[0], ..., z[num_steps - 1]`, drawn from the underlying
    Markov chain, that is most likely to yield those observations.

    It uses the [Viterbi algorithm](
    https://en.wikipedia.org/wiki/Viterbi_algorithm).

    Note: the behavior of this function is undefined if the
    `observations` argument represents impossible observations
    from the model.

    Note: if there isn't a unique most likely sequence then one
    of the equally most likely sequences is chosen.

    Args:
      observations: A tensor representing a batch of observations made on the
        hidden Markov model.  The rightmost dimensions of this tensor correspond
        to the dimensions of the observation distributions of the underlying
        Markov chain.  The next dimension from the right indexes the steps in a
        sequence of observations from a single sample from the hidden Markov
        model.  The size of this dimension should match the `num_steps`
        parameter of the hidden Markov model object.  The other dimensions are
        the dimensions of the batch and these are broadcast with the hidden
        Markov model's parameters.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: "HiddenMarkovModel".

    Returns:
      posterior_mode: A `Tensor` representing the most likely sequence of hidden
        states. The rightmost dimension of this tensor will equal the
        `num_steps` parameter providing one hidden state for each step. The
        other dimensions are those of the batch.

    Raises:
      ValueError: if the `observations` tensor does not consist of
      sequences of `num_steps` observations.

    #### Examples

    ```python
    tfd = tfp.distributions

    # A simple weather model.

    # Represent a cold day with 0 and a hot day with 1.
    # Suppose the first day of a sequence has a 0.8 chance of being cold.

    initial_distribution = tfd.Categorical(probs=[0.8, 0.2])

    # Suppose a cold day has a 30% chance of being followed by a hot day
    # and a hot day has a 20% chance of being followed by a cold day.

    transition_distribution = tfd.Categorical(probs=[[0.7, 0.3],
                                                     [0.2, 0.8]])

    # Suppose additionally that on each day the temperature is
    # normally distributed with mean and standard deviation 0 and 5 on
    # a cold day and mean and standard deviation 15 and 10 on a hot day.

    observation_distribution = tfd.Normal(loc=[0., 15.], scale=[5., 10.])

    # This gives the hidden Markov model:

    model = tfd.HiddenMarkovModel(
        initial_distribution=initial_distribution,
        transition_distribution=transition_distribution,
        observation_distribution=observation_distribution,
        num_steps=7)

    # Suppose we observe gradually rising temperatures over a week:
    temps = [-2., 0., 2., 4., 6., 8., 10.]

    # We can now compute the most probable sequence of hidden states:

    model.posterior_mode(temps)

    # The result is [0 0 0 0 0 1 1] telling us that the transition
    # from "cold" to "hot" most likely happened between the
    # 5th and 6th days.
    ```
    """

    with tf.name_scope(name or "posterior_mode"):
      with tf.control_dependencies(self._runtime_assertions):
        observation_tensor_shape = tf.shape(input=observations)

        with self._observation_shape_preconditions(observation_tensor_shape):
          observation_batch_shape = observation_tensor_shape[
              :-1 - self._underlying_event_rank]
          observation_event_shape = observation_tensor_shape[
              -1 - self._underlying_event_rank:]

          batch_shape = tf.broadcast_dynamic_shape(observation_batch_shape,
                                                   self.batch_shape_tensor())
          log_init = tf.broadcast_to(self._log_init,
                                     tf.concat([batch_shape,
                                                [self._num_states]],
                                               axis=0))

          observations = tf.broadcast_to(observations,
                                         tf.concat([batch_shape,
                                                    observation_event_shape],
                                                   axis=0))
          observation_rank = tf.rank(observations)
          underlying_event_rank = self._underlying_event_rank
          observations = distribution_util.move_dimension(
              observations, observation_rank - underlying_event_rank - 1, 0)

          # We need to compute the probability of each observation for
          # each possible state.
          # This requires inserting an extra index just before the
          # observation event indices that will be broadcast with the
          # last batch index in `observation_distribution`.
          observations = tf.expand_dims(
              observations,
              observation_rank - underlying_event_rank)
          observation_log_probs = self._observation_distribution.log_prob(
              observations)

          log_prob = log_init + observation_log_probs[0]

          if self._num_steps == 1:
            most_likely_end = tf.argmax(input=log_prob, axis=-1)
            return most_likely_end[..., tf.newaxis]

          def forward_step(previous_step_pair, log_prob_observation):
            log_prob_previous = previous_step_pair[0]
            log_prob = (log_prob_previous[..., tf.newaxis] +
                        self._log_trans +
                        log_prob_observation[..., tf.newaxis, :])
            most_likely_given_successor = tf.argmax(input=log_prob, axis=-2)
            max_log_p_given_successor = tf.reduce_max(input_tensor=log_prob,
                                                      axis=-2)
            return (max_log_p_given_successor, most_likely_given_successor)

          forward_log_probs, all_most_likely_given_successor = tf.scan(
              forward_step,
              observation_log_probs[1:],
              initializer=(log_prob,
                           tf.zeros(tf.shape(input=log_init), dtype=tf.int64)),
              name="forward_log_probs")

          most_likely_end = tf.argmax(input=forward_log_probs[-1], axis=-1)

          # We require the operation that gives C from A and B where
          # C[i...j] = A[i...j, B[i...j]]
          # and A = most_likely_given_successor
          #     B = most_likely_successor.
          # tf.gather requires indices of known shape so instead we use
          # reduction with tf.one_hot(B) to pick out elements from B
          def backward_step(most_likely_successor, most_likely_given_successor):
            return tf.reduce_sum(
                input_tensor=(most_likely_given_successor *
                              tf.one_hot(most_likely_successor,
                                         self._num_states,
                                         dtype=tf.int64)),
                axis=-1)

          backward_scan = tf.scan(
              backward_step,
              all_most_likely_given_successor,
              most_likely_end,
              reverse=True)
          most_likely_sequences = tf.concat([backward_scan, [most_likely_end]],
                                            axis=0)
          return distribution_util.move_dimension(most_likely_sequences, 0, -1)

def _log_vector_matrix(vs, ms):
  """Multiply tensor of vectors by matrices assuming values stored are logs."""

  return tf.reduce_logsumexp(input_tensor=vs[..., tf.newaxis] + ms, axis=-2)

def _log_matrix_vector(ms, vs):
  """Multiply tensor of matrices by vectors assuming values stored are logs."""

  return tf.reduce_logsumexp(input_tensor=ms + vs[..., tf.newaxis, :], axis=-1)

def _vector_matrix(vs, ms):
  """Multiply tensor of vectors by matrices."""

  return tf.reduce_sum(input_tensor=vs[..., tf.newaxis] * ms, axis=-2)

def _extract_log_probs(num_states, dist):
  """Tabulate log probabilities from a batch of distributions."""

  states = tf.reshape(tf.range(num_states),
                      tf.concat([[num_states],
                                 tf.ones_like(dist.batch_shape_tensor())],
                                axis=0))
  return distribution_util.move_dimension(dist.log_prob(states), 0, -1)


# Copyright 2018 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Slice sampler transition kernel."""

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 import distributions
from tensorflow_probability.python.internal import distribution_util
from tensorflow_probability.python.mcmc import kernel as kernel_base

from tensorflow_probability.python.mcmc.internal import slice_sampler_utils as ssu
from tensorflow_probability.python.mcmc.internal import util as mcmc_util

__all__ = [
    'SliceSampler',
]

SliceSamplerKernelResults = collections.namedtuple(
    'SliceSamplerKernelResults',
    [
        'target_log_prob',  # For "next_state".
        'bounds_satisfied',  # Were the slice bounds chosen outside the slice.
        'direction',  # The direction in which the slice was sampled.
        'upper_bounds',  # Upper bound of the slice in the sampling direction.
        'lower_bounds',  # Lower bound of the slice in the sampling direction.
    ])

class SliceSampler(kernel_base.TransitionKernel):
  """Runs one step of the slice sampler using a hit and run approach.

  Slice Sampling is a Markov Chain Monte Carlo (MCMC) algorithm based, as stated
  by [Neal (2003)][1], on the observation that "...one can sample from a
  distribution by sampling uniformly from the region under the plot of its
  density function. A Markov chain that converges to this uniform distribution
  can be constructed by alternately uniform sampling in the vertical direction
  with uniform sampling from the horizontal `slice` defined by the current
  vertical position, or more generally, with some update that leaves the uniform
  distribution over this slice invariant". Mathematical details and derivations
  can be found in [Neal (2003)][1]. The one dimensional slice sampler is
  extended to n-dimensions through use of a hit-and-run approach: choose a
  random direction in n-dimensional space and take a step, as determined by the
  one-dimensional slice sampling algorithm, along that direction
  [Belisle at al. 1993][2].

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

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

  ### Examples:

  #### Simple chain with warm-up.

  In this example we sample from a standard univariate normal
  distribution using slice sampling.

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

    tfd = tfp.distributions

    dtype = np.float32

    target = tfd.Normal(loc=dtype(0), scale=dtype(1))

    samples, _ = tfp.mcmc.sample_chain(
        num_results=1000,
        current_state=dtype(1),
        kernel=tfp.mcmc.SliceSampler(
            target.log_prob,
            step_size=1.0,
            max_doublings=5,
            seed=1234),
        num_burnin_steps=500,
        parallel_iterations=1)  # For determinism.

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

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

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

  #### Sample from a Two Dimensional Normal.

  In the following example we sample from a two dimensional Normal
  distribution using slice sampling.

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

    tfd = tfp.distributions

    dtype = np.float32
    true_mean = dtype([0, 0])
    true_cov = dtype([[1, 0.5], [0.5, 1]])
    num_results = 500
    num_chains = 50

    # Target distribution is defined through the Cholesky decomposition
    chol = tf.linalg.cholesky(true_cov)
    target = tfd.MultivariateNormalTriL(loc=true_mean, scale_tril=chol)

    # Assume that the state is passed as a list of 1-d tensors `x` and `y`.
    # Then the target log-density is defined as follows:
    def target_log_prob(x, y):
      # Stack the input tensors together
      z = tf.stack([x, y], axis=-1) - true_mean
      return target.log_prob(z)

    # Initial state of the chain
    init_state = [np.ones([num_chains, 1], dtype=dtype),
                  np.ones([num_chains, 1], dtype=dtype)]

    # Run Slice Samper for `num_results` iterations for `num_chains`
    # independent chains:
    [x, y], _ = tfp.mcmc.sample_chain(
        num_results=num_results,
        current_state=init_state,
        kernel=tfp.mcmc.SliceSampler(
            target_log_prob_fn=target_log_prob,
            step_size=1.0,
            max_doublings=5,
            seed=47),
        num_burnin_steps=200,
        num_steps_between_results=1,
        parallel_iterations=1)

    states = tf.stack([x, y], axis=-1)
    sample_mean = tf.reduce_mean(states, axis=[0, 1])
    z = states - sample_mean
    sample_cov = tf.reduce_mean(tf.matmul(z, z, transpose_a=True),
                                axis=[0, 1])

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

    print "sample_mean: ", sample_mean
    print "sample_cov: ", sample_cov
  ```

  ### References

  [1]: Radford M. Neal. Slice Sampling. The Annals of Statistics. 2003, Vol 31,
       No. 3 , 705-767.
       https://projecteuclid.org/download/pdf_1/euclid.aos/1056562461

  [2]: C.J.P. Belisle, H.E. Romeijn, R.L. Smith. Hit-and-run algorithms for
       generating multivariate distributions. Math. Oper. Res., 18(1993),
       225-266.
       https://www.jstor.org/stable/3690278?seq=1#page_scan_tab_contents
  """

  def __init__(self,
               target_log_prob_fn,
               step_size,
               max_doublings,
               seed=None,
               name=None):
    """Initializes this transition kernel.

    Args:
      target_log_prob_fn: Python callable which takes an argument like
        `current_state` (or `*current_state` if it is a list) and returns its
        (possibly unnormalized) log-density under the target distribution.
      step_size: Scalar or `tf.Tensor` with same dtype as and shape compatible
        with `x_initial`. The size of the initial interval.
      max_doublings: Scalar positive int32 `tf.Tensor`. The maximum number of
      doublings to consider.
      seed: Python integer to seed the random number generator.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., 'slice_sampler_kernel').

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

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

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

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

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

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

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

  @property
  def is_calibrated(self):
    return True

  def one_step(self, current_state, previous_kernel_results):
    """Runs one iteration of Slice Sampler.

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

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

    Raises:
      ValueError: if there isn't one `step_size` or a list with same length as
        `current_state`.
      TypeError: if `not target_log_prob.dtype.is_floating`.
    """
    with tf.compat.v1.name_scope(
        name=mcmc_util.make_name(self.name, 'slice', 'one_step'),
        values=[
            self.step_size, self.max_doublings, self._seed_stream,
            current_state, previous_kernel_results.target_log_prob
        ]):
      with tf.compat.v1.name_scope('initialize'):
        [
            current_state_parts,
            step_sizes,
            current_target_log_prob
        ] = _prepare_args(
            self.target_log_prob_fn,
            current_state,
            self.step_size,
            previous_kernel_results.target_log_prob,
            maybe_expand=True)

        max_doublings = tf.convert_to_tensor(
            value=self.max_doublings, dtype=tf.int32, name='max_doublings')

      independent_chain_ndims = distribution_util.prefer_static_rank(
          current_target_log_prob)

      [
          next_state_parts,
          next_target_log_prob,
          bounds_satisfied,
          direction,
          upper_bounds,
          lower_bounds
      ] = _sample_next(
          self.target_log_prob_fn,
          current_state_parts,
          step_sizes,
          max_doublings,
          current_target_log_prob,
          independent_chain_ndims,
          seed=self._seed_stream()
      )

      def maybe_flatten(x):
        return x if mcmc_util.is_list_like(current_state) else x[0]

      return [
          maybe_flatten(next_state_parts),
          SliceSamplerKernelResults(
              target_log_prob=next_target_log_prob,
              bounds_satisfied=bounds_satisfied,
              direction=direction,
              upper_bounds=upper_bounds,
              lower_bounds=lower_bounds
          ),
      ]

  def bootstrap_results(self, init_state):
    with tf.compat.v1.name_scope(
        name=mcmc_util.make_name(self.name, 'slice', 'bootstrap_results'),
        values=[init_state]):
      if not mcmc_util.is_list_like(init_state):
        init_state = [init_state]
      init_state = [tf.convert_to_tensor(value=x) for x in init_state]
      direction = [tf.zeros_like(x) for x in init_state]
      init_target_log_prob = self.target_log_prob_fn(*init_state)  # pylint:disable=not-callable
      return SliceSamplerKernelResults(
          target_log_prob=init_target_log_prob,
          bounds_satisfied=tf.zeros(
              shape=tf.shape(input=init_target_log_prob), dtype=tf.bool),
          direction=direction,
          upper_bounds=tf.zeros_like(init_target_log_prob),
          lower_bounds=tf.zeros_like(init_target_log_prob))

def _choose_random_direction(current_state_parts, batch_rank, seed=None):
  """Chooses a random direction in the event space."""
  seed_gen = distributions.SeedStream(seed, salt='_choose_random_direction')
  # Chooses the random directions across each of the input components.
  rnd_direction_parts = [
      tf.random.normal(
          tf.shape(input=current_state_part), dtype=tf.float32, seed=seed_gen())
      for current_state_part in current_state_parts
  ]

  # Sum squares over all of the input components. Note this takes all
  # components into account.
  sum_squares = sum(
      tf.reduce_sum(
          input_tensor=rnd_direction**2.,
          axis=tf.range(batch_rank, tf.rank(rnd_direction)),
          keepdims=True) for rnd_direction in rnd_direction_parts)

  # Normalizes the random direction fragments.
  rnd_direction_parts = [rnd_direction / tf.sqrt(sum_squares)
                         for rnd_direction in rnd_direction_parts]

  return rnd_direction_parts

def _sample_next(target_log_prob_fn,
                 current_state_parts,
                 step_sizes,
                 max_doublings,
                 current_target_log_prob,
                 batch_rank,
                 seed=None,
                 name=None):
  """Applies a single iteration of slice sampling update.

  Applies hit and run style slice sampling. Chooses a uniform random direction
  on the unit sphere in the event space. Applies the one dimensional slice
  sampling update along that direction.

  Args:
    target_log_prob_fn: Python callable which takes an argument like
      `*current_state_parts` and returns its (possibly unnormalized) log-density
      under the target distribution.
    current_state_parts: Python `list` of `Tensor`s representing the current
      state(s) of the Markov chain(s). The first `independent_chain_ndims` of
      the `Tensor`(s) index different chains.
    step_sizes: Python `list` of `Tensor`s. Provides a measure of the width
      of the density. Used to find the slice bounds. Must broadcast with the
      shape of `current_state_parts`.
    max_doublings: Integer number of doublings to allow while locating the slice
      boundaries.
    current_target_log_prob: `Tensor` representing the value of
      `target_log_prob_fn(*current_state_parts)`. The only reason to specify
      this argument is to reduce TF graph size.
    batch_rank: Integer. The number of axes in the state that correspond to
      independent batches.
    seed: Python integer to seed random number generators.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'find_slice_bounds').

  Returns:
    proposed_state_parts: Tensor or Python list of `Tensor`s representing the
      state(s) of the Markov chain(s) at each result step. Has same shape as
      input `current_state_parts`.
    proposed_target_log_prob: `Tensor` representing the value of
      `target_log_prob_fn` at `next_state`.
    bounds_satisfied: Boolean `Tensor` of the same shape as the log density.
      True indicates whether the an interval containing the slice for that
      batch was found successfully.
    direction: `Tensor` or Python list of `Tensors`s representing the direction
      along which the slice was sampled. Has the same shape and dtype(s) as
      `current_state_parts`.
    upper_bounds: `Tensor` of batch shape and the dtype of the input state. The
      upper bounds of the slices along the sampling direction.
    lower_bounds: `Tensor` of batch shape and the dtype of the input state. The
      lower bounds of the slices along the sampling direction.
  """
  with tf.compat.v1.name_scope(name, 'sample_next', [
      current_state_parts, step_sizes, max_doublings, current_target_log_prob,
      batch_rank
  ]):
    # First step: Choose a random direction.
    # Direction is a list of tensors. The i'th tensor should have the same shape
    # as the i'th state part.
    direction = _choose_random_direction(current_state_parts,
                                         batch_rank=batch_rank,
                                         seed=seed)

    # Interpolates the step sizes for the chosen direction.
    # Applies an ellipsoidal interpolation to compute the step direction for
    # the chosen direction. Suppose we are given step sizes for each direction.
    # Label these s_1, s_2, ... s_k. These are the step sizes to use if moving
    # in a direction parallel to one of the axes. Consider an ellipsoid which
    # intercepts the i'th axis at s_i. The step size for a direction specified
    # by the unit vector (n_1, n_2 ...n_k) is then defined as the intersection
    # of the line through this vector with this ellipsoid.
    #
    # One can show that the length of the vector from the origin to the
    # intersection point is given by:
    # 1 / sqrt(n_1^2 / s_1^2  + n_2^2 / s_2^2  + ...).
    #
    # Proof:
    # The equation of the ellipsoid is:
    # Sum_i [x_i^2 / s_i^2 ] = 1. Let n be a unit direction vector. Points
    # along the line given by n may be parameterized as alpha*n where alpha is
    # the distance along the vector. Plugging this into the equation for the
    # ellipsoid, we get:
    # alpha^2 ( n_1^2 / s_1^2 + n_2^2 / s_2^2 + ...) = 1
    # so alpha = \sqrt { \frac{1} { ( n_1^2 / s_1^2 + n_2^2 / s_2^2 + ...) } }
    reduce_axes = [tf.range(batch_rank, tf.rank(dirn_part))
                   for dirn_part in direction]

    components = [
        tf.reduce_sum(
            input_tensor=(dirn_part / step_size)**2, axis=reduce_axes[i])
        for i, (step_size, dirn_part) in enumerate(zip(step_sizes, direction))
    ]
    step_size = tf.math.rsqrt(tf.add_n(components))
    # Computes the rank of a tensor. Uses the static rank if possible.
    def _get_rank(x):
      return (len(x.shape.as_list()) if x.shape.dims is not None
              else tf.rank(x))
    state_part_ranks = [_get_rank(part) for part in current_state_parts]

    def _step_along_direction(alpha):
      """Converts the scalar alpha into an n-dim vector with full state info.

      Computes x_0 + alpha * direction where x_0 is the current state and
      direction is the direction chosen above.

      Args:
        alpha: A tensor of shape equal to the batch dimensions of
          `current_state_parts`.

      Returns:
        state_parts: Tensor or Python list of `Tensor`s representing the
          state(s) of the Markov chain(s) for a given alpha and a given chosen
          direction. Has the same shape as `current_state_parts`.
      """
      padded_alphas = [_right_pad(alpha, final_rank=part_rank)
                       for part_rank in state_part_ranks]

      state_parts = [state_part + padded_alpha
                     * direction_part for state_part, direction_part,
                     padded_alpha in
                     zip(current_state_parts, direction, padded_alphas)]
      return state_parts

    def projected_target_log_prob_fn(alpha):
      """The target log density projected along the chosen direction.

      Args:
        alpha: A tensor of shape equal to the batch dimensions of
          `current_state_parts`.

      Returns:
        Target log density evaluated at x_0 + alpha * direction where x_0 is the
        current state and direction is the direction chosen above. Has the same
        shape as `alpha`.
      """
      return target_log_prob_fn(*_step_along_direction(alpha))

    alpha_init = tf.zeros_like(current_target_log_prob,
                               dtype=current_state_parts[0].dtype.base_dtype)
    [
        next_alpha,
        next_target_log_prob,
        bounds_satisfied,
        upper_bounds,
        lower_bounds
    ] = ssu.slice_sampler_one_dim(projected_target_log_prob_fn,
                                  x_initial=alpha_init,
                                  max_doublings=max_doublings,
                                  step_size=step_size, seed=seed)
    return [
        _step_along_direction(next_alpha),
        next_target_log_prob,
        bounds_satisfied,
        direction,
        upper_bounds,
        lower_bounds
    ]

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

def _right_pad(x, final_rank):
  """Pads the shape of x to the right to be of rank final_rank.

  Expands the dims of `x` to the right such that its rank is equal to
  final_rank. For example, if `x` is of shape [1, 5, 7, 2] and `final_rank` is
  7, we return padded_x, which is of shape [1, 5, 7, 2, 1, 1, 1].

  Args:
    x: The tensor whose shape is to be padded.
    final_rank: Scalar int32 `Tensor` or Python `int`. The desired rank of x.

  Returns:
    padded_x: A tensor of rank final_rank.
  """
  padded_shape = tf.concat(
      [tf.shape(input=x),
       tf.ones(final_rank - tf.rank(x), dtype=tf.int32)],
      axis=0)
  static_padded_shape = None
  if x.shape.is_fully_defined() and isinstance(final_rank, int):
    static_padded_shape = x.shape.as_list()
    extra_dims = final_rank - len(static_padded_shape)
    static_padded_shape.extend([1] * extra_dims)

  padded_x = tf.reshape(x, static_padded_shape or padded_shape)
  return padded_x

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

  target_log_prob = _maybe_call_fn(
      target_log_prob_fn,
      state_parts,
      target_log_prob,
      description)
  step_sizes = (list(step_size) if mcmc_util.is_list_like(step_size)
                else [step_size])
  step_sizes = [
      tf.convert_to_tensor(
          value=s, name='step_size', dtype=target_log_prob.dtype)
      for s in step_sizes
  ]
  if len(step_sizes) == 1:
    step_sizes *= len(state_parts)
  if len(state_parts) != len(step_sizes):
    raise ValueError('There should be exactly one `step_size` or it should '
                     'have same length as `current_state`.')
  def maybe_flatten(x):
    return x if maybe_expand or mcmc_util.is_list_like(state) else x[0]
  return [
      maybe_flatten(state_parts),
      maybe_flatten(step_sizes),
      target_log_prob
  ]


# Copyright 2018 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Methods for fitting StructuralTimeSeries models to data."""
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 import distributions as tfd
from tensorflow_probability.python import mcmc
from tensorflow_probability.python.sts.internal import util as sts_util

def sample_uniform_initial_state(parameter,
                                 return_constrained=True,
                                 init_sample_shape=(),
                                 seed=None):
  """Initialize from a uniform [-2, 2] distribution in unconstrained space.

  Args:
    parameter: `sts.Parameter` named tuple instance.
    return_constrained: if `True`, re-applies the constraining bijector
      to return initializations in the original domain. Otherwise, returns
      initializations in the unconstrained space.
      Default value: `True`.
    init_sample_shape: `sample_shape` of the sampled initializations.
      Default value: `[]`.
    seed: Python integer to seed the random number generator.

  Returns:
    uniform_initializer: `Tensor` of shape `concat([init_sample_shape,
    parameter.prior.batch_shape, transformed_event_shape])`, where
    `transformed_event_shape` is `parameter.prior.event_shape`, if
    `return_constrained=True`, and otherwise it is
    `parameter.bijector.inverse_event_shape(parameteter.prior.event_shape)`.
  """
  unconstrained_prior_sample = parameter.bijector.inverse(
      parameter.prior.sample(init_sample_shape, seed=seed))
  uniform_initializer = 4 * tf.random.uniform(
      tf.shape(input=unconstrained_prior_sample),
      dtype=unconstrained_prior_sample.dtype,
      seed=seed) - 2
  if return_constrained:
    return parameter.bijector.forward(uniform_initializer)
  else:
    return uniform_initializer

def _build_trainable_posterior(param, initial_loc_fn):
  """Built a transformed-normal variational dist over a parameter's support."""
  loc = tf.compat.v1.get_variable(
      param.name + '_loc',
      initializer=lambda: initial_loc_fn(param),
      dtype=param.prior.dtype,
      use_resource=True)
  scale = tf.nn.softplus(
      tf.compat.v1.get_variable(
          param.name + '_scale',
          initializer=lambda: -4 * tf.ones_like(initial_loc_fn(param)),
          dtype=param.prior.dtype,
          use_resource=True))

  q = tfd.Normal(loc=loc, scale=scale)

  # Ensure the `event_shape` of the variational distribution matches the
  # parameter.
  if (param.prior.event_shape.ndims is None
      or param.prior.event_shape.ndims > 0):
    q = tfd.Independent(
        q, reinterpreted_batch_ndims=param.prior.event_shape.ndims)

  # Transform to constrained parameter space.
  return tfd.TransformedDistribution(q, param.bijector)

def build_factored_variational_loss(model,
                                    observed_time_series,
                                    init_batch_shape=(),
                                    seed=None,
                                    name=None):
  """Build a loss function for variational inference in STS models.

  Variational inference searches for the distribution within some family of
  approximate posteriors that minimizes a divergence between the approximate
  posterior `q(z)` and true posterior `p(z|observed_time_series)`. By converting
  inference to optimization, it's generally much faster than sampling-based
  inference algorithms such as HMC. The tradeoff is that the approximating
  family rarely contains the true posterior, so it may miss important aspects of
  posterior structure (in particular, dependence between variables) and should
  not be blindly trusted. Results may vary; it's generally wise to compare to
  HMC to evaluate whether inference quality is sufficient for your task at hand.

  This method constructs a loss function for variational inference using the
  Kullback-Liebler divergence `KL[q(z) || p(z|observed_time_series)]`, with an
  approximating family given by independent Normal distributions transformed to
  the appropriate parameter space for each parameter. Minimizing this loss (the
  negative ELBO) maximizes a lower bound on the log model evidence `-log
  p(observed_time_series)`. This is equivalent to the 'mean-field' method
  implemented in [1]. and is a standard approach. The resulting posterior
  approximations are unimodal; they will tend to underestimate posterior
  uncertainty when the true posterior contains multiple modes (the `KL[q||p]`
  divergence encourages choosing a single mode) or dependence between variables.

  Args:
    model: An instance of `StructuralTimeSeries` representing a
      time-series model. This represents a joint distribution over
      time-series and their parameters with batch shape `[b1, ..., bN]`.
    observed_time_series: `float` `Tensor` of shape
      `concat([sample_shape, model.batch_shape, [num_timesteps, 1]]) where
      `sample_shape` corresponds to i.i.d. observations, and the trailing `[1]`
      dimension may (optionally) be omitted if `num_timesteps > 1`. May
      optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
      a mask `Tensor` to specify timesteps with missing observations.
    init_batch_shape: Batch shape (Python `tuple`, `list`, or `int`) of initial
      states to optimize in parallel.
      Default value: `()`. (i.e., just run a single optimization).
    seed: Python integer to seed the random number generator.
    name: Python `str` name prefixed to ops created by this function.
      Default value: `None` (i.e., 'build_factored_variational_loss').

  Returns:
    variational_loss: `float` `Tensor` of shape
      `concat([init_batch_shape, model.batch_shape])`, encoding a stochastic
      estimate of an upper bound on the negative model evidence `-log p(y)`.
      Minimizing this loss performs variational inference; the gap between the
      variational bound and the true (generally unknown) model evidence
      corresponds to the divergence `KL[q||p]` between the approximate and true
      posterior.
    variational_distributions: `collections.OrderedDict` giving
      the approximate posterior for each model parameter. The keys are
      Python `str` parameter names in order, corresponding to
      `[param.name for param in model.parameters]`. The values are
      `tfd.Distribution` instances with batch shape
      `concat([init_batch_shape, model.batch_shape])`; these will typically be
      of the form `tfd.TransformedDistribution(tfd.Normal(...),
      bijector=param.bijector)`.

  #### Examples

  Assume we've built a structural time-series model:

  ```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)
  ```

  To run variational inference, we simply construct the loss and optimize
  it:

  ```python
    (variational_loss,
     variational_distributions) = tfp.sts.build_factored_variational_loss(
       model=model, observed_time_series=observed_time_series)

    train_op = tf.train.AdamOptimizer(0.1).minimize(variational_loss)
    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())

      for step in range(200):
        _, loss_ = sess.run((train_op, variational_loss))

        if step % 20 == 0:
          print("step {} loss {}".format(step, loss_))

      posterior_samples_ = sess.run({
        param_name: q.sample(50)
        for param_name, q in variational_distributions.items()})
  ```

  As a more complex example, we might try to avoid local optima by optimizing
  from multiple initializations in parallel, and selecting the result with the
  lowest loss:

  ```python
    (variational_loss,
     variational_distributions) = tfp.sts.build_factored_variational_loss(
       model=model, observed_time_series=observed_time_series,
       init_batch_shape=[10])

    train_op = tf.train.AdamOptimizer(0.1).minimize(variational_loss)
    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())

      for step in range(200):
        _, loss_ = sess.run((train_op, variational_loss))

        if step % 20 == 0:
          print("step {} losses {}".format(step, loss_))

      # Draw multiple samples to reduce Monte Carlo error in the optimized
      # variational bounds.
      avg_loss = np.mean(
        [sess.run(variational_loss) for _ in range(25)], axis=0)
      best_posterior_idx = np.argmin(avg_loss, axis=0).astype(np.int32)
  ```

  #### References

  [1]: Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, and
       David M. Blei. Automatic Differentiation Variational Inference. In
       _Journal of Machine Learning Research_, 2017.
       https://arxiv.org/abs/1603.00788

  """

  with tf.compat.v1.name_scope(
      name, 'build_factored_variational_loss',
      values=[observed_time_series]) as name:
    seed = tfd.SeedStream(
        seed, salt='StructuralTimeSeries_build_factored_variational_loss')

    variational_distributions = collections.OrderedDict()
    variational_samples = []
    for param in model.parameters:
      def initial_loc_fn(param):
        return sample_uniform_initial_state(
            param, return_constrained=True,
            init_sample_shape=init_batch_shape,
            seed=seed())
      q = _build_trainable_posterior(param, initial_loc_fn=initial_loc_fn)
      variational_distributions[param.name] = q
      variational_samples.append(q.sample(seed=seed()))

    # Multiple initializations (similar to HMC chains) manifest as an extra
    # param batch dimension, so we need to add corresponding batch dimension(s)
    # to `observed_time_series`.
    observed_time_series = sts_util.pad_batch_dimension_for_multiple_chains(
        observed_time_series, model, chain_batch_shape=init_batch_shape)

    # Construct the variational bound.
    log_prob_fn = model.joint_log_prob(observed_time_series)
    expected_log_joint = log_prob_fn(*variational_samples)
    entropy = tf.reduce_sum(
        input_tensor=[
            -q.log_prob(sample) for (q, sample) in zip(
                variational_distributions.values(), variational_samples)
        ],
        axis=0)
    variational_loss = -(expected_log_joint + entropy)  # -ELBO

  return variational_loss, variational_distributions

def _minimize_in_graph(build_loss_fn, num_steps=200, optimizer=None):
  """Run an optimizer within the graph to minimize a loss function."""
  optimizer = tf.compat.v1.train.AdamOptimizer(
      0.1) if optimizer is None else optimizer

  def train_loop_body(step):
    train_op = optimizer.minimize(
        build_loss_fn if tf.executing_eagerly() else build_loss_fn())
    return tf.tuple(tensors=[tf.add(step, 1)], control_inputs=[train_op])

  minimize_op = tf.compat.v1.while_loop(
      cond=lambda step: step < num_steps,
      body=train_loop_body,
      loop_vars=[tf.constant(0)],
      return_same_structure=True)[0]  # Always return a single op.
  return minimize_op

def fit_with_hmc(model,
                 observed_time_series,
                 num_results=100,
                 num_warmup_steps=50,
                 num_leapfrog_steps=15,
                 initial_state=None,
                 initial_step_size=None,
                 chain_batch_shape=(),
                 num_variational_steps=150,
                 variational_optimizer=None,
                 seed=None,
                 name=None):
  """Draw posterior samples using Hamiltonian Monte Carlo (HMC).

  Markov chain Monte Carlo (MCMC) methods are considered the gold standard of
  Bayesian inference; under suitable conditions and in the limit of infinitely
  many draws they generate samples from the true posterior distribution. HMC [1]
  uses gradients of the model's log-density function to propose samples,
  allowing it to exploit posterior geometry. However, it is computationally more
  expensive than variational inference and relatively sensitive to tuning.

  This method attempts to provide a sensible default approach for fitting
  StructuralTimeSeries models using HMC. It first runs variational inference as
  a fast posterior approximation, and initializes the HMC sampler from the
  variational posterior, using the posterior standard deviations to set
  per-variable step sizes (equivalently, a diagonal mass matrix). During the
  warmup phase, it adapts the step size to target an acceptance rate of 0.75,
  which is thought to be in the desirable range for optimal mixing [2].

  Args:
    model: An instance of `StructuralTimeSeries` representing a
      time-series model. This represents a joint distribution over
      time-series and their parameters with batch shape `[b1, ..., bN]`.
    observed_time_series: `float` `Tensor` of shape
      `concat([sample_shape, model.batch_shape, [num_timesteps, 1]]) where
      `sample_shape` corresponds to i.i.d. observations, and the trailing `[1]`
      dimension may (optionally) be omitted if `num_timesteps > 1`. May
      optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
      a mask `Tensor` to specify timesteps with missing observations.
    num_results: Integer number of Markov chain draws.
      Default value: `100`.
    num_warmup_steps: Integer number of steps to take before starting to
      collect results. The warmup steps are also used to adapt the step size
      towards a target acceptance rate of 0.75.
      Default value: `50`.
    num_leapfrog_steps: Integer number of steps to run the leapfrog integrator
      for. Total progress per HMC step is roughly proportional to
      `step_size * num_leapfrog_steps`.
      Default value: `15`.
    initial_state: Optional Python `list` of `Tensor`s, one for each model
      parameter, representing the initial state(s) of the Markov chain(s). These
      should have shape `concat([chain_batch_shape, param.prior.batch_shape,
      param.prior.event_shape])`. If `None`, the initial state is set
      automatically using a sample from a variational posterior.
      Default value: `None`.
    initial_step_size: Python `list` of `Tensor`s, one for each model parameter,
      representing the step size for the leapfrog integrator. Must
      broadcast with the shape of `initial_state`. Larger step sizes lead to
      faster progress, but too-large step sizes make rejection exponentially
      more likely. If `None`, the step size is set automatically using the
      standard deviation of a variational posterior.
      Default value: `None`.
    chain_batch_shape: Batch shape (Python `tuple`, `list`, or `int`) of chains
      to run in parallel.
      Default value: `[]` (i.e., a single chain).
    num_variational_steps: Python `int` number of steps to run the variational
      optimization to determine the initial state and step sizes.
      Default value: `150`.
    variational_optimizer: Optional `tf.train.Optimizer` instance to use in
      the variational optimization. If `None`, defaults to
      `tf.train.AdamOptimizer(0.1)`.
      Default value: `None`.
    seed: Python integer to seed the random number generator.
    name: Python `str` name prefixed to ops created by this function.
      Default value: `None` (i.e., 'fit_with_hmc').

  Returns:
    samples: Python `list` of `Tensors` representing posterior samples of model
      parameters, with shapes `[concat([[num_results], chain_batch_shape,
      param.prior.batch_shape, param.prior.event_shape]) for param in
      model.parameters]`.
    kernel_results: A (possibly nested) `tuple`, `namedtuple` or `list` of
      `Tensor`s representing internal calculations made within the HMC sampler.

  #### Examples

  Assume we've built a structural time-series model:

  ```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)
  ```

  To draw posterior samples using HMC under default settings:

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

  with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    samples_, kernel_results_ = sess.run((samples, kernel_results))

  print("acceptance rate: {}".format(
    np.mean(kernel_results_.inner_results.is_accepted, axis=0)))
  print("posterior means: {}".format(
    {param.name: np.mean(param_draws, axis=0)
     for (param, param_draws) in zip(model.parameters, samples_)}))
  ```

  We can also run multiple chains. This may help diagnose convergence issues
  and allows us to exploit vectorization to draw samples more quickly, although
  warmup still requires the same number of sequential steps.

  ```python
  from matplotlib import pylab as plt

  samples, kernel_results = tfp.sts.fit_with_hmc(
    model, observed_time_series, chain_batch_shape=[10])

  with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    samples_, kernel_results_ = sess.run((samples, kernel_results))

  print("acceptance rate: {}".format(
    np.mean(kernel_results_.inner_results.inner_results.is_accepted, axis=0)))

  # Plot the sampled traces for each parameter. If the chains have mixed, their
  # traces should all cover the same region of state space, frequently crossing
  # over each other.
  for (param, param_draws) in zip(model.parameters, samples_):
    if param.prior.event_shape.ndims > 0:
      print("Only plotting traces for scalar parameters, skipping {}".format(
        param.name))
      continue
    plt.figure(figsize=[10, 4])
    plt.title(param.name)
    plt.plot(param_draws)
    plt.ylabel(param.name)
    plt.xlabel("HMC step")

  # Combining the samples from multiple chains into a single dimension allows
  # us to easily pass sampled parameters to downstream forecasting methods.
  combined_samples_ = [np.reshape(param_draws,
                                  [-1] + list(param_draws.shape[2:]))
                       for param_draws in samples_]
  ```

  For greater flexibility, you may prefer to implement your own sampler using
  the TensorFlow Probability primitives in `tfp.mcmc`. The following recipe
  constructs a basic HMC sampler, using a `TransformedTransitionKernel` to
  incorporate constraints on the parameter space.

  ```python
  transformed_hmc_kernel = mcmc.TransformedTransitionKernel(
      inner_kernel=mcmc.SimpleStepSizeAdaptation(
          inner_kernel=mcmc.HamiltonianMonteCarlo(
              target_log_prob_fn=model.joint_log_prob(observed_time_series),
              step_size=step_size,
              num_leapfrog_steps=num_leapfrog_steps,
              state_gradients_are_stopped=True,
              seed=seed),
          num_adaptation_steps = int(0.8 * num_warmup_steps)),
      bijector=[param.bijector for param in model.parameters])

  # Initialize from a Uniform[-2, 2] distribution in unconstrained space.
  initial_state = [tfp.sts.sample_uniform_initial_state(
    param, return_constrained=True) for param in model.parameters]

  samples, kernel_results = tfp.mcmc.sample_chain(
    kernel=transformed_hmc_kernel,
    num_results=num_results,
    current_state=initial_state,
    num_burnin_steps=num_warmup_steps)
  ```

  #### References

  [1]: Radford Neal. MCMC Using Hamiltonian Dynamics. _Handbook of Markov Chain
       Monte Carlo_, 2011. https://arxiv.org/abs/1206.1901
  [2]  M.J. Betancourt, Simon Byrne, and Mark Girolami. Optimizing The
       Integrator Step Size for Hamiltonian Monte Carlo.
       https://arxiv.org/abs/1411.6669

  """
  with tf.compat.v1.name_scope(
      name, 'fit_with_hmc', values=[observed_time_series]) as name:
    seed = tfd.SeedStream(seed, salt='StructuralTimeSeries_fit_with_hmc')

    # Initialize state and step sizes from a variational posterior if not
    # specified.
    if initial_step_size is None or initial_state is None:

      # To avoid threading variational distributions through the training
      # while loop, we build our own copy here. `make_template` ensures
      # that our variational distributions share the optimized parameters.
      def make_variational():
        return build_factored_variational_loss(
            model, observed_time_series,
            init_batch_shape=chain_batch_shape, seed=seed())

      make_variational = tf.compat.v1.make_template('make_variational',
                                                    make_variational)
      _, variational_distributions = make_variational()
      minimize_op = _minimize_in_graph(
          build_loss_fn=lambda: make_variational()[0],  # return just the loss.
          num_steps=num_variational_steps,
          optimizer=variational_optimizer)

      with tf.control_dependencies([minimize_op]):
        if initial_state is None:
          initial_state = [tf.stop_gradient(d.sample())
                           for d in variational_distributions.values()]

        # Set step sizes using the unconstrained variational distribution.
        if initial_step_size is None:
          initial_step_size = [
              transformed_q.distribution.stddev()
              for transformed_q in variational_distributions.values()]

    # Multiple chains manifest as an extra param batch dimension, so we need to
    # add a corresponding batch dimension to `observed_time_series`.
    observed_time_series = sts_util.pad_batch_dimension_for_multiple_chains(
        observed_time_series, model, chain_batch_shape=chain_batch_shape)

    # Run HMC to sample from the posterior on parameters.
    samples, kernel_results = mcmc.sample_chain(
        num_results=num_results,
        current_state=initial_state,
        num_burnin_steps=num_warmup_steps,
        kernel=mcmc.SimpleStepSizeAdaptation(
            inner_kernel=mcmc.TransformedTransitionKernel(
                inner_kernel=mcmc.HamiltonianMonteCarlo(
                    target_log_prob_fn=model.joint_log_prob(
                        observed_time_series),
                    step_size=initial_step_size,
                    num_leapfrog_steps=num_leapfrog_steps,
                    state_gradients_are_stopped=True,
                    seed=seed()),
                bijector=[param.bijector for param in model.parameters]),
            num_adaptation_steps=int(num_warmup_steps * 0.8),
            adaptation_rate=tf.convert_to_tensor(
                value=0.1, dtype=initial_state[0].dtype)),
        parallel_iterations=1 if seed is not None else 10)

    return samples, kernel_results


# Copyright 2018 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Utilities for time series with missing values."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections

# Dependency imports
import tensorflow as tf

tfl = tf.linalg

class MaskedTimeSeries(collections.namedtuple('MaskedTimeSeries',
                                              ['time_series', 'is_missing'])):
  """Named tuple encoding a time series `Tensor` and optional missingness mask.

  Structural time series models handle missing values naturally, following the
  rules of conditional probability. Posterior inference can be used to impute
  missing values, with uncertainties. Forecasting and posterior decomposition
  are also supported for time series with missing values; the missing values
  will generally lead to corresponding higher forecast uncertainty.

  All methods in the `tfp.sts` API that accept an `observed_time_series`
  `Tensor` should optionally also accept a `MaskedTimeSeries` instance.

  The time series should be a float `Tensor` of shape `[..., num_timesteps]` or
  `[..., num_timesteps, 1]`. The `is_missing` mask must be either a boolean
  `Tensor` of shape `[..., num_timesteps]`, or `None`. `True` values in
  `is_missing` denote missing (masked) observations; `False` denotes observed
  (unmasked) values. Note that these semantics are opposite that of low-level
  TensorFlow methods like `tf.boolean_mask`, but consistent with the behavior
  of Numpy masked arrays.

  The batch dimensions of `is_missing` must broadcast with the batch
  dimensions of `time_series`.

  A `MaskedTimeSeries` is just a `collections.namedtuple` instance, i.e., a dumb
  container. Although the convention for the elements is as described here, it's
  left to downstream methods to validate or convert the elements as required.
  In particular, most downstream methods will call `tf.convert_to_tensor`
  on the components. In order to prevent duplicate `Tensor` creation, you may
  (if memory is an issue) wish to ensure that the components are *already*
  `Tensors`, as opposed to numpy arrays or similar.

  #### Examples

  To construct a simple MaskedTimeSeries instance:

  ```
  observed_time_series = tfp.sts.MaskedTimeSeries(
    time_series=tf.random_normal([3, 4, 5]),
    is_missing=[True, False, False, True, False])
  ```

  Note that the mask we specified will broadcast against the batch dimensions of
  the time series.

  For time series with missing entries specified as NaN 'magic values', you can
  generate a mask using `tf.is_nan`:

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

  time_series_with_nans = [-1., 1., np.nan, 2.4, np.nan, 5]
  observed_time_series = tfp.sts.MaskedTimeSeries(
    time_series=time_series_with_nans,
    is_missing=tf.is_nan(time_series_with_nans))

  # Build model using observed time series to set heuristic priors.
  linear_trend_model = tfp.sts.LocalLinearTrend(
    observed_time_series=observed_time_series)
  model = tfp.sts.Sum([linear_trend_model],
                      observed_time_series=observed_time_series)

  # Fit model to data
  parameter_samples, _ = tfp.sts.fit_with_hmc(model, observed_time_series)

  # Forecast
  forecast_dist = tfp.sts.forecast(
    model, observed_time_series, num_steps_forecast=5)
  ```

  """
  # This class is just a namedtuple, but we use a subclass to override the
  # docstring, following
  # https://stackoverflow.com/questions/1606436/adding-docstrings-to-namedtuples
  # To keep the child class from constructing instance dicts and losing the
  # lightweightness/immutability of a namedtuple, we need the following magic:
  __slots__ = ()

def moments_of_masked_time_series(time_series_tensor, broadcast_mask):
  """Compute mean and variance, accounting for a mask.

  Args:
    time_series_tensor: float `Tensor` time series of shape
      `concat([batch_shape, [num_timesteps]])`.
    broadcast_mask: bool `Tensor` of the same shape as `time_series`.
  Returns:
    mean: float `Tensor` of shape `batch_shape`.
    variance: float `Tensor` of shape `batch_shape`.
  """
  num_unmasked_entries = tf.cast(
      tf.reduce_sum(input_tensor=tf.cast(~broadcast_mask, tf.int32), axis=-1),
      time_series_tensor.dtype)

  # Manually compute mean and variance, excluding masked entries.
  mean = (tf.reduce_sum(input_tensor=tf.where(
      broadcast_mask,
      tf.zeros_like(time_series_tensor),
      time_series_tensor), axis=-1) / num_unmasked_entries)
  variance = (tf.reduce_sum(input_tensor=tf.where(
      broadcast_mask,
      tf.zeros_like(time_series_tensor),
      (time_series_tensor - mean[..., tf.newaxis]) ** 2), axis=-1)
              / num_unmasked_entries)
  return mean, variance

def initial_value_of_masked_time_series(time_series_tensor, broadcast_mask):
  """Get the first unmasked entry of each time series in the batch.

  Args:
    time_series_tensor: float `Tensor` of shape [..., num_timesteps].
    broadcast_mask: bool `Tensor` of same shape as `time_series`.
  """

  num_timesteps = tf.shape(input=time_series_tensor)[-1]

  # Compute the index of the first unmasked entry for each series in the batch.
  unmasked_negindices = (
      tf.cast(~broadcast_mask, tf.int32) *
      tf.range(num_timesteps, 0, -1))
  first_unmasked_indices = num_timesteps - tf.reduce_max(
      input_tensor=unmasked_negindices, axis=-1)

  if first_unmasked_indices.shape.ndims is None:
    raise NotImplementedError(
        'Cannot compute initial values of a masked time series with'
        'dynamic rank.')  # `batch_gather` requires static rank

  # Extract the initial value for each series in the batch.
  return tf.squeeze(tf.compat.v1.batch_gather(
      params=time_series_tensor,
      indices=first_unmasked_indices[..., tf.newaxis]), axis=-1)