###
### The file contains functions related to the paper
###
### Gromov--Witten theory of K3 surfaces and a Kaneko--Zagier equation for Jacobi forms
###    by Jan-Willem van Ittersum, Georg Oberdieck, Aaron Pixton
### 
### Run on: SageMath version 7.4, Release Date: 2016-10-18
### 
### The ring R is defined in Proposition 2.1 where 
###  -- K is the theta function Theta(z)
###  -- WP is the Weierstrass elliptic function
###  -- dWP is its derivative with respect to z
###  -- KInv is 1/Theta(z)
###
### The function phi(m) and phi(m,n) for all m,n in Z returns the main functions of the paper in the ring R.
###
### e.g. the first pairs (m, phi(m)) is given by
### (1, K)
### (2, 2*K^2*A)
### (3, 9/2*K^3*A^2 - 3/2*K^3*WP)
### (4, 32/3*K^4*A^3 - 8*K^4*A*WP - 2/3*K^4*dWP)
###
### The code at the end
###  - expands the functions \mathcal{F}(x,z) appearing in Proof of Theorem 1.2 around x=0 and x=z, and checks the vanishings of the holomorphic anomaly.
###  - expands the function \mathcal{F}(x,y,z) appearing in Proof of Proposition 4.4 / equation (23) around y=0 and y=z, and then each coefficient around x=0 and x=z.
###
### By Georg Oberdieck, July 2020.

from sage.rings.polynomial.polydict import ETuple
R = PolynomialRing(QQ, 'K, A, G2, WP, dWP, G4, KInv')
K, A, G2, WP, dWP, G4, KInv = R.gens()

G6 = Rational((-1,504))*( 6*36*dWP**2 + 18*(240*G4)*WP - 24*36*WP**3 )

E2 = -24*G2
E4 = 240*G4
E6 = -504*G6

#FFF = dtau(dtau(K))/K
FFF = Rational((1,4))*A**4 - Rational((3,2))*A**2*WP - Rational((3,4))*WP**2 - A*dWP + Rational((1,48))*E4

@cached_function
def dtau(f, iterations=1):
  if iterations == 0:
    return f
  return dtau( dtau_once(f), iterations-1)

@cached_function
def dtau_once(f):
  if f in QQ:
    return 0
  return f.derivative(K)  * K*(Rational((1,2))*(A**2 - WP) - Rational((1,12))*E2) \
       + f.derivative(A)  * (A*(Rational((1,12))*E2 - WP) - Rational((1,2))*dWP) \
       + f.derivative(G2) * Rational((-1,24)) * Rational((1,12))*(E2**2 - E4) \
       + f.derivative(WP) * ( dWP*A + Rational((1,6))*E2*WP - Rational((1,36))*E4 + 2*WP**2 ) \
       + f.derivative(dWP) * ( 6*WP**2*A - Rational((1,24))*E4*A + Rational((1,4))*E2*dWP + 3*WP*dWP ) \
       + f.derivative(G4) * Rational((1,240))*Rational((1,3))*(E2*E4 - (6*36*dWP**2 + 18*E4*WP - 24*36*WP**3)) \
       + f.derivative(KInv) * (-KInv)*(Rational((1,2))*(A**2 - WP) - Rational((1,12))*E2)

@cached_function
def dz(f, iterations=1):
  if iterations == 0:
    return f
  return dz( dz_once(f), iterations-1)

@cached_function
def dz_once(f):
  if f in QQ:
    return 0
  return f.derivative(K)  * A*K \
       + f.derivative(A)  * (-WP + Rational((1,12))*E2) \
       + f.derivative(WP) * dWP \
       + f.derivative(dWP) * (6*WP**2 - Rational((1,24))*E4) \
       + f.derivative(KInv) * (-A*KInv)

@cached_function
def G(k):
  if is_odd(k) or k<=0:
    return 0
  if k==2: return G2
  if k==4: return G4
  if k==6: return G6
  n=int(k/2)
  return Rational(( factorial(2*n-2), (n-3)*(2*n+1) ))*6*sum( Rational((1,factorial(2*p-2)*factorial(2*(n-p)-2) ))*G(2*p)*G(2*(n-p)) for p in range(2,n-1))

@cached_function
def Eis(k):
  #return Eisenstein series in 1 + O(q) normalization in terms of G2, G4, G6
  if is_odd(k) or k<=0: return 0
  return (-2*k / bernoulli(k))*G(k)

