CC=ComplexField(200); R = LaurentSeriesRing(CC, 'q');R;q=R.0; CCdisp = ComplexField(25) # for debugging

load 'modeta.sage'

Q= LaurentSeriesRing(QQ, 't')  #t-series expansions always have rational coefficients, q-series expansion always have complex coefficients
t=Q.0

def etagen(N):
    exp = MixedIntegerLinearProgram(maximization=True)
    r = exp.new_variable(integer=True)
    divs = N.divisors()
    numdivs = len(divs)
    primedivs = N.prime_divisors()
    numprimes = len(primedivs)
    AN = matrix([[N*(gcd(d1,d2))^2/(d1*d2*gcd(d2,N/d2)) for d2 in divs] for d1 in divs]).change_ring(ZZ)
    #print AN
    exp.set_objective(sum(AN[numdivs-1,n]*r[n] for n in range(numdivs))) # maximize the order at oo
    exp.add_constraint(sum(AN[numdivs-1,n]*r[n] for n in range(numdivs)), max=-24) # but require it have a pole
    for i in range(numdivs-1):
        exp.add_constraint(sum(AN[i,n]*r[n] for n in range(numdivs)), min=0) # require holomorphic at other cusps
# also require the orders at the cusps to be integers
        exp.add_constraint(sum(AN[i,n]*r[n] for n in range(numdivs))==24*r[numdivs+i])
    exp.add_constraint(sum(AN[numdivs-1,n]*r[n] for n in range(numdivs))==24*r[2*numdivs-1]) 
# now we impose that \prod_{d|N} d^{r_d} = square
# to do this, we figure out the power of each prime divisor of N which divides the product
# this power will be a linear function of the vector r, which is why we can use linear programming
# so first we determine what linear combination it is, then impose the constraint that it be even
# (by making it be 2 times a new integer variable)
    squareconstraints = []
    for p in primedivs:
        pconstraint = []
        for d in divs:
            f = list(d.factor())
            dprimes = [x[0] for x in f]
            if dprimes.count(p) == 0:
                ppower = 0
            else:
                ppower = f[dprimes.index(p)][1]
            pconstraint.append(ppower)
        squareconstraints.append(pconstraint)
    #print squareconstraints
    for i in range(numprimes):
        exp.add_constraint(sum(squareconstraints[i][n]*r[n] for n in range(numdivs))==2*r[2*numdivs+i])
# We allow positive AND negative values for all the variables
    for i in range(len(r.items())):
        exp.set_min(r[i],None)
# Finally we impose the condition \sum r_d = 0
    exp.add_constraint(sum(r[n] for n in range(numdivs))==0)
    opt = exp.solve()
    l = []
    for i in range(numdivs):
        l.append([divs[i],ZZ(exp.get_values(r[i]))])
    #print l
    dic = dict(l)
    e = EtaProduct(N,dic)
    return e

def residue(f):
    exp=f.exponents()
    if 0 in exp:
        i=exp.index(0)
        coef=f.coefficients()
        r=coef[i]
    if 0 not in exp:
        r=0
    return r

def periodg(f,gg):
    R = LaurentSeriesRing(CC, 'q')
    q=R.0
    r=residue(R(f))
    f0=f-r
    w0=f0/q
    W=R(w0.integral())
    a=gg[0,0]
    b=gg[0,1]
    c=gg[1,0]
    d=gg[1,1]
    tau=-(d/c)+(CC(I)/abs(c))
    gammatau=(a*tau+b)/(c*tau+d)
    period=W.substitute({q:exp(2*CC(pi)*CC(I)*gammatau)})-W.substitute({q:exp(2*CC(pi)*CC(I)*tau)})+r*(2*CC(pi)*CC(I))*(gammatau - tau)
    return period

def periodM(f,M):
    per=[]
    G=M.group()
    lM=len(M)
    for a in [0..lM-1]:
        B=M[a].integral_basis()
        lB=len(B)
        for b in [0..lB-1]:
            m=B[b].modular_symbol_rep()
            lm=len(m)
            per.append(sum(m[c][0]*periodg(f,G.are_equivalent(Cusp(m[c][1].alpha()),Cusp(m[c][1].beta()),trans=True)) for c in [0..lm-1]))
    return per

