//
//  Tadic.m
//  10/02/2018.
//
//  T-ADIC CODE, IN LEVEL 1 FOR SIMPLICITY.
//  
//  This code computes eigenforms in the space of ``boundary forms'', following the paper
//      ``Modular eigenforms at the boundary of weight space''
//
//  The algorithm is based on earlier work of Lauder, and contains significant parts 
//  of his code, with adjustments to transpose everything to the T-adic setting. 
//  To give ample credit to Alan, here is the log file for the code I used:

/*

	February 20, 2013: (Alan Lauder)
Computes Hecke series for characters eps of order dividing (p-1), provided 
p >= 5 and Q(eps) has class number ONE.
 
	March 15, 2013: (Alan Lauder)
Code to compute L_p(*,**,***)(1,2,1) using version "b" - recall version "a" is
the faster space efficient one which will fail for small precisions for trivial reasons.
 
	March 27, 2013: (Alan Lauder)
New version has been tested successfully on p = 5, N = 23 example (Eisenstein series) with
precision m = 20 and is *MUCH FASTER* at finding complementary spaces e,g, times
2.3,12.5 versus 1.8, 1.02.

	April 8, 2013: (Alan Lauder)
Just a slightly tidier version 27/03/13.

	December 22, 2014: (Jan Vonk)
Extended to p=2,3 using the algorithm in "Computing overconvergent forms for small primes", LMS J. Comp.

*/


// The main function below is 
// > eigenfunctions(p,n,L);
//
// where
// p >= 3 is a prime
// n > 1 is the T-adic precision to which basis elements should be computed
// L = Sequence of integers, the desired eigenforms, sorted via their T-adic slope. 
//

// As an example, the command
// > E := eigenfunctions(5,20,[1..4]);
//
// Will return a sequence of q-expansions, which are approximations of the 1st, 2nd, 3rd, and 4th T-adic eigenforms 
// for p=5, level 1, and the trivial character. Here, the eigenforms are sorted by slope, and if there are multiple eigenforms
// of the same slope, they are ordered by Magma in some arbitrary fashion.
//
// The algorithm works up to some T-adic precision. In the final step of the algorithm, an internal function
// of magma seems to often lose quite a bit of precision. In a future version this should be circumvented 
// carefully, but for now the level 1 code is sufficiently efficient that one can easily increase the precision
// to obtain more digits of accuracy in the output. We make no detailed precision analysis.
//
// WARNING: The algorithm was written solely to provide a first tool with which to experiment with boundary eigenforms. 
// I have not tried to obtain precision bounds, nor have I done any optimisation of the code. 
//
// There are some convincing tests one can do to verify to which precision the output is correct:
// 1) Check that all the Hecke algebra relations are satisfied for the Fourier coefficients. 
//    This is a convincing check to see how many digits of precision were lost in the computation.
// 2) Check that the entries of the matrix of Up that was computed outgrow the desired precision. 
//    This is a convincing check that the approximate basis overconverges far enough.


// ------------------------------------------------------------------------------------------
// ---                               1)  EISENSTEIN SERIES                                ---
// ------------------------------------------------------------------------------------------



// Returns Eisenstein series E_k over given ring S modulo q^NN,
// normalised to have constant term a_0 = 1.

function ESeries(k,NN,S: chi:=DirichletGroup(1)!1, psi := DirichletGroup(1)!1)

	assert chi(-1)*psi(-1) eq (-1)^k; // Weight character must make sense.
	assert k ge 1; // The Eisenstein series exists in weight k >= 1.

	R<q> := PowerSeriesRing(S,NN);

	if Conductor(chi) ne 1 then
		a1 := 1;
		Ek := a1*q;
	else
		a1 := -S!(2*k/(Bernoulli(k,psi)));
		Ek := 1 + a1*q;
	end if;

	for n := 2 to NN-1 do
		coeffn := S!0;
		divs := Divisors(n);
		for d in divs do
			nd := Integers()!(n/d);
			coeffn := coeffn + S!(chi(nd)*psi(d))*d^(k-1);
		end for;
		Ek := Ek + a1*coeffn*q^n;
	end for;

	return Ek;

end function;



// Returns p-stabilised Eisenstein series E_k* over given ring S modulo q^NN,
// normalised to have constant term a_0 = 1.
//
// (Optional) If characters chi, psi are given (of order dividing p-1) then it returns
//       E_k*(chi,psi) = sum_n  [sum_(d | n, gcd(p,d)=1)   chi(n/d)*psi(d)*d^(k-1)]  q^n.