#=============== phi(m,n) ================
@cached_function
def phi(m,l=None):
  #Return main functions phi(m,l) in Jacobi Ring.
  if m==0 or l==0:
    return 0
  if l is None:
    if m<0:
      return -phi(-m)
    return phi_of_first_kind(m)

  if abs(l)>abs(m):
    return phi(l,m)
  if m<0:
    return phi(-m,-l)

  #We have ensured m>0 and m>=|l|.
  return phi_of_second_kind(m,l)

def psi(m,n):
  return phi(m,n) - (abs(m) if m==-n else 0)

def phi_of_first_kind(m):
  w=PowerSeriesRing(R, 'w').gen()
  A=2*m*sum( G(k)*Rational((1,factorial(k)))*w**k for k in range(2, m)) + O(w**m)
  f=exp(A)
  #f=sum(Rational((1,factorial(k)))*A**k for k in range(0,m)) + O(w**(m))
  g=sum(Rational((1,factorial(k)))*dz(K**m,k)*w**k for k in range(0,m)) + O(w**m)
  return (f*g).coefficients()[m-1]

def phi_of_second_kind(m,n):
  if m != -n:
    return Rational((1,m+n))*(m*phi(m)*dtau(phi(n)) + n*dtau(phi(m))*phi(n) )

  #now m = - n
  n=abs(n)
  B=phi(-n-1,n).derivative(A) - sum( Rational((n+1,j))*phi(-n-1+j,n)*phi(j) for j in range(2,n+1)) - sum( Rational((n,j))*phi(-n-1, n-j)*phi(j) for j in range(1,n) )
  #print B
  return abs(n)+Rational((1,n+1))*R( B/K )


#===============================
# Expansions of differential equations from now on.
#===============================

#Preliminaries:
Rx = LaurentSeriesRing(R, 'x')
x=Rx.gen()
prec=10

#The expansion of f(x) = Theta(x+z)/Theta(x):
Klogx = sum( bernoulli(m)*Rational((1,m*factorial(m)))*Eis(m)*x**m for m in range(1,prec+1) if is_even(m)) + O(x**(prec+1))
Kx = x* sum( Rational((1,factorial(k)))*Klogx**k for k in range(0,prec+1)) + O(x**(prec+1))
f_xexp = Kx**(-1)*sum( Rational((1,factorial(i)))*dz(K,i)*x**i for i in range(0,prec+1)) + O(x**(prec+1))

#The expansion of \tilde{f}(x) = Theta(x)/Theta(x-z) obtained from f(x) by changing variables x -> x-z.
f_xexp_sign = sum( ((-1) if is_odd(e) else 1)*R(c/K)*x**e for (c,e) in zip(f_xexp.coefficients(), f_xexp.exponents()) ) + O(x**(prec+1))
ftilde_xexp = KInv*(1/f_xexp_sign) + O(x**(prec+1))


RX = PolynomialRing(R, 'f,fx,fxx,fz,ft,fzx,ftx, ftt' )
f,fx,fxx,fz,ft,fzx,ftx, ftt = RX.gens()

def RXdAx(F):
  #calculate d/dA(x) F in the ring RX
  return f*F.derivative(fz) + fx*F.derivative(fzx) + fx*F.derivative(ft) + fxx*F.derivative(ftx) + 2*ftx*F.derivative(ftt)

def RXdG2(F):
  #calculate d/dG2 F in the ring RX
  return -2*fx*F.derivative(ftx) - 2*f*F.derivative(fzx) -4*ft*F.derivative(ftt) + F.derivative(G2)

#Evaluation the formal expression as Taylor series near x
Eval=Hom(RX, Rx)([ f_xexp, f_xexp.derivative(x), f_xexp.derivative(x,2), dz(f_xexp), dtau(f_xexp), dz(f_xexp.derivative(x)), dtau(f_xexp.derivative(x)), dtau(f_xexp,2) ])

#Evaluation the formal expression as Taylor series near x = z using ftilde
#We have the chain rule d/dz( f(x-z) ) = (df/dz)(x-z) + (df/dx)(x-z) * (-1), hence (df/dz) gets evaluated to d\tilde{f}/dx + d\tilde{f}/dz, etc.
Eval2=Hom(RX, Rx)([ ftilde_xexp, ftilde_xexp.derivative(x), ftilde_xexp.derivative(x,2), dz(ftilde_xexp) + ftilde_xexp.derivative(x), dtau(ftilde_xexp), dz(ftilde_xexp.derivative(x)) + ftilde_xexp.derivative(x,2), dtau(ftilde_xexp.derivative(x)), dtau(ftilde_xexp,2) ])




