# 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.
# ============================================================================
"""Tests for tensorflow_probability.python.sts.structural_time_series."""

import numpy as np
import tensorflow.compat.v1 as tf1
import tensorflow.compat.v2 as tf
import tensorflow_probability as tfp
from tensorflow_probability.python import distributions as tfd
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import samplers
from tensorflow_probability.python.internal import test_util
from tensorflow_probability.python.sts import Autoregressive
from tensorflow_probability.python.sts import AutoregressiveIntegratedMovingAverage
from tensorflow_probability.python.sts import DynamicLinearRegression
from tensorflow_probability.python.sts import LinearRegression
from tensorflow_probability.python.sts import LocalLevel
from tensorflow_probability.python.sts import LocalLinearTrend
from tensorflow_probability.python.sts import Seasonal
from tensorflow_probability.python.sts import SemiLocalLinearTrend
from tensorflow_probability.python.sts import SmoothSeasonal
from tensorflow_probability.python.sts import SparseLinearRegression
from tensorflow_probability.python.sts import Sum
from tensorflow_probability.python.sts.internal import util as sts_util


class _StructuralTimeSeriesTests(object):

  def test_broadcast_batch_shapes(self):
    seed = test_util.test_seed(sampler_type='stateless')

    batch_shape = [3, 1, 4]
    partial_batch_shape = [2, 1]
    expected_broadcast_batch_shape = [3, 2, 4]

    # Build a model where parameters have different batch shapes.
    partial_batch_loc = self._build_placeholder(
        np.random.randn(*partial_batch_shape))
    full_batch_loc = self._build_placeholder(
        np.random.randn(*batch_shape))

    partial_scale_prior = tfd.LogNormal(
        loc=partial_batch_loc, scale=tf.ones_like(partial_batch_loc))
    full_scale_prior = tfd.LogNormal(
        loc=full_batch_loc, scale=tf.ones_like(full_batch_loc))
    loc_prior = tfd.Normal(loc=partial_batch_loc,
                           scale=tf.ones_like(partial_batch_loc))

    linear_trend = LocalLinearTrend(level_scale_prior=full_scale_prior,
                                    slope_scale_prior=full_scale_prior,
                                    initial_level_prior=loc_prior,
                                    initial_slope_prior=loc_prior)
    seasonal = Seasonal(num_seasons=3,
                        drift_scale_prior=partial_scale_prior,
                        initial_effect_prior=loc_prior)
    model = Sum([linear_trend, seasonal],
                observation_noise_scale_prior=partial_scale_prior)
    param_samples = [p.prior.sample(seed=seed) for p in model.parameters]
    ssm = model.make_state_space_model(num_timesteps=2,
                                       param_vals=param_samples)

    # Test that the model's batch shape matches the SSM's batch shape,
    # and that they both match the expected broadcast shape.
    self.assertAllEqual(model.batch_shape, ssm.batch_shape)

    (model_batch_shape_tensor_,
     ssm_batch_shape_tensor_) = self.evaluate((model.batch_shape_tensor(),
                                               ssm.batch_shape_tensor()))
    self.assertAllEqual(model_batch_shape_tensor_, ssm_batch_shape_tensor_)
    self.assertAllEqual(model_batch_shape_tensor_,
                        expected_broadcast_batch_shape)

  def test_addition_raises_error_with_no_observed_time_series(self):
    c1 = tfp.sts.LocalLevel(level_scale_prior=tfd.Normal(0., 1.),
                            initial_level_prior=tfd.Normal(0., 1.))
    c2 = tfp.sts.LocalLevel(level_scale_prior=tfd.Normal(0., 0.1),
                            initial_level_prior=tfd.Normal(1., 2.))
    with self.assertRaisesRegex(
        ValueError, 'Could not automatically create a `Sum` component'):
      c1 + c2  # pylint: disable=pointless-statement

  def test_adding_two_sums(self):
    observed_time_series = self._build_placeholder([1., 2., 3., 4., 5.])
    s1 = tfp.sts.Sum(
        [tfp.sts.LocalLevel(observed_time_series=observed_time_series)],
        observed_time_series=observed_time_series)
    s2 = tfp.sts.Sum(
        [tfp.sts.LocalLinearTrend(observed_time_series=observed_time_series)],
        observed_time_series=observed_time_series)
    s3 = s1 + s2
    self.assertLen(s3.components, 2)

    seed = test_util.test_seed(sampler_type='stateless')
    def observation_noise_scale_prior_sample(s):
      return s.parameters[0].prior.sample(seed=seed)
    self.assertAllEqual(observation_noise_scale_prior_sample(s3),
                        observation_noise_scale_prior_sample(s1))
    self.assertAllEqual(observation_noise_scale_prior_sample(s3),
                        observation_noise_scale_prior_sample(s2))

    self.assertAllEqual(s3.constant_offset, s1.constant_offset)
    self.assertAllEqual(s3.constant_offset, s2.constant_offset)

    with self.assertRaisesRegex(
        ValueError, 'Cannot add Sum components'):
      s1.copy(observed_time_series=3 * observed_time_series) + s2  # pylint: disable=expression-not-assigned

    with self.assertRaisesRegex(
        ValueError, 'Cannot add Sum components'):
      s1.copy(constant_offset=4.) + s2  # pylint: disable=expression-not-assigned

    with self.assertRaisesRegex(
        ValueError, 'Cannot add Sum components'):
      s1.copy(observation_noise_scale_prior=tfd.Normal(  # pylint: disable=expression-not-assigned
          self._build_placeholder(0.), self._build_placeholder(1.))) + s2

  def _build_placeholder(self, ndarray, dtype=None):
    """Convert a numpy array to a TF placeholder.

    Args:
      ndarray: any object convertible to a numpy array via `np.asarray()`.
      dtype: optional `dtype`; if not specified, defaults to `self.dtype`.

    Returns:
      placeholder: a TensorFlow `placeholder` with default value given by the
      provided `ndarray`, dtype given by `self.dtype`, and shape specified
      statically only if `self.use_static_shape` is `True`.
    """

    if dtype is None:
      dtype = self.dtype

    ndarray = np.asarray(ndarray).astype(dtype)
    return tf1.placeholder_with_default(
        ndarray, shape=ndarray.shape if self.use_static_shape else None)


