# Copyright (c) 2022 PaddlePaddle 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.

from paddle.distribution import distribution


class Independent(distribution.Distribution):
    r"""
    Reinterprets some of the batch dimensions of a distribution as event dimensions.

    This is mainly useful for changing the shape of the result of
    :meth:`log_prob`. 

    Args:
        base (Distribution): The base distribution.
        reinterpreted_batch_rank (int): The number of batch dimensions to 
            reinterpret as event dimensions.

    Examples:

        .. code-block:: python
        
            import paddle
            from paddle.distribution import independent

            beta = paddle.distribution.Beta(paddle.to_tensor([0.5, 0.5]), paddle.to_tensor([0.5, 0.5]))
            print(beta.batch_shape, beta.event_shape)
            # (2,) ()
            print(beta.log_prob(paddle.to_tensor(0.2)))
            # Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True,
            #        [-0.22843921, -0.22843921])
            reinterpreted_beta = independent.Independent(beta, 1)
            print(reinterpreted_beta.batch_shape, reinterpreted_beta.event_shape)
            # () (2,)
            print(reinterpreted_beta.log_prob(paddle.to_tensor([0.2,  0.2])))
            # Tensor(shape=[1], dtype=float32, place=Place(gpu:0), stop_gradient=True,
            #        [-0.45687842])
    """

    def __init__(self, base, reinterpreted_batch_rank):
        if not isinstance(base, distribution.Distribution):
            raise TypeError(
                f"Expected type of 'base' is Distribution, but got {type(base)}"
            )
        if not (0 < reinterpreted_batch_rank <= len(base.batch_shape)):
            raise ValueError(
                f"Expected 0 < reinterpreted_batch_rank <= {len(base.batch_shape)}, but got {reinterpreted_batch_rank}"
            )
        self._base = base
        self._reinterpreted_batch_rank = reinterpreted_batch_rank

        shape = base.batch_shape + base.event_shape
        super(Independent,
              self).__init__(batch_shape=shape[:len(base.batch_shape) -
                                               reinterpreted_batch_rank],
                             event_shape=shape[len(base.batch_shape) -
                                               reinterpreted_batch_rank:])

    @property
    def mean(self):
        return self._base.mean

    @property
    def variance(self):
        return self._base.variance

    def sample(self, shape=()):
        return self._base.sample(shape)

    def log_prob(self, value):
        return self._sum_rightmost(self._base.log_prob(value),
                                   self._reinterpreted_batch_rank)

    def prob(self, value):
        return self.log_prob(value).exp()

    def entropy(self):
        return self._sum_rightmost(self._base.entropy(),
                                   self._reinterpreted_batch_rank)

    def _sum_rightmost(self, value, n):
        return value.sum(list(range(-n, 0))) if n > 0 else value