// Magma code related to "Modular curves of prime power level with infinitely many rational points",
// by Andrew V. Sutherland and David Zywina.

load "g0groups.m";

/*
  Functions for Siegel functions and finding hauptmoduls.
*/


SiegelPower := recformat<a:SeqEnum, m:RngIntElt, c:FldCycElt>;     
// Record type that represents c*g_a^m, where g_a is the Siegel function.


function DedekindFactor(A)
/*  
    For a matrix A in SL(2,Z) with last row (c,d), returns the 12-th root of unity such that
    eta(Az)^2= w*(cz+d)*eta(z)^2, where eta is the Dedekind eta function.
*/
    C<w>:=CyclicGroup(12);
    SL2:=SL(2,Integers(12)); 
    SL2:=sub<SL2|[[1,1,0,1],[0,1,-1,0]]>;
    f:=hom<SL2->C|[w^1,w^3]>;
    i:=[i: i in [0..11] | f( SL2!A) eq w^i][1];
    return RootOfUnity(12)^i;  
end function;


function ReduceKlein(a,c)
/*  Input: A pair a in Q^2-Z^2 and a constant c.
    Output: A constant C and a "reduced" pair b in Q^2-Z^2 such that c*k_a(z) = C*k_b(z),
            where k_a(z) is a Klein form.
*/
    b1:=Floor(a[1]); b2:=Floor(a[2]);
    a1:=a[1] - b1; a2:=a[2] - b2;
    m:=(b2*a1-b1*a2)/2; 
    e:=(-1)^(b1+b2+b1*b2)*RootOfUnity(Denominator(m))^Numerator(m);

    if ((0 lt a1) and (a1 lt 1/2) and (0 le a2) and (a2 lt 1)) or
       ((a1 eq 0) and (0 lt a2) and (a2 le 1/2)) or
       ((a1 eq 1/2) and (0 le a2) and (a2 le 1/2)) then
       return c*e, [a1,a2];
    end if;

    return $$([-a1,-a2],-c*e);
end function;