##
#Check of the differential equation appearing in Proof of Theorem 1.2.
##
DiffEqn1 = fx**2*ftt - 2*fx*ftx*ft + fxx*ft**2 + FFF*( fx**2*f - fxx*f**2 )

print "Differential equation 1"
print RXdG2(DiffEqn1), RXdAx(DiffEqn1), Eval(DiffEqn1), Eval2( DiffEqn1)


##
#Check of the differential equation appearing in Proof of Proposition 4.4 (see equation (23))
##
Ry=LaurentSeriesRing(RX,'y')
y=Ry.gen()

#expansion of f(y) = Theta(y+z)/Theta(y)
Klogy = sum( bernoulli(m)*Rational((1,m*factorial(m)))*Eis(m)*y**m for m in range(1,prec+1) if is_even(m)) + O(y**(prec+1))
Ky = y* sum( Rational((1,factorial(k)))*Klogy**k for k in range(0,prec+1)) + O(y**(prec+1))
fy = Ky**(-1)*sum( Rational((1,factorial(i)))*dz(K,i)*y**i for i in range(0,prec+1)) + O(y**(prec+1))


#Diff eqn at y=0:
C1 = fxx*f*fy.derivative(y)**2 - fy.derivative(y,2)*fy*fx**2
C = fx*fy.derivative(y) - fz*fy.derivative(y) - fx*dz(fy) + f*dtau(fy) + ft*fy

D1 = f*fy.derivative(y) - fx*fy
Da = ( fxx*fz - fx*fzx )*fy.derivative(y)**2 \
   + ( fy.derivative(y,2)*dz(fy) - fy.derivative(y)*dz(fy.derivative(y)) )*fx**2
Db = -( fxx*f - fx**2 )*fy.derivative(y)*dtau(fy) - ( fxx*ft - fx*ftx )*fy*fy.derivative(y) \
     -( fy.derivative(y,2)*fy - fy.derivative(y)**2 )*fx*ft - ( fy.derivative(y,2)*dtau(fy) - fy.derivative(y)*dtau(fy.derivative(y)) )*f*fx

DiffEqn2 = C1*C + D1*(Da+Db)

print "Differential equation 2 at y=0:"
for e,c in zip(DiffEqn2.exponents(),DiffEqn2.coefficients()):
  print e, RXdG2(c), RXdAx(c), Eval(c), Eval2( c )

#Diff eqn at y=z. Again we need to apply the chain rule (d/dz( f(x-z) ) = (df/dz)(x-z) + (df/dx)(x-z) * (-1))
#Expansion of \tilde{f}(y)
f_yexp_sign = sum( ((-1) if is_odd(e) else 1)*R(c/K)*y**e for (c,e) in zip(f_xexp.coefficients(), f_xexp.exponents()) ) + O(y**(prec+1))
fytilde = KInv*(1/f_yexp_sign) + O(y**(prec+1))

#Diff eqn at y=z:
C1tilde = fxx*f*fytilde.derivative(y)**2 - fytilde.derivative(y,2)*fytilde*fx**2
Ctilde = fx*fytilde.derivative(y) - fz*fytilde.derivative(y) - fx*(fytilde.derivative(y) + dz(fytilde)) + f*dtau(fytilde) + ft*fytilde

D1tilde = f*fytilde.derivative(y) - fx*fytilde
Datilde = ( fxx*fz - fx*fzx )*fytilde.derivative(y)**2 \
   + ( fytilde.derivative(y,2)*(fytilde.derivative(y) + dz(fytilde)) - fytilde.derivative(y)*( dz(fytilde.derivative(y)) + fytilde.derivative(y,2) ) )*fx**2
Dbtilde = -( fxx*f - fx**2 )*fytilde.derivative(y)*dtau(fytilde) - ( fxx*ft - fx*ftx )*fytilde*fytilde.derivative(y) \
     -( fytilde.derivative(y,2)*fytilde - fytilde.derivative(y)**2 )*fx*ft - ( fytilde.derivative(y,2)*dtau(fytilde) - fytilde.derivative(y)*dtau(fytilde.derivative(y)) )*f*fx

DiffEqn2tilde = C1tilde*Ctilde + D1tilde*(Datilde+Dbtilde)

print "Differential equation 2 at y=z:"
for e,c in zip(DiffEqn2tilde.exponents(),DiffEqn2tilde.coefficients()):
  print e, RXdG2(c), RXdAx(c), Eval(c), Eval2( c )