def periodv(f,v):
    l=len(v)
    period=sum(v[i][0]*periodg(f,v[i][1]) for i in [0..l-1])
    return period

def periodS(f, m, gens, mat):
    """ Input a differential form f (= omegaf q/dq), a modular symbol m of weight 2 for Gamma0(N), a list of hyperbolic generators for the homology of X0(N), and the change of basis between the corresponding modular symbols and the standard integral basis for the homology... return \int_m f."""
    vec = mat.transpose().solve_right(vector(m.list())).list()
    path = []
    for i in range(len(vec)):
        path.append([vec[i],gens[i]])
    return periodv(f, path)

def poincare(w1,w2):
    W1=(w1/t).integral()
    return(residue(W1*w2))

def pd(w1,w2):
    e1=w1.exponents()
    e2=w2.exponents()
    a=w1.coefficients()
    b=w2.coefficients()
    M=-w2.valuation()
    pd=[]
    for i in [1..M]:
        if i in e1:
            if -i in e2:
                pd.append(a[e1.index(i)]*b[e2.index(-i)]/i)
    return(sum(pd[k] for k in [0..len(pd)-1]))

def pd_matrix(b):
    A = []
    for i in [0..len(b)-1]:
        A.append([])
	for j in [0..i]:
	    A[i].append(0)
	for j in [i+1..len(b)-1]:
	    A[i].append(poincare(b[i],b[j]))
    return Matrix(A) - Matrix(A).transpose()

def poincare_matrix(b,c,mat,gens):
    A = []
    for i in [0..len(b)-1]:
        A.append([])
	for j in [0..len(c)-1]:
	    A[i].append(periodS(b[i],c[j],gens,mat))
    return Matrix(A)

def iteratedintegral(tau0,g,G,etag,H,wgPrimEtag,gamma):
    a=gamma[0,0];
    b=gamma[0,1];
    c=gamma[1,0];
    d=gamma[1,1];
    tau=-(d/c)+(CC(I)/abs(c));
    ptau=periodg(wgPrimEtag,gamma)-periodg(g,gamma)*H.substitute({q:exp(2*CC(pi)*CC(I)*tau0)});
    path=G.substitute({q:exp(2*CC(pi)*CC(I)*tau)})-G.substitute({q:exp(2*CC(pi)*CC(I)*tau0)});
    ptau0=path*periodg(etag,gamma);
    return(ptau-ptau0)

def iteratedintegralv(tau0,g,G,etag,H,wgPrimEtag,v):
    l=len(v);
    period=sum(v[i][0]*iteratedintegral(tau0,g,G,etag,H,wgPrimEtag,v[i][1]) for i in [0..l-1]);
    return period

#The following command searches for a homology basis of X_0(N) consisting of matrices whose lower
#left entries are positive and as small as possible.

def hypgens(N):
    a=1;
    G=Gamma0(N);
    M=ModularSymbols(G);
    d=M.cuspidal_subspace().dimension();
    Gens=[];
    GensM=[];
    Mat=[];
    MatAux=[];
    rnext=matrix(Mat).rank();
    x = 1;
    while rnext<d:
        a=a+1;
        rant=rnext;
        g=G.are_equivalent(Cusp(Infinity),Cusp(a/(x*N)),trans=True);
        if g:
            if (g.a()+g.d()).abs()>2:
                m=M.modular_symbol([Infinity,g.a()/g.c()]);
                c=m.list();
                MatAux.append(c);
                rnext=matrix(MatAux).rank();
                if rnext>rant:
                    if rnext <= d:
                        Mat.append(c);
                        Gens.append(g);
                        GensM.append(m);
        if a > 2*N:
            a = 1
            x = x+1
            print 'Warning: ideal homology basis does not exist'
    return Gens,Matrix(Mat)

