"""
Associated code for:
   Search techniques for root-unitary polynomials
   paper by Kiran S. Kedlaya
   this code by Kiran S. Kedlaya and Andre Wibisono

Last modified: 17 Jan 2007
This version fixes an error in Algorithm 4.2 of the original preprint,
and uses binary search in Algorithm 4.3; thanks to Alan Lauder for comments.

EXAMPLES:
sage: polRing.<x> = PolynomialRing(Rationals())
sage: P0 = 3*x^21 + 5*x^20 + 6*x^19 + 7*x^18 + 5*x^17 + 4*x^16 + 2*x^15 \
     - x^14 - 3*x^13 - 5*x^12 - 5*x^11 - 5*x^10 - 5*x^9 - 3*x^8 - x^7 + 2*x^6 \
     + 4*x^5 + 5*x^4 + 7*x^3 + 6*x^2 + 5*x + 3
sage: Q0, cofactor = convertPtoQ(P0)
sage: tree = []
sage: modulus = 3
sage: terms = 3
sage: ans = findQ(Q0, 3^modulus, terms, tree)
sage: print "Number of terminal nodes:", len(tree)
Number of terminal nodes: 63
sage: ans2 = [convertQtoP(i)*cofactor for i in ans]
sage: print ans2
[3*x^21 + 5*x^20 + 6*x^19 + 7*x^18 + 5*x^17 + 4*x^16 + 2*x^15 - x^14 - 3*x^13
- 5*x^12 - 5*x^11 - 5*x^10 - 5*x^9 - 3*x^8 - x^7 + 2*x^6 + 4*x^5 + 5*x^4 +
7*x^3 + 6*x^2 + 5*x + 3]
"""

# Global variables
prec = 12
pari.set_real_precision(prec+1)
tprec = pari(10^(prec+1))
rectprec = pari(10^(-prec-1))
diff = pari(3) * rectprec

def findQ(Q0, m, numGiven, tree=None):
    """
    Return a list of all polynomials Q with roots in [-2,2] such that Q0 is
    congruent to Q modulo m and shares its top numGiven coefficients with Q.
    The optional argument tree, if provided, should be a mutable sequence
    type; findQ will append to tree all terminal nodes enumerated during the
    search. The second return value is the number of terminal nodes.

    Note: we always assume the leading coefficient of Q0 is positive and that
    numGiven is positive.
    """
    polRing = Q0.parent()
    (x,) = polRing.gens()
    loopAns, count, tree = loopStep(pari(Q0/m), Q0.degree(), pari.vector(numGiven), tree)
    return [(Q0 + m*polRing(flip(list(t)))) for t in loopAns], count

def flip(t):
    """
    Return a copy of the list t with the entries in reverse order.
    """
    return t[::-1]

def convertPtoQ(P):
    """
    In the notation of the paper, convert from the polynomial P to Q. The
    second return value is a cofactor, which is either 1, x-1, x+1, or x^2-1.
    """
    polRing = P.parent()
    (x,) = polRing.gens()
    d = P.degree()
    P0 = x^d * P(1/x)
    if (P == P0):
        sign = +1
    elif (P == -P0):
        sign = -1
    else:
        raise RuntimeError, "Polynomial not self-inversive"
    if P.degree() % 2 == 1:
        if sign == 1:
            cofactor = x+1
        else:
            cofactor = x-1
    elif sign == 1:
        cofactor = polRing(1)
    else:
        cofactor = x^2 - 1
    Q = P // cofactor
    coeffs = []
    m = Q.degree() // 2
    for i in reversed(range(m+1)):
        t = Q // (x^2+1)^i
        coeffs.insert(0, t.constant_coefficient())
        Q = (Q % (x^2+1)^i) // x
    return polRing(coeffs), cofactor

def convertQtoP(Q):
    """
    In the notation of the paper, convert from the polynomial Q to P.
    """
    polRing.<x> = Q.parent()
    return polRing(x^(Q.degree()) * Q(x + 1/x))

# Main iterative step
def loopStep(pQ0, dtotal, pgivens, tree):
    d = len(pgivens)
    if d > dtotal:
        isLeaf = True
        ans = [pgivens]
    else:
        e = pQ0.poldegree() - d
        pcoefQ = pgivens.concat(pari.vector(e+1))
        pR = pQ0 + pcoefQ.Pol()
	for _ in xrange(e):
            pR = pR.deriv()
        pR = pR / pari("factorial(" + str(e) + ")")
        lower, upper = seeker(pR)
        ans = []
        isLeaf = (lower > upper)
    if d <= 3:
        print pgivens
    if isLeaf:
        if tree != None:
            tree.append(pgivens)
        count = 1
    else:
        count = 0
        pgivens2 = pgivens.concat(pari.vector(1))
        for t in xrange(lower, upper+1):
            pgivens2[d] = t
            newans, newcount, tree = \
                    loopStep(pQ0, dtotal, pgivens2, tree)
            count += newcount
            ans += newans
    return ans, count, tree

