/*****************************************************
 tr.m

 AUTHOR -- William A. Stein
 CREATED: 14 January 2001

 DESCRIPTION: Hijikata trace formula for tr(T_n) 
 on S_k(Gamma_0(N)), any n>=1, except if n|N, 
 then n must be prime.

 *****************************************************/

//classgroup_algorithm := "ClassGroup";
//classgroup_algorithm := "Shanks";   // heuristic -- can be *wrong*!!
classgroup_algorithm := "ReducedForms";  

function w(d)
   case d:
      when -4:
         return 2;
      when -3:
         return 3;
   end case;
   return 1;
end function;


function tof(a)
   t := &*[pn[1]^((pn[2]-(pn[2] mod 2)) div 2) : 
                        pn in Factorization(a)];

   if (Integers()!(a/t^2) mod 4) eq 1 then
      return t;
   else
      return Integers()!(t/2);
   end if;
end function;


function mu(N)
   return N * 
           &*[Rationals()|1+1/p[1] : p in Factorization(N)];
end function;


function sig(n, N)
   if N mod n eq 0 then
      return n;
   else
      return &*[ (1 -p[1]^(p[2]+1))/(1-p[1]) : p in Factorization(n)];
   end if;  
end function;


function quadpoints(s,n,p,v)   
   // WARNING: This is a very very stupid algorithm!
   pv := p^v;
   return [x : x in [0..pv-1] | (x^2 - s*x + n) mod pv eq 0];
end function;


function A(s,f,n,p,v)
   rho := Valuation(f,p);
   sr := Set([x mod p^(v+rho) : x in quadpoints(s,n,p,v+2*rho)]);   
   return #[x : x in sr | 
          (2*x-s) mod p^rho eq 0 
          and ((n ne p) or ((n eq p) and x mod p ne 0))];
end function;


function B(s,f,n,p,v)
   rho := Valuation(f,p);
   return #Set([x mod p^(v+rho) : x in quadpoints(s,n,p,v+2*rho+1)]);
end function;


function cp(s,f,n,p,v)
   ans :=  A(s,f,n,p,v);
   if (Integers()!((s^2-4*n)/f^2)) mod p eq 0 then
      ans +:= B(s,f,n,p,v);
   end if;
   return ans;
end function;


function c(s,f,n,N)
   return &*[Integers()| cp(s,f,n,p[1],p[2]) : p in Factorization(N)];
end function;


function type_p(n,k,N)
   if IsSquare(n) then
      s := Floor(Sqrt(4*n));
      ans := 1/4*(s/2)*n^((k div 2)-1)*(c(s,1,n,N) + (-1)^k*c(-s,1,n,N));
      return ans;
   else
      return 0;
   end if;
end function;


function absxy(s,n,k,N)
   t := Floor(Sqrt(s^2-4*n));
   x := (s-t)/2;
   y := (s+t)/2;
   ans :=  (Min([Abs(x), Abs(y)])^(k-1)/Abs(x-y)) * Sign(x)^k *
            &+[1/2*EulerPhi(Integers()!(t/f))*c(s,f,n,N) : f in Divisors(t)];
   return ans;
end function;


function type_h(n,k,N)
   start := Ceiling(2*Sqrt(n));
   if IsSquare(n) then
      start +:= 1;
   end if;
   ans :=  &+[Rationals()|absxy(s,n,k,N) + absxy(-s,n,k,N) :
          s in [start..4*n] | IsSquare(s^2-4*n)];
   return ans;
end function;

function xy(s,n,k,N)
   K<a> := QuadraticField(s^2-4*n);
   D := Discriminant(K);   // a = sqrt(D) or sqrt(D/4).
   if D mod 4 eq 0 then
      D := D div 4;
   end if;
   b := Round(Sqrt((s^2-4*n)/D));
   x := (s+a*b)/2;
   y := (s-a*b)/2;
   ans := 1/2 *  (x^(k-1)-y^(k-1))/(x-y) 
          * &+[ClassNumber(Integers()!((s^2-4*n)/f^2) : Al:=classgroup_algorithm)
             / w((s^2-4*n)/f^2) * c(s,f,n,N) : 
          f in Divisors(tof(s^2-4*n))];
   return Rationals()!ans;
end function;


function type_e(n,k,N)
   r := Floor(2*Sqrt(n));
   if IsSquare(n) then
      r -:= 1;
   end if;

//   return &+[xy(s,n,k,N) : s in [-r..r]];
// WARNING: I *conjecture* that the sum below is the same.

   return xy(0,n,k,N) + 2* &+[xy(s,n,k,N) : s in [1..r]];
end function;


function sum_s(n,k,N)
   return type_p(n,k,N)+type_h(n,k,N)+type_e(n,k,N);
end function;


intrinsic TraceHeckeOperator(n::RngIntElt,  
                             k::RngIntElt) -> RngIntElt
{}
   require n ge 1 : "Argument 1 must be at least 1.";
   require k ge 2 : "Argument 2 must be at least 2.";
   return TraceHeckeOperator(n,k,1);
end intrinsic;

