from sage.structure.element import Element
from sage.combinat.free_module import CombinatorialFreeModule

# TODO: doc and tests for all of those
class SymbolicMatrix(SageObject):
    def __init__(self, name, nrows, ncols, inverted = False, transposed = False):
        #Element.__init__(self, parent)
        self._name = name
        self._nrows = nrows
        self._ncols = ncols
        self._inverted = inverted
        self._transposed = transposed

    def _repr_(self):
        result = self._name
        if self._inverted:
            result += "^-1"
        if self._transposed:
            result += "^t"
        return result

    def transpose(self):
        result = copy(self)
        result._transposed = not self._transposed
        (result._nrows, result._ncols) = (self._ncols, self._nrows)
        return result

    def __invert__(self):
        assert self._nrows == self._ncols, "Can't inverse non square matrix"
        result = copy(self)
        result._inverted = not self._inverted
        return result


class SymbolicMatrixAlgebra(CombinatorialFreeModule):
    r"""

    EXAMPLES::

        sage: Alg = SymbolicMatrixAlgebra(QQ)

        sage: A = Alg.matrix("A",3,2)
        sage: B = Alg.matrix("B",2,3)
        sage: C = Alg.matrix("C",2,3)
        sage: D = Alg.matrix("D",3,3)

    Example 1: Adding/Multiplying matrices of correct size::

        sage: A * (B + C)
        A B + A C

    Example 2: Transposing a sum of matrices::

        sage: (B + C).transpose()
        C^t + B^t

    Example 3: Transposing a product of matrices::

        sage: (A * B).transpose()
        B^t A^t

    Example 4: Inverting a product of matrices::

        sage: (A * B)^-1
        B^-1 A^-1

    Example 5: Multiplying by its inverse::

        sage: D * D^-1 # todo: not implemented
        I

    TODO: decide on the best output; do we want to be able to
    copy-paste back? do we prefer short notations?

    .. warnings::

    The identity does not know it is size, so the following should
    complain, but does not::

        sage: D * D^-1 * A

    TODO: describe all the abuses
    """

    def __init__(self, R):
        """
        EXAMPLES::

            sage: A = AlgebrasWithBasis(QQ).example(); A
            An example of an algebra with basis: the free algebra on the generators ('a', 'b', 'c') over Rational Field
            sage: TestSuite(A).run()

        """
        CombinatorialFreeModule.__init__(self, R, Words(), category = AlgebrasWithBasis(R))

    def matrix(self, name, nrows, ncols):
        """ TODO: doctest"""
        return self.monomial(Word([SymbolicMatrix(name, nrows, ncols)]))

    def _repr_(self):
        """
        EXAMPLES::

            sage: SymbolicMatrixAlgebra(QQ)
            The symbolic matrix algebra over Rational Field
        """
        return "The symbolic matrix algebra over %s"%(self.base_ring())

    @cached_method
    def one_basis(self):
        """
        Returns the empty word, which index the one of this algebra,
        as per :meth:`AlgebrasWithBasis.ParentMethods.one_basis`.

        EXAMPLES::

            sage: Alg = SymbolicMatrixAlgebra(QQ)
            sage: Alg.one_basis()
            word:
            sage: A.one()
            I
        """
        return self.basis().keys()([])

    def product_on_basis(self, w1, w2):
        r"""
        Product of basis elements, as per :meth:`AlgebrasWithBasis.ParentMethods.product_on_basis`.

        EXAMPLES::

            sage: Alg = SymbolicMatrixAlgebra(QQ)
            sage: P = Alg.matrix("P", 3, 2)
            sage: Q = Alg.matrix("Q", 3, 2)
            sage: R = Alg.matrix("R", 2, 2)
            sage: S = Alg.matrix("S", 2, 3)
            sage: P * P
            Traceback (most recent call last):
            ...
            AssertionError: Non-conformable matrices: matrix sizes are incompatible for multiplication
            sage: P * R * S
            P*R*S
            sage: (P+Q) * R
            P*R + Q*R

            sage: (P+Q) * R * S
            P*R*S + Q*R*S
            sage: S * (P+Q) * R
        """
        if len(w1) == 0:
            return self.monomial(w2)
        if len(w2) == 0:
            return self.monomial(w1)
        assert w1[-1]._ncols == w2[0]._nrows, "Non-conformable matrices: matrix sizes are incompatible for multiplication"
        # TODO: handle cancelations between w1[-1] and w2[0]
        return self.monomial(w1 + w2)

    # TODO: define an_element

    def _repr_term(self, t):
        """
        EXAMPLES::

            sage: Alg.one() # indirect doctest
            I
            sage: Alg.an_element() # todo: not implemented
        """
        if len(t) == 0:
            return 'I'
        else:
            return ' '.join(repr(m) for m in t)


    class Element(CombinatorialFreeModule.Element):

        def transpose(self):
            """

            EXAMPLES::

                sage: Alg = SymbolicMatrixAlgebra(QQ)
                sage: A = Alg.matrix("A", 3, 2)
                sage: B = Alg.matrix("B", 3, 2)
                sage: C = Alg.matrix("C", 2, 2)
                sage: D = Alg.matrix("D", 2, 3)
                sage: x = D * (A+B) * C
                sage: x
                D B C + D A C
                sage: x.transpose()
                C^t B^t D^t + C^t A^t D^t

            """
            return self.map_support(lambda w: Word(m.transpose() for m in reversed(w)))

        def __invert__(self):
            """
            EXAMPLES::

                sage: Alg = SymbolicMatrixAlgebra(QQ)
                sage: A = Alg.matrix("A", 2, 2)
                sage: B = Alg.matrix("B", 2, 2)
                sage: C = Alg.matrix("C", 2, 2)
                sage: D = Alg.matrix("D", 3, 2)
                sage: E = Alg.matrix("E", 2, 3)
                sage: ~(A*B*C)
                C^-1 B^-1 A^-1
                sage: ~(A*B^-1*C^-1)
                C B A^-1
                sage: ~(D*C*E)
                Traceback (most recent call last):
                ...
                AssertionError: Can't inverse non square matrix
                sage: ~(A+B)
                Traceback (most recent call last):
                ...
                AssertionError: Can't inverse non trivial linear combinations of matrices
            """
            assert len(self) == 1, "Can't inverse non trivial linear combinations of matrices"
            (w, c) = self.leading_item()
            return self.parent().term(Word(~m for m in reversed(w)), c)