# This routine implements Algorithm 4.3 of the paper. Note that the check on
# the sign of R''(a) has been moved into seekercore for efficiency.
def findextremum(pR, pRpr, interval, diff):
    """
    Given R is a polynomial such that R' has all real roots, and
    interval = (a,b) is an interval containing only one local
    maximum x of R and no root of R'', or if a = b, then find floor(R(x)).
    Exception may be thrown if R'' has a root in [a, x].
    """
    (a,b) = interval
    temp = pR(a)
    t = temp.floor()
    if (a == b):
        return t
    temp = temp + diff*pRpr(a)
    u = temp.floor() 
    if u == t:
        return t
# The following occurs rarely in tested examples, so is not optimized.
    print "Signal: needed to shift (probably because of an integral extremum)"
    polRing.<x> = PolynomialRing(Rationals())
    while t != u:
        v = (t+u)/2
        v = v.ceiling() 
        temp = polRing(pR - v).radical()
# Note that polsturm only looks in the interval (a,b], but temp(a) > 0
# by construction, so it is safe to ignore the left endpoint.
        if pari(temp).polsturm(a,b) > 0:
            t = v
        else:
            u = v - 1
    return t

# This routine implements Algorithm 4.2 of the paper.
def seekercore(pR):
    d = pR.poldegree()
    pRpr = pR.deriv()
    pRprpr = pRpr.deriv()
# The optional argument 1 to polroots uses a faster algorithm, but the
# results are not guaranteed.
    rootlist = pRpr.polroots()
# Reminders:
#    prec = pari.get_real_precision()-1
#    tprec = pari(10^(prec+1))
#    rectprec = pari(10^(-prec-1))
#    diff = pari(3) * rectprec
    sign = +1
    for i in xrange(d+1):
        if (i == 0):
            interval = (pari(2), pari(2))
        elif (i == d):
            interval = (pari(-2), pari(-2))
        else:
            t = rootlist[d-i-1].real() * tprec - 1
            a = t.floor() * rectprec
            b = a + diff
            if b >= interval[0] or pRprpr(a).sign() == sign:
                raise RuntimeError, "Insufficient precision"
            interval = (a,b)
        if (sign == 1):
            t = -findextremum(pR, pRpr, interval, diff)
            if i == 0:
                lower = t
            else:
                if t > upper:
                    return 1,0
                lower = max(lower, t)
        else:
            t = findextremum(-pR, -pRpr, interval, diff)
            if i == 1:
                upper = t
            else:
                if lower > t:
                    return 1,0
                upper = min(upper, t)
        sign = -sign
    return lower, upper

def seeker(pR):
    pRpr = pR.deriv()
    if (not pRpr.issquarefree()) or (pRpr(2) == 0) or (pRpr(-2) == 0):
        print "Signal: needed to treat multiple roots"
	PR.<x> = PolynomialRing(Rationals())
	R = PR(pR)
        Rpr = R.derivative()
        Rprpr = Rpr.derivative()
        gc = Rpr.gcd(Rprpr * (x+2)*(x-2))
        return multiplerootseeker(R, gc)
    return seekercore(pR)

def allinrange(Q,a,b):
    R = Q.radical()
    if (R(a) == 0):
        R = R // (R.parent().gen() - a)
    return (pari(R).polsturm(a,b) == R.degree())

# This routine implements Algorithm 4.1 of the paper.
# This routine has not been optimized because it is rarely invoked.
def multiplerootseeker(R, gc):
    fac = gc.factor()
    for i in range(len(fac)):
        rem = R % fac[i][0]
        if not rem.is_constant():
            return 1,0
        if (i == 0):
            c = -rem.constant_coefficient()
            if c != c.floor():
                return 1,0
        else:
            if (c != -rem.constant_coefficient()):
                return 1,0
    if allinrange(R+c, -2, 2):
        return c,c
    else:
        return 1,0

polRing.<x> = PolynomialRing(Rationals())
P0 = 3*x^21 + 5*x^20 + 6*x^19 + 7*x^18 + 5*x^17 + 4*x^16 + 2*x^15 - x^14 \
     - 3*x^13 - 5*x^12 - 5*x^11 - 5*x^10 - 5*x^9 - 3*x^8 - x^7 + 2*x^6 \
     + 4*x^5 + 5*x^4 + 7*x^3 + 6*x^2 + 5*x + 3
Q0, cofactor = convertPtoQ(P0)
#modulus = 3
#terms = 3
#print "Modulus", modulus, "Terms", terms
#ans, count = findQ(Q0, 3^modulus, terms)
#print "Number of terminal nodes:", count
#ans2 = [convertQtoP(i)*cofactor for i in ans]
#print P0
#print ans2