# 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.,
               initial_step=0,
               validate_args=False,
               allow_nan_stats=True,
               name=None):
    """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.`.
      initial_step: 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: 'SmoothSeasonalStateSpaceModel'.

    """
    parameters = dict(locals())
    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,
          initial_step=initial_step,
          allow_nan_stats=allow_nan_stats,
          validate_args=validate_args,
          name=name)
      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 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: 'SmoothSeasonal'.

    """

    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,
          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,
                              initial_step=0):

    if initial_state_prior is None:
      initial_state_prior = self.initial_state_prior

    if not self.allow_drift:
      param_map['drift_scale'] = 0.

    return SmoothSeasonalStateSpaceModel(
        num_timesteps=num_timesteps,
        period=self.period,
        frequency_multipliers=self.frequency_multipliers,
        initial_state_prior=initial_state_prior,
        initial_step=initial_step,
        **param_map)