function pStabESeries(p,k,NN,S: chi:=DirichletGroup(1)!1, psi := DirichletGroup(1)!1)

	assert chi(-1)*psi(-1) eq (-1)^k; // Weight character must make sense.
	assert k ge 0; // The p-stabilised Eisenstein series exists in weight k >= 2, and is set 1 when k=0.

	R<q> := PowerSeriesRing(S,NN);
	if k eq 0 then
		Ek := R!1;
	else
		if Conductor(chi) ne 1 then
			a1 := 1;
			Ek := a1*q;
		else
			a1 := -S!(2*k/(Bernoulli(k,psi)*(1-p^(k-1))));
			Ek := 1 + a1*q;
		end if;
		for n := 2 to NN-1 do
			coeffn := S!0;
			divs := Divisors(n);
			for d in divs do
				if d mod p ne 0 then
					nd := Integers()!(n/d);
					coeffn := coeffn + S!(chi(nd)*psi(d))*d^(k-1);
				end if;
			end for;
			Ek := Ek + a1*coeffn*q^n;
		end for;
	end if;

	return Ek;

end function;



// Returns the power of the Hasse invariant that will been lifted for the Katz basis.
// This is always 1, except when p = 2,3.

function ap(p)

	a := 1;
	if p eq 2 then
		a := 4;
	end if;
	if p eq 3 then 
		a := 3;
	end if;
	return a;

end function;



// Returns the weight of the Eisenstein series to be used for lifting the Hasse invariant. 
// This is always p-1, unless p = 2,3.

function EisWeight(p)

	return ap(p)*(p-1);

end function;



// Computes T-adic Eisenstein series E(q) and E(p-1)(q) over Z_p[T] mod (p^m, T^n), up to prec q^NN.
// These are the Eisenstein families centered at weights 0 and p-1 respectively.
// Here, T = exp(q)^(wt-w0) - 1 is a parameter of the connected component of weight space of trivial character.
//
// This version is a lot faster for **large** T-adic precision. (See mb.m for small T-adic prec)
//
// p - a prime
// N - tame level
// NN - precision in q-expansion, O(q^NN)
// m - p-adic precision, O(p^m)
// n - T-adic precision, O(T^n)


function TadicE(p,NN,m,n)

	// ---------------------------------------
	// 1) PARAMETERS AND INTERPOLATION POINTS.

	// Set qq such that weight space has parameter T = (1+qq)^wt - 1.
	qq := p;
	if p eq 2 then
		qq := 4;
	end if;

	// Compute the weights that will be used in the finite difference method.
	Ws := [0]; // List of weights as integers, start with wt0 = 0.
	T0 := 0; // Comes from (1+qq)^wt0 - 1 where we take wt0 = 0.
	Ts := [0];
	h := EisWeight(p); // Weight difference for interpolation points.
	par := (1+qq)^h;
	ndash := Ceiling(m/Valuation(h*qq,p) + n)-1; // Number of interpolation points we need, apart from T0.
	j := 0; 
	parj := 1; //par^j
	for i := 1 to ndash do
		j := j+1;
		parj := parj*par;
		Append(~Ws,j*h);
		Append(~Ts,Integers()!(parj-1));
	end for;

	// Compute Valuation (T_1-T_0)*(T_2-T_0)*...*(T_n-T_0).
	// This is the precision we lose due to dividing out in the finite difference procedure.
	loss := 0;
	for i := 2 to ndash+1 do
		loss := loss + Valuation(Ts[i]-Ts[1],p);
	end for;
	mdash := m + loss;
	print "mdash, ndash are ", mdash, ndash;


	// -----------------------------------
	// 2) EISENSTEIN SERIES INTERPOLATION.

	Sdash := IntegerRing(p^mdash);
	Es := [];
	for wt in Ws do
		Ewt := pStabESeries(p,wt,NN,Sdash);
		Append(~Es,Ewt); // List of Gs at all these weights.
	end for;

	// Base rings: q-Expansions will live in B.
	A<T> := LaurentSeriesRing(Integers(p^m),n);
	if m eq 1 then
		A<T> := PowerSeriesRing(GF(p),n);
	end if;
	B<q> := PowerSeriesRing(A,NN);
	E := B!0; // Eisenstein family centered at weight 0.
	Ep1 := B!0; // Eisenstein family centered at weight p-1.

	// Construct the 2-variable series G using finite differences.
	// Initial finite differences are just evaluations:
	Diffs:= [[[Coefficient(Ei,j) : Ei in Es] : j in [0 .. NN-1]]];
	for i := 0 to NN-1 do
		E := E + A!(Diffs[1][i+1][1])*q^i;
	end for;

	// Construct the finite difference set and E simultaneously.
	for d := 1 to ndash do
		Diffd := [[Integers()!(x[j+1]-x[j])/(Ts[j+d]-Ts[j]): j in [1 .. ndash-d+1]] : x in Diffs[d]];
		Append(~Diffs, Diffd);
		for i := 1 to NN-1 do
			Kap := 1;
			for a := 1 to d do
				Kap := Kap*(T-Ts[a]);
			end for;
			E := E + A!(Diffd[i+1][1])*q^i*Kap;
		end for;
	end for;

	return E;

