# Compute sum of witt vectors of lenth 3 using
# the fcts from sec_coord.sage

load 'sec_coord.sage'

# sum two vectors using recursion
def WittSum2(v1,v2):

    p=v1[0].parent().characteristic()
    
    c1= v1[0] + v2[0]
    c2= v1[1] + v2[1] + fSpcFct(v1[0],v2[0])
    c3= v1[2] + v2[2] + fSpcFct(v1[1],v2[1]) + fSpcFct(v1[1]+v2[1],fSpcFct(v1[0],v2[0])) + fSpc4Fct(v1[0],v2[0])

    return [c1,c2,c3]
    

# multiply two vectors using recursion
def WittProd2(v1,v2):

    p=v1[0].parent().characteristic()
    
    c1= v1[0]*v2[0]
    c2= v1[1]*v2[0]^p + v2[1]*v1[0]^p
    c3= v1[2]*v2[0]^(p^2) + v1[1]^p*v2[1]^p + v2[2]*v1[0]^(p^2) + fSpcFct(v2[0]^p*v1[1],v1[0]^p*v2[1])

    return [c1,c2,c3]
    

# compute the polynomials for the sum using the recursion
def WittSum2Pol(p):

    PR.<X0,X1,X2,Y0,Y1,Y2>=PolynomialRing(GF(p),6)
 
    c1= X0 + Y0
    c2= X1 + Y1 + fSpcFct(X0,Y0)
    c3= X2 + Y2 + fSpcFct(X1,Y1) + fSpcFct(X1+Y1,fSpcFct(X0,Y0)) + fSpc4Fct(X0,Y0)

    return [c1,c2,c3]
  


# compute the polynomials for the product using the recursion
def WittProd2Pol(p):
    
    PR.<X0,X1,X2,Y0,Y1,Y2>=PolynomialRing(GF(p),6)
 
    
    c1= X0*Y0
    c2= X1*Y0^p + Y1*X0^p
    c3= X2*Y0^(p^2) + X1^p*Y1^p + Y2*X0^(p^2) + fSpcFct(X0^p*Y1,Y0^p*X1)

    return [c1,c2,c3]
    


# these will compute in Z/p^n

def Divp(f):

    P=f.parent()
    p=P.characteristic().factor()[0][0]

    R.<var1,var2,var3,var4,var5,var6>=PolynomialRing(Integers(),6)

    return P(R(f)//p)




def Divp2(f):

    P=f.parent()
    p=P.characteristic().factor()[0][0]

    R.<var1,var2,var3,var4,var5,var6>=PolynomialRing(Integers(),6)

    return P(R(f)//p^2)




def WittSum2Pol2(p):

    PR.<X0,X1,X2,Y0,Y1,Y2>=PolynomialRing(GF(p),6)

    c1=X0+Y0

    Rp2=Zmod(p^2)
    RR.<XX0,XX1,XX2,YY0,YY1,YY2>=PolynomialRing(Rp2,6)

    c2=X1+Y1 + PR(Divp(XX0^p+YY0^p - (XX0+YY0)^p))
    
    Rp3=Zmod(p^3)
    RR.<XX0,XX1,XX2,YY0,YY1,YY2>=PolynomialRing(Rp3,6)

    c3=X2 + Y2 + PR(Divp2(p*(XX1^p+YY1^p-(XX1+YY1+Divp(XX0^p+YY0^p-(XX0+YY0)^p))^p) + (XX0^(p^2)+YY0^(p^2)-(XX0+YY0)^(p^2)))) 
        

    return [c1,c2,c3]
  


def WittProd2Pol2(p):

    PR.<X0,X1,X2,Y0,Y1,Y2>=PolynomialRing(GF(p),6)

    c1=X0*Y0

    c2=X0^p*Y1+Y0^p*X1

    Rp2=Zmod(p^2)
    RR.<XX0,XX1,XX2,YY0,YY1,YY2>=PolynomialRing(Rp2,6)

    c3=X0^(p^2)*Y2 + X1^p*Y1^p + Y0^(p^2)*X2 + PR(Divp(XX0^(p^2)*YY1^p+YY0^(p^2)*XX1^p - (XX0^p*YY1+YY0^p*XX1)^p))

    return [c1,c2,c3]




# now compute in Z

def WittSum2Pol3(p):

    PR.<X0,X1,X2,Y0,Y1,Y2>=PolynomialRing(GF(p),6)

    c1=X0+Y0

    RR.<XX0,XX1,XX2,YY0,YY1,YY2>=PolynomialRing(Integers(),6)

    c2=X1+Y1 + PR((XX0^p+YY0^p - (XX0+YY0)^p)//p)
    
    c3=X2 + Y2 + PR((p*(XX1^p+YY1^p-(XX1+YY1+(XX0^p+YY0^p-(XX0+YY0)^p)//p)^p) + (XX0^(p^2)+YY0^(p^2)-(XX0+YY0)^(p^2)))//p^2) 
        

    return [c1,c2,c3]
  

def WittProd2Pol3(p):

    PR.<X0,X1,X2,Y0,Y1,Y2>=PolynomialRing(GF(p),6)

    c1=X0*Y0

    c2=X0^p*Y1+Y0^p*X1

    RR.<XX0,XX1,XX2,YY0,YY1,YY2>=PolynomialRing(Integers(),6)

    c3=X0^(p^2)*Y2 + X1^p*Y1^p + Y0^(p^2)*X2 + PR((XX0^(p^2)*YY1^p+YY0^(p^2)*XX1^p - (XX0^p*YY1+YY0^p*XX1)^p)//p)

    return [c1,c2,c3]





# evaluate the vector of polynomials in aux
def EvalVec(vec,aux):
    P=vec[0].parent()
    dic={} # create dictionary to use with subs
    for i in range(6):
        dic[P.gens()[i]]=aux[i]
    return [ term.subs(dic) for term in vec ]






# compute sum and product by precomputing the formulas
# for sum and product and then evaluating
def WittSum22(v1,v2, vecS=[]):

    P=v1[0].parent()
    p=P.characteristic()
    if vecS == []:
	vecS=WittSum2Pol(p)

    return EvalVec(vecS,v1+v2)



def WittProd22(v1,v2, vecP=[]):

    P=v1[0].parent()
    p=P.characteristic()
    if vecP == []:
	vecP=WittProd2Pol(p)
    
    return EvalVec(vecP,v1+v2)