# Copyright 2016 The TensorFlow Authors. All Rights Reserved.

#

# 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.

# ==========================================================================

# ======================================================================== #

#

# Copyright (c) 2017 - 2020 scVAE 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 ZeroInflated distribution class."""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

from tensorflow import where

from tensorflow.python.framework import ops

from tensorflow.python.ops import array_ops

from tensorflow.python.ops import check_ops

from tensorflow.python.ops import math_ops

from tensorflow_probability.python.distributions import distribution

from tensorflow_probability.python.internal import reparameterization

class ZeroInflated(distribution.Distribution):

    """zero-inflated distribution.

    The `zero-inflated` object implements batched zero-inflated distributions.

    The zero-inflated model is defined by a zero-inflation rate

    and a python list of `Distribution` objects.

    Methods supported include `log_prob`, `prob`, `mean`, `sample`, and

    `entropy_lower_bound`.

    """

    def __init__(self,

                 dist,

                 pi,

                 validate_args=False,

                 allow_nan_stats=True,

                 name="ZeroInflated"):

        """Initialise a zero-inflated distribution.

        A `ZeroInflated` is defined by a zero-inflation rate (`pi`,

        representing the probabilities of excess zeroes) and a `Distribution`

        object having matching dtype, batch shape, event shape, and continuity

        properties (the dist).

        Args:

            pi: A zero-inflation rate, representing the probabilities of excess

                zeroes.

            dist: A `Distribution` instance.

                The instance must have `batch_shape` matching the

                zero-inflation rate.

            validate_args: Python `bool`, default `False`. If `True`, raise a

                runtime error if batch or event ranks are inconsistent between

                pi and any of the distributions. This is only checked if the

                ranks cannot be determined statically at graph construction

                time.

            allow_nan_stats: Boolean, default `True`. If `False`, raise an

                 exception if a statistic (e.g. mean/mode/etc...) is undefined

                for any batch member. If `True`, batch members with valid

                parameters leading to undefined statistics will return NaN for

                this statistic.

            name: A name for this distribution (optional).

        Raises:

            TypeError: If pi is not a zero-inflation rate, or `dist` is not

                `Distibution` are not instances of `Distribution`, or do not

                 have matching `dtype`.

            ValueError: If `dist` is an empty list or tuple, or its

                elements do not have a statically known event rank.

                If `pi.num_classes` cannot be inferred at graph creation time,

                or the constant value of `pi.num_classes` is not equal to

                `len(dist)`, or all `dist` and `pi` do not have

                matching static batch shapes, or all dist do not

                have matching static event shapes.

        """

        parameters = locals()

        if not dist:

            raise ValueError("dist must be non-empty")

        if not isinstance(dist, distribution.Distribution):

            raise TypeError(

                "dist must be a Distribution instance"

                " but saw: %s" % dist)

        dtype = dist.dtype

        static_event_shape = dist.event_shape

        static_batch_shape = pi.get_shape()

        if static_event_shape.ndims is None:

            raise ValueError(

                "Expected to know rank(event_shape) from dist, but "

                "the distribution does not provide a static number of ndims")

        # Ensure that all batch and event ndims are consistent.

        with ops.name_scope(name, values=[pi]):

            with ops.control_dependencies([check_ops.assert_positive(pi)] if

                                          validate_args else []):

                pi_batch_shape = array_ops.shape(pi)

                pi_batch_rank = array_ops.size(pi_batch_shape)

                if validate_args:

                    dist_batch_shape = dist.batch_shape_tensor()

                    dist_batch_rank = array_ops.size(dist_batch_shape)

                    check_message = (

                        "dist batch shape must match pi batch shape")

                    self._assertions = [check_ops.assert_equal(

                        pi_batch_rank, dist_batch_rank, message=check_message)]

                    self._assertions += [

                        check_ops.assert_equal(

                            pi_batch_shape, dist_batch_shape,

                            message=check_message)]

                else:

                    self._assertions = []

                self._pi = pi

                self._dist = dist

                self._static_event_shape = static_event_shape

                self._static_batch_shape = static_batch_shape

        # We let the zero-inflated distribution access _graph_parents since its

        # arguably more like a baseclass.

        graph_parents = [self._pi]

        graph_parents += self._dist._graph_parents

        super(ZeroInflated, self).__init__(

                dtype=dtype,

                reparameterization_type=reparameterization.NOT_REPARAMETERIZED,

                validate_args=validate_args,

                allow_nan_stats=allow_nan_stats,

                parameters=parameters,

                graph_parents=graph_parents,

                name=name)

    @property

    def pi(self):

        return self._pi

    @property

    def dist(self):

        return self._dist

    def _batch_shape_tensor(self):

        return array_ops.shape(self._pi)

    def _batch_shape(self):

        return self._static_batch_shape

    def _event_shape_tensor(self):

        return self._dist.event_shape_tensor()

    def _event_shape(self):

        return self._static_event_shape

    def _mean(self):

        with ops.control_dependencies(self._assertions):

            # These should all be the same shape by virtue of matching

            # batch_shape and event_shape.

            return (1-self._pi) * self._dist.mean()

    def _variance(self):

        with ops.control_dependencies(self._assertions):

            # These should all be the same shape by virtue of matching

            # batch_shape and event_shape.

            return ((1-self._pi) * (self._dist.variance()

                                    + math_ops.square(self._dist.mean()))

                    - math_ops.square(self._mean()))

    def _log_prob(self, x):

        with ops.control_dependencies(self._assertions):

            x = ops.convert_to_tensor(x, name="x")

            y_0 = math_ops.log(self.pi + (1 - self.pi) * self._dist.prob(x))

            y_1 = math_ops.log(1 - self.pi) + self._dist.log_prob(x)

            return where(x > 0, y_1, y_0)

    def _prob(self, x):

        return math_ops.exp(self._log_prob(x))