intrinsic TraceHeckeOperator(n::RngIntElt,  
                             k::RngIntElt, 
                             N::RngIntElt) -> RngIntElt
{The trace of the Hecke operator T_n on S_k(Gamma_0(N))).  
The only constraint on n is that if GCD(n,N) ne 1, 
then n must be prime.}
   require n ge 1 : "Argument 1 must be at least 1.";
   require k ge 2 : "Argument 2 must be at least 2.";
   require N ge 1 : "Argument 3 must be at least 1.";
 
   if GCD(n,N) ne 1 then
      require IsPrime(n) : 
        "Argument 1 must be prime when the GCD of arguments 1 and 3 is not 1.";
   end if;
 
   if k mod 2 ne 0 then
      return 0;
   end if;
   if k eq 2 then 
      t := sig(n,N);
   else
      t := 0;
   end if;
   t +:= -sum_s(n,k,N);
   if IsSquare(n) then
      t +:= (k-1)*mu(N)/12*n^((k div 2)-1);
   end if;
   return t;
end intrinsic;


intrinsic TraceModularForm(k::RngIntElt, prec::RngIntElt)
      -> RngSerPowElt
{The trace modular form sum Tr(T_n)*q^n to absolute precision prec,
where T_n is the nth Hecke operator on S_k(SL_2(Z)).  The complexity
is almost entirely a function of prec, but *not* of k.}

   R<q>:=PowerSeriesRing(Rationals());
   return &+[TraceHeckeOperator(n,k,1)*q^n : n in [1..prec-1]] + O(q^(prec));

end intrinsic;

intrinsic TraceToFile(k::RngIntElt, 
                      start::RngIntElt, stop::RngIntElt,
                      file::File)
{}
  for n in [start..stop] do
     t := Cputime();
     tr := TraceHeckeOperator(n,k,1);
     fprintf file, "+%o*q^%o // time = %o s\n", tr, n, Cputime(t);
     Flush(file);
  end for;
end intrinsic;

function Tp(p, r, k, f)
   assert Type(p) eq RngIntElt;
   assert Type(r) eq RngIntElt;
   assert Type(f) eq RngSerPowElt;

   if r gt 1 then
      return Tp(p,1,k, Tp(p,r-1,k,f)) - p^(k-1)*Tp(p,r-2,k, f);
   end if;
   if r eq 1 then
      R<q> := Parent(f);
      prec := Floor(((AbsolutePrecision(f)-1)/p) + 1);
      return &+[R|Coefficient(f,n*p)*q^n + p^(k-1)*Coefficient(f,n)*q^(n*p) :
                    n in [1..prec-1]] + O(q^prec);
   end if;
   if r eq 0 then
      return f;
   end if;
end function;


intrinsic HeckeOperator(n::RngIntElt, k::RngIntElt, f::RngSerPowElt)
      -> RngSerPowElt
{The image T_n(f) of f under the Hecke operator T_n on level-1 
weight-k modular forms.}
   require n ge 1 : "Argument 1 must be at least 1.";
   require k ge 2 : "Argument 2 must be at least 2.";

   for p in Factorization(n) do
      f := Tp(p[1],p[2],k,f);
   end for;
   return f;

end intrinsic;

intrinsic TraceFormulaBasis(k::RngIntElt, prec::RngIntElt) -> SeqEnum
{Basis for S_k(SL_2(Z)) computed using the trace formula.}

   f := TraceModularForm(k, prec);
   d := Integers()!Coefficient(f,1);  // the dimension
   B := [HeckeOperator(n,k,f) : n in [1..d]];
   return B;
end intrinsic;

function BasisMatrix(B)
   d := #B;
   I := MatrixAlgebra(Rationals(),d)!0;
   for r in [1..d] do 
      for c in [1..d] do
         I[r,c] := Coefficient(B[r],c);
      end for;
   end for;
   return I;
end function;

intrinsic HeckeOperatorMatrix(k::RngIntElt, n::RngIntElt) -> AlgMatElt
{Matrix representing the Hecke operator T_n (acting on the
 right, as usual in MAGMA) on weight-k level-1 modular forms
 with respect to the trace basis.}
   d := DimensionCuspFormsGamma0(1,k);
   B := TraceFormulaBasis(k, n*d^2+1);
   return HeckeOperatorMatrix(k,n,B);
end intrinsic;


intrinsic HeckeOperatorMatrix(k::RngIntElt, 
                              n::RngIntElt, 
                              B::SeqEnum ) -> AlgMatElt

{Matrix representing the Hecke operator T_n (acting on the
right, as usual in MAGMA) on weight-k level-1 modular forms
with respect to the basis B.}
   d := #B;
   I := BasisMatrix(B)^(-1);
   A := Parent(I)!&cat[ [Coefficient(g,i) : i in [1..d]]    
             where g is HeckeOperator(n,k,B[j]) : j in [1..d]];
   return I*A;
end intrinsic;


/*
Here's a good test that the program is working:

  for k in [i : i in [48..60] | IsEven(i)] do
     B:=TraceFormBasis(k,k);  
     I := BasisMatrix(B);
     D := DiscriminantOfHeckeAlgebra(CS(MS(1,k,+1)));
     k, Determinant(I)-D;
  end for;

*/