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

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

import tensorflow.compat.v2 as tf

from tensorflow_probability.python.distributions import gaussian_process
from tensorflow_probability.python.distributions import independent
from tensorflow_probability.python.distributions import kullback_leibler
from tensorflow_probability.python.distributions import mvn_linear_operator
from tensorflow_probability.python.distributions import normal
from tensorflow_probability.python.internal import dtype_util

__all__ = [
    'VariationalGaussianProcess',
]

def _add_diagonal_shift(matrix, shift):
  return tf.linalg.set_diag(
      matrix, tf.linalg.diag_part(matrix) + shift, name='add_diagonal_shift')

def _solve_cholesky_factored_system(
    cholesky_factor, rhs, name=None):
  with tf.name_scope(
      name or '_solve_cholesky_factored_system') as scope:
    cholesky_factor = tf.convert_to_tensor(value=cholesky_factor,
                                           name='cholesky_factor')
    rhs = tf.convert_to_tensor(value=rhs, name='rhs')
    lin_op = tf.linalg.LinearOperatorLowerTriangular(
        cholesky_factor, name=scope)
    return lin_op.solve(lin_op.solve(rhs), adjoint=True)

def _solve_cholesky_factored_system_vec(cholesky_factor, rhs, name=None):
  with tf.name_scope(
      name or '_solve_cholesky_factored_system') as scope:
    cholesky_factor = tf.convert_to_tensor(
        value=cholesky_factor, name='cholesky_factor')
    rhs = tf.convert_to_tensor(value=rhs, name='rhs')
    lin_op = tf.linalg.LinearOperatorLowerTriangular(
        cholesky_factor, name=scope)
    return lin_op.solvevec(lin_op.solvevec(rhs), adjoint=True)