@test_util.test_all_tf_execution_regimes
class StructuralTimeSeriesTestsStaticShape32(
    _StructuralTimeSeriesTests, test_util.TestCase):
  dtype = np.float32
  use_static_shape = True


@test_util.test_all_tf_execution_regimes
class StructuralTimeSeriesTestsDynamicShape32(
    _StructuralTimeSeriesTests, test_util.TestCase):
  dtype = np.float32
  use_static_shape = False


@test_util.test_all_tf_execution_regimes
class StructuralTimeSeriesTestsStaticShape64(
    _StructuralTimeSeriesTests, test_util.TestCase):
  dtype = np.float64
  use_static_shape = True


class _StsTestHarness(object):

  def test_state_space_model(self):
    seed = test_util.test_seed(sampler_type='stateless')
    model = self._build_sts()

    dummy_param_vals = [p.prior.sample(seed=seed) for p in model.parameters]
    initial_state_prior = tfd.MultivariateNormalDiag(
        loc=-2. + tf.zeros([model.latent_size]),
        scale_diag=3. * tf.ones([model.latent_size]))

    mask = tf.convert_to_tensor(
        [False, True, True, False, False, False, False, True, False, False],
        dtype=tf.bool)

    # Verify we build the LGSSM without errors.
    ssm = model.make_state_space_model(
        num_timesteps=10,
        param_vals=dummy_param_vals,
        initial_state_prior=initial_state_prior,
        initial_step=1,
        mask=mask)

    # Verify that the child class passes the initial step, prior, and mask
    # arguments through to the SSM.
    self.assertEqual(self.evaluate(ssm.initial_step), 1)
    self.assertEqual(ssm.initial_state_prior, initial_state_prior)
    self.assertAllEqual(ssm.mask, mask)

    # Verify the model has the correct latent size.
    self.assertEqual(
        self.evaluate(
            tf.convert_to_tensor(
                ssm.latent_size_tensor())),
        model.latent_size)

    # Verify that the SSM tracks its parameters.
    seed = test_util.test_seed(sampler_type='stateless')
    observed_time_series = self.evaluate(
        samplers.normal([10, 1], seed=seed))
    ssm_copy = ssm.copy(name='copied_ssm')
    self.assertAllClose(*self.evaluate((
        ssm.log_prob(observed_time_series),
        ssm_copy.log_prob(observed_time_series))))

  def test_log_joint(self):
    seed = test_util.test_seed(sampler_type='stateless')
    model = self._build_sts()
    num_timesteps = 5

    # simple case: single observation, and all params unbatched
    log_joint_fn = model.joint_log_prob(
        observed_time_series=np.float32(
            np.random.standard_normal([num_timesteps, 1])))
    lp_seed1, lp_seed2 = tfp.random.split_seed(seed, n=2)
    seeds = tfp.random.split_seed(lp_seed1, n=len(model.parameters))
    lp = self.evaluate(
        log_joint_fn(*[p.prior.sample(seed=seed) for seed, p in zip(
            seeds, model.parameters)]))
    self.assertEqual(tf.TensorShape([]), lp.shape)

    # more complex case: y has sample and batch shapes, some parameters
    # have partial batch shape.
    full_batch_shape = [2, 3]
    partial_batch_shape = [3]
    sample_shape = [4]
    log_joint_fn = model.joint_log_prob(
        observed_time_series=np.float32(
            np.random.standard_normal(sample_shape + full_batch_shape +
                                      [num_timesteps, 1])))

    # We alternate full_batch_shape, partial_batch_shape in sequence so that in
    # a model with only one parameter, that parameter is constructed with full
    # batch shape.
    seeds = tfp.random.split_seed(lp_seed2, n=len(model.parameters))
    batch_shaped_parameters_ = self.evaluate([
        p.prior.sample(sample_shape=full_batch_shape if (i % 2 == 0)
                       else partial_batch_shape, seed=seed)
        for (i, (seed, p)) in enumerate(zip(seeds, model.parameters))])

    lp = self.evaluate(log_joint_fn(*batch_shaped_parameters_))
    self.assertEqual(tf.TensorShape(full_batch_shape), lp.shape)

    # Check that the log joint function also supports parameters passed
    # as kwargs.
    parameters_by_name_ = {
        p.name: v for (p, v) in zip(model.parameters, batch_shaped_parameters_)}
    lp_with_kwargs = self.evaluate(log_joint_fn(**parameters_by_name_))
    self.assertAllClose(lp, lp_with_kwargs)

  def test_constant_series_does_not_induce_constant_prior(self):
    observed_time_series = np.array([1.0, 1.0, 1.0]).astype(np.float32)
    model = self._build_sts(observed_time_series=observed_time_series)
    for parameter in model.parameters:
      param_samples = self.evaluate(
          parameter.prior.sample(
              30, seed=test_util.test_seed(sampler_type='stateless')))
      self.assertAllGreater(tf.math.reduce_std(param_samples), 0.)

  def test_log_joint_with_missing_observations(self):
    # Test that this component accepts MaskedTimeSeries inputs. In most
    # cases, it is sufficient that the component accesses only
    # `empirical_statistics(observed_time_series)`.
    # TODO(b/139483802): De-flake this test when run with --vary_seed.
    seed = test_util.test_seed(hardcoded_seed=123, sampler_type='stateless')
    observed_time_series = np.array(
        [1.0, 2.0, -1000., 0.4, np.nan, 1000., 4.2, np.inf]).astype(np.float32)
    observation_mask = np.array(
        [False, False, True, False, True, True, False, True]).astype(np.bool_)
    masked_time_series = tfp.sts.MaskedTimeSeries(observed_time_series,
                                                  is_missing=observation_mask)

    model = self._build_sts(observed_time_series=masked_time_series)

    log_joint_fn = model.joint_log_prob(
        observed_time_series=masked_time_series)
    seeds = tfp.random.split_seed(seed, n=len(model.parameters))
    lp = self.evaluate(
        log_joint_fn(*[p.prior.sample(seed=seed) for seed, p in zip(
            seeds, model.parameters)]))

    self.assertEqual(tf.TensorShape([]), lp.shape)
    self.assertTrue(np.isfinite(lp))

  def test_prior_sample(self):
    model = self._build_sts()
    ys, param_samples = model.prior_sample(
        num_timesteps=5, params_sample_shape=[2], trajectories_sample_shape=[3],
        seed=test_util.test_seed(sampler_type='stateless'))

    self.assertAllEqual(ys.shape, [3, 2, 5, 1])
    self.assertEqual(len(param_samples), len(model.parameters))
    for i in range(len(param_samples)):
      sampled = param_samples[i]
      param = model.parameters[i]
      self.assertAllEqual(sampled.shape, [
          2,
      ] + param.prior.batch_shape.as_list() + param.prior.event_shape.as_list())

  def test_joint_distribution_name(self):
    model = self._build_sts(name='foo')
    jd = model.joint_distribution(num_timesteps=5)
    self.assertEqual('foo', jd.name)

  def test_joint_distribution_log_prob(self):
    model = self._build_sts(
        # Dummy series to build the model with float64 priors. Working in
        # float64 minimizes numeric inconsistencies between log-prob
        # implementations.
        observed_time_series=np.float64([1., 0., 1., 0.]))

    jd_no_trajectory_shape = model.joint_distribution(num_timesteps=11)