# Copyright 2019 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.
# ============================================================================
"""Scalar Tests."""

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


# Dependency imports
from absl.testing import parameterized
import numpy as np
import tensorflow.compat.v2 as tf

import tensorflow_probability as tfp
from tensorflow_probability.python import bijectors as tfb
from tensorflow_probability.python.bijectors import bijector_test_util
from tensorflow_probability.python.internal import test_util


@test_util.test_all_tf_execution_regimes
class ScaleBijectorTest(test_util.TestCase, parameterized.TestCase):
  """Tests correctness of the Y = scale @ x transformation."""

  def testName(self):
    # scale corresponds to 1.
    bijector = tfb.Scale(scale=-1.)
    self.assertStartsWith(bijector.name, 'scale')

  @parameterized.named_parameters(
      dict(testcase_name='float32', dtype=np.float32),
      dict(testcase_name='float64', dtype=np.float64),
  )
  def testTinyScale(self, dtype):
    log_scale = tf.cast(-2000., dtype)
    x = tf.cast(1., dtype)
    scale = tf.math.exp(log_scale)
    fldj_linear = tfb.Scale(scale=scale).forward_log_det_jacobian(
        x, event_ndims=0)
    fldj_log = tfb.Scale(log_scale=log_scale).forward_log_det_jacobian(
        x, event_ndims=0)
    fldj_linear_, fldj_log_ = self.evaluate([fldj_linear, fldj_log])
    # Using the linear scale will saturate to 0, and produce bad log-det
    # Jacobians.
    self.assertNotEqual(fldj_linear_, fldj_log_)
    self.assertAllClose(-2000., fldj_log_)

  @parameterized.named_parameters(
      dict(testcase_name='static_float32', is_static=True, dtype=np.float32),
      dict(testcase_name='static_float64', is_static=True, dtype=np.float64),
      dict(testcase_name='dynamic_float32', is_static=False, dtype=np.float32),
      dict(testcase_name='dynamic_float64', is_static=False, dtype=np.float64),
  )
  def testNoBatchScale(self, is_static, dtype):
    bijector = tfb.Scale(scale=dtype(2.))
    x = self.maybe_static(np.array([1., 2, 3], dtype), is_static)
    self.assertAllClose([2., 4, 6], bijector.forward(x))
    self.assertAllClose([.5, 1, 1.5], bijector.inverse(x))
    self.assertAllClose(
        -np.log(2.),
        bijector.inverse_log_det_jacobian(x, event_ndims=0))

  @parameterized.named_parameters(
      dict(testcase_name='static_float32', is_static=True, dtype=np.float32),
      dict(testcase_name='static_float64', is_static=True, dtype=np.float64),
      dict(testcase_name='dynamic_float32', is_static=False, dtype=np.float32),
      dict(testcase_name='dynamic_float64', is_static=False, dtype=np.float64),
  )
  def testBatchScale(self, is_static, dtype):
    # Batched scale
    bijector = tfb.Scale(scale=dtype([2., 3.]))
    x = self.maybe_static(np.array([1.], dtype=dtype), is_static)
    self.assertAllClose([2., 3.], bijector.forward(x))
    self.assertAllClose([0.5, 1./3.], bijector.inverse(x))
    self.assertAllClose(
        [-np.log(2.), -np.log(3.)],
        bijector.inverse_log_det_jacobian(x, event_ndims=0))

  @parameterized.named_parameters(
      dict(testcase_name='float32', dtype=np.float32),
      dict(testcase_name='float64', dtype=np.float64),
  )
  def testScalarCongruency(self, dtype):
    bijector = tfb.Scale(scale=dtype(0.42))
    bijector_test_util.assert_scalar_congruency(
        bijector,
        lower_x=dtype(-2.),
        upper_x=dtype(2.),
        eval_func=self.evaluate)

  @parameterized.named_parameters(
      dict(testcase_name='float32', dtype=np.float32),
      dict(testcase_name='float64', dtype=np.float64),
  )
  def testScalarCongruencyLogScale(self, dtype):
    bijector = tfb.Scale(log_scale=dtype(np.log(0.42)))
    bijector_test_util.assert_scalar_congruency(
        bijector,
        lower_x=dtype(-2.),
        upper_x=dtype(2.),
        eval_func=self.evaluate)

  @test_util.jax_disable_variable_test
  @test_util.numpy_disable_gradient_test
  def testVariableGradients(self):
    b = tfb.Scale(scale=tf.Variable(2.))

    with tf.GradientTape() as tape:
      y = b.forward(.1)
    self.assertAllNotNone(tape.gradient(y, b.trainable_variables))

  @test_util.numpy_disable_gradient_test
  def testNonVariableGradients(self):
    def _func(scale):
      b = tfb.Scale(scale=scale)
      return b.forward(.1)

    value, grad = tfp.math.value_and_gradient(_func, [2.])
    self.assertAllNotNone([value, grad])
    self.assertNotAllZero(grad)

  def testImmutableScaleAssertion(self):
    with self.assertRaisesOpError('Argument `scale` must be non-zero'):
      b = tfb.Scale(scale=0., validate_args=True)
      _ = self.evaluate(b.forward(1.))

  def testVariableScaleAssertion(self):
    v = tf.Variable(0.)
    self.evaluate(v.initializer)
    with self.assertRaisesOpError('Argument `scale` must be non-zero'):
      b = tfb.Scale(scale=v, validate_args=True)
      _ = self.evaluate(b.forward(1.))

  def testModifiedVariableScaleAssertion(self):
    v = tf.Variable(1.)
    self.evaluate(v.initializer)
    b = tfb.Scale(scale=v, validate_args=True)
    with self.assertRaisesOpError('Argument `scale` must be non-zero'):
      with tf.control_dependencies([v.assign(0.)]):
        _ = self.evaluate(b.forward(1.))


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