#This function writes a cusp form w of weight 2 and level N with exact algebraic coefficients
#in terms of an echelon basis for S_2(N, QQ) and returns the vector of coefficients
def echeloncoeffs(w,N,new=False):
    if new:
        S = CuspForms(N,2).new_subspace()    
    else:
        S = CuspForms(N,2)
    B = S.echelon_basis()
    g = len(B)
    basis = [f.q_expansion(3*g) for f in B]
    orders = [f.valuation() for f in basis]
    coeffs = []
    for i in range(g):
        if w.valuation() == orders[i]:
            lead = w[orders[i]]
            coeffs.append(lead)
            w = w - lead*basis[i]
        else:
            coeffs.append(0)
    if w != 0: # for debugging
        print R(w).change_ring(ComplexField(10))
    assert w == 0
    return coeffs

def identitycharmodN(n,N):
    if gcd(n,N)==1:
        return 1
    else:
	return 0

def nonzerocoef(n,v,w):
    if n in v:
        return w[v.index(n)]
    else:
	return 0

#This next command takes as input a meromorphic differential of the second kind on X_0(N) for
#any N, and positive integers n and m. The output is the q expansion of f|T_n up to the term q^m.

def hecke_q_action(f,n,N,m):
    expf = f.exponents()
    coef = f.coefficients()
    A = f.base_ring()
    t = f.parent().0
    Tnf = A(0)
    for i in [min(1,expf[0]*n)..m]:
        Tnf = Tnf+A(sum([identitycharmodN(d,N)*d*nonzerocoef(i*n/d^2,expf,coef) for d in divisors(gcd(n,i))]))*t^(i)
    return Tnf

#The input is a set b of meromorphic differentials of the second kind forming a basis of H^1_dR(X_0(N))
#and a prime n not dividing N. The output is the matrix giving the action of T_n on H^1_dR(X_0(N)) with
#respect to this basis. The last two inputs are optional; the code automatically computes the Poincare
#pairing matrix with respect to b, but this can be entered as input A if the matrix has been previously
#computed. In this case, findA should be set to False.

def hecke_q_matrix(b,n,N,A=0,findA=True):
    m = max([-x.exponents()[0] for x in b])
    Tnb = [hecke_q_action(x,n,N,m+1) for x in b]
    if findA==True:
    	A = pd_matrix(b)
    Ainv = A^(-1)
    Tn = []
    for i in [0..len(b)-1]:
        Tn.append([])
        Tnv = []
        for j in [0..len(b)-1]:
        	Tnv.append(poincare(b[j].truncate(m*n+1),Tnb[i]))
        W=Ainv*vector(Tnv)
        for j in [0..len(b)-1]:
            Tn[i].append(W[j])
    return Matrix(Tn)

def hecke_CC_matrix(A,n,M):
    B=M.hecke_matrix(n)
    return A*B.transpose()*A^(-1)

#This function takes in a matrix defining an isotypic component of H^1_dR in terms of a chosen basis
#for the entire space, and outputs a new choice of basis such that the first half lies in the subspace
#generated by the first half of the vectors in the chosen basis. In particular, if the first half of
#the chosen basis is holomorphic, then so is the first half of the output.

def holomorphic_basis(M,A):
	Mcols = M.columns()
	N = []
	for i in [len(Mcols)/2..len(Mcols)-1]:
		N.append(Mcols[i])
	N = Matrix(N).transpose()
	N1 = N.kernel().matrix()
	N1rows = N1.rows()
	N2 = N.column_space().matrix()
	N2rows = N2.rows()
	N3 = []
	for i in [0..len(N1rows)-1]:
		N3.append(N1rows[i])
	for i in [0..len(N2rows)-1]:
		N3.append(N2rows[i])
	N3 = Matrix(N3)
	return N3*M, N3*A*N3.transpose()

def symplectic_basis(M,A):
	N = []
	n = len(A.rows())
	for i in [0..n/2-1]:
		N.append([])
		for j in [0..n-1]:
			if j==i:
				N[i].append(1)
			if j!=i:
				N[i].append(0)
	Ainvt = (A^(-1)).transpose()
	for i in [0..n/2-1]:
		N.append(vector(N[i])*Ainvt)
	N = Matrix(N)
	Nrows = N.rows()
	Mtemp = N*M
	Atemp = N*A*N.transpose()
	P = []
	for i in [0..n/2-1]:
		P.append(Nrows[i])
	for i in [n/2..n-1]:
		P.append([])
		for j in [0..n-1]:
			if j < i-n/2:
				P[i].append(0)
			if j > i-n/2-1 and j < n/2:
				P[i].append(-Atemp[i][j+n/2])
			if j==i:
				P[i].append(1)
			if j > n/2-1 and j!=i:
				P[i].append(0)
	P = Matrix(P)
	Mnew = P*Mtemp
	Anew = P*Atemp*P.transpose()
	return Mnew,Anew

