"""
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 is_symmetric_psd(A):
    """
    Determine whether or not the matrix ``A`` is symmetric
    positive-semidefinite.

    INPUT:

    - ``A`` - The matrix in question

    OUTPUT:

    Either ``True`` if ``A`` is symmetric positive-semidefinite, or
    ``False`` otherwise.

    EXAMPLES:

    Every completely positive matrix is symmetric
    positive-semidefinite::

        sage: v = vector(map(abs, random_vector(ZZ, 10)))
        sage: A = v.column() * v.row()
        sage: is_symmetric_psd(A)
        True

    The following matrix is symmetric but not positive semidefinite::

        sage: A = matrix(ZZ, [[1, 2], [2, 1]])
        sage: is_symmetric_psd(A)
        False

    This matrix isn't even symmetric::

        sage: A = matrix(ZZ, [[1, 2], [3, 4]])
        sage: is_symmetric_psd(A)
        False

    """

    if A.base_ring() == SR:
        msg = 'The matrix ``A`` cannot be symbolic.'
        raise ValueError.new(msg)

    # First make sure that ``A`` is symmetric.
    if not A.is_symmetric():
        return False

    # If ``A`` is symmetric, we only need to check that it is positive
    # semidefinite. For that we can consult its minimum eigenvalue,
    # which should be zero or greater. Since ``A`` is symmetric, its
    # eigenvalues are guaranteed to be real.
    return min(A.eigenvalues()) >= 0


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. The base ring of ``A`` must be either
              exact or the symbolic ring (to compute eigenvalues), and it
              must be a field so that we can take its algebraic closure
              (necessary to e.g. take square roots).

    OUTPUT:

    A matrix ``X`` such that `A = XX^{T}`. The base field of ``X`` will
    be the algebraic closure of the base field of ``A``.

    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:

    Create a symmetric positive-semidefinite matrix over the symbolic
    ring and factor it::

        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

    Attempt to factor the same matrix over ``RR`` which won't work
    because ``RR`` isn't exact::

        sage: A = matrix(RR, [[0, 2, 3], [2, 0, 0], [3, 0, 0]])
        sage: factor_psd(A)
        Traceback (most recent call last):
        ...
        ValueError: The base ring of ``A`` must be either exact or symbolic.

    Attempt to factor the same matrix over ``ZZ`` which won't work
    because ``ZZ`` isn't a field::

        sage: A = matrix(ZZ, [[0, 2, 3], [2, 0, 0], [3, 0, 0]])
        sage: factor_psd(A)
        Traceback (most recent call last):
        ...
        ValueError: The base ring of ``A`` must be a field.

    """

    if not A.base_ring().is_exact() and not A.base_ring() is SR:
        msg = 'The base ring of ``A`` must be either exact or symbolic.'
        raise ValueError(msg)

    if not A.base_ring().is_field():
        raise ValueError('The base ring of ``A`` must be a field.')

    if not A.base_ring() is SR:
        # Change the base field of ``A`` so that we are sure we can take
        # roots. The symbolic ring has no algebraic_closure method.
        A = A.change_ring(A.base_ring().algebraic_closure())


    # 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()