end function;


// Given a q-expansion whose coefficients are T-adic series, return image under T |-> fT.

function SubEis(E,fT)

	NN := AbsolutePrecision(E);
	R<q> := PowerSeriesRing(Parent(fT),NN);
	ES := R!0;
	for i := 0 to NN-1 do
		ES := ES + Evaluate(Coefficient(E,i),fT)*q^i;
	end for;

	return ES;

end function;



// ------------------------------------------------------------------------------------------
// ---                              2)  KATZ EXPANSIONS                                   ---
// ------------------------------------------------------------------------------------------


// Returns dimension of space of modular forms of weight k and level 1.

function DimensionMk(k)

	if ((k mod 2) eq 1) or (k lt 0) then
		return 0;
	end if;
	kmod12 := k mod 12;
	if (kmod12 eq 2) then
		return Floor(k/12);
	else
		return Floor(k/12) + 1;
	end if;

end function;

// Returns basis for W_i (those elements in the Miller basis of weight k which 
// are not in E_(p-1) x {weight (k - (p-1)) forms} and a power of delta
// which can be reused on the next call of function.

function ComputeWi(k,p,delta,deltaj,E4,E6)
	
	// Define a and b
	a := k mod 3;
	b := (k div 2) mod 2;
		
	// Compute dimensions required for Miller basis
	d := DimensionMk(k) - 1;
	e := DimensionMk(k - EisWeight(p)) - 1;
	
	// Construct basis for Wi
	Wi := [];
	for j := e+1 to d do
		// compute aj = delta^j*E6^(2*(d-j) + b)*E4^a
		aj := deltaj*E6^(2*(d-j) + b)*E4^a;
		deltaj := deltaj*delta;
		Wi := Append(Wi,aj);
	end for;

	return Wi,deltaj;

end function;

// Returns list e of Katz expansions e_(i,s) from Step 3 of Algorithm 1
// and E_(p-1) for use in Step 4.

function KatzExpansions(k0,p,ellp,m,n,kdiv)
	
	E := TadicE(p,ellp,m,n);
	q := Parent(E).1;
	T := Parent(Coefficient(E,0)).1;
	//G := E*Evaluate(E,q^p)^-1;

	R := Parent(Coefficient(E,0));

	Ep1cl:= pStabESeries(p,EisWeight(p),ellp,R);
	E4 := ESeries(4,ellp,R);
	E6 := ESeries(6,ellp,R);
	delta := Delta(q);
	
	deltaj := Parent(delta)!1;
	e := [];	
	for i := 0 to kdiv do
		Wi,deltaj := ComputeWi(k0 + i*EisWeight(p),p,delta,deltaj,E4,E6);
		for bis in Wi do
			//eis := p^Floor(i*ap(p)/(p+1))*Ep1^(-i)*bis; // ORIGINAL.
			//eis := T^Floor(i/(p+1))*bis*E;  // WORKS MOD P.
			//eis := G*(Ep1)^(i)*bis;
			//eis := pStabESeries(p,EisWeight(p)*i,ellp,A)^(-1)*bis;
			//eis := (p+T)^Floor(i*ap(p)/p^3)*(Ep1cl)^(-i)*bis*E; // Perhaps p+3?
			eis := (T)^Floor(i*ap(p)/(p+3))*(Ep1cl)^(-i)*bis*E; // Perhaps p+3?
			e := Append(e,eis);
		end for;
	end for;
	
	return e,R;

end function;




// ------------------------------------------------------------------------------------------
// ---                             3)  THE U_p OPERATOR                                   ---
// ------------------------------------------------------------------------------------------


// Computes the Up operator on
// Throughout, the p-adic precision is m, and the T-adic precision is n.