def isotypic_bases(H1dR,M,pdmat=0,findA=True):
    if findA==True:
        minpolord=max([-x.exponents()[0] for x in H1dR])
        H1dRtrunc=[x.truncate(minpolord+2) for x in H1dR]
        pdmat=pd_matrix(H1dRtrunc)
    N=M.level()
    numiso = len(M)
    n=1
    Iso1=[[] for i in [0..numiso-1]]
    Iso=[]
    check = [[] for i in [0..numiso-1]]
    while sum([len(check[i]) for i in [0..numiso-1]])<numiso:
    	n=n+1
    	if gcd(n,N)==1:
	    Tn=hecke_q_matrix(H1dR,n,N,pdmat,False)
	    for i in [0..numiso-1]:
	    	if len(check[i]) == 0:
	    	    Tncharpoly=M[i].hecke_polynomial(n)
    		    if Tncharpoly.subs(Tn).nullity()==M[i].dimension():
    		        Iso1[i].append(Tncharpoly.subs(Tn).kernel().matrix())
    			check[i].append(0)
    for i in [0..len(Iso1)-1]:
    	hbasis = holomorphic_basis(Iso1[i][0],Iso1[i][0]*pdmat*Iso1[i][0].transpose())
    	sbasis = symplectic_basis(hbasis[0],hbasis[1])
    	Iso.append(sbasis[0])
    omegas = []
    etas = []
    for i in [0..len(Iso)-1]:
    	iso_rows = Iso[i].rows()
	omegas.append([])
	etas.append([])
	for j in [0..len(iso_rows)/2-1]:
	    omegas[i].append(Q(dotprod(iso_rows[j],H1dR)))
	    etas[i].append(Q(dotprod(iso_rows[j+len(iso_rows)/2],H1dR)))
    return omegas,etas

def dotprod(v,w):
	n = min(len(v),len(w))
	prod = 0
	for i in [0..n-1]:
		prod = prod + v[i]*w[i]
	return prod

#The following command computes the degree of X_0(N) over X(1).

def mu(N):
    mu = N
    for p in N.prime_divisors():
    	mu = mu*(1+1/p)
    return mu

def TtensorQgens(M):
    d=M.dimension()
    heckegens=[]
    heckegens1=[]
    heckegens2=[]
    nums=[]
    n=0
    while rank(Matrix(heckegens2))<d/2:
        n=n+1
	A=M.hecke_matrix(n)
        heckegens1.append(A.list())
        if rank(Matrix(heckegens2)) < rank(Matrix(heckegens1)):
	    heckegens2.append(A.list())
            heckegens.append(A)
            nums.append(n)
    return nums,heckegens

def HeckeAlgebra(M):
    N=M.level()
    plus=[M[i].plus_submodule() for i in [0..len(M)-1]]
    d=sum([(M[i].dimension()/2)^2 for i in [0..len(M)-1]])
    R=ZZ^d
    gens=[]
    r = mu(N)/6 - (mu(N)-1)/N
    for n in [1..ceil(r)]:
        x=[]
	for i in [0..len(M)-1]:
	    x+=plus[i].integral_hecke_matrix(n).list()
	gens.append(x)
    return R.submodule(gens)

def hecke_denominator(n,M,i,T):
    plus=M[i].plus_submodule()
    dims=[(M[j].dimension()/2)^2 for j in [0..len(M)-1]]
    R=ZZ^(sum(dims))
    x=[]
    for j in [0..len(M)-1]:
    	if j==i:
	    x+=plus.integral_hecke_matrix(n).list()
	else:
	    x+=[0 for k in [1..dims[j]]]
    y = R(x) in T
    r=1
    while y==False:
        r=r+1
	y = r*R(x) in T
    return r

def roundmatrix(A,n):
    B=[]
    for row in A.rows():
    	B.append([round(real(x),n) for x in row])
    return Matrix(B)