"""Module providing a user-friendly interface to the polynomial algorithms."""

from sympy import S, Add, Symbol, Mul
from sympy.polynomials.base import coeff_ring, get_numbers, Polynomial
from sympy.polynomials import div_
from sympy.polynomials import factor_
from sympy.polynomials import groebner_
from sympy.polynomials import roots_

def div(f, g, var=None, order=None, coeff=None):
    """Division with remainder.

    A thin wrapper that returns SymPy expressions coming from L{div_.div}.

    """

    q, r = div_.div(f, g, var, order, coeff)

    if isinstance(q, list):
        return map(lambda x: x.sympy_expr, q), r.sympy_expr
    else:
        return q.sympy_expr, r.sympy_expr

def quo(*args, **kwargs):
    return div(*args, **kwargs)[0]

def rem(*args, **kwargs):
    return div(*args, **kwargs)[1]

def factor(f, var=None, order=None):
    """Factorization of polynomials over the rationals.

    A thin wrapper that returns a SymPy expression by multiplying the
    different factors coming from L{factor_.factor}

    """

    # Re-assemble factors:
    return Mul(*[ff.sympy_expr for ff in factor_.factor(f, var, order)])


def gcd(f, g, var=None, order=None, coeff=None):
    """Greatest common divisor of two polynomials.

    A thin wrapper returning a SymPy expression from L{div_.gcd}'s
    result after checking the coefficients' type.

    """

    # First check if the polynomials have integer coefficients.
    if coeff is None:
        coeff_f = coeff_ring(get_numbers(f))
        coeff_g = coeff_ring(get_numbers(g))
        if coeff_f == 'int' and coeff_g == 'int':
            coeff = 'int'
        if var is not None:
            if isinstance(var, Symbol):
                var = [var]
            if [v for v in f.atoms(Symbol).union(g.atoms(Symbol))
                if v not in var]:
                coeff = 'sym'

    return div_.gcd(f, g, var, order, coeff).sympy_expr

def groebner(f, var=None, order=None, reduced=True):
    """Computes the (reduced) Groebner base of given polynomials.

    Thin wrapper returning SymPy expressions from L{groebner_.groebner}.

    """

    g = groebner_.groebner(f, var, order, reduced)
    return map(lambda p: p.sympy_expr, g)

def lcm(f, g, var=None, order=None, coeff=None):
    """Least common divisor of two polynomials.

    A thin wrapper returning a SymPy expression from L{div_.lcm}'s result.

    """

    # First check if the polynomials have integer coefficients.
    if coeff is None:
        coeff_f = coeff_ring(get_numbers(f))
        coeff_g = coeff_ring(get_numbers(g))
        if coeff_f == 'int' and coeff_g == 'int':
            coeff = 'int'
        if var is not None:
            if isinstance(var, Symbol):
                var = [var]
            if [v for v in f.atoms(Symbol).union(g.atoms(Symbol))
                if v not in var]:
                coeff = 'sym'

    return div_.lcm(f, g, var, order, coeff).sympy_expr

def resultant(f, g, x, method='bezout'):
    """Computes resultant of two polynomials in one variable. This
       method can be used to verify if selected polynomials have
       common root withot factoring them or computing any GCD's.

       It can be also be used for variable elemination in case of
       bivariate polynomials. Just assume that one of the var
       is a parameter and compute resultant with respect to the other
       one and you will get univariate polynomial in single variable.

       For now two algorithm have been implemented, based on
       Sylvester and Bezout matrices. The default is Bezout.

       TODO: Make Bezout approach run in O(s**2). Currently
             it is O(s**3), where s = max(deg(f), deg(g)).

    """

    from sympy.matrices import zeros

    fp = Polynomial(f, var=x).coeffs
    gp = Polynomial(g, var=x).coeffs

    m, n = int(fp[0][1]), int(gp[0][1])

    if method is 'sylvester':
        M = zeros(m+n)

        for i in range(n):
            for coeff, j in fp:
                M[i, m-int(j)+i] = coeff

        for i in range(m):
            for coeff, j in gp:
                M[i+n, n-int(j)+i] = coeff

        return M.det()
    elif method is 'bezout':
        if n > m:
            s, fp, gp = n, gp, fp
        else:
            s = m

        p = [S.Zero] * (s+1)
        q = [S.Zero] * (s+1)

        for coeff, j in fp:
            p[int(j)] = coeff

        for coeff, j in gp:
            q[int(j)] = coeff

        M = zeros(s)

        for i in range(s):
            for j in range(i, s):
                z = 1 + min(i, s-1-j)
                terms = [S.Zero] * z

                for k in range(z):
                    terms[k] = p[j+k+1]*q[i-k] - p[i-k]*q[j+k+1]

                M[i, j] = M[j, i] = Add(*terms)

        if (s*(s-1)/2) % 2 == 0:
            b = p[-1]**abs(n-m)
        else:
            b = -p[-1]**abs(n-m)

        return (1/b) * M.det()
    else:
        raise ValueError("Invalid method: '%s'" % method)


def sqf(f, var=None, order=None, coeff=None):
    """Square-free decomposition.

    Thin wrapper that multiplies the factors of L{factor_.sqf} as
    SymPy expressions.

    """

    a = factor_.sqf(f, var, order, coeff)
    result = S.One
    for i, p in enumerate(a):
        result *= (p.sympy_expr)**(i+1)
    return result


def sqf_part(f, var=None, order=None):
    """Square-free part.

    Thin wrapper that returns the result of L{factor_.sqf_part} as a
    SymPy expression.

    """

    return factor_.sqf_part(f, var=var, order=order).sympy_expr


def egcd(p, q, x):
    U = (p, S.One, S.Zero)
    V = (q, S.Zero, S.One)

    while V[0] is not S.Zero:
        u, v = U[0], V[0]

        if u.has(x):
            if v.has(x):
                q = quo(u, v, x)
            else:
                q = u / v
        else:
            if v.has(x):
                q = S.Zero
            else:
                q = u / v

        U, V = V, [ (a - q*b).expand() for a, b in zip(U, V) ]

    return U
