# Copyright (c) 2021 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 __future__ import annotations

import math
from collections.abc import Sequence
from typing import TYPE_CHECKING

import paddle
from paddle.base.data_feeder import convert_dtype
from paddle.distribution import constraint, distribution
from paddle.framework import in_dynamic_mode
from paddle.utils.decorator_utils import param_one_alias

if TYPE_CHECKING:
    from paddle import Tensor
    from paddle._typing.dtype_like import _DTypeLiteral


class MultivariateNormal(distribution.Distribution):
    r"""The Multivariate Normal distribution is a type multivariate continuous distribution defined on the real set, with parameter: `loc` and any one
    of the following parameters characterizing the variance: `covariance_matrix`, `precision_matrix`, `scale_tril`.

    Mathematical details

    The probability density function (pdf) is

    .. math::

        p(X ;\mu, \Sigma) = \frac{1}{\sqrt{(2\pi)^k |\Sigma|}} \exp(-\frac{1}{2}(X - \mu)^{\intercal} \Sigma^{-1} (X - \mu))

    In the above equation:

    * :math:`X`: is a k-dim random vector.
    * :math:`loc = \mu`: is the k-dim mean vector.
    * :math:`covariance_matrix = \Sigma`: is the k-by-k covariance matrix.

    Args:
        loc(int|float|Tensor): The mean of Multivariate Normal distribution. If the input data type is int or float, the data type of `loc` will be
            convert to a 1-D Tensor the paddle global default dtype.
        covariance_matrix(Tensor|None): The covariance matrix of Multivariate Normal distribution. The data type of `covariance_matrix` will be convert
            to be the same as the type of loc.
        precision_matrix(Tensor|None): The inverse of the covariance matrix. The data type of `precision_matrix` will be convert to be the same as the
            type of loc.
        scale_tril(Tensor|None): The cholesky decomposition (lower triangular matrix) of the covariance matrix. The data type of `scale_tril` will be
            convert to be the same as the type of loc.

    Examples:
        .. code-block:: pycon

            >>> import paddle
            >>> from paddle.distribution import MultivariateNormal
            >>> paddle.set_device("cpu")
            >>> paddle.seed(100)

            >>> rv = MultivariateNormal(
            ...     loc=paddle.to_tensor([2.0, 5.0]),
            ...     covariance_matrix=paddle.to_tensor([[2.0, 1.0], [1.0, 2.0]]),
            ... )

            >>> print(rv.sample([3, 2]))
            Tensor(shape=[3, 2, 2], dtype=float32, place=Place(cpu), stop_gradient=True,
            [[[-0.00339603,  4.31556797],
              [ 2.01385283,  4.63553190]],
             [[ 0.10132277,  3.11323833],
              [ 2.37435842,  3.56635118]],
             [[ 2.89701366,  5.10602522],
              [-0.46329355,  3.14768648]]])

            >>> print(rv.mean)
            Tensor(shape=[2], dtype=float32, place=Place(cpu), stop_gradient=True,
            [2., 5.])

            >>> print(rv.variance)
            Tensor(shape=[2], dtype=float32, place=Place(cpu), stop_gradient=True,
            [1.99999988, 2.        ])

            >>> print(rv.entropy())
            Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True,
            3.38718319)

            >>> rv1 = MultivariateNormal(
            ...     loc=paddle.to_tensor([2.0, 5.0]),
            ...     covariance_matrix=paddle.to_tensor([[2.0, 1.0], [1.0, 2.0]]),
            ... )
            >>> rv2 = MultivariateNormal(
            ...     loc=paddle.to_tensor([-1.0, 3.0]), covariance_matrix=paddle.to_tensor([[3.0, 2.0], [2.0, 3.0]])
            ... )
            >>> print(rv1.kl_divergence(rv2))
            Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True,
            1.55541301)
    """

    loc: Tensor
    covariance_matrix: Tensor | None
    precision_matrix: Tensor | None
    scale_tril: Tensor | None
    dtype: _DTypeLiteral

    arg_constraints = {
        "loc": constraint.real_vector,
        "covariance_matrix": constraint.positive_definite,
        "precision_matrix": constraint.positive_definite,
        "scale_tril": constraint.lower_cholesky,
    }
    support = constraint.real_vector
    has_rsample = True

    def __init__(
        self,
        loc: float | Tensor,
        covariance_matrix: Tensor | None = None,
        precision_matrix: Tensor | None = None,
        scale_tril: Tensor | None = None,
        validate_args: bool | None = None,
    ):
        self.dtype = paddle.get_default_dtype()
        if isinstance(loc, (float, int)):
            loc = paddle.to_tensor([loc], dtype=self.dtype)
        else:
            self.dtype = convert_dtype(loc.dtype)
        if loc.dim() < 1:
            raise ValueError("loc must be at least one-dimensional.")
        if (covariance_matrix is not None) + (scale_tril is not None) + (
            precision_matrix is not None
        ) != 1:
            raise ValueError(
                "Exactly one of covariance_matrix or precision_matrix or scale_tril may be specified."
            )

        if scale_tril is not None:
            if scale_tril.dim() < 2:
                raise ValueError(
                    "scale_tril matrix must be at least two-dimensional, "
                    "with optional leading batch dimensions"
                )
            scale_tril = paddle.cast(scale_tril, dtype=self.dtype)
            batch_shape = paddle.broadcast_shape(
                scale_tril.shape[:-2], loc.shape[:-1]
            )
            self.scale_tril = scale_tril.expand(
                [*batch_shape, scale_tril.shape[-2], scale_tril.shape[-1]]
            )
        elif covariance_matrix is not None:
            if covariance_matrix.dim() < 2:
                raise ValueError(
                    "covariance_matrix must be at least two-dimensional, "
                    "with optional leading batch dimensions"
                )
            covariance_matrix = paddle.cast(covariance_matrix, dtype=self.dtype)
            batch_shape = paddle.broadcast_shape(
                covariance_matrix.shape[:-2], loc.shape[:-1]
            )
            self.covariance_matrix = covariance_matrix.expand(
                [
                    *batch_shape,
                    covariance_matrix.shape[-2],
                    covariance_matrix.shape[-1],
                ]
            )
        else:
            if precision_matrix.dim() < 2:
                raise ValueError(
                    "precision_matrix must be at least two-dimensional, "
                    "with optional leading batch dimensions"
                )
            precision_matrix = paddle.cast(precision_matrix, dtype=self.dtype)
            batch_shape = paddle.broadcast_shape(
                precision_matrix.shape[:-2], loc.shape[:-1]
            )
            self.precision_matrix = precision_matrix.expand(
                [
                    *batch_shape,
                    precision_matrix.shape[-2],
                    precision_matrix.shape[-1],
                ]
            )
        self.loc = loc.expand([*batch_shape, -1])
        event_shape = self.loc.shape[-1:]
        super().__init__(batch_shape, event_shape, validate_args=validate_args)
        if in_dynamic_mode() and self._validate_args_enabled:
            self._validate_parameters(
                scale_tril=scale_tril,
                covariance_matrix=covariance_matrix,
                precision_matrix=precision_matrix,
            )

        if scale_tril is not None:
            self._unbroadcasted_scale_tril = scale_tril
        elif covariance_matrix is not None:
            self._unbroadcasted_scale_tril = paddle.linalg.cholesky(
                covariance_matrix
            )
        else:
            self._unbroadcasted_scale_tril = precision_to_scale_tril(
                precision_matrix
            )

    def _validate_parameters(
        self,
        *,
        scale_tril: Tensor | None = None,
        covariance_matrix: Tensor | None = None,
        precision_matrix: Tensor | None = None,
    ) -> None:
        if scale_tril is not None:
            matrix_name = "scale_tril"
            matrix_value = self.scale_tril
        elif covariance_matrix is not None:
            matrix_name = "covariance_matrix"
            matrix_value = self.covariance_matrix
        else:
            matrix_name = "precision_matrix"
            matrix_value = self.precision_matrix

        for param, value in (
            ("loc", self.loc),
            (matrix_name, matrix_value),
        ):
            constraint_ = self.arg_constraints[param]
            valid = constraint_.check(value)
            if not bool(valid.all()):
                raise ValueError(
                    f"Expected parameter {param} "
                    f"({type(value).__name__} of shape {tuple(value.shape)}) "
                    f"of distribution {self!r} "
                    f"to satisfy the constraint {constraint_!r}, "
                    f"but found invalid values:\n{value}"
                )

    def expand(self, batch_shape, _instance=None):
        new = (
            self.__class__.__new__(self.__class__)
            if _instance is None
            else _instance
        )
        batch_shape = tuple(batch_shape)
        loc_shape = batch_shape + self.event_shape
        cov_shape = batch_shape + self.event_shape + self.event_shape
        new.loc = self.loc.expand(loc_shape)
        new._unbroadcasted_scale_tril = self._unbroadcasted_scale_tril
        if "_covariance_matrix" in self.__dict__:
            new._covariance_matrix = self.covariance_matrix.expand(cov_shape)
        if "_scale_tril" in self.__dict__:
            new._scale_tril = self.scale_tril.expand(cov_shape)
        if "_precision_matrix" in self.__dict__:
            new._precision_matrix = self.precision_matrix.expand(cov_shape)
        super(MultivariateNormal, new).__init__(
            batch_shape, self.event_shape, validate_args=False
        )
        new._validate_args_enabled = self._validate_args_enabled
        return new

    @property
    def scale_tril(self) -> Tensor:
        if "_scale_tril" not in self.__dict__:
            self._scale_tril = self._unbroadcasted_scale_tril.expand(
                self._batch_shape + self._event_shape + self._event_shape
            )
        return self._scale_tril

    @scale_tril.setter
    def scale_tril(self, value: Tensor) -> None:
        self._scale_tril = value

    @property
    def covariance_matrix(self) -> Tensor:
        if "_covariance_matrix" not in self.__dict__:
            new_perm = list(range(len(self._unbroadcasted_scale_tril.shape)))
            new_perm[-1], new_perm[-2] = new_perm[-2], new_perm[-1]
            self._covariance_matrix = paddle.matmul(
                self._unbroadcasted_scale_tril,
                self._unbroadcasted_scale_tril.transpose(new_perm),
            ).expand(self._batch_shape + self._event_shape + self._event_shape)
        return self._covariance_matrix

    @covariance_matrix.setter
    def covariance_matrix(self, value: Tensor) -> None:
        self._covariance_matrix = value

    @property
    def precision_matrix(self) -> Tensor:
        if "_precision_matrix" not in self.__dict__:
            self._precision_matrix = paddle.linalg.cholesky_inverse(
                self._unbroadcasted_scale_tril
            ).expand(self._batch_shape + self._event_shape + self._event_shape)
        return self._precision_matrix

    @precision_matrix.setter
    def precision_matrix(self, value: Tensor) -> None:
        self._precision_matrix = value

    @property
    def mean(self) -> Tensor:
        """Mean of Multivariate Normal distribution.

        Returns:
            Tensor: mean value.
        """
        return self.loc

    @property
    def variance(self) -> Tensor:
        """Variance of Multivariate Normal distribution.

        Returns:
            Tensor: variance value.
        """
        return (
            paddle.square(self._unbroadcasted_scale_tril)
            .sum(-1)
            .expand(self._batch_shape + self._event_shape)
        )

    @property
    def mode(self) -> Tensor:
        return self.loc

    @mode.setter
    def mode(self, value: Tensor) -> None:
        self.loc = value

    @param_one_alias(["shape", "sample_shape"])
    def sample(self, shape: Sequence[int] = []) -> Tensor:
        """Generate Multivariate Normal samples of the specified shape. The final shape would be ``sample_shape + batch_shape + event_shape``.

        Args:
            shape (Sequence[int], optional): Prepended shape of the generated samples.

        Returns:
            Tensor, Sampled data with shape `sample_shape` + `batch_shape` + `event_shape`. The data type is the same as `self.loc`.
        """
        with paddle.no_grad():
            return self.rsample(shape)

    @param_one_alias(["shape", "sample_shape"])
    def rsample(self, shape: Sequence[int] = []) -> Tensor:
        """Generate Multivariate Normal samples of the specified shape. The final shape would be ``sample_shape + batch_shape + event_shape``.

        Args:
            shape (Sequence[int], optional): Prepended shape of the generated samples.

        Returns:
            Tensor, Sampled data with shape `sample_shape` + `batch_shape` + `event_shape`. The data type is the same as `self.loc`.
        """
        if not isinstance(shape, Sequence):
            raise TypeError('sample shape must be Sequence object.')
        output_shape = self._extend_shape(shape)
        eps = paddle.cast(paddle.normal(shape=output_shape), dtype=self.dtype)
        return self.loc + paddle.matmul(
            self._unbroadcasted_scale_tril, eps.unsqueeze(-1)
        ).squeeze(-1)

    def log_prob(self, value: Tensor) -> Tensor:
        """Log probability density function.

        Args:
          value (Tensor): The input tensor.

        Returns:
          Tensor: log probability. The data type is the same as `self.loc`.
        """
        value = paddle.cast(value, dtype=self.dtype)
        if in_dynamic_mode() and self._validate_args_enabled:
            self._validate_sample(value)

        diff = value - self.loc
        M = batch_mahalanobis(self._unbroadcasted_scale_tril, diff)
        half_log_det = (
            self._unbroadcasted_scale_tril.diagonal(axis1=-2, axis2=-1)
            .log()
            .sum(-1)
        )
        return (
            -0.5 * (self._event_shape[0] * math.log(2 * math.pi) + M)
            - half_log_det
        )

    def prob(self, value: Tensor) -> Tensor:
        """Probability density function.

        Args:
            value (Tensor): The input tensor.

        Returns:
            Tensor: probability. The data type is the same as `self.loc`.
        """
        return paddle.exp(self.log_prob(value))

    def entropy(self) -> Tensor:
        r"""Shannon entropy in nats.

        The entropy is

        .. math::

            \mathcal{H}(X) = \frac{n}{2} \log(2\pi) + \log {\det A} + \frac{n}{2}

        In the above equation:

        * :math:`\Omega`: is the support of the distribution.

        Returns:
            Tensor, Shannon entropy of Multivariate Normal distribution. The data type is the same as `self.loc`.
        """
        half_log_det = (
            self._unbroadcasted_scale_tril.diagonal(axis1=-2, axis2=-1)
            .log()
            .sum(-1)
        )
        H = (
            0.5 * self._event_shape[0] * (1.0 + math.log(2 * math.pi))
            + half_log_det
        )
        if len(self._batch_shape) == 0:
            return H
        else:
            return H.expand(self._batch_shape)

    def kl_divergence(self, other: MultivariateNormal) -> Tensor:
        r"""The KL-divergence between two poisson distributions with the same `batch_shape` and `event_shape`.

        The probability density function (pdf) is

        .. math::

            KL\_divergence(\lambda_1, \lambda_2) = \log(\det A_2) - \log(\det A_1) -\frac{n}{2} +\frac{1}{2}[tr [\Sigma_2^{-1} \Sigma_1] + (\mu_1 - \mu_2)^{\intercal} \Sigma_2^{-1}  (\mu_1 - \mu_2)]

        Args:
            other (MultivariateNormal): instance of Multivariate Normal.

        Returns:
            Tensor, kl-divergence between two Multivariate Normal distributions. The data type is the same as `self.loc`.

        """
        if (
            self._batch_shape != other._batch_shape
            and self._event_shape != other._event_shape
        ):
            raise ValueError(
                "KL divergence of two Multivariate Normal distributions should share the same `batch_shape` and `event_shape`."
            )
        half_log_det_1 = (
            self._unbroadcasted_scale_tril.diagonal(axis1=-2, axis2=-1)
            .log()
            .sum(-1)
        )
        half_log_det_2 = (
            other._unbroadcasted_scale_tril.diagonal(axis1=-2, axis2=-1)
            .log()
            .sum(-1)
        )
        new_perm = list(range(len(self._unbroadcasted_scale_tril.shape)))
        new_perm[-1], new_perm[-2] = new_perm[-2], new_perm[-1]
        cov_mat_1 = paddle.matmul(
            self._unbroadcasted_scale_tril,
            self._unbroadcasted_scale_tril.transpose(new_perm),
        )
        cov_mat_2 = paddle.matmul(
            other._unbroadcasted_scale_tril,
            other._unbroadcasted_scale_tril.transpose(new_perm),
        )
        expectation = (
            paddle.linalg.solve(cov_mat_2, cov_mat_1)
            .diagonal(axis1=-2, axis2=-1)
            .sum(-1)
        )
        expectation += batch_mahalanobis(
            other._unbroadcasted_scale_tril, self.loc - other.loc
        )
        return (
            half_log_det_2
            - half_log_det_1
            + 0.5 * (expectation - self._event_shape[0])
        )


