# Computes the canonical lifting and elliptic Teichmuller lift of points
# of an elliptic curve or the minimal degree lifting.

load 'sec_coord.sage'

################

# need to define an integration function
def my_int(f,x):
    # x is the variable!
    P=f.parent()
    p=P.characteristic()

    if p== 0:
        F=QQ()
    else:
        F=GF(p)
    
    v=f.monomials()
    return sum([ f.monomial_coefficient(mon)*mon*x/F((mon.degree(x)+1)) for mon in v ])


# need to extract coefficients with respect to one variable
def my_coef(f,x):
    # x is the variable!
    d=f.degree(x)
    return [ f.coefficient({x : i}) for i in range(d+1) ]


"""
# powers of powers of p
def pol_p_power(f,p,k=1):
    # computes f^(p^k)
    pp=parent(f).characteristic()
    if pp== p:
        # right call!
        res=0
        for mon in f.monomials():
            res+=f.monomial_coefficient(mon)^(p^k)*mon^(p^k)
        return res
    else:
        # wrong p -- regular power
        f^(p^k)
"""

#########################################

# main function
def lift(a0,b0, tm=False, test=False, ncoord=2, minimal=False, choice=2):
    # tm = time it?
    # ncoord = no. of coordinates to compute (1 -> x1,P1; 2 -> x2,P2)
    # test = test result
    # minimal = compute minimal instead of canonical lift?
    # choice = make a_1=a_2=0 (choice=2) or b_1=b_2=0 (choice=1)
    
    if tm:
        import time
        ttime=time.time()
        ptime=ttime
        print "Testing input."

    F1=a0.parent() # field of definition

    if F1 != b0.parent():
        print "ERROR: ",a0, " and ", b0, " are not in the same field"
        return 0


    p=F1.characteristic() # characteristic

    if p == 0:
        print "ERROR: Char. of the field is zero!"
        return 0


    if 4*a0^3+27*b0^2 == 0:
        print "ERROR: singular curve"
        return 0

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Create polynomials."

    if choice == 1:
        # we will assume b1=0
        PR1.<x0,c0,c1,a1>=PolynomialRing(F1,4)
        b1=0
    else:
        # we will assume a1=0;
        PR1.<x0,c0,c1,b1>=PolynomialRing(F1,4)
        a1=0

    f=PR1(x0^3+a0*x0+b0) # cubic
    sf=SpcFct(f) # sfct (psi_1) of f

    # y0^(p-1)
    y0pm1=f^(((p-1)//2))

    # y0^(p+1)
    y0pp1=y0pm1*f


    # Hasse Invariant
    HI=(y0pm1.coefficient(x0^(p-1))).constant_coefficient()

    if HI == 0:
        print "ERROR: hasse inv. = 0!"
        return 0

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Compute x1"


    # dx1/dx0
    tmpx1=PR1((HI^(-1))*y0pm1-x0^(p-1))

    x1=my_int(tmpx1,x0)+c0+c1*x0^p

    del tmpx1

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Compute RHS."

    # RHS - 2nd coord
    rhs1=PR1((3*x0^(2*p)+a0^p)*x1+a1*x0^p+b1) + sf


    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Long division of RHS."

    # Find canditate to P1 (y1=y0*P1)
    # rem1 must be zero!
    rhs1*=F1(1/2)
    P1=0
    deg1=rhs1.degree(x0)
    deg2=(3*((p+1)//2))

    while deg1 >= deg2:
        lterm=rhs1.coefficient({x0:deg1})*x0^(deg1-deg2)
        rhs1-=lterm*y0pp1
        P1+=lterm
        deg1=rhs1.degree(x0)

    rem1=rhs1

    del rhs1

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Separate coefficients that must be zero."

    # separate the coefficients of rem1 -- they must be zero
    vcoef=my_coef(rem1,x0)

    len=rem1.degree(x0)+1

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Create matrix and vector."

    # 3 by len matrix -- 3 variables: (c0,c1,a1) or (c0,c1,b1)
    M1=matrix(F1,3,len,[[(vcoef[j].coefficient({PR1.gens()[i+1] : 1})).constant_coefficient() for j in range(len)] for i in range(3)])
    # print M1

    # We now solve the linear system.

    # term independent of the variables
    # (note: the 4th variable is x0, which does not show up in the coef.)
    v1=matrix(F1,1,len,[-vcoef[i].constant_coefficient() for i in range(len)])


    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Solve the linear system."
        
    # atime=time.time()
    # find ONE solution!
    vsol=M1.solve_left(v1)
    # ftime=time.time()
    # print "Time: ", ftime-atime

    # IMPROVE?
    # We could find all...  

    # Now, compute the solutions

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Find x1 from solution."

    # x1 only involves c0 and c1, but not b1
    x1=x1.subs({PR1.1: vsol[0,0], PR1.2 : vsol[0,1]})

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Find P1 from solution."

    P1=P1.subs({PR1.1 : vsol[0,0], PR1.2 : vsol[0,1], PR1.3 : vsol[0,2]})


    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
    
    # ********** test ***************

    if test:
        print "%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%"
        if rem1.subs({PR1.1 : vsol[0,0], PR1.2 : vsol[0,1], PR1.3 : vsol[0,2]}) == 0:
            print "test1 is OK"
        else:
            print "test 1 in bad!"
            print "rem1 = ", rem1.subs({PR1.1 : vsol[0,0], PR1.2 : vsol[0,1], PR1.3 : vsol[0,2]})

        print "%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%"

    del M1, rem1, vcoef, len, v1


    if tm:
        ftime=time.time()
        print "Done with 1st coord."
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"


    # return a1, b1, x1, P1
    # if ncoord=1, stops here!
    if ncoord == 1:
        if choice == 1:
            return vsol[0,2], 0, x1, P1
        else:
            return 0, vsol[0,2], x1, P1

    """
    //////////////////////////////////////////////////
    ////////////// SECOND COORDINATE /////////////////
    //////////////////////////////////////////////////
    """

    #  Now to red. mod p^3

    if tm:
        ptime=time.time()
        print "Creating new structures."


    # introduce the new variables: a2 (and b2=0) or b2 (and a2=0)
    #  and the coeff. of x0^(pk) in x2
    # NOTE: need one less var. than if a2!=0 and b2!=0!!!

    if choice == 1:
        vars=['x0','a2']+['d%s' %i for i in range((3*p+1)//2) ]        
        PR2=PolynomialRing(F1,vars)
        # PR2.inject_variables(verbose=False) # define the names for interactive use
        # the above does not seem to work inside a function...
        # we use PR2.x instead
        b1=0
        b2=0
        a2=PR2.1
    else:
        vars=['x0','b2']+['d%s' %i for i in range((3*p+1)//2) ]        
        PR2=PolynomialRing(F1,vars)
        # PR2.inject_variables(verbose=False) # define the names for interactive use
        # the above does not seem to work inside a function...
        # we use PR2.x instead
        a1=0
        a2=0
        b2=PR2.1
  
    # We need to convert the previous results to this new ring.
    def Convert(P,M):
        # converts P to new structure M
        PP=P.parent()
        tmp={ PP.0 : M.0, PP.1 : 0, PP.2 : 0, PP.3 : 0}
        return P.subs(tmp)

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Converte polynomials."

    # convert!
    P1=Convert(P1,PR2)
    x1=Convert(x1,PR2)
    y0pm1=Convert(y0pm1,PR2)
    y0pp1=Convert(y0pp1,PR2)
    f=Convert(f,PR2)
    sf=Convert(sf,PR2)

    # set the values of coefficients
    if choice == 1:
        va1=vsol[0,2] # a1
        vb1=0 # b1
    else:
        vb1=vsol[0,2] # b1
        va1=0 # a1

    del vsol

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Compute new polynomials."

    # compute new powers of y0:
    # y0^(p^2-1) and y0^(p^2+1)
    # y0psm1=y0pm1^p*y0pm1
    # sage is not smart with p powers
    # y0pm1p=sum([ m^p*y0pm1.monomial_coefficient(m)^p for m in y0pm1.monomials() ])
    # y0psm1=y0pm1p*y0pm1  # y0psm1=y0pm1^p*y0pm1
    # del y0pm1p
    # y0psp1=f*y0psm1
    y0psm1=pol_p_power(y0pm1,p)*y0pm1
    y0psp1=f*y0psm1


    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Compute x2."


    # x0, despite "inject_variables" above, seems to still belong to PR1
    # print x0.parent()

    # dx2/dx0
    # compute x1^(p-1) as "pol_p_power(x1,p)/x1"???
    # tmpx2=(HI^(-1-p))*y0psm1-PR2.0^(p^2-1)-x1^(p-1)*(HI^(-1)*y0pm1-PR2.0^(p-1))
    # this HI^(-1-p) is NOT computed in an efficient way!!
    tmpx2=(HI^(-1-p))*y0psm1-PR2.0^(p^2-1)-x1^(p-1)*(HI^(-1)*y0pm1-PR2.0^(p-1))

    # we ARE adding the term in x0^(p^2)!!
    x2=my_int(tmpx2,PR2.0)+ sum([PR2.gens()[i+2]*PR2.0^(i*p) for i in range(((3*p-1)//2)+1)])


    del tmpx2, HI

    # add extra terms (known)
    if not minimal:
        tmp= F1(3/4)*x1^2
        x2+= sum([pol_p_power(tmp.coefficient({ PR2.0 : ((i//p)+p) }),p,1)*PR2.0^i for i in range(((3*p^2+p)//2),2*p^2-p+1,p)])
        del tmp


    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Create vector for RHS."


    # find the necessary vector for the call of SpcFctv
    if choice == 1: # b1=0
        if a0 == 0:
            ev=[ 3*PR2.0^(2*p)*x1 , va1*PR2.0^p ]
        else:
            ev=[ 3*PR2.0^(2*p)*x1 , a0^p*x1, va1*PR2.0^p ]
    else: # a1=0
        if a0 == 0:
            ev=[ 3*PR2.0^(2*p)*x1 , vb1 ]
        else:
            ev=[ 3*PR2.0^(2*p)*x1 , a0^p*x1, vb1 ]

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "First part of RHS."


    # RHS
    # va1^p is NOT computed in an efficient way!!
    rhs2= (3*PR2.0^(2*p^2)+a0^(p^2))*x2 + 3*PR2.0^(p^2)*pol_p_power(x1^2,p) + (-NextCoord(3,p)*PR2.0^(2*p^2)+va1^p)*pol_p_power(x1,p) + a2*PR2.0^(p^2) + b2

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Second part (SpcFct(ev))."

    rhs2+= SpcFctv(ev)

    del ev


    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Third part (fSpcFct)."

    rhs2+= fSpcFct( (3*PR2.0^(2*p)+a0^p)*x1 + va1*PR2.0^p + vb1 , sf )

    del sf

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Fourth part (Spc4Fct)."

    rhs2+= Spc4Fct(f)

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Subtract known terms."
    
    # now subtract parts of the LHS that are known!
    # what is left is 2*y0^(p^2+1)*P2!
    rhs2-=(pol_power(f,p)*pol_power(P1^2,p)+F1((2-2^p)//p)*pol_power(y0pp1,p)*pol_power(P1,p)) # NOT really rhs2...

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Long division of RHS."
        
    # rem2 has to be zero
    rhs2*=F1(1/2)
    P2=0
    deg1=rhs2.degree(PR2.0)
    deg2=(3*((p^2+1)//2))

    while deg1 >= deg2:
        lterm=rhs2.coefficient({PR2.0 : deg1})*PR2.0^(deg1-deg2)
        rhs2-=lterm*y0psp1
        P2+=lterm
        deg1=rhs2.degree(PR2.0)


    rem2=rhs2

    del rhs2, y0pm1, y0pp1, y0psm1, y0psp1

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Collect terms that must be zero."

    vcoef=[ rem2.coefficient({PR2.0 : i}) for i in range(rem2.degree(PR2.0)+1) ]

    len=rem2.degree(PR2.0)+1

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Compute matrix and vector."

    r=((3*p+3)//2)  # number of variables (except x0) = rank!



    # CAREFUL HERE!!!
    # Sage (and MAGMA) solve system like x*A=b, not A*x=b

    # full matrix and vector
    M2=matrix(F1,r,len,[[(vcoef[j].coefficient({PR2.gens()[i+1] : 1})).constant_coefficient() for j in range(len)] for i in range(r)])
    
    v2=matrix(F1,1,len,[-vcoef[j].constant_coefficient() for j in range(len)])


    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Solve the system."

    vsol=M2.solve_left(v2)


    # IMPROVE?
    # Find all solutions??

    del M2, v2

    # no. of varibales - 1 (for x0)
    # r=((3*p+3) div 2);

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Create dictionary to evaluate solutions."


    # we now create a dictionary to evaluate the solutions
    # it must have [x0,a2 or b2, ..coefs of x2.. ]
    tp={}
    for i in range(r):
        tp[PR2.gens()[i+1]]=vsol[0,i]

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()


    # ********** test ***************
    if test:
        print "%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%"
        if rem2.subs(tp) == 0:
            print "test2 is OK"
        else:
            print "test 2 in bad!"
            print "rem2 = ",  rem2(tp)
        print "%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%";

    del rem2
    
    if tm:
        ptime=time.time()
        print "Find x2 from solution."
        
    x2=x2.subs(tp)

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"
        ptime=time.time()
        print "Find P2 from solution."

    P2=P2.subs(tp)

    del tp, r

    if tm:
        ftime=time.time()
        print "Done!"
        print "Partial time = ", ftime - ptime
        print "Total time = ", ftime - ttime
        print "**************************************************"


    if choice == 1:
        return va1, vb1, x1, P1, vsol[0,0], 0, x2, P2
    else:
        return va1, vb1, x1, P1, 0, vsol[0,0], x2, P2