from sage.structure.element import is_Matrix

def _scale(x, alpha):
    r"""
    Scale the vector, matrix, or cartesian-product-of-those-things
    ``x`` by ``alpha``.

    This works around the inability to scale certain elements of
    Cartesian product spaces, as reported in

      https://trac.sagemath.org/ticket/31435

    ..WARNING:

        This will do the wrong thing if you feed it a tuple or list.

    SETUP::

        sage: from mjo.eja.eja_utils import _scale

    EXAMPLES::

        sage: v = vector(QQ, (1,2,3))
        sage: _scale(v,2)
        (2, 4, 6)
        sage: m = matrix(QQ, [[1,2],[3,4]])
        sage: M = cartesian_product([m.parent(), m.parent()])
        sage: _scale(M((m,m)), 2)
        ([2 4]
        [6 8], [2 4]
        [6 8])

    """
    if hasattr(x, 'cartesian_factors'):
        P = x.parent()
        return P(tuple( _scale(x_i, alpha)
                        for x_i in x.cartesian_factors() ))
    else:
        return x*alpha


def _all2list(x):
    r"""
    Flatten a vector, matrix, or cartesian product of those things
    into a long list.

    If the entries of the matrix themselves belong to a real vector
    space (such as the complex numbers which can be thought of as
    pairs of real numbers), they will also be expanded in vector form
    and flattened into the list.

    SETUP::

        sage: from mjo.eja.eja_utils import _all2list
        sage: from mjo.hurwitz import (QuaternionMatrixAlgebra,
        ....:                          Octonions,
        ....:                          OctonionMatrixAlgebra)

    EXAMPLES::

        sage: _all2list([[1]])
        [1]

    ::

        sage: V1 = VectorSpace(QQ,2)
        sage: V2 = MatrixSpace(QQ,2)
        sage: x1 = V1([1,1])
        sage: x2 = V1([1,-1])
        sage: y1 = V2.one()
        sage: y2 = V2([0,1,1,0])
        sage: _all2list((x1,y1))
        [1, 1, 1, 0, 0, 1]
        sage: _all2list((x2,y2))
        [1, -1, 0, 1, 1, 0]
        sage: M = cartesian_product([V1,V2])
        sage: _all2list(M((x1,y1)))
        [1, 1, 1, 0, 0, 1]
        sage: _all2list(M((x2,y2)))
        [1, -1, 0, 1, 1, 0]

    ::

        sage: _all2list(Octonions().one())
        [1, 0, 0, 0, 0, 0, 0, 0]
        sage: _all2list(OctonionMatrixAlgebra(1).one())
        [1, 0, 0, 0, 0, 0, 0, 0]

    ::

        sage: _all2list(QuaternionAlgebra(QQ, -1, -1).one())
        [1, 0, 0, 0]
        sage: _all2list(QuaternionMatrixAlgebra(1).one())
        [1, 0, 0, 0]

    ::

        sage: V1 = VectorSpace(QQ,2)
        sage: V2 = OctonionMatrixAlgebra(1,field=QQ)
        sage: C = cartesian_product([V1,V2])
        sage: x1 = V1([3,4])
        sage: y1 = V2.one()
        sage: _all2list(C( (x1,y1) ))
        [3, 4, 1, 0, 0, 0, 0, 0, 0, 0]

    """
    if hasattr(x, 'to_vector'):
        # This works on matrices of e.g. octonions directly, without
        # first needing to convert them to a list of octonions and
        # then recursing down into the list. It also avoids the wonky
        # list(x) when x is an element of a CFM. I don't know what it
        # returns but it aint the coordinates. We don't recurse
        # because vectors can only contain ring elements as entries.
        return x.to_vector().list()

    if is_Matrix(x):
        # This sucks, but for performance reasons we don't want to
        # call _all2list recursively on the contents of a matrix
        # when we don't have to (they only contain ring elements
        # as entries)
        return x.list()

    try:
        xl = list(x)
    except TypeError: # x is not iterable
        return [x]

    if xl == [x]:
        # Avoid the retardation of list(QQ(1)) == [1].
        return [x]

    return sum( map(_all2list, xl) , [])