def precision_to_scale_tril(P: Tensor) -> Tensor:
    """Convert precision matrix to scale tril matrix

    Args:
        P (Tensor): input precision matrix

    Returns:
        Tensor: scale tril matrix
    """
    Lf = paddle.linalg.cholesky(paddle.flip(P, (-2, -1)))
    tmp = paddle.flip(Lf, (-2, -1))
    new_perm = list(range(len(tmp.shape)))
    new_perm[-2], new_perm[-1] = new_perm[-1], new_perm[-2]
    L_inv = paddle.transpose(tmp, new_perm)
    Id = paddle.eye(P.shape[-1], dtype=P.dtype)
    L = paddle.linalg.triangular_solve(L_inv, Id, upper=False)
    return L


def batch_mahalanobis(bL: Tensor, bx: Tensor) -> Tensor:
    r"""
    Computes the squared Mahalanobis distance of the Multivariate Normal distribution with cholesky decomposition of the covariance matrix.
    Accepts batches for both bL and bx.

    Args:
        bL (Tensor): scale trial matrix (batched)
        bx (Tensor): difference vector(batched)

    Returns:
        Tensor: squared Mahalanobis distance
    """
    n = bx.shape[-1]
    bx_batch_shape = bx.shape[:-1]

    # Assume that bL.shape = (i, 1, n, n), bx.shape = (..., i, j, n),
    # we are going to make bx have shape (..., 1, j, i, 1, n) to apply batched tri.solve
    bx_batch_dims = len(bx_batch_shape)
    bL_batch_dims = bL.dim() - 2
    outer_batch_dims = bx_batch_dims - bL_batch_dims
    old_batch_dims = outer_batch_dims + bL_batch_dims
    new_batch_dims = outer_batch_dims + 2 * bL_batch_dims
    # Reshape bx with the shape (..., 1, i, j, 1, n)
    bx_new_shape = bx.shape[:outer_batch_dims]
    for sL, sx in zip(bL.shape[:-2], bx.shape[outer_batch_dims:-1]):
        bx_new_shape += (sx // sL, sL)
    bx_new_shape += (n,)
    bx = bx.reshape(bx_new_shape)
    # Permute bx to make it have shape (..., 1, j, i, 1, n)
    permute_dims = (
        list(range(outer_batch_dims))
        + list(range(outer_batch_dims, new_batch_dims, 2))
        + list(range(outer_batch_dims + 1, new_batch_dims, 2))
        + [new_batch_dims]
    )
    bx = bx.transpose(permute_dims)

    flat_L = bL.reshape((-1, n, n))  # shape = b x n x n
    flat_x = bx.reshape((-1, flat_L.shape[0], n))  # shape = c x b x n
    flat_x_swap = flat_x.transpose((1, 2, 0))  # shape = b x n x c
    M_swap = (
        paddle.linalg.triangular_solve(flat_L, flat_x_swap, upper=False)
        .pow(2)
        .sum(-2)
    )  # shape = b x c
    M = M_swap.t()  # shape = c x b

    # Now we revert the above reshape and permute operators.
    permuted_M = M.reshape(bx.shape[:-1])  # shape = (..., 1, j, i, 1)
    permute_inv_dims = list(range(outer_batch_dims))
    for i in range(bL_batch_dims):
        permute_inv_dims += [outer_batch_dims + i, old_batch_dims + i]
    reshaped_M = permuted_M.transpose(
        permute_inv_dims
    )  # shape = (..., 1, i, j, 1)
    return reshaped_M.reshape(bx_batch_shape)