"""Calculates the Von neumann entropy metric of a quantum state."""

import numpy as np

from toqito.matrix_props import is_density


def von_neumann_entropy(rho: np.ndarray) -> float:
    r"""Compute the von Neumann entropy of a density matrix [@wikipediavonneumann].

    Let \(P \in \text{Pos}(\mathcal{X})\) be a positive semidefinite operator, for a complex
    Euclidean space \(\mathcal{X}\). Then one defines the *von Neumann entropy* as

    \[
        H(P) = H(\lambda(P)),
    \]

    where \(\lambda(P)\) is the vector of eigenvalues of \(P\) and where the function
    \(H(\cdot)\) is the Shannon entropy function defined as

    \[
        H(u) = -\sum_{\substack{a \in \Sigma \\ u(a) > 0}} u(a) \text{log}(u(a)),
    \]

    where the \(\text{log}\) function is assumed to be the base-2 logarithm, and where
    \(\Sigma\) is an alphabet where \(u \in [0, \infty)^{\Sigma}\) is a vector of
    nonnegative real numbers indexed by \(\Sigma\).

    Further information for computing the von Neumann entropy of a density matrix can be found in Section: "Definitions
    Of Quantum Entropic Functions" from [@watrous2018theory]).

    Examples:
        Consider the following Bell state:

        \[
            u = \frac{1}{\sqrt{2}} \left(|00 \rangle + |11 \rangle \right) \in \mathcal{X}.
        \]

        The corresponding density matrix of \(u\) may be calculated by:

        \[
            \rho = u u^* = \frac{1}{2} \begin{pmatrix}
                             1 & 0 & 0 & 1 \\
                             0 & 0 & 0 & 0 \\
                             0 & 0 & 0 & 0 \\
                             1 & 0 & 0 & 1
                           \end{pmatrix} \in \text{D}(\mathcal{X}).
        \]

        Calculating the von Neumann entropy of \(\rho\) in `|toqito⟩` can be done as follows.

        ```python exec="1" source="above"
        from toqito.state_props import von_neumann_entropy
        import numpy as np
        test_input_mat = np.array(
                [[1 / 2, 0, 0, 1 / 2], [0, 0, 0, 0],
                [0, 0, 0, 0], [1 / 2, 0, 0, 1 / 2]]
            )
        print(von_neumann_entropy(test_input_mat))
        ```

        Consider the density operator corresponding to the maximally mixed state of dimension two

        \[
            \rho = \frac{1}{2}
            \begin{pmatrix}
                1 & 0 \\
                0 & 1
            \end{pmatrix}.
        \]

        As this state is maximally mixed, the von Neumann entropy of \(\rho\) is
        equal to one. We can see this in `|toqito⟩` as follows.

        ```python exec="1" source="above"
        from toqito.state_props import von_neumann_entropy
        import numpy as np
        rho = 1/2 * np.identity(2)
        print(von_neumann_entropy(rho))
        ```

    Args:
        rho: Density operator.

    Returns:
        The von Neumann entropy of `rho`.

    """
    if not is_density(rho):
        raise ValueError("Von Neumann entropy is only defined for density operators.")
    eigs, _ = np.linalg.eig(rho)
    eigs = [eig for eig in eigs if eig > 0]
    return -np.sum(np.real(eigs * np.log2(eigs)))