function Level1Up(p,m,n)
 
	// Step 1

	k0 := 0; // Currently we always take the component of weight space of trivial character.
	//kdiv := Floor(((p+1)/EisWeight(p))*(n+1+(ap(p)-1)/(p+1))); // Perhaps p+3?
	kdiv := Floor((p+2)*(m+n+2)/(2*EisWeight(p)));
	ell := DimensionMk(k0 + kdiv*EisWeight(p));
	ellp := ell*p;
	ndash := n + Ceiling(kdiv*ap(p)/p);

	print "kdiv is ", kdiv;
	print "ellp is ", ellp;

	// Steps 2 and 3

	t0 := Cputime();
	S := IntegerRing(p^m);
	e,R:= KatzExpansions(k0,p,ellp,m,ndash,kdiv);

	print "Katz expansions done, time: ", Cputime() - t0;

	// Step 5 and computation of T in Step 6
	T := ZeroMatrix(R,ell,ell);

	for i := 1 to ell do
		for j := 1 to ell do
			T[i,j] := Coefficient(e[i],p*(j-1));
		end for;
	end for;

	// Step 6: solve T = AE using fact E is upper triangular.
	// Warning: assumes that T = AE (rather than pT = AE) has 
	// a solution over Z/(p^mdash). This has always been the case in 
	// examples computed by the author, see Note 3.1.

	//Quot := ChangeRing(Parent(T[1][1]),Integers()); // Ring needed to take quotient later.
	//T := ChangeRing(T,RQ);

	RInt := ChangeRing(R,Integers());
	A := ZeroMatrix(RInt,ell,ell);

	for i := 1 to ell do
		Ti := T[i];
		for j := 1 to ell do
			ej := Parent(Ti)![Coefficient(e[j],l-1): l in [1 .. ell]];
			lj := RInt!(ej[j]);
			//print i,j, RInt!Ti[j]/lj;
			A[i,j] := RInt!Ti[j]/lj;
			Ti := Ti - A[i,j]*ej;
		end for;
	end for;

	//Now cut down to correct precision.
	R<T> := LaurentSeriesRing(Integers(p^m),n);
	if m eq 1 then
		R<T> := PowerSeriesRing(GF(p),n);
	end if;
	return MatrixRing(R,ell)!A, e;

end function;





// ------------------------------------------------------------------------------------------
// ---                             4)  T-ADIC EIGENFORMS                                  ---
// ------------------------------------------------------------------------------------------



// Auxiliary function to turn a polynomial over T-power series (where coefficients have precision)
// into polynomial over T-polynomials (infinite precision).

function AbsolutePolynomial(f,p)

	A<T> := LaurentSeriesRing(GF(p));
	B<x> := PolynomialRing(A);
	F := B!0;

	for i := 0 to Degree(f) do
		C := Coefficient(f,i);
		for j := 0 to Precision(C)-1 do
			F := F + Coefficient(C,j)*T^j*x^i;
		end for;
	end for;

	return F;

end function;



// Returns reciprocal of the characteristic series of a matrix.

function ReverseCharacteristicPolynomial(A)
	
	P:=CharacteristicPolynomial(A);
	return ReciprocalPolynomial(P);

end function;



// Sorts a sequence of T-adic eigenvalues by slope.

function SortSlopes(A)
	
	valf := func<x,y | Valuation(x[1])-Valuation(y[1])>;
	return Sort(A,valf);

end function;


// Computes T-adic valuation of a power series: used for checking in qExpToKatz

function ValPS(f)

	qprec:=Precision(Parent(f));
	val:=Infinity();

	for i:=0 to qprec - 1 do
		val:=Minimum(val,Valuation(Coefficient(f,i)));
	end for;

	return val;

end function;

// Returns position of the leading entry of a NON-ZERO vector.

function LeadingPosition(v)

	d:=Degree(Parent(v));
	i:=1;
	vi:=v[i];
	while vi eq 0 do
		i:=i+1;
		vi:=v[i];
	end while;
	
	return i;

end function;


// Increases the T-adic precision of the entries in a matrix.

function PrecIncrease(A,nn)

	p := Characteristic(Parent(A[1,1]));
	T := Parent(A[1,1]).1;
	n := Precision(Parent(A[1,1]));
	d := Degree(Parent(A[1]));
	R<T> := PowerSeriesRing(GF(p),nn);
	A0 := ZeroMatrix(R,d,d);
	for i := 1 to d do
		for j := 1 to d do
			Aij := A[i,j];
			for k := 0 to n-1 do
				A0[i,j] := A0[i,j] + Coefficient(Aij,k)*T^k;
			end for;
		end for;
	end for;
	
	return A0;				