def gram_schmidt(v, inner_product=None):
    """
    Perform Gram-Schmidt on the list ``v`` which are assumed to be
    vectors over the same base ring. Returns a list of orthonormalized
    vectors over the same base ring, which means that your base ring
    needs to contain the appropriate roots.

    SETUP::

        sage: from mjo.eja.eja_utils import gram_schmidt

    EXAMPLES:

    If you start with an orthonormal set, you get it back. We can use
    the rationals here because we don't need any square roots::

        sage: v1 = vector(QQ, (1,0,0))
        sage: v2 = vector(QQ, (0,1,0))
        sage: v3 = vector(QQ, (0,0,1))
        sage: v = [v1,v2,v3]
        sage: gram_schmidt(v) == v
        True

    The usual inner-product and norm are default::

        sage: v1 = vector(AA,(1,2,3))
        sage: v2 = vector(AA,(1,-1,6))
        sage: v3 = vector(AA,(2,1,-1))
        sage: v = [v1,v2,v3]
        sage: u = gram_schmidt(v)
        sage: all( u_i.inner_product(u_i).sqrt() == 1 for u_i in u )
        True
        sage: bool(u[0].inner_product(u[1]) == 0)
        True
        sage: bool(u[0].inner_product(u[2]) == 0)
        True
        sage: bool(u[1].inner_product(u[2]) == 0)
        True


    But if you supply a custom inner product, the result is
    orthonormal with respect to that (and not the usual inner
    product)::

        sage: v1 = vector(AA,(1,2,3))
        sage: v2 = vector(AA,(1,-1,6))
        sage: v3 = vector(AA,(2,1,-1))
        sage: v = [v1,v2,v3]
        sage: B = matrix(AA, [ [6, 4, 2],
        ....:                  [4, 5, 4],
        ....:                  [2, 4, 9] ])
        sage: ip = lambda x,y: (B*x).inner_product(y)
        sage: norm = lambda x: ip(x,x)
        sage: u = gram_schmidt(v,ip)
        sage: all( norm(u_i) == 1 for u_i in u )
        True
        sage: ip(u[0],u[1]).is_zero()
        True
        sage: ip(u[0],u[2]).is_zero()
        True
        sage: ip(u[1],u[2]).is_zero()
        True

    This Gram-Schmidt routine can be used on matrices as well, so long
    as an appropriate inner-product is provided::

        sage: E11 = matrix(AA, [ [1,0],
        ....:                    [0,0] ])
        sage: E12 = matrix(AA, [ [0,1],
        ....:                    [1,0] ])
        sage: E22 = matrix(AA, [ [0,0],
        ....:                    [0,1] ])
        sage: I = matrix.identity(AA,2)
        sage: trace_ip = lambda X,Y: (X*Y).trace()
        sage: gram_schmidt([E11,E12,I,E22], inner_product=trace_ip)
        [
        [1 0]  [                  0 0.7071067811865475?]  [0 0]
        [0 0], [0.7071067811865475?                   0], [0 1]
        ]

    It even works on Cartesian product spaces whose factors are vector
    or matrix spaces::

        sage: V1 = VectorSpace(AA,2)
        sage: V2 = MatrixSpace(AA,2)
        sage: M = cartesian_product([V1,V2])
        sage: x1 = V1([1,1])
        sage: x2 = V1([1,-1])
        sage: y1 = V2.one()
        sage: y2 = V2([0,1,1,0])
        sage: z1 = M((x1,y1))
        sage: z2 = M((x2,y2))
        sage: def ip(a,b):
        ....:     return a[0].inner_product(b[0]) + (a[1]*b[1]).trace()
        sage: U = gram_schmidt([z1,z2], inner_product=ip)
        sage: ip(U[0],U[1])
        0
        sage: ip(U[0],U[0])
        1
        sage: ip(U[1],U[1])
        1

    TESTS:

    Ensure that zero vectors don't get in the way::

        sage: v1 = vector(AA,(1,2,3))
        sage: v2 = vector(AA,(1,-1,6))
        sage: v3 = vector(AA,(0,0,0))
        sage: v = [v1,v2,v3]
        sage: len(gram_schmidt(v)) == 2
        True

    """
    if len(v) == 0:
        # cool
        return v

    V = v[0].parent()

    if inner_product is None:
        inner_product = lambda x,y: x.inner_product(y)

    sc = lambda x,a: a*x
    if hasattr(V, 'cartesian_factors'):
        # Only use the slow implementation if necessary.
        sc = _scale

    def proj(x,y):
        # project y onto the span of {x}
        return sc(x, (inner_product(x,y)/inner_product(x,x)))

    def normalize(x):
        # Don't extend the given field with the necessary
        # square roots. This will probably throw weird
        # errors about the symbolic ring if you e.g. try
        # to use it on a set of rational vectors that isn't
        # already orthonormalized.
        return sc(x, ~inner_product(x,x).sqrt())

    v_out = [] # make a copy, don't clobber the input

    for (i, v_i) in enumerate(v):
        ortho_v_i = v_i - V.sum( proj(v_out[j],v_i) for j in range(i) )
        if not ortho_v_i.is_zero():
            v_out.append(normalize(ortho_v_i))

    return v_out