class VariationalGaussianProcess(
    mvn_linear_operator.MultivariateNormalLinearOperator):
  """Posterior predictive of a variational Gaussian process.

  This distribution implements the variational Gaussian process (VGP), as
  described in [Titsias, 2009][1] and [Hensman, 2013][2]. The VGP is an
  inducing point-based approximation of an exact GP posterior
  (see Mathematical Details, below). Ultimately, this Distribution class
  represents a marginal distrbution over function values at a
  collection of `index_points`. It is parameterized by

    - a kernel function,
    - a mean function,
    - the (scalar) observation noise variance of the normal likelihood,
    - a set of index points,
    - a set of inducing index points, and
    - the parameters of the (full-rank, Gaussian) variational posterior
      distribution over function values at the inducing points, conditional on
      some observations.

  A VGP is "trained" by selecting any kernel parameters, the locations of the
  inducing index points, and the variational parameters. [Titsias, 2009][1] and
  [Hensman, 2013][2] describe a variational lower bound on the marginal log
  likelihood of observed data, which this class offers through the
  `variational_loss` method (this is the negative lower bound, for convenience
  when plugging into a TF Optimizer's `minimize` function).
  Training may be done in minibatches.

  [Titsias, 2009][1] describes a closed form for the optimal variational
  parameters, in the case of sufficiently small observational data (ie,
  small enough to fit in memory but big enough to warrant approximating the GP
  posterior). A method to compute these optimal parameters in terms of the full
  observational data set is provided as a staticmethod,
  `optimal_variational_posterior`. It returns a
  `MultivariateNormalLinearOperator` instance with optimal location and
  scale parameters.

  #### Mathematical Details

  ##### Notation

  We will in general be concerned about three collections of index points, and
  it'll be good to give them names:

    * `x[1], ..., x[N]`: observation index points -- locations of our observed
      data.
    * `z[1], ..., z[M]`: inducing index points  -- locations of the
      "summarizing" inducing points
    * `t[1], ..., t[P]`: predictive index points -- locations where we are
      making posterior predictions based on observations and the variational
      parameters.

  To lighten notation, we'll use `X, Z, T` to denote the above collections.
  Similarly, we'll denote by `f(X)` the collection of function values at each of
  the `x[i]`, and by `Y`, the collection of (noisy) observed data at each `x[i].
  We'll denote kernel matrices generated from pairs of index points as `K_tt`,
  `K_xt`, `K_tz`, etc, e.g.,

  ```none
           | k(t[1], z[1])    k(t[1], z[2])  ...  k(t[1], z[M]) |
    K_tz = | k(t[2], z[1])    k(t[2], z[2])  ...  k(t[2], z[M]) |
           |      ...              ...                 ...      |
           | k(t[P], z[1])    k(t[P], z[2])  ...  k(t[P], z[M]) |
  ```

  ##### Preliminaries

  A Gaussian process is an indexed collection of random variables, any finite
  collection of which are jointly Gaussian. Typically, the index set is some
  finite-dimensional, real vector space, and indeed we make this assumption in
  what follows. The GP may then be thought of as a distribution over functions
  on the index set. Samples from the GP are functions *on the whole index set*;
  these can't be represented in finite compute memory, so one typically works
  with the marginals at a finite collection of index points. The properties of
  the GP are entirely determined by its mean function `m` and covariance
  function `k`. The generative process, assuming a mean-zero normal likelihood
  with stddev `sigma`, is

  ```none
    f ~ GP(m, k)

    Y | f(X) ~ Normal(f(X), sigma),   i = 1, ... , N
  ```

  In finite terms (ie, marginalizing out all but a finite number of f(X)'sigma),
  we can write

  ```none
    f(X) ~ MVN(loc=m(X), cov=K_xx)

    Y | f(X) ~ Normal(f(X), sigma),   i = 1, ... , N
  ```

  Posterior inference is possible in analytical closed form but becomes
  intractible as data sizes get large. See [Rasmussen, 2006][3] for details.

  ##### The VGP

  The VGP is an inducing point-based approximation of an exact GP posterior,
  where two approximating assumptions have been made:

    1. function values at non-inducing points are mutually independent
       conditioned on function values at the inducing points,
    2. the (expensive) posterior over function values at inducing points
       conditional on obseravtions is replaced with an arbitrary (learnable)
       full-rank Gaussian distribution,

       ```none
         q(f(Z)) = MVN(loc=m, scale=S),
       ```

       where `m` and `S` are parameters to be chosen by optimizing an evidence
       lower bound (ELBO).

  The posterior predictive distribution becomes

  ```none
    q(f(T)) = integral df(Z) p(f(T) | f(Z)) q(f(Z))
            = MVN(loc = A @ m, scale = B^(1/2))
  ```

  where

  ```none
    A = K_tz @ K_zz^-1
    B = K_tt - A @ (K_zz - S S^T) A^T
  ```

  ***The approximate posterior predictive distribution `q(f(T))` is what the
  `VariationalGaussianProcess` class represents.***

  Model selection in this framework entails choosing the kernel parameters,
  inducing point locations, and variational parameters. We do this by optimizing
  a variational lower bound on the marginal log likelihood of observed data. The
  lower bound takes the following form (see [Titsias, 2009][1] and
  [Hensman, 2013][2] for details on the derivation):

  ```none
    L(Z, m, S, Y) = (
        MVN(loc=(K_zx @ K_zz^-1) @ m, scale_diag=sigma).log_prob(Y) -
        (Tr(K_xx - K_zx @ K_zz^-1 @ K_xz) +
         Tr(S @ S^T @ K_zz^1 @ K_zx @ K_xz @ K_zz^-1)) / (2 * sigma^2) -
        KL(q(f(Z)) || p(f(Z))))
  ```

  where in the final KL term, `p(f(Z))` is the GP prior on inducing point
  function values. This variational lower bound can be computed on minibatches
  of the full data set `(X, Y)`. A method to compute the *negative* variational
  lower bound is implemented as `VariationalGaussianProcess.variational_loss`.

  ##### Optimal variational parameters

  As described in [Titsias, 2009][1], a closed form optimum for the variational
  location and scale parameters, `m` and `S`, can be computed when the
  observational data are not prohibitively voluminous. The
  `optimal_variational_posterior` function to computes the optimal variational
  posterior distribution over inducing point function values in terms of the GP
  parameters (mean and kernel functions), inducing point locations, observation
  index points, and observations. Note that the inducing index point locations
  must still be optimized even when these parameters are known functions of the
  inducing index points. The optimal parameters are computed as follows:

  ```none
    C = sigma^-2 (K_zz + K_zx @ K_xz)^-1

    optimal Gaussian covariance: K_zz @ C @ K_zz
    optimal Gaussian location: sigma^-2 K_zz @ C @ K_zx @ Y
  ```

  #### Usage Examples

  Here's an example of defining and training a VariationalGaussianProcess on
  some toy generated data.

  ```python
  # We'll use double precision throughout for better numerics.
  dtype = np.float64

  # Generate noisy data from a known function.
  f = lambda x: np.exp(-x[..., 0]**2 / 20.) * np.sin(1. * x[..., 0])
  true_observation_noise_variance_ = dtype(1e-1) ** 2

  num_training_points_ = 100
  x_train_ = np.stack(
      [np.random.uniform(-6., 0., [num_training_points_/ 2 , 1]),
       np.random.uniform(1., 10., [num_training_points_/ 2 , 1])],
      axis=0).astype(dtype)
  y_train_ = (f(x_train_) +
              np.random.normal(
                  0., np.sqrt(true_observation_noise_variance_),
                  [num_training_points_]).astype(dtype))

  # Create kernel with trainable parameters, and trainable observation noise
  # variance variable. Each of these is constrained to be positive.
  amplitude = (tf.nn.softplus(tf.Variable(-1., dtype=dtype, name='amplitude')))
  length_scale = (1e-5 +
                  tf.nn.softplus(
                      tf.Variable(-3., dtype=dtype, name='length_scale')))
  kernel = tfk.ExponentiatedQuadratic(
      amplitude=amplitude,
      length_scale=length_scale)

  observation_noise_variance = tf.nn.softplus(
      tf.Variable(0, dtype=dtype, name='observation_noise_variance'))

  # Create trainable inducing point locations and variational parameters.
  num_inducing_points_ = 20

  inducing_index_points = tf.Variable(
      initial_inducing_points_, dtype=dtype,
      name='inducing_index_points')
  variational_inducing_observations_loc = tf.Variable(
      np.zeros([num_inducing_points_], dtype=dtype),
      name='variational_inducing_observations_loc')
  variational_inducing_observations_scale = tf.Variable(
      np.eye(num_inducing_points_, dtype=dtype),
      name='variational_inducing_observations_scale')

  # These are the index point locations over which we'll construct the
  # (approximate) posterior predictive distribution.
  num_predictive_index_points_ = 500
  index_points_ = np.linspace(-13, 13,
                              num_predictive_index_points_,
                              dtype=dtype)[..., np.newaxis]

  # Construct our variational GP Distribution instance.
  vgp = tfd.VariationalGaussianProcess(
      kernel,
      index_points=index_points_,
      inducing_index_points=inducing_index_points,
      variational_inducing_observations_loc=variational_inducing_observations_loc,
      variational_inducing_observations_scale=variational_inducing_observations_scale,
      observation_noise_variance=observation_noise_variance)

  # For training, we use some simplistic numpy-based minibatching.
  batch_size = 64
  x_train_batch = tf.placeholder(dtype, [batch_size, 1], name='x_train_batch')
  y_train_batch = tf.placeholder(dtype, [batch_size], name='y_train_batch')

  # Create the loss function we want to optimize.
  loss = vgp.variational_loss(
      observations=y_train_batch,
      observation_index_points=x_train_batch,
      kl_weight=float(batch_size) / float(num_training_points_))

  optimizer = tf.train.AdamOptimizer()
  train_op = optimizer.minimize(loss)

  num_iters = 10000
  num_logs = 10
  with tf.Session() as sess:
    for i in range(num_iters):
      batch_idxs = np.random.randint(num_training_points_, size=[batch_size])
      x_train_batch_ = x_train_[batch_idxs, ...]
      y_train_batch_ = y_train_[batch_idxs]

      [_, loss_] = sess.run([train_op, loss],
                            feed_dict={x_train_batch: x_train_batch_,
                                       y_train_batch: y_train_batch_})
      if i % (num_iters / num_logs) == 0 or i + 1 == num_iters:
        print(i, loss_)

  # Generate a plot with
  #   - the posterior predictive mean
  #   - training data
  #   - inducing index points (plotted vertically at the mean of the variational
  #     posterior over inducing point function values)
  #   - 50 posterior predictive samples

  num_samples = 50
  [
      samples_,
      mean_,
      inducing_index_points_,
      variational_loc_,
  ] = sess.run([
      vgp.sample(num_samples),
      vgp.mean(),
      inducing_index_points,
      variational_inducing_observations_loc
  ])
  plt.figure(figsize=(15, 5))
  plt.scatter(inducing_index_points_[..., 0], variational_loc_
              marker='x', s=50, color='k', zorder=10)
  plt.scatter(x_train_[..., 0], y_train_, color='#00ff00', zorder=9)
  plt.plot(np.tile(index_points_[..., 0], num_samples),
           samples_.T, color='r', alpha=.1)
  plt.plot(index_points_, mean_, color='k')
  plt.plot(index_points_, f(index_points_), color='b')
  ```

  # Here we use the same data setup, but compute the optimal variational
  # parameters instead of training them.
  ```python
  # We'll use double precision throughout for better numerics.
  dtype = np.float64

  # Generate noisy data from a known function.
  f = lambda x: np.exp(-x[..., 0]**2 / 20.) * np.sin(1. * x[..., 0])
  true_observation_noise_variance_ = dtype(1e-1) ** 2

  num_training_points_ = 1000
  x_train_ = np.random.uniform(-10., 10., [num_training_points_, 1])
  y_train_ = (f(x_train_) +
              np.random.normal(
                  0., np.sqrt(true_observation_noise_variance_),
                  [num_training_points_]))

  # Create kernel with trainable parameters, and trainable observation noise
  # variance variable. Each of these is constrained to be positive.
  amplitude = (tf.nn.softplus(
    tf.Variable(.54, dtype=dtype, name='amplitude', use_resource=True)))
  length_scale = (
    1e-5 +
    tf.nn.softplus(
      tf.Variable(.54, dtype=dtype, name='length_scale', use_resource=True)))
  kernel = tfk.ExponentiatedQuadratic(
      amplitude=amplitude,
      length_scale=length_scale)

  observation_noise_variance = tf.nn.softplus(
      tf.Variable(
        .54, dtype=dtype, name='observation_noise_variance', use_resource=True))

  # Create trainable inducing point locations and variational parameters.
  num_inducing_points_ = 10

  inducing_index_points = tf.Variable(
      np.linspace(-10., 10., num_inducing_points_)[..., np.newaxis],
      dtype=dtype, name='inducing_index_points', use_resource=True)

  variational_loc, variational_scale = (
      tfd.VariationalGaussianProcess.optimal_variational_posterior(
          kernel=kernel,
          inducing_index_points=inducing_index_points,
          observation_index_points=x_train_,
          observations=y_train_,
          observation_noise_variance=observation_noise_variance))

  # These are the index point locations over which we'll construct the
  # (approximate) posterior predictive distribution.
  num_predictive_index_points_ = 500
  index_points_ = np.linspace(-13, 13,
                              num_predictive_index_points_,
                              dtype=dtype)[..., np.newaxis]

  # Construct our variational GP Distribution instance.
  vgp = tfd.VariationalGaussianProcess(
      kernel,
      index_points=index_points_,
      inducing_index_points=inducing_index_points,
      variational_inducing_observations_loc=variational_loc,
      variational_inducing_observations_scale=variational_scale,
      observation_noise_variance=observation_noise_variance)

  # For training, we use some simplistic numpy-based minibatching.
  batch_size = 64
  x_train_batch = tf.placeholder(dtype, [batch_size, 1], name='x_train_batch')
  y_train_batch = tf.placeholder(dtype, [batch_size], name='y_train_batch')

  # Create the loss function we want to optimize.
  loss = vgp.variational_loss(
      observations=y_train_batch,
      observation_index_points=x_train_batch,
      kl_weight=float(batch_size) / float(num_training_points_))

  optimizer = tf.train.AdamOptimizer(learning_rate=.01)
  train_op = optimizer.minimize(loss)

  num_iters = 300
  num_logs = 10
  with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for i in range(num_iters):
      batch_idxs = np.random.randint(num_training_points_, size=[batch_size])
      x_train_batch_ = x_train_[batch_idxs, ...]
      y_train_batch_ = y_train_[batch_idxs]

      [_, loss_] = sess.run([train_op, loss],
                            feed_dict={x_train_batch: x_train_batch_,
                                       y_train_batch: y_train_batch_})
      if i % (num_iters / num_logs) == 0 or i + 1 == num_iters:
        print(i, loss_)

    # Generate a plot with
    #   - the posterior predictive mean
    #   - training data
    #   - inducing index points (plotted vertically at the mean of the
    #     variational posterior over inducing point function values)
    #   - 50 posterior predictive samples

    num_samples = 50
    [
        samples_,
        mean_,
        inducing_index_points_,
        variational_loc_,
    ] = sess.run([
        vgp.sample(num_samples),
        vgp.mean(),
        inducing_index_points,
        variational_loc
    ])
    plt.figure(figsize=(15, 5))
    plt.scatter(inducing_index_points_[..., 0], variational_loc_,
                marker='x', s=50, color='k', zorder=10)
    plt.scatter(x_train_[..., 0], y_train_, color='#00ff00', alpha=.1, zorder=9)
    plt.plot(np.tile(index_points_, num_samples),
             samples_.T, color='r', alpha=.1)
    plt.plot(index_points_, mean_, color='k')
    plt.plot(index_points_, f(index_points_), color='b')

  ```

  #### References

  [1]: Titsias, M. "Variational Model Selection for Sparse Gaussian Process
       Regression", 2009.
       http://proceedings.mlr.press/v5/titsias09a/titsias09a.pdf
  [2]: Hensman, J., Lawrence, N. "Gaussian Processes for Big Data", 2013
       https://arxiv.org/abs/1309.6835
  [3]: Carl Rasmussen, Chris Williams. Gaussian Processes For Machine Learning,
       2006. http://www.gaussianprocess.org/gpml/
  """

  def __init__(self,
               kernel,
               index_points,
               inducing_index_points,
               variational_inducing_observations_loc,
               variational_inducing_observations_scale,
               mean_fn=None,
               observation_noise_variance=0.,
               predictive_noise_variance=0.,
               jitter=1e-6,
               validate_args=False,
               allow_nan_stats=False,
               name='VariataionalGaussianProcess'):
    """Instantiate a VariationalGaussianProcess Distribution.

    Args:
      kernel: `PositiveSemidefiniteKernel`-like instance representing the
        GP's covariance function.
      index_points: `float` `Tensor` representing finite (batch of) vector(s) of
        points in the index set over which the VGP is defined. Shape has the
        form `[b1, ..., bB, e1, f1, ..., fF]` where `F` is the number of feature
        dimensions and must equal `kernel.feature_ndims` and `e1` is the number
        (size) of index points in each batch (we denote it `e1` to distinguish
        it from the numer of inducing index points, denoted `e2` below).
        Ultimately the VariationalGaussianProcess distribution corresponds to an
        `e1`-dimensional multivariate normal. The batch shape must be
        broadcastable with `kernel.batch_shape`, the batch shape of
        `inducing_index_points`, and any batch dims yielded by `mean_fn`.
      inducing_index_points: `float` `Tensor` of locations of inducing points in
        the index set. Shape has the form `[b1, ..., bB, e2, f1, ..., fF]`, just
        like `index_points`. The batch shape components needn't be identical to
        those of `index_points`, but must be broadcast compatible with them.
      variational_inducing_observations_loc: `float` `Tensor`; the mean of the
        (full-rank Gaussian) variational posterior over function values at the
        inducing points, conditional on observed data. Shape has the form `[b1,
        ..., bB, e2]`, where `b1, ..., bB` is broadcast compatible with other
        parameters' batch shapes, and `e2` is the number of inducing points.
      variational_inducing_observations_scale: `float` `Tensor`; the scale
        matrix of the (full-rank Gaussian) variational posterior over function
        values at the inducing points, conditional on observed data. Shape has
        the form `[b1, ..., bB, e2, e2]`, where `b1, ..., bB` is broadcast
        compatible with other parameters and `e2` is the number of inducing
        points.
      mean_fn: Python `callable` that acts on index points to produce a (batch
        of) vector(s) of mean values at those index points. Takes a `Tensor` of
        shape `[b1, ..., bB, f1, ..., fF]` and returns a `Tensor` whose shape is
        (broadcastable with) `[b1, ..., bB]`. Default value: `None` implies
        constant zero function.
      observation_noise_variance: `float` `Tensor` representing the variance
        of the noise in the Normal likelihood distribution of the model. May be
        batched, in which case the batch shape must be broadcastable with the
        shapes of all other batched parameters (`kernel.batch_shape`,
        `index_points`, etc.).
        Default value: `0.`
      predictive_noise_variance: `float` `Tensor` representing additional
        variance in the posterior predictive model. If `None`, we simply re-use
        `observation_noise_variance` for the posterior predictive noise. If set
        explicitly, however, we use the given value. This allows us, for
        example, to omit predictive noise variance (by setting this to zero) to
        obtain noiseless posterior predictions of function values, conditioned
        on noisy observations.
      jitter: `float` scalar `Tensor` added to the diagonal of the covariance
        matrix to ensure positive definiteness of the covariance matrix.
        Default value: `1e-6`.
      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: `False`.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: "VariationalGaussianProcess".

    Raises:
      ValueError: if `mean_fn` is not `None` and is not callable.
    """
    parameters = dict(locals())
    with tf.name_scope(
        name or 'VariationalGaussianProcess') as name:
      dtype = dtype_util.common_dtype(
          [kernel,
           index_points,
           inducing_index_points,
           variational_inducing_observations_loc,
           variational_inducing_observations_scale,
           observation_noise_variance,
           predictive_noise_variance,
           jitter], tf.float32)

      index_points = tf.convert_to_tensor(
          value=index_points, dtype=dtype, name='index_points')
      inducing_index_points = tf.convert_to_tensor(
          value=inducing_index_points, dtype=dtype,
          name='inducing_index_points')
      variational_inducing_observations_loc = tf.convert_to_tensor(
          value=variational_inducing_observations_loc, dtype=dtype,
          name='variational_inducing_observations_loc')
      variational_inducing_observations_scale = tf.convert_to_tensor(
          value=variational_inducing_observations_scale, dtype=dtype,
          name='variational_inducing_observations_scale')
      observation_noise_variance = tf.convert_to_tensor(
          value=observation_noise_variance,
          dtype=dtype,
          name='observation_noise_variance')
      if predictive_noise_variance is None:
        predictive_noise_variance = observation_noise_variance
      else:
        predictive_noise_variance = tf.convert_to_tensor(
            value=predictive_noise_variance, dtype=dtype,
            name='predictive_noise_variance')
      jitter = tf.convert_to_tensor(
          value=jitter, dtype=dtype, name='jitter')

      self._kernel = kernel
      self._index_points = index_points
      self._inducing_index_points = inducing_index_points
      self._variational_inducing_observations_posterior = (
          mvn_linear_operator.MultivariateNormalLinearOperator(
              loc=variational_inducing_observations_loc,
              scale=tf.linalg.LinearOperatorFullMatrix(
                  variational_inducing_observations_scale),
              name='variational_inducing_observations_posterior'))

      # Default to a constant zero function, borrowing the dtype from
      # index_points to ensure consistency.
      if mean_fn is None:
        mean_fn = lambda x: tf.zeros([1], dtype=dtype)
      else:
        if not callable(mean_fn):
          raise ValueError('`mean_fn` must be a Python callable')
      self._mean_fn = mean_fn
      self._observation_noise_variance = observation_noise_variance
      self._predictive_noise_variance = predictive_noise_variance
      self._jitter = jitter

      with tf.name_scope('init'):
        # We let t and z denote predictive and inducing index points, resp.
        kzz = _add_diagonal_shift(
            kernel.matrix(inducing_index_points, inducing_index_points),
            jitter)

        self._chol_kzz = tf.linalg.cholesky(kzz)
        self._kzz_inv_varloc = _solve_cholesky_factored_system_vec(
            self._chol_kzz,
            (variational_inducing_observations_loc -
             mean_fn(inducing_index_points)),
            name='kzz_inv_varloc')

        loc, scale = self._compute_posterior_predictive_params()

        super(VariationalGaussianProcess, self).__init__(
            loc=loc,
            scale=scale,
            validate_args=validate_args,
            allow_nan_stats=allow_nan_stats,
            name=name)
        self._parameters = parameters
        self._graph_parents = [
            index_points,
            inducing_index_points,
            variational_inducing_observations_loc,
            variational_inducing_observations_scale,
            observation_noise_variance,
            jitter]

  def _compute_posterior_predictive_params(self):
    ktt = _add_diagonal_shift(
        self._kernel.matrix(self._index_points, self._index_points),
        self._jitter)
    kzt = tf.linalg.LinearOperatorFullMatrix(
        self._kernel.matrix(self._inducing_index_points, self._index_points))

    kzz_inv_kzt = tf.linalg.LinearOperatorFullMatrix(
        _solve_cholesky_factored_system(
            self._chol_kzz, kzt.to_dense(), name='kzz_inv_kzt'))

    var_cov = tf.linalg.LinearOperatorFullMatrix(
        self._variational_inducing_observations_posterior.covariance())
    posterior_predictive_cov = (
        ktt -
        kzt.matmul(kzz_inv_kzt.to_dense(), adjoint=True) +
        kzz_inv_kzt.matmul(var_cov.matmul(kzz_inv_kzt.to_dense()),
                           adjoint=True))

    # Add predictive_noise_variance
    posterior_predictive_cov = _add_diagonal_shift(
        posterior_predictive_cov, self._predictive_noise_variance)

    scale = tf.linalg.LinearOperatorLowerTriangular(
        tf.linalg.cholesky(posterior_predictive_cov))

    loc = (self._mean_fn(self._index_points) +
           kzt.matvec(self._kzz_inv_varloc, adjoint=True))

    return loc, scale

  @property
  def mean_fn(self):
    return self._mean_fn

  @property
  def kernel(self):
    return self._kernel

  @property
  def index_points(self):
    return self._index_points

  @property
  def inducing_index_points(self):
    return self._inducing_index_points

  @property
  def variational_inducing_observations_loc(self):
    return self._variational_inducing_observations_posterior.loc

  @property
  def variational_inducing_observations_scale(self):
    return self._variational_inducing_observations_posterior.scale

  @property
  def observation_noise_variance(self):
    return self._observation_noise_variance

  @property
  def predictive_noise_variance(self):
    return self._predictive_noise_variance

  @property
  def jitter(self):
    return self._jitter

  def _covariance(self):
    return self._covariance_matrix

  def variational_loss(self,
                       observations,
                       observation_index_points=None,
                       kl_weight=1.,
                       name='variational_loss'):
    """Variational loss for the VGP.

    Given `observations` and `observation_index_points`, compute the
    negative variational lower bound as specified in [Hensman, 2013][1].

    Args:
      observations: `float` `Tensor` representing collection, or batch of
        collections, of observations corresponding to
        `observation_index_points`. Shape has the form `[b1, ..., bB, e]`, which
        must be brodcastable with the batch and example shapes of
        `observation_index_points`. The batch shape `[b1, ..., bB]` must be
        broadcastable with the shapes of all other batched parameters
        (`kernel.batch_shape`, `observation_index_points`, etc.).
      observation_index_points: `float` `Tensor` representing finite (batch of)
        vector(s) of points where observations are defined. Shape has the
        form `[b1, ..., bB, e1, f1, ..., fF]` where `F` is the number of feature
        dimensions and must equal `kernel.feature_ndims` and `e1` is the number
        (size) of index points in each batch (we denote it `e1` to distinguish
        it from the numer of inducing index points, denoted `e2` below). If
        set to `None` uses `index_points` as the origin for observations.
        Default value: None.
      kl_weight: Amount by which to scale the KL divergence loss between prior
        and posterior.
        Default value: 1.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: "GaussianProcess".
    Returns:
      loss: Scalar tensor representing the negative variational lower bound.
        Can be directly used in a `tf.Optimizer`.
    Raises:
      ValueError: if `mean_fn` is not `None` and is not callable.

    #### References

    [1]: Hensman, J., Lawrence, N. "Gaussian Processes for Big Data", 2013
         https://arxiv.org/abs/1309.6835
    """

    with tf.name_scope(name or 'variational_gp_loss'):
      if observation_index_points is None:
        observation_index_points = self._index_points
      observation_index_points = tf.convert_to_tensor(
          value=observation_index_points, dtype=self._dtype,
          name='observation_index_points')
      observations = tf.convert_to_tensor(
          value=observations, dtype=self._dtype, name='observations')
      kl_weight = tf.convert_to_tensor(
          value=kl_weight, dtype=self._dtype,
          name='kl_weight')

      # The variational loss is a negative ELBO. The ELBO can be broken down
      # into three terms:
      #  1. a likelihood term
      #  2. a trace term arising from the covariance of the posterior predictive

      kzx = self.kernel.matrix(self._inducing_index_points,
                               observation_index_points)

      kzx_linop = tf.linalg.LinearOperatorFullMatrix(kzx)
      loc = (self._mean_fn(observation_index_points) +
             kzx_linop.matvec(self._kzz_inv_varloc, adjoint=True))

      likelihood = independent.Independent(
          normal.Normal(
              loc=loc,
              scale=tf.sqrt(self._observation_noise_variance + self._jitter),
              name='NormalLikelihood'),
          reinterpreted_batch_ndims=1)
      obs_ll = likelihood.log_prob(observations)

      chol_kzz_linop = tf.linalg.LinearOperatorLowerTriangular(self._chol_kzz)
      chol_kzz_inv_kzx = chol_kzz_linop.solve(kzx)
      kzz_inv_kzx = chol_kzz_linop.solve(chol_kzz_inv_kzx, adjoint=True)

      kxx_diag = tf.linalg.diag_part(
          self.kernel.matrix(
              observation_index_points, observation_index_points))
      ktilde_trace_term = (
          tf.reduce_sum(input_tensor=kxx_diag, axis=-1) -
          tf.reduce_sum(input_tensor=chol_kzz_inv_kzx ** 2, axis=[-2, -1]))

      # Tr(SB)
      # where S = A A.T, A = variational_inducing_observations_scale
      # and B = Kzz^-1 Kzx Kzx.T Kzz^-1
      #
      # Now Tr(SB) = Tr(A A.T Kzz^-1 Kzx Kzx.T Kzz^-1)
      #            = Tr(A.T Kzz^-1 Kzx Kzx.T Kzz^-1 A)
      #            = sum_ij (A.T Kzz^-1 Kzx)_{ij}^2
      other_trace_term = tf.reduce_sum(
          input_tensor=(
              self._variational_inducing_observations_posterior.scale.matmul(
                  kzz_inv_kzx) ** 2),
          axis=[-2, -1])

      trace_term = (.5 * (ktilde_trace_term + other_trace_term) /
                    self._observation_noise_variance)

      inducing_prior = gaussian_process.GaussianProcess(
          kernel=self._kernel,
          mean_fn=self._mean_fn,
          index_points=self._inducing_index_points,
          observation_noise_variance=self._observation_noise_variance)

      kl_term = kl_weight * kullback_leibler.kl_divergence(
          self._variational_inducing_observations_posterior,
          inducing_prior)

      lower_bound = (obs_ll - trace_term - kl_term)

      return -tf.reduce_mean(input_tensor=lower_bound)

  @staticmethod
  def optimal_variational_posterior(
      kernel,
      inducing_index_points,
      observation_index_points,
      observations,
      observation_noise_variance,
      mean_fn=None,
      jitter=1e-6,
      name=None):
    """Model selection for optimal variational hyperparameters.

    Given the full training set (parameterized by `observations` and
    `observation_index_points`), compute the optimal variational
    location and scale for the VGP. This is based of the method suggested
    in [Titsias, 2009][1].

    Args:
      kernel: `PositiveSemidefiniteKernel`-like instance representing the
        GP's covariance function.
      inducing_index_points: `float` `Tensor` of locations of inducing points in
        the index set. Shape has the form `[b1, ..., bB, e2, f1, ..., fF]`, just
        like `observation_index_points`. The batch shape components needn't be
        identical to those of `observation_index_points`, but must be broadcast
        compatible with them.
      observation_index_points: `float` `Tensor` representing finite (batch of)
        vector(s) of points where observations are defined. Shape has the
        form `[b1, ..., bB, e1, f1, ..., fF]` where `F` is the number of feature
        dimensions and must equal `kernel.feature_ndims` and `e1` is the number
        (size) of index points in each batch (we denote it `e1` to distinguish
        it from the numer of inducing index points, denoted `e2` below).
      observations: `float` `Tensor` representing collection, or batch of
        collections, of observations corresponding to
        `observation_index_points`. Shape has the form `[b1, ..., bB, e]`, which
        must be brodcastable with the batch and example shapes of
        `observation_index_points`. The batch shape `[b1, ..., bB]` must be
        broadcastable with the shapes of all other batched parameters
        (`kernel.batch_shape`, `observation_index_points`, etc.).
      observation_noise_variance: `float` `Tensor` representing the variance
        of the noise in the Normal likelihood distribution of the model. May be
        batched, in which case the batch shape must be broadcastable with the
        shapes of all other batched parameters (`kernel.batch_shape`,
        `index_points`, etc.).
        Default value: `0.`
      mean_fn: Python `callable` that acts on index points to produce a (batch
        of) vector(s) of mean values at those index points. Takes a `Tensor` of
        shape `[b1, ..., bB, f1, ..., fF]` and returns a `Tensor` whose shape is
        (broadcastable with) `[b1, ..., bB]`. Default value: `None` implies
        constant zero function.
      jitter: `float` scalar `Tensor` added to the diagonal of the covariance
        matrix to ensure positive definiteness of the covariance matrix.
        Default value: `1e-6`.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: "optimal_variational_posterior".
    Returns:
      loc, scale: Tuple representing the variational location and scale.
    Raises:
      ValueError: if `mean_fn` is not `None` and is not callable.

    #### References

    [1]: Titsias, M. "Variational Model Selection for Sparse Gaussian Process
         Regression", 2009.
         http://proceedings.mlr.press/v5/titsias09a/titsias09a.pdf
    """

    with tf.name_scope(name or 'optimal_variational_posterior'):
      dtype = dtype_util.common_dtype(
          [inducing_index_points,
           observation_index_points,
           observations,
           observation_noise_variance,
           jitter], tf.float32)

      inducing_index_points = tf.convert_to_tensor(
          value=inducing_index_points,
          dtype=dtype, name='inducing_index_points')
      observation_index_points = tf.convert_to_tensor(
          value=observation_index_points, dtype=dtype,
          name='observation_index_points')
      observations = tf.convert_to_tensor(
          value=observations, dtype=dtype, name='observations')
      observation_noise_variance = tf.convert_to_tensor(
          value=observation_noise_variance,
          dtype=dtype,
          name='observation_noise_variance')
      jitter = tf.convert_to_tensor(
          value=jitter, dtype=dtype, name='jitter')

      # Default to a constant zero function.
      if mean_fn is None:
        mean_fn = lambda x: tf.zeros([1], dtype=dtype)
      else:
        if not callable(mean_fn):
          raise ValueError('`mean_fn` must be a Python callable')

      # z are the inducing points and x are the observation index points.
      kzz = kernel.matrix(inducing_index_points, inducing_index_points)
      kzx = kernel.matrix(inducing_index_points, observation_index_points)

      noise_var_inv = tf.math.reciprocal(observation_noise_variance)

      sigma_inv = _add_diagonal_shift(
          kzz + noise_var_inv * tf.matmul(kzx, kzx, adjoint_b=True),
          jitter)

      chol_sigma_inv = tf.linalg.cholesky(sigma_inv)

      kzx_lin_op = tf.linalg.LinearOperatorFullMatrix(kzx)
      kzx_obs = kzx_lin_op.matvec(
          observations - mean_fn(observation_index_points))
      kzz_lin_op = tf.linalg.LinearOperatorFullMatrix(kzz)
      loc = (mean_fn(inducing_index_points) +
             noise_var_inv * kzz_lin_op.matvec(
                 _solve_cholesky_factored_system_vec(chol_sigma_inv, kzx_obs)))

      chol_sigma_inv_lin_op = tf.linalg.LinearOperatorLowerTriangular(
          chol_sigma_inv)
      scale = chol_sigma_inv_lin_op.solve(kzz)

      return loc, scale