end function;



// Converts form in Katz Basis E to q-expansion.

function KatzToqExp(E,v)

	ell:=#E;
	qprec:=Precision(Parent(E[1]));
	S:=Parent(v[1]);
	PW<q>:=PowerSeriesRing(S,qprec);
	
	w:=PW!0;
	for i:=1 to ell do
		w:=w + v[i]*(PW!E[i]);
	end for;

	return w;

end function;



// Input: E is the Katz basis modulo T^mdash, and f a modular form as a POWER SERIES
// Output: f in the Katz basis modulo T^m.

function qExpToKatz(f,E,p,m)

	print "f,E,p,m", f,E;
	print "f-E", f - Parent(f)!E[1];
	ell := #E;
	PS := Parent(E[1]);
	Smdash := Parent(Coefficient(f,0)); // IntegerRing(p^mdash)
	Sm:=Parent(Coefficient(f,0));

	elldashp:=Precision(PS);
	elldash:=Minimum(Integers()!(elldashp/p),Precision(Parent(f)));

	Rmdashelldash:=RSpace(Smdash,elldash);
	w := Rmdashelldash![Coefficient(f,i-1): i in [1..elldash]]; // first elldash coefficients
	Rmdashell:=RSpace(Smdash,ell);
	v := Rmdashell![0: i in [1..ell]];


	for j:=1 to ell do
		Ej:=Rmdashelldash![Coefficient(E[j],l-1): l in [1 .. elldash]];
		Ejleadpos:=LeadingPosition(Ej);
		lj:=Sm!(Ej[Ejleadpos]);
		if Valuation(lj) le m then
			v[j]:=Sm!((w[Ejleadpos])/lj);
			w:=w - v[j]*Ej;
			f:=f - v[j]*(Parent(f)!E[j]); // for checking only //WARNING: This coercion does NOT work properly for T,T2
		end if;
	end for;

	if ValPS(f) lt m then // check
		print "Error in qExpToKatz: form not l.c. of katz basis!", f;
	end if;
	
	Rm:=RSpace(Sm,ell);

	print "dashdash: ", Rm!v;

	return Rm!v; // loss of precision mdash - m.

end function;


// This function calculates the q-expansion of the n-th T-adic eigenform (for all n in L),
// where the eigenforms are ordered by slope.
// WARNING: It assumes that every slope occurs with multiplicity one.

// INPUT: p - a prime
// n - the T-adic precision to which we compute U_p
// L - sequence of integers
//
// OUTPUT: The q-expansions of the l-th T-adic eigenforms, for all l in L.

function eigenfunctions(p,n,L)

	A,e := Level1Up(p,1,n);
	print "Size of Up Matrix: ", Degree(Parent(A[1]));
	dim := Degree(Parent(A[1])); // size of matrix A.
	R<T> := PowerSeriesRing(GF(p),n);
	S<q> := PowerSeriesRing(R,dim*p);
	
	
	// Print slopes.
	print "Calculating characteristicSeries of Up... ";
	f := CharacteristicPolynomial(PrecIncrease(A,n^2));
	F := AbsolutePolynomial(f,p); 
	print "Slopes, with multiplicities, are: ";
	N := NewtonPolygon(F);
	V := Vertices(N);
	for i := 1 to #V-1 do
		print (V[i+1][2]-V[i][2])/(V[i][1]-V[i+1][1]), V[i+1][1] - V[i][1];
	end for;
	print 0, Degree(f)-V[#V][1];


	// Find all eigenvalues for sequence L.
	RtsChaos := Roots(F);
	Rts := SortSlopes(RtsChaos); // Sort roots according to their slopes.
	print Rts;
	lambda := [];
	for l in L do
		Append(~lambda,R!(Rts[l][1])); // The l-th eigenvalue of U_p, sorted by T-adic slope.
	end for;
	//print lambda;

	// Now project onto
	G := [];
	for l in L do
		K := Basis(Kernel(A-lambda[l]));
		g := S!0;
		for i:=1 to dim do
			Ki := K[1,i]; // We assume multiplicity one of slopes!
			ei := S!e[i];
			for j := 0 to dim*p-1 do
				g := g + Ki*Coefficient(ei,j)*q^j;
			end for;
		end for;

		Nrm := 1; // Normalisation of coefficient of q for cusp forms.
		if Coefficient(g,0) eq 0 then
			Nrm := Coefficient(g,1);
			print Nrm;
		end if;
		Append(~G, S!(g/Nrm));
	end for;

	return G;

end function;

