"""
The positive semidefinite cone `$S^{n}_{+}$` is the cone consisting of
all symmetric positive-semidefinite matrices (as a subset of
`$\mathbb{R}^{n \times n}$`
"""

from sage.all import *

# Sage doesn't load ~/.sage/init.sage during testing (sage -t), so we
# have to explicitly mangle our sitedir here so that "mjo.symbolic"
# resolves.
from os.path import abspath
from site import addsitedir
addsitedir(abspath('../../'))
from mjo.symbolic import matrix_simplify_full


def unit_eigenvectors(A):
    """
    Return the unit eigenvectors of a symmetric positive-definite matrix.

    INPUT:

      - ``A`` - The matrix whose eigenvectors we want to compute.

    OUTPUT:

    A list of (eigenvalue, eigenvector) pairs where each eigenvector is
    associated with its paired eigenvalue of ``A`` and has norm `1`.

    EXAMPLES::

        sage: A = matrix(QQ, [[0, 2, 3], [2, 0, 0], [3, 0, 0]])
        sage: unit_evs = unit_eigenvectors(A)
        sage: bool(unit_evs[0][1].norm() == 1)
        True
        sage: bool(unit_evs[1][1].norm() == 1)
        True
        sage: bool(unit_evs[2][1].norm() == 1)
        True

    """
    # This will give us a list of lists whose elements are the
    # eigenvectors we want.
    ev_lists = [ (val,vecs) for (val,vecs,multiplicity)
                            in A.eigenvectors_right() ]

    # Pair each eigenvector with its eigenvalue and normalize it.
    evs = [ [(l, vec/vec.norm()) for vec in vecs] for (l,vecs) in ev_lists ]

    # Flatten the list, abusing the fact that "+" is overloaded on lists.
    return sum(evs, [])




def factor_psd(A):
    """
    Factor a symmetric positive-semidefinite matrix ``A`` into
    `XX^{T}`.

    INPUT:

      - ``A`` - The matrix to factor.

    OUTPUT:

    A matrix ``X`` such that `A = XX^{T}`.

    ALGORITHM:

    Since ``A`` is symmetric and positive-semidefinite, we can
    diagonalize it by some matrix `$Q$` whose columns are orthogonal
    eigenvectors of ``A``. Then,

      `$A = QDQ^{T}$`

    From this representation we can take the square root of `$D$`
    (since all eigenvalues of ``A`` are nonnegative). If we then let
    `$X = Q*sqrt(D)*Q^{T}$`, we have,

      `$XX^{T} = Q*sqrt(D)*Q^{T}Q*sqrt(D)*Q^{T} = Q*D*Q^{T} = A$`

    as desired.

    In principle, this is the algorithm used, although we ignore the
    eigenvectors corresponding to the eigenvalue zero. Thus if `$rank(A)
    = k$`, the matrix `$Q$` will have dimention `$n \times k$`, and
    `$D$` will have dimension `$k \times k$`. In the end everything
    works out the same.

    EXAMPLES::

        sage: A = matrix(SR, [[0, 2, 3], [2, 0, 0], [3, 0, 0]])
        sage: X = factor_psd(A)
        sage: A2 = matrix_simplify_full(X*X.transpose())
        sage: A == A2
        True

    """

    # Get the eigenvectors, and filter out the ones that correspond to
    # the eigenvalue zero.
    all_evs = unit_eigenvectors(A)
    evs = [ (val,vec) for (val,vec) in all_evs if not val == 0 ]

    d = [ sqrt(val) for (val,vec) in evs ]
    root_D = diagonal_matrix(d).change_ring(A.base_ring())

    Q = matrix(A.base_ring(), [ vec for (val,vec) in evs ]).transpose()

    return Q*root_D*Q.transpose()