# Copyright 2022 The GPJax Contributors. 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 beartype.typing import (
    Optional,
)
from jax import vmap
import jax.numpy as jnp
import jax.random as jr
from jaxtyping import Float
from numpyro.distributions import constraints
from numpyro.distributions.distribution import Distribution
from numpyro.distributions.util import is_prng_key

from gpjax.linalg.operations import (
    diag,
    logdet,
    lower_cholesky,
    solve,
)
from gpjax.linalg.operators import LinearOperator
from gpjax.linalg.utils import psd
from gpjax.typing import (
    Array,
    ScalarFloat,
)


class GaussianDistribution(Distribution):
    r"""Multivariate Gaussian distribution for GP predictions.

    This is the return type of all ``predict()`` methods in GPJax. It wraps a
    mean vector and a covariance :class:`~gpjax.linalg.operators.LinearOperator`,
    providing methods for sampling, computing log-probabilities, and evaluating
    KL divergences.

    The distribution is parameterised as

    .. math::

        p(\mathbf{x}) = \mathcal{N}(\mathbf{x}; \boldsymbol{\mu}, \mathbf{\Sigma})

    where :math:`\boldsymbol{\mu}` is the ``loc`` (mean) vector and
    :math:`\mathbf{\Sigma}` is represented by the ``scale``
    :class:`~gpjax.linalg.operators.LinearOperator`. The ``scale`` is
    automatically annotated as positive semi-definite on construction.

    Parameters
    ----------
    loc : Float[Array, " N"]
        Mean vector of the distribution.
    scale : LinearOperator
        Covariance matrix represented as a
        :class:`~gpjax.linalg.operators.LinearOperator` (e.g.
        :class:`~gpjax.linalg.operators.Dense` or
        :class:`~gpjax.linalg.operators.Diagonal`).

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from gpjax.distributions import GaussianDistribution
    >>> from gpjax.linalg.operators import Dense
    >>> mu = jnp.array([0.0, 1.0])
    >>> cov = Dense(jnp.eye(2))
    >>> dist = GaussianDistribution(loc=mu, scale=cov)
    >>> dist.mean
    Array([0., 1.], dtype=float32)
    >>> dist.variance
    Array([1., 1.], dtype=float32)
    """

    support = constraints.real_vector

    def __init__(
        self,
        loc: Optional[Float[Array, " N"]],
        scale: Optional[LinearOperator],
        validate_args=None,
    ):
        self.loc = loc
        self.scale = psd(scale)
        batch_shape = ()
        event_shape = jnp.shape(self.loc)
        super().__init__(batch_shape, event_shape, validate_args=validate_args)

    def sample(self, key, sample_shape=()):
        r"""Draw samples from the distribution.

        Generates samples via the reparameterisation trick:

        .. math::

            \mathbf{x} = \boldsymbol{\mu} + \mathbf{L}\mathbf{z},
            \quad \mathbf{z} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})

        where :math:`\mathbf{L}` is the lower Cholesky factor of the covariance.

        Parameters
        ----------
        key : KeyArray
            JAX PRNG key.
        sample_shape : tuple of int, optional
            Leading batch dimensions for the samples. Defaults to ``()``,
            returning a single sample.

        Returns
        -------
        Float[Array, "... N"]
            Array of samples with shape ``(*sample_shape, N)``.
        """
        assert is_prng_key(key)
        # Obtain covariance root.
        covariance_root = lower_cholesky(self.scale)

        # Gather n samples from standard normal distribution Z = [z₁, ..., zₙ]ᵀ.
        white_noise = jr.normal(
            key, shape=sample_shape + self.batch_shape + self.event_shape
        )

        # xᵢ ~ N(loc, cov) <=> xᵢ = loc + sqrt zᵢ, where zᵢ ~ N(0, I).
        def affine_transformation(_x):
            return self.loc + covariance_root @ _x

        if not sample_shape:
            return affine_transformation(white_noise)

        return vmap(affine_transformation)(white_noise)

    @property
    def mean(self) -> Float[Array, " N"]:
        r"""Calculates the mean."""
        return self.loc

    @property
    def variance(self) -> Float[Array, " N"]:
        r"""Calculates the marginal variance (diagonal of the covariance)."""
        return diag(self.scale)

    def entropy(self) -> ScalarFloat:
        r"""Calculates the differential entropy of the distribution.

        .. math::

            H[p] = \tfrac{1}{2}\bigl(N(1 + \ln 2\pi) + \ln|\mathbf{\Sigma}|\bigr)

        Returns
        -------
        ScalarFloat
            Entropy in nats.
        """
        return 0.5 * (
            self.event_shape[0] * (1.0 + jnp.log(2.0 * jnp.pi)) + logdet(self.scale)
        )

    def median(self) -> Float[Array, " N"]:
        r"""Calculates the median (equal to the mean for a Gaussian)."""
        return self.loc

    def mode(self) -> Float[Array, " N"]:
        r"""Calculates the mode (equal to the mean for a Gaussian)."""
        return self.loc

    def covariance(self) -> Float[Array, "N N"]:
        r"""Materialises the full covariance matrix as a dense array.

        Returns
        -------
        Float[Array, "N N"]
            Dense covariance matrix.
        """
        return self.scale.to_dense()

    @property
    def covariance_matrix(self) -> Float[Array, "N N"]:
        r"""Property alias for :meth:`covariance`."""
        return self.covariance()

    def stddev(self) -> Float[Array, " N"]:
        r"""Calculates the marginal standard deviation."""
        return jnp.sqrt(diag(self.scale))

    def log_prob(self, y: Float[Array, " N"]) -> ScalarFloat:
        r"""Calculates the log pdf of the multivariate Gaussian.

        .. math::

            \log p(\mathbf{y}) = -\tfrac{1}{2}\bigl[
                N\ln 2\pi + \ln|\mathbf{\Sigma}|
                + (\mathbf{y} - \boldsymbol{\mu})^\top
                  \mathbf{\Sigma}^{-1}
                  (\mathbf{y} - \boldsymbol{\mu})
            \bigr]

        Parameters
        ----------
        y : Float[Array, " N"]
            Point at which to evaluate the log-density.

        Returns
        -------
        ScalarFloat
            Log probability.
        """
        mu = self.loc
        sigma = self.scale
        n = mu.shape[-1]

        # diff, y - µ
        diff = y - mu

        # compute the pdf, -1/2[ n log(2π) + log|Σ| + (y - µ)ᵀΣ⁻¹(y - µ) ]
        return -0.5 * (
            n * jnp.log(2.0 * jnp.pi) + logdet(sigma) + diff.T @ solve(sigma, diff)
        )

    def kl_divergence(self, other: "GaussianDistribution") -> ScalarFloat:
        r"""KL divergence from ``self`` to ``other``.

        Computes :math:`\operatorname{KL}[q \| p]` where ``self`` is *q* and
        ``other`` is *p*.

        Parameters
        ----------
        other : GaussianDistribution
            The reference distribution *p*.

        Returns
        -------
        ScalarFloat
            KL divergence in nats.
        """
        return _kl_divergence(self, other)


