from sympy import *
from sympy.abc import q,t
init_printing()

def integrate_matrix(m):
    newm = zeros(n)
    for i in range(0,n):
        for j in range(0,n):
            p = m[i,j]
            pnew = integrate(p,q)
            newm[i,j] = pnew
    return(newm)

def differentiate_matrix(m):
    newm = zeros(n)
    for i in range(0,n):
        for j in range(0,n):
            p = m[i,j]
            pnew = diff(p)
            newm[i,j] = pnew
    return(newm)

def truncate_matrix(m):
    newm = zeros(n)
    for i in range(0,n):
        for j in range(0,n):
            p = m[i,j]
            pnew=(p+O(q**3)).removeO()
            newm[i,j] = pnew
    return(newm)
    
def poly_mod3(r):
    s = Poly(simplify(r*4*4*4*4+q**27), modulus=3)-q**27
    t = s.as_expr()
    return(t)
     
def matrix_mod3(m):
    newm = zeros(n)
    for i in range(0,n):
        for j in range(0,n):
            p = m[i,j]
            pnew=poly_mod3(p)
            newm[i,j] = pnew
    return(newm)

print()
print('This is a partial computation of the Q\Sigma_b operations for a cubic surface')
print('Using the prime p = 3. We do two cases, b = c_1 the first Chern class, and')
print('b = the homology class of a Lagrangian sphere. Everything is written in a')
print('basis of cohomology which consists of 1 in H^0; l in H^2 the class of a')
print('generic line in CP^2, e_1,...,e_6 in H^2 the exceptional curves obtained from blowups,')
print('and finally the Poincare dual of a point in H^4')
print('Throughout, the equivariant variable is set to t = 1.')

print()

n = 9
""" 
Here, the matrix a is the ordinary and quantum product with c_1.
Generators are in the order 1, (line), (exceptionals), (point)
 """
a0 = Matrix([
        [0,0,0,0,0,0,0,0,0],
        [3,0,0,0,0,0,0,0,0],
        [-1,0,0,0,0,0,0,0,0],
        [-1,0,0,0,0,0,0,0,0],
        [-1,0,0,0,0,0,0,0,0],
        [-1,0,0,0,0,0,0,0,0],
        [-1,0,0,0,0,0,0,0,0],
        [-1,0,0,0,0,0,0,0,0],
        [0,3,1,1,1,1,1,1,0]])
              
a1 = Matrix([
        [0,108*q**2,36*q**2,36*q**2,36*q**2,36*q**2,36*q**2,36*q**2,252*q**3],
        [0,39*q,15*q,15*q,15*q,15*q,15*q,15*q,108*q**2],
        [0,-15*q,-11*q,-5*q,-5*q,-5*q,-5*q,-5*q,-36*q**2],
        [0,-15*q,-5*q,-11*q,-5*q,-5*q,-5*q,-5*q,-36*q**2],
        [0,-15*q,-5*q,-5*q,-11*q,-5*q,-5*q,-5*q,-36*q**2],
        [0,-15*q,-5*q,-5*q,-5*q,-11*q,-5*q,-5*q,-36*q**2],
        [0,-15*q,-5*q,-5*q,-5*q,-5*q,-11*q,-5*q,-36*q**2],
        [0,-15*q,-5*q,-5*q,-5*q,-5*q,-5*q,-11*q,-36*q**2],
        [0,0,0,0,0,0,0,0,0]]
        )
a = a0+a1
print("Checking the eigenvalues: set q = 1, the char. poly. of quantum\
multiplication with c_1 is")
set1 = a.subs(q,1)
print(factor_list(set1.charpoly(t)))
print()


"""
Cup product with the class of a Lagrangian sphere, e_1-e_2
"""


b0 = Matrix([
        [0,0,0,0,0,0,0,0,0],
        [0,0,0,0,0,0,0,0,0],
        [1,0,0,0,0,0,0,0,0],
        [-1,0,0,0,0,0,0,0,0],
        [0,0,0,0,0,0,0,0,0],
        [0,0,0,0,0,0,0,0,0],
        [0,0,0,0,0,0,0,0,0],
        [0,0,0,0,0,0,0,0,0],
        [0,0,-1,1,0,0,0,0,0]])

iterations = 24

theta0 = -a0
theta = theta0
for m in range(0,iterations):
    c = theta*a-a*theta
    newtheta = theta0 + integrate_matrix(c/q)
    theta = newtheta
theta = truncate_matrix(theta)

print('Covariantly constant endomorphism whose constant term is quantum product with -c_1, modulo q^3:')
print()
print(theta)

print()
print('Reduction mod 3, after which -c_1 = Sq(c_1):')
print()

thetamod3 = matrix_mod3(theta)

print(thetamod3)

print()
print('Check that it solves the covariant-constancy differential equation (the outcome should be zero mod q^3):')
print()
covariant = matrix_mod3(q*differentiate_matrix(thetamod3) + a*thetamod3-thetamod3*a)
print(covariant)

theta0 = -b0
theta = theta0
for m in range(0,iterations):
    c = theta*a-a*theta
    newtheta = theta0 + integrate_matrix(c/q)
    theta = newtheta
theta = truncate_matrix(theta)

print()
print('Covariantly constant endomorphism whose constant term is quantum product with -s,')
print('s = e_1-e_2 spherical class, again modulo q^3:')
print()
print(theta)

print()
print('Reduction mod 3:')
print()

thetamod3 = matrix_mod3(theta)

print(thetamod3)

print()
print('Check that it solves the differential equation (the outcome should be zero mod q^3):')
print()
covariant = matrix_mod3(q*differentiate_matrix(thetamod3) + a*thetamod3-thetamod3*a)
print(covariant)