from sage.all import *
from sage.interfaces.maxima_lib import maxima_lib
from sage.symbolic.expression import Expression


def set_simplification_domain(d):
    """
    Set Maxima's simplification domain.

    INPUT:

      - ``d`` -- The domain, either 'real' or 'complex'.

    TESTS:

    With the default 'complex' domain, we don't simplify this::

        sage: (abs(x)^2).simplify()
        abs(x)^2

    But in the 'real' domain, we do::

        sage: set_simplification_domain('real')
        'real'
        sage: (abs(x)^2).simplify()
        x^2
        sage: set_simplification_domain('complex')
        'complex'

    """
    cmd = 'domain: %s;' % d
    result = maxima_lib._eval_line(cmd)
    return result


def safe_simplify(expr):
    """
    What should be a totally safe simplification operation that works
    a little better than the plain simplify().

    Uses a top-level function because we can't monkey-patch Cython
    classes.
    """
    expr = expr.simplify_factorial()
    expr = expr.simplify_log()
    return expr


def matrix_subs_expr(m, *equations):
    """
    Symbolic matrices have a `subs()` method, but no `subs_expr()`.
    This makes it diffucult to substitute in a list of solutions obtained
    with `solve()`.

    INPUT:

      - ``m`` -- A symbolic matrix.

      - ``equations`` - One or more symbolic equations, presumably for
        the entries of `m`.

    OUTPUT:

    The result of substituting each equation into `m`, one after another.

    EXAMPLES::

    sage: w,x,y,z = SR.var('w,x,y,z')
    sage: A = matrix(SR, [[w,x],[y,z]])
    sage: matrix_subs_expr(A, w == 1, x == 2, y == 3, z == 4)
    [1 2]
    [3 4]

    """
    from sage.symbolic.expression import is_SymbolicEquation

    if not m.base_ring() == SR:
        raise TypeError, 'the matrix "m" must be symbolic'

    if isinstance(equations[0], dict):
        eq_dict = equations[0]
        equations = [ x == eq_dict[x] for x in eq_dict.keys() ]

    if not all([is_SymbolicEquation(eq) for eq in equations]):
            raise TypeError, "each expression must be an equation"

    d = dict([(eq.lhs(), eq.rhs()) for eq in equations])
    return m.subs(d)


def matrix_simplify_full(A):
    """
    Simplify each entry of a symbolic matrix using the
    Expression.simplify_full() method.

    INPUT:

      - ``A`` - The matrix whose entries we should simplify.

    OUTPUT:

    A copy of ``A`` with all of its entries simplified.

    EXAMPLES:

    Symbolic matrices (examples stolen from Expression.simplify_full())
    will have their entries simplified::

        sage: a,n,k = SR.var('a,n,k')
        sage: f1 = sin(x)^2 + cos(x)^2
        sage: f2 = sin(x/(x^2 + x))
        sage: f3 = binomial(n,k)*factorial(k)*factorial(n-k)
        sage: f4 = x*sin(2)/(x^a)
        sage: A = matrix(SR, [[f1,f2],[f3,f4]])
        sage: matrix_simplify_full(A)
        [                1    sin(1/(x + 1))]
        [     factorial(n) x^(-a + 1)*sin(2)]

    But an exception will be raised if ``A`` is not symbolic::

        sage: A = matrix(QQ, [[1,2],[3,4]])
        sage: matrix_simplify_full(A)
        Traceback (most recent call last):
        ...
        ValueError: The base ring of `A` must be the Symbolic Ring.

    """
    if not A.base_ring() == SR:
        raise ValueError('The base ring of `A` must be the Symbolic Ring.')

    M = A.matrix_space()
    return M(map(lambda x: x.simplify_full(), A))