from sympy import *
from sympy.abc import q
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 the intersection of two quadrics')
print('in CP^5, using the prime p = 3. Everything is written in a')
print('basis of even degree cohomology which consists of one generator in each degree 0,2,4,6')
print('Throughout, the equivariant variable is set to t = 1.')
print()

n = 4

a0 = Matrix([[0,0,0,0],[1,0,0,0],[0,4,0,0],[0,0,1,0]])
a1 = Matrix([[0,q,0,q**2],[0,0,2*q,0],[0,0,0,4*q],[0,0,0,0]])
a = a0+a1

iterations = 12

theta0 = -a0+a0*a0*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('Let b be the degree 2 generator (half of the first Chern class)')
print('Covariantly constant endomorphism whose constant term is quantum product with -b+b^3,') 
print('modulo q^3 (obtained by solving the ODE):')
print()
print(theta)

print()
print('Reduction mod 3, and add the missing q^3 term by hand:')
print()

thetamod3 = matrix_mod3(theta) + Matrix([[0,0,0,q**3],[0,0,0,0],[0,0,0,0],[0,0,0,0]])

print(thetamod3)

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

print()
print('Take its square:')
print()    
print(matrix_mod3(thetamod3*thetamod3))

print()
print('And its cube:')
p = print
print(matrix_mod3(thetamod3*thetamod3*thetamod3))