function ActOnSiegel(g,A :d:=1)  
/*  A is a matrix in GL(2,Z) and d an integer.
    If g is a record of type SiegelPower, representing a function c*g_a^m with g_a a Siegel function, 
    then the function returns a record of the same type representing sigma_d(g)|_A.  
    [the return function c'*g_{a'}^m has "reduced" a', i.e., a' lies in some set A_N (as in the paper).]
    If g is a sequence, then returns sequence with "ActOnSiegel( ,A,d:=d)" applied to its elements.   

    Warning: It is assumed that d is relatively prime to the denominators of a=(a1,a2) and to 2 and 3.
*/

    if Type(A) eq SeqEnum then
       R:=Parent(A[1]); 
       A:=GL(2,R)!A;
    end if;
    R:=BaseRing(Parent(A));
    if Type(R) ne RngInt then
       N:=#R;
       d:=Integers()!Determinant(A); while GCD(d,6) ne 1 do d:=d+N; end while;
       B:=(GL(2,Integers(N))![1,0,0,d])^(-1) * A;
       B:=SL(2,Integers())!SL2ZLift(B,N);
       return $$(g,B :d:=d);
    end if;

    if Type(g) eq SeqEnum then return [$$(s,A: d:=d): s in g]; end if;
   
    if d ne 1 then 
       g:=rec<SiegelPower|a:=[g`a[1],d*g`a[2]], m:=g`m, c:=Conjugate(g`c,d)>;
       return $$(g,A: d:=1);
    end if;

    A    :=SL(2,Integers())!A;
    a    :=g`a; 
    c_,a_:=ReduceKlein([a[1]*A[1,1]+a[2]*A[2,1],a[1]*A[1,2]+a[2]*A[2,2]],1); 
    C    :=g`c * (DedekindFactor(A) * c_)^(g`m);
    return rec<SiegelPower|a:=a_, m:=g`m, c:=C>;
end function;



function SiegelExpansion(g,Precision)
/*  Computes the q-expansion of g.   We assume that g is given in the format as described above for "hauptmodul" for
    the record format "CongSubgroup".  The quantity Precision determines how many terms are used in the product 
    giving the Siegel functions that occur.

    Warning: We will assume that a is "reduced" for each Siegel function g_a, i.e., it occurs in some set A_N from the paper.
*/
        if Type(g) eq SeqEnum and Type(g[1]) eq SeqEnum then
           return &+[$$(b,Precision) : b in g];
        end if;

        e:=&+[r`a[2]*(r`a[1]-1)/2 * r`m : r in g];
        NN:=LCM( [Denominator(r`a[2]): r in g] cat [Denominator(e)]);

        c:=&*[r`c: r in g];
        K:=Compositum(CyclotomicField(NN),Parent(c));    
        c:=K!c;

        R<q>:=PuiseuxSeriesRing(K);
        gg:= R!(RootOfUnity(Denominator(e))^Numerator(e)); 
        gg:=gg* &*[ (1-q^r`a[1]*RootOfUnity(Denominator(r`a[2]))^Numerator(r`a[2]) +O(q^Precision))^r`m : r in g];  

        S:=[1] cat 
           &cat[ [(1-q^n * q^r`a[1]*RootOfUnity(Denominator(r`a[2]))^Numerator(r`a[2]) +O(q^Precision))^r`m : n in [1..Floor(Precision-r`a[1])]] : r in g];  
        S:=S cat 
           &cat[ [(1-q^n * q^(-r`a[1])*RootOfUnity(Denominator(r`a[2]))^Numerator(-r`a[2]) +O(q^Precision))^r`m : n in [1..Floor(Precision+r`a[1])]] : r in g];            
        gg:=gg*&*S;

        v:=&+[1/2*(r`a[1]^2-r`a[1]+1/6)*r`m : r in g]; 
        gg:=gg* &*[(-1)^r`m: r in g] * q^v;

        return R!(c*gg);
end function;


function LeadingTermSiegel(g)  
/*
   Suppose g is a function of the form c*g_a^m (given as a record of type SiegelPower) where
   g_a is a Siegel function with a in A_N for some N.
   The function returns C and v, where the leading term of the q-expansion of the function
   c*g_a^m is of the form C*q^v.

   Suppose g is a sequence.  If g is a sequence of functions of the form c*g_a^m as above, then 
   returns the C and v as above of their product.  Otherwise, applies LeadingTermSiegel() to 
   each term.
*/

    if Type(g) eq SeqEnum and Type(g[1]) eq SeqEnum then
       return [$$(r): r in g];
    end if;

    C:=&*[r`c:r in g];
    C:=C*(-1)^(&+[r`m: r in g]);
    e:=&+[r`m*r`a[2]*(r`a[1]-1)/2 : r in g];
    C:=C*RootOfUnity(Denominator(e))^Numerator(e);
    C:=C * &*([1] cat [ ( 1 - RootOfUnity(Denominator(r`a[2]))^Numerator(r`a[2]))^r`m : r in g | r`a[1] eq 0]);

    return [* C, &+[1/2*((r`a[1])^2-r`a[1]+1/6)*r`m : r in g] *];
end function;


//  FUNCTIONS FOR FINDING THE HAUPTMODUL OF A GENUS 0 CURVE OF PRIME POWER LEVEL

function CuspData(H)
/*
     For the given subgroup H of SL(2,Z/N) with N>1, let Gamma be the congruence subgroup of level N
     whose image modulo N is H.   
     The function returns a sequence of matrices A and a sequence of integers w.   
     The matrices A lie in SL(2,Z) and the values A*infty represent the cusps of X_Gamma.   
     The integers w gives the widths of the cusps.  
     (The first cusp will always be the one at infinity)
*/
     N:=#BaseRing(H);  SL2:=SL(2,Integers(N));
     if SL2![-1,0,0,-1] notin H then  H:=sub<SL2| Generators(H) join {SL2![-1,0,0,-1]}>; end if;

     Ht:=sub<SL2|{Transpose(g):g in Generators(H)}>;
     cusps:=[ Sort([[Integers()!a[1],Integers()!a[2]]: a in O]) : O in Orbits(Ht)];
     cusps:=Sort([O: O in cusps | GCD([O[1][1],O[1][2],N]) eq 1]);

     matrices:=[ SL(2,Integers())![1,0,0,1] ];
     for c in cusps do
     if [1,0] notin c then
        c_:=c[1];
        a:=Integers()!c_[1]; b:=Integers()!c_[2];
        while GCD(a,b) ne 1 do 
             if b ne 0 then a:=a+N; else b:=b+N; end if;
        end while;
        g,x,y:=XGCD(a,b);
        matrices:=matrices cat [SL(2,Integers())![a,-y,b,x]]; 
     end if;       
     end for;

     SL2:=SL(2,Integers(N));
     widths:=[ Minimum([i:i in [1..N]| (SL2!A)*(SL2![1,i,0,1])*(SL2!A)^(-1) in H]) :  A in matrices];
     return matrices, widths;
end function;


function SiegelOrbits(H)
/*
    Given a subgroup H of SL(2,Z/NZ) with N>1, let Gamma be the congruence subgroup of level N whose image
    in SL(2,Z/NZ) is H.
    With notation as in the paper, computes the sequence of orbits of \pm Gamma on the set A_N.
*/
    N:=#BaseRing(H); 
    AA:={[a/N,b/N] : a, b in [0..N-1] | [a,b] ne [0,0]};
    AA:={a: a in AA |  (0 lt a[1] and a[1] lt 1/2) or (0 eq a[1] and a[2] le 1/2) or (a[1] eq 1/2 and a[2] le 1/2) };

    SL2:=SL(2,Integers(N));
    if SL2![-1,0,0,-1] notin H then  H:=sub<SL2| Generators(H) join {SL2![-1,0,0,-1]}>;  end if;

    orbits:=[ {[1/N*(Integers()!a[1] mod N),1/N*(Integers()!a[2] mod N)] : a in O} meet  AA : O in Orbits(H) ];
    orbits:=Sort([ Sort([a:a in O]) : O in orbits | #O ne 0]);    
    return orbits;
end function;


function SiegelDivisor(H,cusps,widths,orbit)   
//  Input: "orbit" is an orbit of H on A_N, where H is the image modulo the level N of a congruence subgroup
//  with a list of "cusps" and their "widths".
//  Output: the divisor (as a sequence matching the order of "cusps") of g_orbit;  see the paper for details.
    N:=#BaseRing(H); 
    M:=RModule(Rationals(),#cusps);
    v:=M!0;
    for a in orbit do
        D:=[];
        for i in [1..#cusps] do  
            A:=cusps[i];  w:=widths[i]; 

            b:=[a[1]*A[1,1]+a[2]*A[2,1],a[1]*A[1,2]+a[2]*A[2,2]];
            b1:=b[1]-Floor(b[1]);
            D:=D cat [6*N*w*(b1^2-b1+1/6)];
        end for;
        v:=v+M!D;
    end for;  
    return [Integers()!a: a in Eltseq(v)];
end function;


function HasGenusZero(H)
// Given a subgroup H of GL(2,Z/NZ), determines whether the congruence subgroup consisting of A in SL(2,Z)
// with A mod N in H has genus 0.

    N:=#BaseRing(H); 
    SL2:=SL(2,Integers(N)); 
    GL2:=GL(2,Integers(N)); 
    if SL2![-1,0,0,-1] notin H then  H:=sub<SL2|Generators(H) join {SL2![-1,0,0,-1]}>;  end if;

    // Reduce to the case where Gamma has level N.
    if not IsPrime(N) then
       for p in PrimeDivisors(N) do
           N_:=N div p;  SL2_:=SL(2,Integers(N_));
           r:=hom<SL2->SL2_|[SL2_! SL2.i: i in [1..#Generators(SL2)]]>;
           if Kernel(r) subset H then return HasGenusZero(r(H)); end if;
       end for;
    end if;     

    for k in Keys(sl2tab) do
        if N eq sl2tab[k]`level and IsConjugate(GL2,H,sub<SL2|sl2tab[k]`gens>) then
           return true, H;
        end if;          
    end for;

    return false, H;
end function;


function FindHauptmodul(H)
/*  Let H be the image modulo the level N of a congruence subgroup Gamma of genus 0; assume that N is not 1 or 11.
    Returns a hauptmodul for H as described in the paper; output format described in "ComputeHauptmoduls.txt".
*/

    b,H:=HasGenusZero(H);
    assert b;   // stops if X_H does not have genus 0

    N:=#BaseRing(H); 
    SL2:=SL(2,Integers(N)); 

    // Compute cusps and the Siegel divisors of X_Gamma
    cusps, widths := CuspData(H);      
    // Compute the orbits of Gamma on A_N
    orbits:=SiegelOrbits(H);    
    divisors:=[SiegelDivisor(H,cusps,widths,O):O in orbits];

    if #cusps ge 2 then
       A:=Transpose( Matrix([ Eltseq(D) : D in divisors]) );      
       // Transposes needed because of how Magma does linear algebra.
       b0:=Matrix([[0]: j in [1..#cusps]]);
       b:=b0; b[1][1]:=-12*N;  b[2][1]:=+12*N;
       _,m,W:=IsConsistent(Transpose(A),Transpose(b));
       m:=Eltseq(m);

       haupt:=[];
       for j in [1..#orbits] do
           if m[j] ne 0 then
           for a in orbits[j] do
               haupt:=haupt cat [rec<SiegelPower|a:=a,m:=m[j],c:=CyclotomicField(1)!1>];
           end for;
           end if;
       end for;

       e:=&+[m[j] * &+[a[2]*(a[1]-1)/2: a in orbits[j]]   : j in [1..#orbits]];
       c:=RootOfUnity(Denominator(e))^(-Numerator(e)); if Type(c) eq FldRatElt then c:=CyclotomicField(2)!c; end if;
       haupt[1]`c:=c;
    end if;


    if #cusps eq 1 then
       nsub:=[G`subgroup: G in NormalSubgroups(H)];
       nsub:=[H2: H2 in nsub | SL2![-1,0,0,-1] in H2 and HasGenusZero(H2)];
       nsub:=[H2: H2 in nsub | #CuspData(H2) eq Index(H,H2)];
       H_:=nsub[1];
   
       haupt_:=FindHauptmodul(H_);
       T:=Transversal(H,H_);
       T:=[SL2ZLift(A,N): A in T];

       haupt:=[];
       for A in T do
          h_:=[];
          for r in haupt_ do
              h_:=h_ cat [ActOnSiegel(r,A)];
          end for;
          haupt:=haupt cat [h_];          
       end for;
    end if;

    return haupt;
end function;