def _check_and_return_dimension(
    q: GaussianDistribution, p: GaussianDistribution
) -> int:
    r"""Checks that the dimensions of the distributions are compatible."""
    if q.event_shape != p.event_shape:
        raise ValueError(
            "Distribution event shapes are not compatible: `q.event_shape ="
            f" {q.event_shape}` and `p.event_shape = {p.event_shape}`. Please check"
            " your mean and covariance shapes."
        )

    return q.event_shape[-1]


def _frobenius_norm_squared(matrix: Float[Array, "N N"]) -> ScalarFloat:
    r"""Calculates the squared Frobenius norm of a matrix."""
    return jnp.sum(jnp.square(matrix))


def _kl_divergence(q: GaussianDistribution, p: GaussianDistribution) -> ScalarFloat:
    r"""KL-divergence between two Gaussians.

    Computes the KL divergence, $\operatorname{KL}[q\mid\mid p]$, between two
    multivariate Gaussian distributions $q(x) = \mathcal{N}(x; \mu_q, \Sigma_q)$
    and $p(x) = \mathcal{N}(x; \mu_p, \Sigma_p)$.

    Args:
        q: a multivariate Gaussian distribution.
        p: another multivariate Gaussian distribution.

    Returns:
        ScalarFloat: The KL divergence between q and p.
    """
    n_dim = _check_and_return_dimension(q, p)

    # Extract q mean and covariance.
    mu_q = q.loc
    sigma_q = q.scale

    # Extract p mean and covariance.
    mu_p = p.loc
    sigma_p = p.scale

    # Find covariance roots.
    sqrt_p = lower_cholesky(sigma_p)
    sqrt_q = lower_cholesky(sigma_q)

    # diff, μp - μq
    diff = mu_p - mu_q

    # trace term, tr[Σp⁻¹ Σq] = tr[(LpLpᵀ)⁻¹(LqLqᵀ)] = tr[(Lp⁻¹Lq)(Lp⁻¹Lq)ᵀ] = (fr[LqLp⁻¹])²
    trace = _frobenius_norm_squared(
        solve(sqrt_p, sqrt_q.to_dense())
    )  # TODO: Not most efficient, given the `to_dense()` call (e.g., consider diagonal p and q). Need to abstract solving linear operator against another linear operator.

    # Mahalanobis term, (μp - μq)ᵀ Σp⁻¹ (μp - μq) = tr [(μp - μq)ᵀ [LpLpᵀ]⁻¹ (μp - μq)] = (fr[Lp⁻¹(μp - μq)])²
    mahalanobis = jnp.sum(jnp.square(solve(sqrt_p, diff)))

    # KL[q(x)||p(x)] = [ [(μp - μq)ᵀ Σp⁻¹ (μp - μq)] - n - log|Σq| + log|Σp| + tr[Σp⁻¹ Σq] ] / 2
    return (mahalanobis - n_dim - logdet(sigma_q) + logdet(sigma_p) + trace) / 2.0


__all__ = [
    "GaussianDistribution",
]
