# 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.
# ============================================================================
"""Smooth Seasonal model."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# Dependency imports
import numpy as np
import tensorflow.compat.v2 as tf

from tensorflow_probability.python import bijectors as tfb
from tensorflow_probability.python import distributions as tfd
from tensorflow_probability.python.internal import dtype_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


class SmoothSeasonalStateSpaceModel(tfd.LinearGaussianStateSpaceModel):
  """State space model for a smooth 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 smooth seasonal effect model is a special case of a linear Gaussian SSM. It
  is the sum of a set of "cyclic" components, with one component for each
  frequency:

  ```python
  frequencies[j] = 2. * pi * frequency_multipliers[j] / period
  ```

  Each cyclic component contains two latent states which we denote `effect` and
  `auxiliary`. The two latent states for component `j` drift over time via:

  ```python
  effect[t] = (effect[t - 1] * cos(frequencies[j]) +
               auxiliary[t - 1] * sin(frequencies[j]) +
               Normal(0., drift_scale))

  auxiliary[t] = (-effect[t - 1] * sin(frequencies[j]) +
                  auxiliary[t - 1] * cos(frequencies[j]) +
                  Normal(0., drift_scale))
  ```

  The `auxiliary` latent state only appears as a matter of construction and thus
  its interpretation is not particularly important. The total smooth seasonal
  effect is the sum of the `effect` values from each of the cyclic components.

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

  #### Mathematical Details

  The smooth seasonal effect model implements a
  `tfp.distributions.LinearGaussianStateSpaceModel` with `latent_size = 2 *
  len(frequency_multipliers)` and `observation_size = 1`. The latent state is
  the concatenation of the cyclic latent states which themselves comprise an
  `effect` and an `auxiliary` state. The transition matrix is a block diagonal
  matrix where block `j` is:

  ```python
  transition_matrix[j] =  [[cos(frequencies[j]), sin(frequencies[j])],
                           [-sin(frequencies[j]), cos(frequencies[j])]]
  ```

  The observation model picks out the cyclic `effect` values from the latent
  state:

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

  For further mathematical details please see [1].

  #### Examples

  A state space model with smooth daily seasonality on hourly data. In other
  words, each day there is a pattern which broadly repeats itself over the
  course of the day and doesn't change too much from one hour to the next. Four
  random samples from such a model can be obtained via:

  ```python
  from matplotlib import pylab as plt

  ssm = SmoothSeasonalStateSpaceModel(
      num_timesteps=100,
      period=24,
      frequency_multipliers=[1, 4],
      drift_scale=0.1,
      initial_state_prior=tfd.MultivariateNormalDiag(
          scale_diag=tf.fill([4], 2.0)),
  )

  fig, axes = plt.subplots(4)

  series = ssm.sample(4)

  for series, ax in zip(series[..., 0], axes):
    ax.set_xticks(tf.range(ssm.num_timesteps, delta=ssm.period))
    ax.grid()
    ax.plot(series)

  plt.show()
  ```

  A comparison of the above with a comparable `Seasonal` component gives an
  example of the difference between these two components:

  ```python
  ssm = SeasonalStateSpaceModel(
      num_timesteps=100,
      num_seasons=24,
      num_steps_per_season=1,
      drift_scale=0.1,
      initial_state_prior=tfd.MultivariateNormalDiag(
          scale_diag=tf.fill([24], 2.0)),
  )
  ```

  #### References

  [1]: Harvey, A. Forecasting, Structural Time Series Models and the Kalman
    Filter. Cambridge: Cambridge University Press, 1990.

  """

  def __init__(self,
               num_timesteps,
               period,
               frequency_multipliers,
               drift_scale,
               initial_state_prior,
               observation_noise_scale=0.,
               name=None,
               **linear_gaussian_ssm_kwargs):
    """Build a smooth seasonal state space model.

    Args:
      num_timesteps: Scalar `int` `Tensor` number of timesteps to model
        with this distribution.
      period: positive scalar `float` `Tensor` giving the number of timesteps
        required for the longest cyclic effect to repeat.
      frequency_multipliers: One-dimensional `float` `Tensor` listing the
        frequencies (cyclic components) included in the model, as multipliers of
        the base/fundamental frequency `2. * pi / period`. Each component is
        specified by the number of times it repeats per period, and adds two
        latent dimensions to the model. A smooth seasonal model that can
        represent any periodic function is given by `frequency_multipliers = [1,
        2, ..., floor(period / 2)]`. However, it is often desirable to enforce a
        smoothness assumption (and reduce the computational burden) by dropping
        some of the higher frequencies.
      drift_scale: Scalar (any additional dimensions are treated as batch
        dimensions) `float` `Tensor` indicating the standard deviation of the
        latent state transitions.
      initial_state_prior: instance of `tfd.MultivariateNormal`
        representing the prior distribution on latent states.  Must have
        event shape `[num_features]`.
      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.`.
      name: Python `str` name prefixed to ops created by this class.
        Default value: 'SmoothSeasonalStateSpaceModel'.
      **linear_gaussian_ssm_kwargs: Optional additional keyword arguments to
        to the base `tfd.LinearGaussianStateSpaceModel` constructor.
    """
    parameters = dict(locals())
    parameters.update(linear_gaussian_ssm_kwargs)
    del parameters['linear_gaussian_ssm_kwargs']
    with tf.name_scope(name or 'SmoothSeasonalStateSpaceModel') as name:

      dtype = dtype_util.common_dtype(
          [period, frequency_multipliers, drift_scale, initial_state_prior])

      period = tf.convert_to_tensor(
          value=period, name='period', dtype=dtype)

      frequency_multipliers = tf.convert_to_tensor(
          value=frequency_multipliers,
          name='frequency_multipliers',
          dtype=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)

      num_frequencies = static_num_frequencies(frequency_multipliers)

      observation_matrix = tf.tile(
          tf.constant([[1., 0.]], dtype=dtype),
          multiples=[1, num_frequencies])

      transition_matrix = build_smooth_seasonal_transition_matrix(
          period=period,
          frequency_multipliers=frequency_multipliers,
          dtype=dtype)

      self._drift_scale = drift_scale
      self._observation_noise_scale = observation_noise_scale
      self._period = period
      self._frequency_multipliers = frequency_multipliers

      super(SmoothSeasonalStateSpaceModel, self).__init__(
          num_timesteps=num_timesteps,
          transition_matrix=transition_matrix,
          transition_noise=tfd.MultivariateNormalDiag(
              scale_diag=(drift_scale[..., tf.newaxis] *
                          tf.ones([2 * num_frequencies], dtype=dtype)),
              name='transition_noise'),
          observation_matrix=observation_matrix,
          observation_noise=tfd.MultivariateNormalDiag(
              scale_diag=observation_noise_scale[..., tf.newaxis],
              name='observation_noise'),
          initial_state_prior=initial_state_prior,
          name=name,
          **linear_gaussian_ssm_kwargs)
      self._parameters = parameters

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

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

  @property
  def period(self):
    """The seasonal period."""
    return self._period

  @property
  def frequency_multipliers(self):
    """Multipliers of the fundamental frequency."""
    return self._frequency_multipliers


def build_smooth_seasonal_transition_matrix(period,
                                            frequency_multipliers,
                                            dtype):
  """Build the transition matrix for a SmoothSeasonalStateSpaceModel."""

  two_pi = tf.constant(2. * np.pi, dtype=dtype)
  frequencies = two_pi * frequency_multipliers / period
  num_frequencies = static_num_frequencies(frequency_multipliers)

  sin_frequencies = tf.sin(frequencies)
  cos_frequencies = tf.cos(frequencies)

  trigonometric_values = tf.stack(
      [cos_frequencies, sin_frequencies, -sin_frequencies, cos_frequencies],
      axis=-1)

  transition_matrix = tf.linalg.LinearOperatorBlockDiag(
      [tf.linalg.LinearOperatorFullMatrix(
          matrix=tf.reshape(trigonometric_values[i], [2, 2]),
          is_square=True) for i in range(num_frequencies)]
  )

  return transition_matrix


def static_num_frequencies(frequency_multipliers):
  """Statically known number of frequencies. Raises if not possible."""

  frequency_multipliers = tf.convert_to_tensor(
      frequency_multipliers, name='frequency_multipliers')

  num_frequencies = tf.compat.dimension_value(
      dimension=frequency_multipliers.shape[0])

  if num_frequencies is None:
    raise ValueError('The number of frequencies must be statically known. Saw '
                     '`frequency_multipliers` with shape {}'.format(
                         frequency_multipliers.shape))

  return num_frequencies


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

  The smooth seasonal model uses a set of trigonometric terms in order to
  capture a recurring pattern whereby adjacent (in time) effects are
  similar. The model uses `frequencies` calculated via:

  ```python
  frequencies[j] = 2. * pi * frequency_multipliers[j] / period
  ```

  and then posits two latent states for each `frequency`. The two latent states
  associated with frequency `j` drift over time via:

  ```python
  effect[t] = (effect[t - 1] * cos(frequencies[j]) +
               auxiliary[t - 1] * sin(frequencies[j]) +
               Normal(0., drift_scale))

  auxiliary[t] = (-effect[t - 1] * sin(frequencies[j]) +
                  auxiliary[t - 1] * cos(frequencies[j]) +
                  Normal(0., drift_scale))
  ```

  where `effect` is the smooth seasonal effect and `auxiliary` only appears as a
  matter of construction. The interpretation of `auxiliary` is thus not
  particularly important.

  #### Examples

  A smooth seasonal effect model representing smooth weekly seasonality on daily
  data:

  ```python
  component = SmoothSeasonal(
      period=7,
      frequency_multipliers=[1, 2, 3],
      initial_state_prior=tfd.MultivariateNormalDiag(scale_diag=tf.ones([6])),
  )
  ```

  """

  def __init__(self,
               period,
               frequency_multipliers,
               allow_drift=True,
               drift_scale_prior=None,
               initial_state_prior=None,
               observed_time_series=None,
               name=None):
    """Specify a smooth seasonal effects model.

    Args:
      period: positive scalar `float` `Tensor` giving the number of timesteps
        required for the longest cyclic effect to repeat.
      frequency_multipliers: One-dimensional `float` `Tensor` listing the
        frequencies (cyclic components) included in the model, as multipliers of
        the base/fundamental frequency `2. * pi / period`. Each component is
        specified by the number of times it repeats per period, and adds two
        latent dimensions to the model. A smooth seasonal model that can
        represent any periodic function is given by `frequency_multipliers = [1,
        2, ..., floor(period / 2)]`. However, it is often desirable to enforce a
        smoothness assumption (and reduce the computational burden) by dropping
        some of the higher frequencies.
      allow_drift: optional Python `bool` specifying whether the seasonal
        effects can drift over time.  Setting this to `False`
        removes the `drift_scale` parameter from the model. This is
        mathematically equivalent to
        `drift_scale_prior = tfd.Deterministic(0.)`, but removing drift
        directly is preferred because it avoids the use of a degenerate prior.
        Default value: `True`.
      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_state_prior: instance of `tfd.MultivariateNormal` representing
        the prior distribution on the latent states. Must have event shape
        `[2 * len(frequency_multipliers)]`. If `None`, a heuristic default prior
        is constructed based on the provided `observed_time_series`.
      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 `NaN`s
        are interpreted as missing observations; missingness may be also be
        explicitly specified by passing a `tfp.sts.MaskedTimeSeries` instance.
        Any priors not explicitly set will be given default values according to
        the scale of the observed time series (or batch of time series).
        Default value: `None`.
      name: the name of this model component.
        Default value: 'SmoothSeasonal'.

    """
    init_parameters = dict(locals())
    with tf.name_scope(name or 'SmoothSeasonal') as name:

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

      latent_size = 2 * static_num_frequencies(frequency_multipliers)

      # 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_state_prior is None:
        initial_state_scale = (
            tf.abs(observed_initial) + observed_stddev)[..., tf.newaxis]
        ones = tf.ones([latent_size], dtype=drift_scale_prior.dtype)
        initial_state_prior = tfd.MultivariateNormalDiag(
            scale_diag=initial_state_scale * ones)

      self._initial_state_prior = initial_state_prior
      self._period = period
      self._frequency_multipliers = frequency_multipliers

      parameters = []
      if allow_drift:
        parameters.append(Parameter(
            'drift_scale', drift_scale_prior,
            tfb.Chain([tfb.Scale(scale=observed_stddev),
                       tfb.Softplus()])))
      self._allow_drift = allow_drift

      super(SmoothSeasonal, self).__init__(
          parameters=parameters,
          latent_size=latent_size,
          init_parameters=init_parameters,
          name=name)

  @property
  def allow_drift(self):
    """Whether the seasonal effects are allowed to drift over time."""
    return self._allow_drift

  @property
  def period(self):
    """The seasonal period."""
    return self._period

  @property
  def frequency_multipliers(self):
    """Multipliers of the fundamental frequency."""
    return self._frequency_multipliers

  @property
  def initial_state_prior(self):
    """Prior distribution on the initial latent states."""
    return self._initial_state_prior

  def _make_state_space_model(self,
                              num_timesteps,
                              param_map,
                              initial_state_prior=None,
                              **linear_gaussian_ssm_kwargs):

    if initial_state_prior is None:
      initial_state_prior = self.initial_state_prior

    if not self.allow_drift:
      param_map['drift_scale'] = 0.
    linear_gaussian_ssm_kwargs.update(param_map)
    return SmoothSeasonalStateSpaceModel(
        num_timesteps=num_timesteps,
        period=self.period,
        frequency_multipliers=self.frequency_multipliers,
        initial_state_prior=initial_state_prior,
        **linear_gaussian_ssm_kwargs)
