from sympy import MatMul, MatAdd, ZeroMatrix, MatrixBase, Identity, ShapeError, MatrixExpr, S, Number
from sympy.core.decorators import call_highest_priority
from sympy.strategies import (rm_id, unpack, typed, flatten, sort, condition, exhaust,
        do_one, new, glom)

from .supermatbase import SuperMatBase

class SuperMatMul(SuperMatBase, MatMul):
    """
    Redefines some methods of MatMul so as to make them amenable to our application
    """

    _op_priority = 10000

    def __new__(cls, *args, **kwargs):
        return MatMul.__new__(cls, *args, **kwargs)
    
    def as_coeff_mmul(self):
        coeff, matrices = self.as_coeff_matrices()
        return coeff, SuperMatMul(*matrices)
    
    def _eval_transpose(self):
        return SuperMatMul(*[arg.T if isinstance(arg, MatrixExpr) else arg
                             for arg in self.args[::-1]]).doit()
    
    def _eval_inverse(self):
        try:
            return SuperMatMul(*[
                arg.inverse() if isinstance(arg, MatrixExpr) else arg**-1
                    for arg in self.args[::-1]]).doit()
        except ShapeError:
            from .supermatexpr import SuperMatInverse
            return SuperMatInverse(self)

    #def transpose(self):
    #    from .supermatexpr import SuperMatTranspose
    #    return SuperMatTranspose(self).doit()
    
    def doit(self, **kwargs):
        deep = kwargs.get('deep', True)
        if deep:
            args = [arg.doit(**kwargs) for arg in self.args]
        else:
            args = self.args
        return canonicalize(SuperMatMul(*args))

def newmul(*args):
    if args[0] == 1:
        args = args[1:]
    return new(SuperMatMul, *args)

def any_zeros(mul):
    if any([arg.is_zero or (arg.is_Matrix and arg.is_ZeroMatrix)
                       for arg in mul.args]):
        matrices = [arg for arg in mul.args if arg.is_Matrix]
        return ZeroMatrix(matrices[0].rows, matrices[-1].cols)
    return mul

def merge_explicit(matmul):
    """ Merge explicit MatrixBase arguments
    >>> from sympy import MatrixSymbol, eye, Matrix, MatMul, pprint
    >>> from sympy.matrices.expressions.matmul import merge_explicit
    >>> A = MatrixSymbol('A', 2, 2)
    >>> B = Matrix([[1, 1], [1, 1]])
    >>> C = Matrix([[1, 2], [3, 4]])
    >>> X = MatMul(A, B, C)
    >>> pprint(X)
      [1  1] [1  2]
    A*[    ]*[    ]
      [1  1] [3  4]
    >>> pprint(merge_explicit(X))
      [4  6]
    A*[    ]
      [4  6]
    >>> X = MatMul(B, A, C)
    >>> pprint(X)
    [1  1]   [1  2]
    [    ]*A*[    ]
    [1  1]   [3  4]
    >>> pprint(merge_explicit(X))
    [1  1]   [1  2]
    [    ]*A*[    ]
    [1  1]   [3  4]
    """
    if not any(isinstance(arg, MatrixBase) for arg in matmul.args):
        return matmul
    newargs = []
    last = matmul.args[0]
    for arg in matmul.args[1:]:
        if isinstance(arg, (MatrixBase, Number)) and isinstance(last, (MatrixBase, Number)):
            last = last * arg
        else:
            newargs.append(last)
            last = arg
    newargs.append(last)

    return SuperMatMul(*newargs)

def xxinv(mul):
    """ Y * X * X.I -> Y """
    factor, matrices = mul.as_coeff_matrices()
    for i, (X, Y) in enumerate(zip(matrices[:-1], matrices[1:])):
        try:
            if X.is_square and Y.is_square and X == Y.inverse():
                I = Identity(X.rows)
                return newmul(factor, *(matrices[:i] + [I] + matrices[i+2:]))
        except ValueError:  # Y might not be invertible
            pass

    return mul

def remove_ids(mul):
    """ Remove Identities from a MatMul
    This is a modified version of sympy.strategies.rm_id.
    This is necesssary because MatMul may contain both MatrixExprs and Exprs
    as args.
    See Also
    --------
        sympy.strategies.rm_id
    """
    # Separate Exprs from MatrixExprs in args
    factor, mmul = mul.as_coeff_mmul()
    # Apply standard rm_id for MatMuls
    result = rm_id(lambda x: x.is_Identity is True)(mmul)
    if result != mmul:
        return newmul(factor, *result.args)  # Recombine and return
    else:
        return mul

def factor_in_front(mul):
    factor, matrices = mul.as_coeff_matrices()
    if factor != 1:
        return newmul(factor, *matrices)
    return mul

rules = (any_zeros, remove_ids, xxinv, unpack, rm_id(lambda x: x == 1),
         merge_explicit, factor_in_front, flatten)

canonicalize = exhaust(typed({SuperMatMul: do_one(*rules)}))

def only_squares(*matrices):
    """ factor matrices only if they are square """
    if matrices[0].rows != matrices[-1].cols:
        raise RuntimeError("Invalid matrices being multiplied")
    out = []
    start = 0
    for i, M in enumerate(matrices):
        if M.cols == matrices[start].rows:
            out.append(SuperMatMul(*matrices[start:i+1]).doit())
            start = i+1
    return out


from sympy.assumptions.ask import ask, Q
from sympy.assumptions.refine import handlers_dict


def refine_SuperMatMul(expr, assumptions):
    """
    >>> from sympy import MatrixSymbol, Q, assuming, refine
    >>> X = MatrixSymbol('X', 2, 2)
    >>> expr = X * X.T
    >>> print(expr)
    X*X.T
    >>> with assuming(Q.orthogonal(X)):
    ...     print(refine(expr))
    I
    """
    newargs = []
    exprargs = []

    for args in expr.args:
        if args.is_Matrix:
            exprargs.append(args)
        else:
            newargs.append(args)

    last = exprargs[0]
    for arg in exprargs[1:]:
        if arg == last.T and ask(Q.orthogonal(arg), assumptions):
            last = Identity(arg.shape[0])
        elif arg == last.conjugate() and ask(Q.unitary(arg), assumptions):
            last = Identity(arg.shape[0])
        else:
            newargs.append(last)
            last = arg
    newargs.append(last)

    return SuperMatMul(*newargs)


handlers_dict['SuperMatMul'] = refine_SuperMatMul

from .supermatadd import SuperMatAdd