# Fast functions to preform computation with Witt vectors of
# Length 3.
#
# See comp.tex

# Coef's are some divisions of binomial coefficients that
# we need to have explicitly as integer (or in GF(p))
def Coef1(n,p,i):
    return GF(p)(binomial(n*p,i)//p)


def Coef(p,i):
    return Coef1(1,p,i)


def Coef2(n,p,i):
    return GF(p)((binomial(n,i)-binomial(n*p,i*p))//p)


# psi_1 (\sfct in comp.tex)
def SpcFct(f):

    P=f.parent()

    v=f.monomials()
    n=len(v)
    if n <= 1:
        return P(0)

    p=P.characteristic()
 
    # split in half-sized polynomials
    g1=sum([ f.monomial_coefficient(v[i])*v[i] for i in range(n//2) ])
    g2=sum([ f.monomial_coefficient(v[i])*v[i] for i in range((n//2),n) ])

    del v

    # faster to use previous powers
    # will store powers of g2
    pwrs2={ 0 : 1, 1 : g2 }
    for i in range(2,p):
        pwrs2[i]=pwrs2[i-1]*g2
    
    tmp=0
    tmp1=1 # powers of g1 (computed in the loop)

    for i in range(1,p):
        tmp1*=g1 # next power of g1
        tmp+=-Coef(p,i)*tmp1*pwrs2[p-i]
        del pwrs2[p-i] # save memory

    del tmp1,pwrs2 # not needed anymore!
    
    return SpcFct(g1) + SpcFct(g2) + tmp


# eta_1(x,y) (\afctt from comp.tex)
def fSpcFct(x,y):

    if (x == 0) or (y == 0):
        return 0
    else:
	p=x.parent().characteristic()

        pwrs2={ 0 : 1, 1 : y }
        for i in range(2,p):
            pwrs2[i]=pwrs2[i-1]*y
    
        tmp=0
        tmp1=1 # powers of x (computed in the loop)

        for i in range(1,p):
            tmp1*=x # next power of g1
            tmp+=-Coef(p,i)*tmp1*pwrs2[p-i]
            del pwrs2[p-i] # save memory
        
	return tmp



# eta_1 (\afctt from comp.tex)
def SpcFctv(v):

    r=len(v)

    if (r == 1) or (r == 0):
        return 0

    p=v[0].parent().characteristic()

    # again, we split in two vectors
    w1=v[:r//2]
    w2=v[r//2:]

    return fSpcFct(sum(w1),sum(w2)) + SpcFctv(w1) + SpcFctv(w2)


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


# second coordinate of an integer as a Witt vector

def NextCoord(a,p):
    return GF(p)((a^p - a)//p)


# more coeffs.

def Coef3(p,i):
    return GF(p)(NextCoord(binomial(p,i)//p,p))


def Coef4(p,i):
    return GF(p)(binomial(p^2,i)//p^2)
    # F=GF(p)
    # num=prod([ F(j) for j in range(p^2-i+1,p^2) if ((j%p)!=0) ])
    # den=prod([ F(j) for j in range(1,i) if ((j%p)!=0) ])
    # return num/den
    
"""
The number of terms is relevant on how to compute this function.
Picked a random number of terms (namely, 5) to divide.
************************************************************
SHOULD CHECK WHEN ONE METHOD BECOMES FASTER THAN THE OTHER!!!
************************************************************
"""

# psi_2 (\ssfct in comp.tex) = Spc2Fct + Spc3Fct (but simplified)

# split in 2 of same size:
def Spc4Fct1(f):
    # fast for large # of terms
    # break into two pols of aprox. same number of terms

    w=f.monomials()
    n=len(w)

    if n <= 1:
        return 0

    P=f.parent()
    p=P.characteristic()

    g1=sum([ w[i]*f.monomial_coefficient(w[i]) for i in range(n//2)])
    g2=sum([ w[i]*f.monomial_coefficient(w[i]) for i in range(n//2,n) ])

    del w

    # res1
    res= -sum([ Coef4(p,i)*g1^i*g2^(p^2-i) for i in range(1,p^2) if (i%p)!=0 ])
    # res2
    res+= sum([ Coef3(p,i)*g1^(p*i)*g2^(p*(p-i)) for i in range(1,p) ])

    v=[ -Coef(p,i)*g1^i*g2^(p-i) for i in range(1,p) ]

    # res3
    res+= SpcFctv(v)

    # sf=fSpcFct(g1,g2)
    sf = sum(v)

    del v

    sg1=SpcFct(g1)

    sg2=SpcFct(g2)

    # res5
    res+=fSpcFct(sg1+sg2,sf)

    del sf

    # res4
    res+=fSpcFct(sg1,sg2)

    del sg1, sg2

    return Spc4Fct(g1) + Spc4Fct(g2) + res


# split in monomial + rest
def Spc4Fct2(f):
    # smaller number of terms
    # break into a single term + rest

    w=f.monomials()
    n=len(w)

    if n <= 1:
        return 0

    P=f.parent()
    p=P.characteristic()

    g1= w[0]*f.monomial_coefficient(w[0]) # just one term
    g2= f-g1 # rest

    del w

    # res1
    res= -sum([ Coef4(p,i)*g1^i*g2^(p^2-i) for i in range(1,p^2) if (i%p)!=0 ])

    # res2
    res+= sum([ Coef3(p,i)*g1^(i*p)*g2^((p-i)*p) for i in range(1,p) ])

    v=[ -Coef(p,i)*g1^i*g2^(p-i) for i in range(1,p) ]

    # res3
    res+= SpcFctv(v)

    # sf=fSpcFct(g1,g2)
    sf = sum(v)

    del v

    # zero in this case
    # sg1=SpcFct(g1)

    sg2=SpcFct(g2)

    # res5
    # res+=fSpcFct(sg1+sg2,sf)
    res+=fSpcFct(sg2,sf)

    del sf, sg2

    # return Spc4Fct(g1) + Spc4Fct(g2) + res1 + res2 + res3 + res4 + res5
    return Spc4Fct(g2) + res


def Spc4Fct(f):
    # decides which to use
    # NEEDS MORE TESTING! 5 is arbritary
    w=f.monomials()
    n=len(w)

    if n > 5:
        return Spc4Fct1(f)
        
    else:
        return Spc4Fct2(f)        




# eta_2 (\aafct in comp.tex)
def fSpc4Fct(x,y):

    P=x.parent()
    p=P.characteristic()

    res= -sum([ Coef4(p,i)*x^i*y^(p^2-i) for i in range(1,p^2) if (i%p)!=0 ])

    res+= sum([ Coef3(p,i)*x^(p*i)*y^(p*(p-i)) for i in range(1,p) ])
    
    v=[ -Coef1(1,p,i)*x^i*y^(p-i) for i in range(1,p) ]
    
    return res + SpcFctv(v)