# 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.
# ============================================================================
"""The Truncated Normal distribution class."""

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

# Dependency imports
import numpy as np

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

from tensorflow_probability.python.distributions import distribution
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import distribution_util
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import reparameterization
from tensorflow_probability.python.internal import special_math
from tensorflow.python.ops import random_ops  # pylint: disable=g-direct-tensorflow-import


__all__ = [
    "TruncatedNormal",
]


class TruncatedNormal(distribution.Distribution):
  """The Truncated Normal distribution.

  #### Mathematical details

  The truncated normal is a normal distribution bounded between `low`
  and `high` (the pdf is 0 outside these bounds and renormalized).

  Samples from this distribution are differentiable with respect to `loc`,
  `scale` as well as the bounds, `low` and `high`, i.e., this
  implementation is fully reparameterizeable.

  For more details, see [here](
  https://en.wikipedia.org/wiki/Truncated_normal_distribution).

  ### Mathematical Details

  The probability density function (pdf) of this distribution is:
  ```none
    pdf(x; loc, scale, low, high) =
        { (2 pi)**(-0.5) exp(-0.5 y**2) / (scale * z) for low <= x <= high
        { 0                                    otherwise
    y = (x - loc)/scale
    z = NormalCDF((high - loc) / scale) - NormalCDF((lower - loc) / scale)
  ```

  where:

  * `NormalCDF` is the cumulative density function of the Normal distribution
    with 0 mean and unit variance.

  This is a scalar distribution so the event shape is always scalar and the
  dimensions of the parameters defined the batch_shape.

  #### Examples
  ```python

  tfd = tfp.distributions
  # Define a batch of two scalar TruncatedNormals which modes at 0. and 1.0
  dist = tfd.TruncatedNormal(loc=[0., 1.], scale=1.0,
                             low=[-1., 0.],
                             high=[1., 1.])

  # Evaluate the pdf of the distributions at 0.5 and 0.8 respectively returning
  # a 2-vector tensor.
  dist.prob([0.5, 0.8])

  # Get 3 samples, returning a 3 x 2 tensor.
  dist.sample([3])
  ```
  """

  def __init__(self,
               loc,
               scale,
               low,
               high,
               validate_args=False,
               allow_nan_stats=True,
               name="TruncatedNormal"):
    """Construct TruncatedNormal.

    All parameters of the distribution will be broadcast to the same shape,
    so the resulting distribution will have a batch_shape of the broadcast
    shape of all parameters.

    Args:
      loc: Floating point tensor; the mean of the normal distribution(s) (
        note that the mean of the resulting distribution will be different
        since it is modified by the bounds).
      scale: Floating point tensor; the std deviation of the normal
        distribution(s).
      low: `float` `Tensor` representing lower bound of the distribution's
        support. Must be such that `low < high`.
      high: `float` `Tensor` representing upper bound of the distribution's
        support. Must be such that `low < high`.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked at run-time.
      allow_nan_stats: Python `bool`, default `True`. When `True`,
        statistics (e.g., mean, mode, variance) use the value "`NaN`" to
        indicate the result is undefined. When `False`, an exception is raised
        if one or more of the statistic's batch members are undefined.
      name: Python `str` name prefixed to Ops created by this class.
    """
    parameters = dict(locals())
    with tf.name_scope(name) as name:
      dtype = dtype_util.common_dtype([loc, scale, low, high], tf.float32)
      loc = tf.convert_to_tensor(value=loc, name="loc", dtype=dtype)
      scale = tf.convert_to_tensor(value=scale, name="scale", dtype=dtype)
      low = tf.convert_to_tensor(value=low, name="low", dtype=dtype)
      high = tf.convert_to_tensor(value=high, name="high", dtype=dtype)
      dtype_util.assert_same_float_dtype([loc, scale, low, high])

      self._broadcast_batch_shape = distribution_util.get_broadcast_shape(
          loc, scale, low, high)

      # Broadcast all parameters to the same shape
      broadcast_ones = tf.ones(shape=self._broadcast_batch_shape,
                               dtype=scale.dtype)
      self._scale = scale * broadcast_ones
      self._loc = loc * broadcast_ones
      self._low = low * broadcast_ones
      self._high = high * broadcast_ones

      with tf.control_dependencies([self._validate()] if validate_args else []):
        self._loc = tf.identity(self._loc)

    super(TruncatedNormal, self).__init__(
        dtype=dtype,
        # This distribution is fully reparameterized. loc, scale have straight
        # through gradients. The gradients for the bounds are implemented using
        # custom derived expressions based on implicit gradients.
        # For the special case of lower bound zero and a positive upper bound
        # an equivalent expression can also be found in Sec 9.1.1.
        # of https://arxiv.org/pdf/1806.01851.pdf. The implementation here
        # handles arbitrary bounds.
        reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        parameters=parameters,
        graph_parents=[loc, scale, low, high],
        name=name)

  def _validate(self):
    vops = [
        assert_util.assert_positive(self._scale),
        assert_util.assert_positive(self._high - self._low),
        assert_util.assert_finite(self._low, message="Lower bound not finite"),
        assert_util.assert_finite(self._high, message="Upper bound not finite"),
        assert_util.assert_finite(self._loc, message="Loc not finite"),
        assert_util.assert_finite(self._scale, message="scale not finite"),
    ]
    return tf.group(*vops, name="ValidationOps")

  @property
  def _standardized_low(self):
    return (self._low - self._loc) / self._scale

  @property
  def _standardized_high(self):
    return (self._high - self._loc) / self._scale

  @property
  def _normalizer(self):
    return (special_math.ndtr(self._standardized_high) -
            special_math.ndtr(self._standardized_low))

  def _normal_pdf(self, x):
    return 1. / np.sqrt(2 * np.pi) * tf.exp(-0.5 * tf.square(x))

  @staticmethod
  def _param_shapes(sample_shape):
    # All parameters are of the same shape
    shape = tf.convert_to_tensor(value=sample_shape, dtype=tf.int32)
    return {"loc": shape,
            "scale": shape,
            "high": shape,
            "low": shape}

  @classmethod
  def _params_event_ndims(cls):
    return dict(loc=0, scale=0, low=0, high=0)

  @property
  def loc(self):
    return self._loc

  @property
  def scale(self):
    """Distribution parameter for the scale."""
    return self._scale

  @property
  def low(self):
    return self._low

  @property
  def high(self):
    return self._high

  def _batch_shape_tensor(self):
    # All the parameters are broadcast the same shape during construction.
    return tf.shape(input=self.loc)

  def _batch_shape(self):
    # All the parameters are broadcast the same shape during construction.
    return self.loc.shape

  def _event_shape_tensor(self):
    return tf.constant([], dtype=tf.int32)

  def _event_shape(self):
    return tf.TensorShape([])

  def _sample_n(self, n, seed=None):
    sample_and_batch_shape = tf.concat([[n], self.batch_shape_tensor()], 0)
    flat_batch_and_sample_shape = tf.stack(
        [tf.reduce_prod(input_tensor=self.batch_shape_tensor()), n])

    # In order to be reparameterizable we sample on the truncated_normal of
    # unit variance and mean and scale (but with the standardized
    # truncation bounds).

    @tf.custom_gradient
    def _std_samples_with_gradients(lower, upper):
      """Standard truncated Normal with gradient support for low, high."""
      # Note: Unlike the convention in tf_probability,
      # parameterized_truncated_normal returns a tensor with the final dimension
      # being the sample dimension.
      std_samples = random_ops.parameterized_truncated_normal(
          shape=flat_batch_and_sample_shape,
          means=0.0,
          stddevs=1.0,
          minvals=lower,
          maxvals=upper,
          dtype=self.dtype,
          seed=seed)

      def grad(dy):
        """Computes a derivative for the min and max parameters.

        This function implements the derivative wrt the truncation bounds, which
        get blocked by the sampler. We use a custom expression for numerical
        stability instead of automatic differentiation on CDF for implicit
        gradients.

        Args:
          dy: output gradients

        Returns:
           The standard normal samples and the gradients wrt the upper
           bound and lower bound.
        """
        # std_samples has an extra dimension (the sample dimension), expand
        # lower and upper so they broadcast along this dimension.
        # See note above regarding parameterized_truncated_normal, the sample
        # dimension is the final dimension.
        lower_broadcast = lower[..., tf.newaxis]
        upper_broadcast = upper[..., tf.newaxis]

        cdf_samples = ((special_math.ndtr(std_samples) -
                        special_math.ndtr(lower_broadcast)) /
                       (special_math.ndtr(upper_broadcast)
                        - special_math.ndtr(lower_broadcast)))

        # tiny, eps are tolerance parameters to ensure we stay away from giving
        # a zero arg to the log CDF expression.

        tiny = np.finfo(dtype_util.as_numpy_dtype(self.dtype)).tiny
        eps = np.finfo(dtype_util.as_numpy_dtype(self.dtype)).eps
        cdf_samples = tf.clip_by_value(cdf_samples, tiny, 1 - eps)

        du = tf.exp(0.5 * (std_samples**2 - upper_broadcast**2) +
                    tf.math.log(cdf_samples))
        dl = tf.exp(0.5 * (std_samples**2 - lower_broadcast**2) +
                    tf.math.log1p(-cdf_samples))

        # Reduce the gradient across the samples
        grad_u = tf.reduce_sum(input_tensor=dy * du, axis=-1)
        grad_l = tf.reduce_sum(input_tensor=dy * dl, axis=-1)
        return [grad_l, grad_u]

      return std_samples, grad

    std_samples = _std_samples_with_gradients(
        tf.reshape(self._standardized_low, [-1]),
        tf.reshape(self._standardized_high, [-1]))

    # The returned shape is [flat_batch x n]
    std_samples = tf.transpose(a=std_samples, perm=[1, 0])

    std_samples = tf.reshape(std_samples, sample_and_batch_shape)
    samples = (std_samples * tf.expand_dims(self._scale, axis=0) +
               tf.expand_dims(self._loc, axis=0))

    return samples

  def _log_prob(self, x):
    log_prob = -(0.5 *
                 ((x - self.loc) / self.scale)**2 + 0.5 * np.log(2. * np.pi) +
                 tf.math.log(self.scale * self._normalizer))
    # p(x) is 0 outside the bounds.
    neg_inf = tf.math.log(tf.zeros_like(log_prob))
    bounded_log_prob = tf1.where(
        tf.logical_or(x > self._high, x < self._low), neg_inf, log_prob)
    return bounded_log_prob

  def _cdf(self, x):
    cdf_in_support = ((special_math.ndtr((x - self.loc) / self.scale)
                       -  special_math.ndtr(self._standardized_low))
                      / self._normalizer)
    return tf.clip_by_value(cdf_in_support, 0., 1.)

  def _entropy(self):
    return (
        tf.math.log(
            np.sqrt(2. * np.pi * np.e) * self.scale * self._normalizer) +
        (self._standardized_low * self._normal_pdf(self._standardized_low) -
         self._standardized_high * self._normal_pdf(self._standardized_high)) /
        (2. * self._normalizer))

  def _mean(self):
    return (self.loc +
            self._scale * ((self._normal_pdf(self._standardized_low) -
                            self._normal_pdf(self._standardized_high))
                           / self._normalizer))

  def _mode(self):
    # mode = { loc:         for low <= loc <= high
    #          low: for loc < low
    #          high: for loc > high
    #        }
    return tf.clip_by_value(self.loc, self.low, self.high)

  def _variance(self):
    var = (tf.square(self.scale) *
           (1. + (self._standardized_low * self._normal_pdf(
               self._standardized_low) -
                  self._standardized_high * self._normal_pdf(
                      self._standardized_high)) / self._normalizer -
            tf.square((self._normal_pdf(self._standardized_low) -
                       self._normal_pdf(self._standardized_high))
                      / self._normalizer)))
    return var
