# 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 Negative Binomial distribution class."""

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

from tensorflow_probability.python.internal.backend.numpy.compat import v2 as tf

from tensorflow_probability.substrates.numpy import math as tfp_math
from tensorflow_probability.substrates.numpy.bijectors import sigmoid as sigmoid_bijector
from tensorflow_probability.substrates.numpy.distributions import distribution
from tensorflow_probability.substrates.numpy.internal import assert_util
from tensorflow_probability.substrates.numpy.internal import distribution_util
from tensorflow_probability.substrates.numpy.internal import dtype_util
from tensorflow_probability.substrates.numpy.internal import parameter_properties
from tensorflow_probability.substrates.numpy.internal import prefer_static as ps
from tensorflow_probability.python.internal import reparameterization
from tensorflow_probability.substrates.numpy.internal import samplers
from tensorflow_probability.substrates.numpy.internal import tensor_util
from tensorflow_probability.substrates.numpy.util.deferred_tensor import DeferredTensor


class NegativeBinomial(distribution.Distribution):
  """NegativeBinomial distribution.

  The NegativeBinomial distribution is related to the experiment of performing
  Bernoulli trials in sequence. Given a Bernoulli trial with probability `p` of
  success, the NegativeBinomial distribution represents the distribution over
  the number of successes `s` that occur until we observe `f` failures.

  The probability mass function (pmf) is,

  ```none
  pmf(s; f, p) = p**s (1 - p)**f / Z
  Z = s! (f - 1)! / (s + f - 1)!
  ```

  where:
  * `total_count = f`,
  * `probs = p`,
  * `Z` is the normalizaing constant, and,
  * `n!` is the factorial of `n`.
  """

  def __init__(self,
               total_count,
               logits=None,
               probs=None,
               validate_args=False,
               allow_nan_stats=True,
               name='NegativeBinomial'):
    """Construct NegativeBinomial distributions.

    Args:
      total_count: Positive floating-point `Tensor` with shape broadcastable to
        `[B1,..., Bb]` with `b >= 0` and the same dtype as `probs` or
        `logits`. Defines this as a batch of `N1 x ... x Nm` different Negative
        Binomial distributions. In practice, this represents the number of
        negative Bernoulli trials to stop at (the `total_count` of
        failures). Its components should be equal to integer values.
      logits: Floating-point `Tensor` with shape broadcastable to
        `[B1, ..., Bb]` where `b >= 0` indicates the number of batch dimensions.
        Each entry represents logits for the probability of success for
        independent Negative Binomial distributions and must be in the half-open
        interval `[-inf, inf)`. Only one of `logits` or `probs` should be
        specified.
      probs: Positive floating-point `Tensor` with shape broadcastable to
        `[B1, ..., Bb]` where `b >= 0` indicates the number of batch dimensions.
        Each entry represents the probability of success for independent
        Negative Binomial distributions and must be in the half-open interval
        `[0, 1)`. Only one of `logits` or `probs` should be specified.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      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())
    if (probs is None) == (logits is None):
      raise ValueError(
          'Construct `NegativeBinomial` with `probs` or `logits` but not both.')
    with tf.name_scope(name) as name:
      dtype = dtype_util.common_dtype([total_count, logits, probs],
                                      dtype_hint=tf.float32)
      self._probs = tensor_util.convert_nonref_to_tensor(
          probs, dtype=dtype, name='probs')
      self._logits = tensor_util.convert_nonref_to_tensor(
          logits, dtype=dtype, name='logits')
      self._total_count = tensor_util.convert_nonref_to_tensor(
          total_count, dtype=dtype, name='total_count')

      super(NegativeBinomial, self).__init__(
          dtype=dtype,
          reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
          validate_args=validate_args,
          allow_nan_stats=allow_nan_stats,
          parameters=parameters,
          name=name)

  @classmethod
  def experimental_from_mean_dispersion(cls, mean, dispersion, **kwargs):
    """Constructs a NegativeBinomial from its mean and dispersion.

    **Experimental: Naming, location of this API may change.**

    In this parameterization, the dispersion is defined as the reciprocal of the
    total count of failures, i.e. `dispersion = 1 / total_count`.

    Args:
      mean: The mean of the constructed distribution.
      dispersion: The reciprocal of the total_count of the constructed
        distribution.
      **kwargs: Other keyword arguments passed directly to `__init__`, e.g.
        `validate_args`.

    Returns:
      neg_bin: A distribution with the given parameterization.
    """
    dtype = dtype_util.common_dtype([mean, dispersion], dtype_hint=tf.float32)
    dispersion = tensor_util.convert_nonref_to_tensor(dispersion, dtype=dtype)
    mean = tensor_util.convert_nonref_to_tensor(mean, dtype=dtype)

    total_count = DeferredTensor(dispersion, tf.math.reciprocal, dtype=dtype)

    return cls(
        total_count=total_count,
        probs=DeferredTensor(mean, lambda mean: mean / (mean + total_count)),
        **kwargs)

  @classmethod
  def _parameter_properties(cls, dtype, num_classes=None):
    return dict(
        total_count=parameter_properties.ParameterProperties(
            default_constraining_bijector_fn=parameter_properties
            .BIJECTOR_NOT_IMPLEMENTED),
        logits=parameter_properties.ParameterProperties(),
        probs=parameter_properties.ParameterProperties(
            default_constraining_bijector_fn=sigmoid_bijector.Sigmoid,
            is_preferred=False))

  @property
  def total_count(self):
    """Number of negative trials."""
    return self._total_count

  @property
  def logits(self):
    """Input argument `logits`."""
    return self._logits

  @property
  def probs(self):
    """Input argument `probs`."""
    return self._probs

  def _batch_shape_tensor(self, logits_or_probs=None, total_count=None):
    if logits_or_probs is None:
      logits_or_probs = self._logits if self._probs is None else self._logits
    total_count = self._total_count if total_count is None else total_count
    return ps.broadcast_shape(
        ps.shape(logits_or_probs), ps.shape(total_count))

  def _batch_shape(self):
    x = self._probs if self._logits is None else self._logits
    return tf.broadcast_static_shape(self._total_count.shape, x.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):
    # Here we use the fact that if:
    # lam ~ Gamma(concentration=total_count, rate=(1-probs)/probs)
    # then X ~ Poisson(lam) is Negative Binomially distributed.
    logits = self._logits_parameter_no_checks()
    gamma_seed, poisson_seed = samplers.split_seed(
        seed, salt='NegativeBinomial')
    # TODO(b/152785714): For some reason switching to gamma_lib.random_gamma
    # makes tests time out. Note: observed similar in jax_transformation_test.
    rate = samplers.gamma(
        shape=[n],
        alpha=self.total_count,
        beta=tf.math.exp(-logits),
        dtype=self.dtype,
        seed=gamma_seed)
    return samplers.poisson(
        shape=[], lam=rate, dtype=self.dtype, seed=poisson_seed)

  def _cdf(self, x):
    logits = self._logits_parameter_no_checks()
    total_count = tf.convert_to_tensor(self.total_count)
    shape = self._batch_shape_tensor(
        logits_or_probs=logits, total_count=total_count)
    safe_x = tf.where(x >= 0, x, 0.)
    answer = tf.math.betainc(
        tf.broadcast_to(total_count, shape),
        tf.broadcast_to(1. + safe_x, shape),
        tf.broadcast_to(tf.sigmoid(-logits), shape))
    return distribution_util.extend_cdf_outside_support(x, answer, low=0)

  def _log_prob(self, x):
    total_count = tf.convert_to_tensor(self.total_count)
    logits = self._logits_parameter_no_checks()
    log_unnormalized_prob = (
        total_count * tf.math.log_sigmoid(-logits) +
        tf.math.multiply_no_nan(tf.math.log_sigmoid(logits), x))
    log_normalization = (tfp_math.lbeta(1. + x, total_count) +
                         tf.math.log(total_count + x))
    return log_unnormalized_prob - log_normalization

  def _mean(self, logits=None):
    logits = self._logits_parameter_no_checks() if logits is None else logits
    return self.total_count * tf.exp(logits)

  def _mode(self):
    total_count = tf.convert_to_tensor(self.total_count)
    adjusted_count = tf.where(1. < total_count, total_count - 1.,
                              tf.zeros_like(total_count))
    return tf.floor(adjusted_count * tf.exp(self._logits_parameter_no_checks()))

  def _variance(self):
    logits = self._logits_parameter_no_checks()
    return self._mean(logits=logits) / tf.sigmoid(-logits)

  def _logits_parameter_no_checks(self, name=None):
    """Logits computed from non-`None` input arg (`probs` or `logits`)."""
    if self._logits is None:
      probs = tf.convert_to_tensor(self._probs)
      return tf.math.log(probs) - tf.math.log1p(-probs)
    return tf.identity(self._logits)

  def logits_parameter(self, name=None):
    """Logits computed from non-`None` input arg (`probs` or `logits`)."""
    with self._name_and_control_scope(name or 'logits_parameter'):
      return self._logits_parameter_no_checks()

  def _probs_parameter_no_checks(self, name=None):
    """Probs computed from non-`None` input arg (`probs` or `logits`)."""
    if self._logits is None:
      return tf.identity(self._probs)
    return tf.math.sigmoid(self._logits)

  def probs_parameter(self, name=None):
    """Probs computed from non-`None` input arg (`probs` or `logits`)."""
    with self._name_and_control_scope(name or 'probs_parameter'):
      return self._probs_parameter_no_checks()

  def _default_event_space_bijector(self):
    return

  def _parameter_control_dependencies(self, is_init):
    return maybe_assert_negative_binomial_param_correctness(
        is_init, self.validate_args, self._total_count, self._probs,
        self._logits)

  def _sample_control_dependencies(self, x):
    """Check counts for proper shape and values, then return tensor version."""
    assertions = []
    if not self.validate_args:
      return assertions
    assertions.extend(distribution_util.assert_nonnegative_integer_form(x))
    return assertions


def maybe_assert_negative_binomial_param_correctness(
    is_init, validate_args, total_count, probs, logits):
  """Return assertions for `NegativeBinomial`-type distributions."""
  if is_init:
    x, name = (probs, 'probs') if logits is None else (logits, 'logits')
    if not dtype_util.is_floating(x.dtype):
      raise TypeError(
          'Argument `{}` must having floating type.'.format(name))

  if not validate_args:
    return []

  assertions = []
  if is_init != tensor_util.is_ref(total_count):
    total_count = tf.convert_to_tensor(total_count)
    assertions.extend([
        assert_util.assert_positive(
            total_count,
            message='`total_count` has components less than or equal to 0.'),
        distribution_util.assert_integer_form(
            total_count,
            message='`total_count` has fractional components.')
    ])
  if probs is not None:
    if is_init != tensor_util.is_ref(probs):
      probs = tf.convert_to_tensor(probs)
      one = tf.constant(1., probs.dtype)
      assertions.extend([
          assert_util.assert_non_negative(
              probs, message='`probs` has components less than 0.'),
          assert_util.assert_less_equal(
              probs, one, message='`probs` has components greater than 1.')
      ])

  return assertions

# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
# This file is auto-generated by substrates/meta/rewrite.py
# It will be surfaced by the build system as a symlink at:
#   `tensorflow_probability/substrates/numpy/distributions/negative_binomial.py`
# For more info, see substrate_runfiles_symlinks in build_defs.bzl
# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

# (This notice adds 10 to line numbering.)