# 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.
# ============================================================================
"""Local Linear Trend model."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# Dependency imports
import numpy as np
import tensorflow as tf

from tensorflow_probability.python import bijectors as tfb
from tensorflow_probability.python import distributions as tfd

from tensorflow_probability.python.internal import distribution_util as dist_util
from tensorflow_probability.python.sts.internal import util as sts_util
from tensorflow_probability.python.sts.structural_time_series import Parameter
from tensorflow_probability.python.sts.structural_time_series import StructuralTimeSeries
from tensorflow_probability.python.util import docstring as docstring_util

seasonal_init_args = """
    Args:
      num_timesteps: Scalar `int` `Tensor` number of timesteps to model
        with this distribution.
      num_seasons: Scalar Python `int` number of seasons.
      drift_scale: Scalar (any additional dimensions are treated as batch
        dimensions) `float` `Tensor` indicating the standard deviation of the
        change in effect between consecutive occurrences of a given season.
        This is assumed to be the same for all seasons.
      initial_state_prior: instance of `tfd.MultivariateNormal`
        representing the prior distribution on latent states; must
        have event shape `[num_seasons]`.
      observation_noise_scale: Scalar (any additional dimensions are
        treated as batch dimensions) `float` `Tensor` indicating the standard
        deviation of the observation noise.
        Default value: 0.
      num_steps_per_season: Python `int` number of steps in each
        season. This may be either a scalar (shape `[]`), in which case all
        seasons have the same length, or a NumPy array of shape `[num_seasons]`,
        in which seasons have different length, but remain constant around
        different cycles, or a NumPy array of shape `[num_cycles, num_seasons]`,
        in which num_steps_per_season for each season also varies in different
        cycle (e.g., a 4 years cycle with leap day).
        Default value: 1.
      initial_step: Optional scalar `int` `Tensor` specifying the starting
        timestep.
        Default value: 0.
      validate_args: Python `bool`. Whether to validate input
        with asserts. If `validate_args` is `False`, and the inputs are
        invalid, correct behavior is not guaranteed.
        Default value: `False`.
      allow_nan_stats: Python `bool`. If `False`, raise an
        exception if a statistic (e.g. mean/mode/etc...) is undefined for any
        batch member. If `True`, batch members with valid parameters leading to
        undefined statistics will return NaN for this statistic.
        Default value: `True`.
      name: Python `str` name prefixed to ops created by this class.
        Default value: "SeasonalStateSpaceModel".

    Raises:
      ValueError: if `num_steps_per_season` has invalid shape (neither
        scalar nor `[num_seasons]`).
    """

class SeasonalStateSpaceModel(tfd.LinearGaussianStateSpaceModel):
  """State space model for a seasonal effect.

    A state space model (SSM) posits a set of latent (unobserved) variables that
    evolve over time with dynamics specified by a probabilistic transition model
    `p(z[t+1] | z[t])`. At each timestep, we observe a value sampled from an
    observation model conditioned on the current state, `p(x[t] | z[t])`. The
    special case where both the transition and observation models are Gaussians
    with mean specified as a linear function of the inputs, is known as a linear
    Gaussian state space model and supports tractable exact probabilistic
    calculations; see `tfp.distributions.LinearGaussianStateSpaceModel` for
    details.

    A seasonal effect model is a special case of a linear Gaussian SSM. The
    latent states represent an unknown effect from each of several 'seasons';
    these are generally not meteorological seasons, but represent regular
    recurring patterns such as hour-of-day or day-of-week effects. The effect of
    each season drifts from one occurrence to the next, following a Gaussian
    random walk:

    ```python
    effects[season, occurrence[i]] = (
      effects[season, occurrence[i-1]] + Normal(loc=0., scale=drift_scale))
    ```

    The latent state has dimension `num_seasons`, containing one effect for each
    seasonal component. The parameters `drift_scale` and
    `observation_noise_scale` are each (a batch of) scalars. The batch shape of
    this `Distribution` is the broadcast batch shape of these parameters and of
    the `initial_state_prior`.

    Note: there is no requirement that the effects sum to zero.

    #### Mathematical Details

    The seasonal effect model implements a
    `tfp.distributions.LinearGaussianStateSpaceModel` with
    `latent_size = num_seasons` and `observation_size = 1`. The latent state
    is organized so that the *current* seasonal effect is always in the first
    (zeroth) dimension. The transition model rotates the latent state to shift
    to a new effect at the end of each season:

    ```
    transition_matrix[t] = (permutation_matrix([1, 2, ..., num_seasons-1, 0])
                            if season_is_changing(t)
                            else eye(num_seasons)
    transition_noise[t] ~ Normal(loc=0., scale_diag=(
                                 [drift_scale, 0, ..., 0]
                                 if season_is_changing(t)
                                 else [0, 0, ..., 0]))
    ```

    where `season_is_changing(t)` is `True` if ``t `mod`
    sum(num_steps_per_season)`` is in the set of final days for each season,
    given by `cumsum(num_steps_per_season) - 1`. The observation model always
    picks out the effect for the current season, i.e., the first element of
    the latent state:

    ```
    observation_matrix = [[1., 0., ..., 0.]]
    observation_noise ~ Normal(loc=0, scale=observation_noise_scale)
    ```

    #### Examples

    A state-space model with day-of-week seasonality on hourly data:

    ```python
    day_of_week = SeasonalStateSpaceModel(
      num_timesteps=30,
      num_seasons=7,
      drift_scale=0.1,
      initial_state_prior=tfd.MultivariateNormalDiag(
        scale_diag=tf.ones([7], dtype=tf.float32),
      num_steps_per_season=24)
    ```

    A model with basic month-of-year seasonality on daily data, demonstrating
    seasons of varying length:

    ```python
    month_of_year = SeasonalStateSpaceModel(
      num_timesteps=2 * 365,  # 2 years
      num_seasons=12,
      drift_scale=0.1,
      initial_state_prior=tfd.MultivariateNormalDiag(
        scale_diag=tf.ones([12], dtype=tf.float32)),
      num_steps_per_season=[31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],
      initial_step=22)
    ```

    Note that we've used `initial_step=22` to denote that the model begins
    on January 23 (steps are zero-indexed). This version works over time periods
    not involving a leap year. A general implementation of month-of-year
    seasonality would require additional logic:

    ```python
    num_days_per_month = np.array(
      [[31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],
       [31, 29, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],  # year with leap day
       [31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],
       [31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31]])

    month_of_year = SeasonalStateSpaceModel(
      num_timesteps=4 * 365 + 2,  # 8 years with leap days
      num_seasons=12,
      drift_scale=0.1,
      initial_state_prior=tfd.MultivariateNormalDiag(
        scale_diag=tf.ones([12], dtype=tf.float32)),
      num_steps_per_season=num_days_per_month,
      initial_step=22)
    ```

  """

  @docstring_util.expand_docstring(seasonal_init_args=seasonal_init_args)
  def __init__(self,
               num_timesteps,
               num_seasons,
               drift_scale,
               initial_state_prior,
               observation_noise_scale=0.,
               num_steps_per_season=1,
               initial_step=0,
               validate_args=False,
               allow_nan_stats=True,
               name=None):  # pylint: disable=g-doc-args
    """Build a seasonal effect state space model.

    {seasonal_init_args}
    """

    with tf.compat.v1.name_scope(
        name, 'SeasonalStateSpaceModel',
        values=[drift_scale, observation_noise_scale]) as name:

      # The initial state prior determines the dtype of sampled values.
      # Other model parameters must have the same dtype.
      dtype = initial_state_prior.dtype
      drift_scale = tf.convert_to_tensor(
          value=drift_scale, name='drift_scale', dtype=dtype)
      observation_noise_scale = tf.convert_to_tensor(
          value=observation_noise_scale,
          name='observation_noise_scale',
          dtype=dtype)

      # Coerce `num_steps_per_season` to a canonical form, an array of
      # `num_seasons` integers.
      num_steps_per_season = np.squeeze(np.asarray(num_steps_per_season))
      if num_steps_per_season.ndim == 0:  # scalar case
        num_steps_per_season = np.tile(num_steps_per_season, num_seasons)
      elif ((num_steps_per_season.ndim <= 2)  # 1D and 2D case
            and (num_steps_per_season.shape[-1] != num_seasons)):
        raise ValueError('num_steps_per_season must either be scalar (shape [])'
                         ' or have the last dimension equal to [num_seasons] = '
                         '[{}] (saw: shape {})'.format(
                             num_seasons, num_steps_per_season.shape))

      is_last_day_of_season = build_is_last_day_of_season(num_steps_per_season)

      seasonal_transition_matrix = build_seasonal_transition_matrix(
          num_seasons=num_seasons,
          is_last_day_of_season=is_last_day_of_season,
          dtype=dtype)

      seasonal_transition_noise = build_seasonal_transition_noise(
          drift_scale, num_seasons, is_last_day_of_season)

      observation_matrix = tf.concat([
          tf.ones([1, 1], dtype=dtype),
          tf.zeros([1, num_seasons-1], dtype=dtype)], axis=-1)

      self._drift_scale = drift_scale
      self._observation_noise_scale = observation_noise_scale
      self._num_seasons = num_seasons
      self._num_steps_per_season = num_steps_per_season

      super(SeasonalStateSpaceModel, self).__init__(
          num_timesteps=num_timesteps,
          transition_matrix=seasonal_transition_matrix,
          transition_noise=seasonal_transition_noise,
          observation_matrix=observation_matrix,
          observation_noise=tfd.MultivariateNormalDiag(
              scale_diag=observation_noise_scale[..., tf.newaxis]),
          initial_state_prior=initial_state_prior,
          initial_step=initial_step,
          allow_nan_stats=allow_nan_stats,
          validate_args=validate_args,
          name=name)

  @property
  def drift_scale(self):
    """Standard deviation of the drift in effects between seasonal cycles."""
    return self._drift_scale

  @property
  def observation_noise_scale(self):
    """Standard deviation of the observation noise."""
    return self._observation_noise_scale

  @property
  def num_seasons(self):
    """Number of seasons."""
    return self._num_seasons

  @property
  def num_steps_per_season(self):
    """Number of steps in each season."""
    return self._num_steps_per_season

class ConstrainedSeasonalStateSpaceModel(tfd.LinearGaussianStateSpaceModel):
  """Seasonal state space model with effects constrained to sum to zero.

    See `SeasonalStateSpaceModel` for background.

    #### Mathematical details

    The constrained model implements a reparameterization of the
    naive `SeasonalStateSpaceModel`. Instead of directly representing the
    seasonal effects in the latent space, the latent space of the constrained
    model represents the difference between each effect and the mean effect.
    The following discussion assumes familiarity with the mathematical details
    of `SeasonalStateSpaceModel`.

    *Reparameterization and constraints*: let the seasonal effects at a given
    timestep be `E = [e_1, ..., e_N]`. The difference between each effect `e_i`
    and the mean effect is `z_i = e_i - sum_i(e_i)/N`. By itself, this
    transformation is not invertible because recovering the absolute effects
    requires that we know the mean as well. To fix this, we'll define
    `z_N = sum_i(e_i)/N` as the mean effect. It's easy to see that this is
    invertible: given the mean effect and the differences of the first `N - 1`
    effects from the mean, it's easy to solve for all `N` effects. Formally,
    we've defined the invertible linear reparameterization `Z = R E`, where

    ```
    R = [1 - 1/N, -1/N,    ..., -1/N
         -1/N,    1 - 1/N, ..., -1/N,
         ...
         1/N,     1/N,     ...,  1/N]
    ```

    represents the change of basis from 'effect coordinates' E to
    'residual coordinates' Z. The `Z`s form the latent space of the
    `ConstrainedSeasonalStateSpaceModel`.

    To constrain the mean effect `z_N` to zero, we fix the prior to zero,
    `p(z_N) ~ N(0., 0)`, and after the transition at each timestep we project
    `z_N` back to zero. Note that this projection is linear: to set the Nth
    dimension to zero, we simply multiply by the identity matrix with a missing
    element in the bottom right, i.e., `Z_constrained = P Z`,
    where `P = eye(N) - scatter((N-1, N-1), 1)`.

    *Model*: concretely, suppose a naive seasonal effect model has initial state
    prior `N(m, S)`, transition matrix `F` and noise covariance
    `Q`, and observation matrix `H`. Then the corresponding constrained seasonal
    effect model has initial state prior `N(P R m, P R S R' P')`,
    transition matrix `P R F R^-1` and noise covariance `F R Q R' F'`, and
    observation matrix `H R^-1`, where the change-of-basis matrix `R` and
    constraint projection matrix `P` are as defined above. This follows
    directly from applying the reparameterization `Z = R E`, and then enforcing
    the zero-sum constraint on the prior and transition noise covariances.

    In practice, because the sum of effects `z_N` is constrained to be zero, it
    will never contribute a term to any linear operation on the latent space,
    so we can drop that dimension from the model entirely.
    `ConstrainedSeasonalStateSpaceModel` does this, so that it implements the
    `N - 1` dimension latent space `z_1, ..., z_[N-1]`.

    Note that since we constrained the mean effect to be zero, the latent
    `z_i`'s now recover their interpretation as the *actual* effects,
    `z_i = e_i` for `i = `1, ..., N - 1`, even though they were originally
    defined as residuals. The `N`th effect is represented only implicitly, as
    the nonzero mean of the first `N - 1` effects. Although the computational
    represention is not symmetric across all `N` effects, we derived the
    `ConstrainedSeasonalStateSpaceModel` by starting with a symmetric
    representation and imposing only a symmetric constraint (the zero-sum
    constraint), so the probability model remains symmetric over all `N`
    seasonal effects.

    #### Examples

    A constrained state-space model with day-of-week seasonality on hourly data:

    ```python
    day_of_week = ConstrainedSeasonalStateSpaceModel(
      num_timesteps=30,
      num_seasons=7,
      drift_scale=0.1,
      initial_state_prior=tfd.MultivariateNormalDiag(
        scale_diag=tf.ones([7-1], dtype=tf.float32)),
      num_steps_per_season=24)
    ```

    A model with basic month-of-year seasonality on daily data, demonstrating
    seasons of varying length:

    ```python
    month_of_year = ConstrainedSeasonalStateSpaceModel(
      num_timesteps=2 * 365,  # 2 years
      num_seasons=12,
      drift_scale=0.1,
      initial_state_prior=tfd.MultivariateNormalDiag(
        scale_diag=tf.ones([12-1], dtype=tf.float32)),
      num_steps_per_season=[31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],
      initial_step=22)
    ```

    Note that we've used `initial_step=22` to denote that the model begins
    on January 23 (steps are zero-indexed). This version works over time periods
    not involving a leap year. A general implementation of month-of-year
    seasonality would require additional logic:

    ```python
    num_days_per_month = np.array(
      [[31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],
       [31, 29, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],  # year with leap day
       [31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],
       [31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31]])

    month_of_year = ConstrainedSeasonalStateSpaceModel(
      num_timesteps=4 * 365 + 2,  # 8 years with leap days
      num_seasons=12,
      drift_scale=0.1,
      initial_state_prior=tfd.MultivariateNormalDiag(
        scale_diag=tf.ones([12-1], dtype=tf.float32)),
      num_steps_per_season=num_days_per_month,
      initial_step=22)
    ```
  """

  def __init__(self,
               num_timesteps,
               num_seasons,
               drift_scale,
               initial_state_prior,
               observation_noise_scale=1e-4,  # Avoid degeneracy.
               num_steps_per_season=1,
               initial_step=0,
               validate_args=False,
               allow_nan_stats=True,
               name=None):  # pylint: disable=g-doc-args
    """Build a seasonal effect state space model with a zero-sum constraint.

    {seasonal_init_args}
    """

    with tf.compat.v1.name_scope(
        name, 'ConstrainedSeasonalStateSpaceModel',
        values=[drift_scale, observation_noise_scale]) as name:

      # The initial state prior determines the dtype of sampled values.
      # Other model parameters must have the same dtype.
      dtype = initial_state_prior.dtype
      drift_scale = tf.convert_to_tensor(
          value=drift_scale, name='drift_scale', dtype=dtype)
      observation_noise_scale = tf.convert_to_tensor(
          value=observation_noise_scale,
          name='observation_noise_scale',
          dtype=dtype)

      # Coerce `num_steps_per_season` to a canonical form, an array of
      # `num_seasons` integers.
      num_steps_per_season = np.squeeze(np.asarray(num_steps_per_season))
      if num_steps_per_season.ndim == 0:  # scalar case
        num_steps_per_season = np.tile(num_steps_per_season, num_seasons)
      elif ((num_steps_per_season.ndim <= 2)  # 1D and 2D case
            and (num_steps_per_season.shape[-1] != num_seasons)):
        raise ValueError('num_steps_per_season must either be scalar (shape [])'
                         ' or have the last dimension equal to [num_seasons] = '
                         '[{}] (saw: shape {})'.format(
                             num_seasons, num_steps_per_season.shape))

      is_last_day_of_season = build_is_last_day_of_season(num_steps_per_season)

      [
          effects_to_residuals,
          residuals_to_effects
      ] = build_effects_to_residuals_matrix(num_seasons, dtype=dtype)

      seasonal_transition_matrix = build_seasonal_transition_matrix(
          num_seasons=num_seasons,
          is_last_day_of_season=is_last_day_of_season,
          dtype=dtype,
          basis_change_matrix=effects_to_residuals,
          basis_change_matrix_inv=residuals_to_effects)

      seasonal_transition_noise = build_constrained_seasonal_transition_noise(
          drift_scale, num_seasons, is_last_day_of_season)

      observation_matrix = tf.concat(
          [tf.ones([1, 1], dtype=dtype),
           tf.zeros([1, num_seasons-1], dtype=dtype)], axis=-1)
      observation_matrix = tf.matmul(observation_matrix, residuals_to_effects)

      self._drift_scale = drift_scale
      self._observation_noise_scale = observation_noise_scale
      self._num_seasons = num_seasons
      self._num_steps_per_season = num_steps_per_season

      super(ConstrainedSeasonalStateSpaceModel, self).__init__(
          num_timesteps=num_timesteps,
          transition_matrix=seasonal_transition_matrix,
          transition_noise=seasonal_transition_noise,
          observation_matrix=observation_matrix,
          observation_noise=tfd.MultivariateNormalDiag(
              scale_diag=observation_noise_scale[..., tf.newaxis]),
          initial_state_prior=initial_state_prior,
          initial_step=initial_step,
          allow_nan_stats=allow_nan_stats,
          validate_args=validate_args,
          name=name)

  @property
  def drift_scale(self):
    """Standard deviation of the drift in effects between seasonal cycles."""
    return self._drift_scale

  @property
  def observation_noise_scale(self):
    """Standard deviation of the observation noise."""
    return self._observation_noise_scale

  @property
  def num_seasons(self):
    """Number of seasons."""
    return self._num_seasons

  @property
  def num_steps_per_season(self):
    """Number of steps in each season."""
    return self._num_steps_per_season

def build_is_last_day_of_season(num_steps_per_season):
  """Build utility method to compute whether the season is changing."""
  num_steps_per_cycle = np.sum(num_steps_per_season)
  changepoints = np.cumsum(np.ravel(num_steps_per_season)) - 1
  def is_last_day_of_season(t):
    t_ = dist_util.maybe_get_static_value(t)
    if t_ is not None:  # static case
      step_in_cycle = t_ % num_steps_per_cycle
      return any(step_in_cycle == changepoints)
    else:
      step_in_cycle = tf.math.floormod(t, num_steps_per_cycle)
      return tf.reduce_any(
          input_tensor=tf.equal(step_in_cycle, changepoints))
  return is_last_day_of_season

def build_effects_to_residuals_matrix(num_seasons, dtype):
  """Build change-of-basis matrices for constrained seasonal effects.

  This method builds the matrix that transforms seasonal effects into
  effect residuals (differences from the mean effect), and additionally
  projects these residuals onto the subspace where the mean effect is zero.

  See `ConstrainedSeasonalStateSpaceModel` for mathematical details.

  Args:
    num_seasons: scalar `int` number of seasons.
    dtype: TensorFlow `dtype` for the returned values.
  Returns:
    effects_to_residuals: `Tensor` of shape
      `[num_seasons-1, num_seasons]`, such that `differences_from_mean_effect =
      matmul(effects_to_residuals, seasonal_effects)`.  In the
      notation of `ConstrainedSeasonalStateSpaceModel`, this is
      `effects_to_residuals = P * R`.
    residuals_to_effects: the (pseudo)-inverse of the above; a
      `Tensor` of shape `[num_seasons, num_seasons-1]`. In the
      notation of `ConstrainedSeasonalStateSpaceModel`, this is
      `residuals_to_effects = R^{-1} * P'`.
  """

  # Build the matrix that converts effects `e_i` into differences from the mean
  # effect `(e_i - sum(e_i)) / num_seasons`, with the mean effect in the last
  # row so that the transformation is invertible.
  effects_to_residuals_fullrank = np.eye(num_seasons) - 1./num_seasons
  effects_to_residuals_fullrank[-1, :] = 1./num_seasons  # compute mean effect
  residuals_to_effects_fullrank = np.linalg.inv(effects_to_residuals_fullrank)

  # Drop the final dimension, effectively setting the mean effect to zero.
  effects_to_residuals = effects_to_residuals_fullrank[:-1, :]
  residuals_to_effects = residuals_to_effects_fullrank[:, :-1]

  # Return Tensor values of the specified dtype.
  effects_to_residuals = tf.cast(
      effects_to_residuals, dtype=dtype, name='effects_to_residuals')
  residuals_to_effects = tf.cast(
      residuals_to_effects, dtype=dtype, name='residuals_to_effects')

  return effects_to_residuals, residuals_to_effects

def build_seasonal_transition_matrix(
    num_seasons, is_last_day_of_season, dtype,
    basis_change_matrix=None, basis_change_matrix_inv=None):
  """Build a function computing transitions for a seasonal effect model."""

  with tf.compat.v1.name_scope('build_seasonal_transition_matrix'):
    # If the season is changing, the transition matrix permutes the latent
    # state to shift all seasons up by a dimension, and sends the current
    # season's effect to the bottom.
    seasonal_permutation = np.concatenate(
        [np.arange(1, num_seasons), [0]], axis=0)
    seasonal_permutation_matrix = tf.constant(
        np.eye(num_seasons)[seasonal_permutation], dtype=dtype)

    # Optionally transform the transition matrix into a reparameterized space,
    # enforcing the zero-sum constraint for ConstrainedSeasonalStateSpaceModel.
    if basis_change_matrix is not None:
      seasonal_permutation_matrix = tf.matmul(
          basis_change_matrix,
          tf.matmul(seasonal_permutation_matrix, basis_change_matrix_inv))

    identity_matrix = tf.eye(
        tf.shape(input=seasonal_permutation_matrix)[-1], dtype=dtype)

    def seasonal_transition_matrix(t):
      return tf.linalg.LinearOperatorFullMatrix(
          matrix=dist_util.pick_scalar_condition(
              is_last_day_of_season(t),
              seasonal_permutation_matrix,
              identity_matrix))

  return seasonal_transition_matrix

def build_seasonal_transition_noise(
    drift_scale, num_seasons, is_last_day_of_season):
  """Build the transition noise model for a SeasonalStateSpaceModel."""

  # If the current season has just ended, increase the variance of its effect
  # following drift_scale. (the just-ended seasonal effect will always be the
  # bottom element of the vector). Otherwise, do nothing.
  drift_scale_diag = tf.stack(
      [tf.zeros_like(drift_scale)] * (num_seasons - 1) + [drift_scale],
      axis=-1)
  def seasonal_transition_noise(t):
    noise_scale_diag = dist_util.pick_scalar_condition(
        is_last_day_of_season(t),
        drift_scale_diag,
        tf.zeros_like(drift_scale_diag))
    return tfd.MultivariateNormalDiag(
        loc=tf.zeros(num_seasons, dtype=drift_scale.dtype),
        scale_diag=noise_scale_diag)
  return seasonal_transition_noise

def build_constrained_seasonal_transition_noise(
    drift_scale, num_seasons, is_last_day_of_season):
  """Build transition noise distribution for a ConstrainedSeasonalSSM."""

  # Conceptually, this method takes the noise covariance on effects L @ L'
  # computed by `build_seasonal_transition_noise`, with scale factor
  #       L = [ 0, 0, ..., 0
  #             ...
  #             0, 0, ..., drift_scale],
  # and transforms it to act on the constrained-residual representation.
  #
  # The resulting noise covariance M @ M' is equivalent to
  #    M @ M' = effects_to_residuals @ LL' @ residuals_to_effects
  # where `@` is matrix multiplication. However because this matrix is
  # rank-deficient, we can't take its Cholesky decomposition directly, so we'll
  # construct its lower-triangular scale factor `M` by hand instead.
  #
  # Concretely, let `M = P @ R @ L` be the scale factor in the
  # transformed space, with matrices `R`, `P` applying the reparameterization
  # and zero-mean constraint respectively as defined in the
  # "Mathematical Details" section of `ConstrainedSeasonalStateSpaceModel`. It's
  # easy to see (*) that the implied covariance
  # `M @ M' = P @ R @ L @ L' @ R' @ P'` is just the constant matrix
  #  `M @ M' = [ 1, 1, ..., 1, 0
  #              1, 1, ..., 1, 0
  #              ...
  #              1, 1, ..., 1, 0
  #              0, 0, ..., 0, 0] * (drift_scale / num_seasons)**2`
  # with zeros in the final row and column. So we can directly construct
  # the lower-triangular factor
  #  `Q = [ 1, 0, ...  0
  #         1, 0, ..., 0
  #         ...
  #         1, 0, ..., 0
  #         0, 0, ..., 0 ] * drift_scale/num_seasons`
  # such that Q @ Q' = M @ M'. In practice, we don't reify the final row and
  # column full of zeroes, i.e., we construct
  # `Q[:num_seasons-1, :num_seasons-1]` as the scale-TriL covariance factor.
  #
  # (*) Argument: `L` is zero everywhere but the last column, so `R @ L` will be
  # too. Since the last column of `R` is the constant `-1/num_seasons`, `R @ L`
  # is simply the matrix with constant `-drift_scale/num_seasons` in the final
  # column (except the final row, which is negated) and zero in all other
  # columns, and `M = P @ R @ L` additionally zeroes out the final row. Then
  # M @ M' is just the outer product of that final column with itself (since all
  # other columns are zero), which gives the matrix shown above.

  drift_scale_tril_nonzeros = tf.concat([
      tf.ones([num_seasons - 1, 1], dtype=drift_scale.dtype),
      tf.zeros([num_seasons - 1, num_seasons - 2], dtype=drift_scale.dtype)],
                                        axis=-1)
  drift_scale_tril = (drift_scale_tril_nonzeros *
                      drift_scale[..., tf.newaxis, tf.newaxis] / num_seasons)

  # Inject transition noise iff it is the last day of the season.
  def seasonal_transition_noise(t):
    noise_scale_tril = dist_util.pick_scalar_condition(
        is_last_day_of_season(t),
        drift_scale_tril,
        tf.zeros_like(drift_scale_tril))
    return tfd.MultivariateNormalTriL(
        loc=tf.zeros(num_seasons-1, dtype=drift_scale.dtype),
        scale_tril=noise_scale_tril)
  return seasonal_transition_noise

class Seasonal(StructuralTimeSeries):
  """Formal representation of a seasonal effect model.

  A seasonal effect model posits a fixed set of recurring, discrete 'seasons',
  each of which is active for a fixed number of timesteps and, while active,
  contributes a different effect to the time series. These are generally not
  meteorological seasons, but represent regular recurring patterns such as
  hour-of-day or day-of-week effects. Each season lasts for a fixed number of
  timesteps. The effect of each season drifts from one occurrence to the next
  following a Gaussian random walk:

  ```python
  effects[season, occurrence[i]] = (
    effects[season, occurrence[i-1]] + Normal(loc=0., scale=drift_scale))
  ```

  The `drift_scale` parameter governs the standard deviation of the random walk;
  for example, in a day-of-week model it governs the change in effect from this
  Monday to next Monday.

  #### Examples

  A seasonal effect model representing day-of-week seasonality on hourly data:

  ```python
  day_of_week = tfp.sts.Seasonal(num_seasons=7,
                                 num_steps_per_season=24,
                                 observed_time_series=y,
                                 name='day_of_week')
  ```

  A seasonal effect model representing month-of-year seasonality on daily data,
  with explicit priors:

  ```python
  month_of_year = tfp.sts.Seasonal(
    num_seasons=12,
    num_steps_per_season=[31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],
    drift_scale_prior=tfd.LogNormal(loc=-1., scale=0.1),
    initial_effect_prior=tfd.Normal(loc=0., scale=5.),
    name='month_of_year')
  ```

  Note that this version works over time periods not involving a leap year. A
  general implementation of month-of-year seasonality would require additional
  logic:

  ```python
  num_days_per_month = np.array(
    [[31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],
     [31, 29, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],  # year with leap day
     [31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],
     [31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31]])

  month_of_year = tfp.sts.Seasonal(
    num_seasons=12,
    num_steps_per_season=num_days_per_month,
    drift_scale_prior=tfd.LogNormal(loc=-1., scale=0.1),
    initial_effect_prior=tfd.Normal(loc=0., scale=5.),
    name='month_of_year')
  ```

  A model representing both day-of-week and hour-of-day seasonality, on hourly
  data:

  ```
  day_of_week = tfp.sts.Seasonal(num_seasons=7,
                                 num_steps_per_season=24,
                                 observed_time_series=y,
                                 name='day_of_week')
  hour_of_day = tfp.sts.Seasonal(num_seasons=24,
                                 num_steps_per_season=1,
                                 observed_time_series=y,
                                 name='hour_of_day')
  model = tfp.sts.Sum(components=[day_of_week, hour_of_day],
                      observed_time_series=y)
  ```

  """

  def __init__(self,
               num_seasons,
               num_steps_per_season=1,
               drift_scale_prior=None,
               initial_effect_prior=None,
               constrain_mean_effect_to_zero=True,
               observed_time_series=None,
               name=None):
    """Specify a seasonal effects model.

    Args:
      num_seasons: Scalar Python `int` number of seasons.
      num_steps_per_season: Python `int` number of steps in each
        season. This may be either a scalar (shape `[]`), in which case all
        seasons have the same length, or a NumPy array of shape `[num_seasons]`,
        in which seasons have different length, but remain constant around
        different cycles, or a NumPy array of shape `[num_cycles, num_seasons]`,
        in which num_steps_per_season for each season also varies in different
        cycle (e.g., a 4 years cycle with leap day).
        Default value: 1.
      drift_scale_prior: optional `tfd.Distribution` instance specifying a prior
        on the `drift_scale` parameter. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        Default value: `None`.
      initial_effect_prior: optional `tfd.Distribution` instance specifying a
        normal prior on the initial effect of each season. This may be either
        a scalar `tfd.Normal` prior, in which case it applies independently to
        every season, or it may be multivariate normal (e.g.,
        `tfd.MultivariateNormalDiag`) with event shape `[num_seasons]`, in
        which case it specifies a joint prior across all seasons. If `None`, a
        heuristic default prior is constructed based on the provided
        `observed_time_series`.
        Default value: `None`.
      constrain_mean_effect_to_zero: if `True`, use a model parameterization
        that constrains the mean effect across all seasons to be zero. This
        constraint is generally helpful in identifying the contributions of
        different model components and can lead to more interpretable
        posterior decompositions. It may be undesirable if you plan to directly
        examine the latent space of the underlying state space model.
        Default value: `True`.
      observed_time_series: optional `float` `Tensor` of shape
        `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
        supported when `T > 1`), specifying an observed time series.
        Any priors not explicitly set will be given default values according to
        the scale of the observed time series (or batch of time series). May
        optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
        a mask `Tensor` to specify timesteps with missing observations.
        Default value: `None`.
      name: the name of this model component.
        Default value: 'Seasonal'.
    """

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

      _, observed_stddev, observed_initial = (
          sts_util.empirical_statistics(observed_time_series)
          if observed_time_series is not None else (0., 1., 0.))

      # Heuristic default priors. Overriding these may dramatically
      # change inference performance and results.
      if drift_scale_prior is None:
        drift_scale_prior = tfd.LogNormal(
            loc=tf.math.log(.01 * observed_stddev), scale=3.)
      if initial_effect_prior is None:
        initial_effect_prior = tfd.Normal(
            loc=observed_initial,
            scale=tf.abs(observed_initial) + observed_stddev)

      dtype = tf.debugging.assert_same_float_dtype(
          [drift_scale_prior, initial_effect_prior])

      if isinstance(initial_effect_prior, tfd.Normal):
        initial_state_prior = tfd.MultivariateNormalDiag(
            loc=tf.stack([initial_effect_prior.mean()] * num_seasons, axis=-1),
            scale_diag=tf.stack([initial_effect_prior.stddev()] * num_seasons,
                                axis=-1))
      else:
        initial_state_prior = initial_effect_prior

      if constrain_mean_effect_to_zero:
        # Transform the prior to the residual parameterization used by
        # `ConstrainedSeasonalStateSpaceModel`, imposing a zero-sum constraint.
        # This doesn't change the marginal prior on individual effects, but
        # does introduce dependence between the effects.
        (effects_to_residuals, _) = build_effects_to_residuals_matrix(
            num_seasons, dtype=dtype)
        effects_to_residuals_linop = tf.linalg.LinearOperatorFullMatrix(
            effects_to_residuals)  # Use linop so that matmul broadcasts.
        initial_state_prior_loc = effects_to_residuals_linop.matvec(
            initial_state_prior.mean())
        initial_state_prior_scale_linop = effects_to_residuals_linop.matmul(
            initial_state_prior.scale)  # returns LinearOperator
        initial_state_prior = tfd.MultivariateNormalFullCovariance(
            loc=initial_state_prior_loc,
            covariance_matrix=initial_state_prior_scale_linop.matmul(
                initial_state_prior_scale_linop.to_dense(), adjoint_arg=True))

      self._constrain_mean_effect_to_zero = constrain_mean_effect_to_zero
      self._initial_state_prior = initial_state_prior
      self._num_seasons = num_seasons
      self._num_steps_per_season = num_steps_per_season

      super(Seasonal, self).__init__(
          parameters=[
              Parameter('drift_scale', drift_scale_prior, tfb.Softplus()),
          ],
          latent_size=(num_seasons - 1
                       if self.constrain_mean_effect_to_zero else num_seasons),
          name=name)

  @property
  def constrain_mean_effect_to_zero(self):
    """Whether to constrain the mean effect to zero."""
    return self._constrain_mean_effect_to_zero

  @property
  def num_seasons(self):
    """Number of seasons."""
    return self._num_seasons

  @property
  def num_steps_per_season(self):
    """Number of steps per season."""
    return self._num_steps_per_season

  @property
  def initial_state_prior(self):
    """Prior distribution on the initial latent state (level and scale)."""
    return self._initial_state_prior

  def _make_state_space_model(self,
                              num_timesteps,
                              param_map,
                              initial_state_prior=None,
                              initial_step=0):

    if initial_state_prior is None:
      initial_state_prior = self.initial_state_prior

    if self.constrain_mean_effect_to_zero:
      return ConstrainedSeasonalStateSpaceModel(
          num_timesteps=num_timesteps,
          num_seasons=self.num_seasons,
          num_steps_per_season=self.num_steps_per_season,
          initial_state_prior=initial_state_prior,
          initial_step=initial_step,
          **param_map)
    else:
      return SeasonalStateSpaceModel(
          num_timesteps=num_timesteps,
          num_seasons=self.num_seasons,
          num_steps_per_season=self.num_steps_per_season,
          initial_state_prior=initial_state_prior,
          initial_step=initial_step,
          **param_map)