# Copyright 2021 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 trainable distributions and bijectors."""

from absl.testing import parameterized

import tensorflow.compat.v1 as tf1
import tensorflow.compat.v2 as tf
import tensorflow_probability as tfp

from tensorflow_probability.python.internal import samplers
from tensorflow_probability.python.internal import test_util

tfb = tfp.bijectors
tfd = tfp.distributions


@test_util.test_graph_and_eager_modes
class TestTrainableDistributionsAndBijectors(test_util.TestCase):

  @parameterized.named_parameters(
      ('Normal', tfd.Normal, []),
      ('Gamma', tfd.Gamma, [2, 3]),
      ('Dirichlet', tfd.Dirichlet, [5]),
      ('MultivariateNormalTriL', tfd.MultivariateNormalTriL, [4]),
      ('Kumaraswamy', tfd.Kumaraswamy, [2, 1]),
      ('Poisson', tfd.Poisson, []),
      ('Geometric', tfd.Geometric, [1, 1]))
  def test_trainable_distributions(self, cls, batch_and_event_shape):
    distribution = tfp.experimental.util.make_trainable(
        cls,
        batch_and_event_shape=batch_and_event_shape,
        seed=test_util.test_seed(),
        validate_args=True)
    self.evaluate([v.initializer for v in distribution.trainable_variables])

    x = self.evaluate(distribution.sample())
    self.assertAllEqual(x.shape, batch_and_event_shape)

    # Verify expected number of trainable variables.
    self.assertLen(distribution.trainable_variables,
                   len([k for k, p in cls.parameter_properties().items()
                        if p.is_preferred]))

    # Verify gradients to all parameters.
    with tf.GradientTape() as tape:
      lp = distribution.log_prob(x)
    grad = tape.gradient(lp, distribution.trainable_variables)
    self.assertAllNotNone(grad)

  @parameterized.named_parameters(
      ('Scale', tfb.Scale, [2]),
      ('ScaleMatvecTriL', tfb.ScaleMatvecTriL, [4]),
      ('FillScaleTriL', tfb.FillScaleTriL, [2, 6]),
      ('Softplus', tfb.Softplus, [2]),
      ('Identity', tfb.Identity, []))
  def test_trainable_bijectors(self, cls, batch_and_event_shape):
    bijector = tfp.experimental.util.make_trainable(
        cls,
        batch_and_event_shape=batch_and_event_shape,
        seed=test_util.test_seed(),
        validate_args=True)
    if bijector.trainable_variables:
      self.evaluate([v.initializer for v in bijector.trainable_variables])

    # Verify expected number of trainable variables.
    self.assertLen(bijector.trainable_variables,
                   len([k for k, p in cls.parameter_properties().items()
                        if p.is_tensor and p.is_preferred]))

    # Verify gradients to all parameters.
    x = self.evaluate(samplers.normal(batch_and_event_shape,
                                      seed=test_util.test_seed()))
    with tf.GradientTape() as tape:
      y = bijector.forward(x)
    grad = tape.gradient(y, bijector.trainable_variables)
    self.assertAllNotNone(grad)

    # Verify that the round trip doesn't broadcast, i.e., that it preserves
    # batch_and_event_shape.
    self.assertAllCloseNested(
        x,
        bijector.inverse(tf.identity(y)),  # Disable bijector cache.
        atol=1e-2)

  def test_can_specify_initial_values(self):
    distribution = tfp.experimental.util.make_trainable(
        tfd.Normal,
        initial_parameters={'scale': 1e-4},
        batch_and_event_shape=[3],
        seed=test_util.test_seed(),
        validate_args=True)
    self.evaluate([v.initializer for v in distribution.trainable_variables])
    self.assertAllClose(
        # Convert to tensor since we can't eval a TransformedVariable directly.
        tf.convert_to_tensor(distribution.scale),
        [1e-4, 1e-4, 1e-4])

  def test_can_specify_callable_initializer(self):

    def uniform_initializer(_, shape, dtype, seed,
                            constraining_bijector):
      return constraining_bijector.forward(
          tf.random.stateless_uniform(
              constraining_bijector.inverse_event_shape_tensor(shape),
              dtype=dtype, seed=seed))

    distribution = tfp.experimental.util.make_trainable(
        tfd.Normal,
        initial_parameters=uniform_initializer,
        batch_and_event_shape=[3, 4],
        seed=test_util.test_seed(),
        validate_args=True)
    for v in distribution.trainable_variables:
      self.evaluate(v.initializer)
      # Unconstrained variables should be uniform in [0, 1].
      self.assertAllGreater(v, 0.)
      self.assertAllLess(v, 1.)

  def test_initialization_is_deterministic_with_seed(self):
    distribution1 = tfp.experimental.util.make_trainable(
        tfd.Normal,
        seed=test_util.test_seed(sampler_type='stateless'),
        validate_args=True)
    self.evaluate([v.initializer for v in distribution1.trainable_variables])

    distribution2 = tfp.experimental.util.make_trainable(
        tfd.Normal,
        seed=test_util.test_seed(sampler_type='stateless'),
        validate_args=True)
    self.evaluate([v.initializer for v in distribution2.trainable_variables])

    self.assertAllEqual(distribution1.mean(), distribution2.mean())
    self.assertAllEqual(distribution1.stddev(), distribution2.stddev())

  def test_can_specify_parameter_dtype(self):
    distribution = tfp.experimental.util.make_trainable(
        tfd.Normal,
        initial_parameters={'loc': 17.},
        parameter_dtype=tf.float64,
        seed=test_util.test_seed(sampler_type='stateless'),
        validate_args=True)
    self.evaluate([v.initializer for v in distribution.trainable_variables])
    self.assertEqual(distribution.loc.dtype, tf.float64)
    self.assertEqual(distribution.scale.dtype, tf.float64)
    self.assertEqual(distribution.sample().dtype, tf.float64)

  def test_can_specify_fixed_values(self):
    distribution = tfp.experimental.util.make_trainable(
        tfd.WishartTriL,
        batch_and_event_shape=[2, 2],
        seed=test_util.test_seed(sampler_type='stateless'),
        validate_args=True,
        df=3)
    self.evaluate([v.initializer for v in distribution.trainable_variables])
    self.assertAllClose(distribution.df, 3.)
    self.assertLen(distribution.trainable_variables, 1)
    self.assertAllEqual(distribution.sample().shape, [2, 2])

  def test_dynamic_shape(self):
    batch_and_event_shape = tf1.placeholder_with_default(
        [4, 3, 2], shape=None)
    distribution = tfp.experimental.util.make_trainable(
        tfd.Normal,
        batch_and_event_shape=batch_and_event_shape,
        seed=test_util.test_seed(),
        validate_args=True)
    self.evaluate([v.initializer for v in distribution.trainable_variables])
    x = self.evaluate(distribution.sample())
    self.assertAllEqual(x.shape, batch_and_event_shape)

  def test_docstring_example_normal(self):
    samples = [4.57, 6.37, 5.93, 7.98, 2.03, 3.59, 8.55, 3.45, 5.06, 6.44]
    model = tfp.experimental.util.make_trainable(
        tfd.Normal,
        seed=test_util.test_seed(sampler_type='stateless'))
    losses = tfp.math.minimize(
        lambda: -model.log_prob(samples),
        optimizer=tf.optimizers.Adam(0.1),
        num_steps=200)
    self.evaluate(tf1.global_variables_initializer())
    self.evaluate(losses)
    self.assertAllClose(tf.reduce_mean(samples), model.mean(), atol=2.0)
    self.assertAllClose(tf.math.reduce_std(samples), model.stddev(), atol=2.0)


if __name__ == '__main__':
  test_util.main()
