// Code common to both Fermat and Klein twists

// twist record format is <n, H0 generators, h, ID(H), ID(H0), f(x,y,z), abs(disc(f)), g(t), h(t), ST, ST_k(a))>
//   n = # twist, f(x,y,z)=0 is the curve equation, g(t) defines the base field k (either Q or a quadratic field), where g(t)=t for Q
//   the splitting field of h(t) is the field K=L over which all the endomorphisms/homomorphisms are defined
//   ST identifies the ST group over k, ST_k(i) identifies it over k(a)

// In the functions below, phi is a map from the free group on generators s,t,u,z to permutation group G (<196,956> for Fermat, <336,208> for Klein)
// G0 is the index 2 subgroup of G generated by the images of s,t,u (<96,64> for Fermat, <168,42> for Klein)
// H is always a subgroup of G and H0 is an index 2 subgroup of H (except for 3 Klein twists where H=H0)

function H0_group (twist,phi) return sub<Codomain(phi)|phi(twist[2])>; end function;
function H_group (twist,phi) return sub<Codomain(phi)|phi(twist[2]) cat [phi(twist[3])*phi(Domain(phi).4)]>; end function;

function o(H0,n) return  #[g:g in H0|Order(g) eq n]; end function;
function obar(H,H0,n) return #[g:g in H|not g in H0 and Order(g) eq n]; end function;

function groups (twist,phi)
    G0 := sub<Codomain(phi)|[phi(Domain(phi).i):i in [1..3]]>;
    H := H_group (twist,phi);
    H0 := H0_group (twist,phi);    
    return G0,H,H0;
end function;


// functions to compute moments, these all depend on a function exp(G0,H,H0,j,v) that is specific to the Fermat and Klein cases
// the function exp effectively constructs the ST group corresponding to H,H0,j and sums a1^v[1]*a2^v[2]*a3^v[3] over the components
    
function a1Moments (G0,H,H0,j,n,exp)
    return [Integers()!exp(G0,H,H0,j,[i,0,0]):i in [1..n]];
end function;
    
function a2Moments (G0,H,H0,j,n,exp)
    return [Integers()!exp(G0,H,H0,j,[0,i,0]):i in [1..n]];
end function;
    
function a3Moments (G0,H,H0,j,n,exp)
    return [Integers()!exp(G0,H,H0,j,[0,0,i]):i in [1..n]];
end function;

function vExp (G0,H,H0,j,v,exp)
    return Integers()!exp(G0,H,H0,j,v);
end function;

function CorrelationSimplex (G0,H,H0,j,n,exp)
    A:=AssociativeArray();
    for ii in [0..n] do for jj in [0..n-ii] do for kk in [0..n-ii-jj] do
        A[[ii,jj,kk]] := Rationals()!exp(G0,H,H0,j,[ii,jj,kk]);
    end for; end for; end for;
    return A;
end function;

function z1Exp (H,H0,j) assert j in [1,2]; return j eq 1 select 1/2+1/2*o(H0,3)/#H0 else o(H0,3)/#H0; end function;
    
function z2Exp (H,H0,j)
    assert j in [1,2];
    if j eq 2 then return [0,o(H0,3)/#H0,0,0,0]; end if;
    return [(1/2)*(obar(H,H0,4)/#H0),(1/2)*(o(H0,3)/#H0)+(1/2)*(obar(H,H0,6)/#H0),(1/2)*(obar(H,H0,8)/#H0),(1/2)*(obar(H,H0,12)/#H0),(1/2)*(obar(H,H0,2)/#H0)];
end function;

function z3Exp (H,H0,j) assert j in [1,2]; return j eq 1 select 1/2 else 0; end function;

function stKey(G0,H,H0,j,exp);    
    m101 :=Integers()!exp(G0,H,H0,j,[1,0,1]);
    m030 :=Integers()!exp(G0,H,H0,j,[0,3,0]);
    m202 :=Integers()!exp(G0,H,H0,j,[2,0,2]);
    return <m101,m030,m202>;
end function;

function tstKey(twist,j,phi,exp)
    G0,H,H0 := groups(twist,phi);
    return stKey(G0,H,H0,j,exp);
end function;

function STdictionary(families) // families = [<twists,phi,exp>]
    STgroups:=[];
    STdict := AssociativeArray();

    for r in families do 
        twists := r[1]; phi := r[2]; exp := r[3];
        for i in [1..#twists] do
            G0,H,H0:=groups(twists[i],phi);
            for j in [1,2] do
                k := tstKey(twists[i],j,phi,exp);
                if not k in Keys(STdict) then
                    G0,H,H0 := groups(twists[i],phi);
                    Append(~STgroups,<G0,H,H0,j,exp>);
                    STdict[k] := #STgroups;
                end if;
            end for;
        end for;
    end for;

    function STcmp (a,b)
        ida := IdentifyGroup(a[4] eq 1 select a[2] else a[3]);
        idb := IdentifyGroup(b[4] eq 1 select b[2] else b[3]);
        if ida lt idb then return -1; end if;
        if ida gt idb then return 1; end if;
        ka := stKey(a[1],a[2],a[3],a[4],a[5]);
        kb := stKey(b[1],b[2],b[3],b[4],b[5]);
        if ka lt kb then return -1; end if;
        if ka gt kb then return 1; end if;
        return 0;
    end function;

    Sort(~STgroups,STcmp);
    for i:= 1 to #STgroups do r:=STgroups[i]; STdict[stKey(r[1],r[2],r[3],r[4],r[5])]:=i; end for;
    return STdict, STgroups;
end function;

function stGroups(twist,phi,exp,STdict)
    return [STdict[tstKey(twist,j,phi,exp)]:j in [1..2]];
end function;

// verification
    
function verify_groups(twists,phi,M) // M is the CM field Q(i) or Q(sqrt(-7))
    for r in twists do
        if IdentifyGroup(H0_group(r,phi)) ne r[5] then print "Bad H0_group for twist", r[1]; return false; end if;
        if IdentifyGroup(H_group(r,phi)) ne r[4] then print "Bad H_group for twist", r[1]; return false; end if;
        k := SplittingField(r[8]);
        if IdentifyGroup(GaloisGroup(ChangeRing(r[9],k))) ne r[4] then print "Bad field K for twist",r[1]; return false; end if;
        kM := Compositum(k,M);
        if IdentifyGroup(GaloisGroup(ChangeRing(r[9],kM))) ne r[5] then print "Bad field K for twist",r[1]; return false; end if;
        print r[1], true;
    end for;
    return true;
end function;

function verify_split(twist,E,n)
    p := 1;
    kf := twist[8];
    f := twist[9];
    C := Curve(ProjectiveSpace(Parent(twist[6])),twist[6]);    
    D := Integers()!(twist[7]*Discriminant(E)*Discriminant(f));
    i := 0; j:= 0;
    while i lt n do
        p:= NextPrime(p);
        while IsDivisibleBy(D,p) or #Roots(ChangeRing(kf,GF(p))) ne Degree(kf) do p := NextPrime(p); end while;
        P:=Decomposition(BaseRing(C),p)[1][1]; assert #ResidueClassField(P) eq p;
        if #Roots(ChangeRing(f,GF(p))) eq Degree(f) then
            assert 3*TraceOfFrobenius(E,p) eq p+1-#RationalPoints(Reduction(C,P));
            printf ".";
            i +:= 1;
        else
            if 10*j lt n then
                t := TraceOfFrobenius(E,p);
                if t ne 0 then assert 3*TraceOfFrobenius(E,p) ne p+1-#RationalPoints(Reduction(C,P)); end if;
                j +:= 1;
            end if;
        end if;
    end while;
    return true;
end function;
    
function verify_splits(twists,ec)
    for i:=1 to #twists do print i,verify_split(twists[i],ec,10); end for;
    return true;
end function;

function verify_isogenies(twists,phi,exp,STdict)
    T := Sort([<stGroups(twists[i],phi,exp,STdict),i>:i in [1..#twists]]);
    i := 1;
    while i lt #T do
        Ci := Curve(ProjectiveSpace(Parent(twists[T[i][2]][6])), twists[T[i][2]][6]);
        j := i+1;
        while j le #T and T[j][1] eq T[i][1] do
            Cj := Curve(ProjectiveSpace(Parent(twists[T[j][2]][6])), twists[T[j][2]][6]);
            assert BaseRing(Ci) eq BaseRing(Cj);
            p := 1;
            for k := 1 to 100 do
                p := NextPrime(p);
                while IsDivisibleBy(twists[T[i][2]][7],p) or IsDivisibleBy(twists[T[j][2]][7],p) do p := NextPrime(p); end while;
                Ps:=[r[1]:r in Decomposition(BaseRing(Ci),p)];
                assert Set([#RationalPoints(Reduction(Ci,P)):P in Ps]) eq Set([#RationalPoints(Reduction(Cj,P)):P in Ps]);
            end for;
            printf "%o(%o) ~ %o(%o)\n", T[i][2],twists[T[i][2]][1], T[j][2],twists[T[j][2]][1];
            j +:= 1;
        end while;
        i := j;
    end while;
    return true;
end function;

function verify_no_real_embedding (H,H0)
    if H[1] eq H0[1] then return true; end if;
    H := SmallGroup(H[1],H[2]);
    H0 := SmallGroup(H0[1],H0[2]);
    I := [g:g in H|Order(g) eq 2];
    for G in [K`subgroup : K in Subgroups(H:IndexEqual:=2) | IsIsomorphic(K`subgroup,H0)] do
        K,pi:=quo<H|G>;
        if #[g:g in I|Order(pi(g)) eq 2] gt 0 then return false; end if;
    end for;
    return true;
end function;

function verify_minimal_fields_of_definition (twists)
    printf "Verifying minimality of fields of definition...";
    for r in twists do
        if Degree(r[8]) eq 1 then return true; end if;
        if r[4] eq r[5] then assert Degree(r[8]) eq 2; return true; end if;
        assert r[4][1] eq r[5][1]*2;
        assert verify_no_real_embedding (r[4],r[5]);
    end for;
    return true;
end function;
    
// functions for formatting tables in the paper
    
function sprint_id (id)
    return Sprintf("$\\langle%o,%o\\rangle$",id[1],id[2]);
end function;

function sprint_rat (r)
    if r eq 0 then return "0"; end if;
    return Sprintf("$\\nicefrac{%o}{%o}$",Numerator(r),Denominator(r));
end function;

function sprint_stgroup (i,G0,H,H0,j,exp)
    assert j in [1,2];
    a1 := a1Moments(G0,H,H0,j,4,exp);
    a2 := a2Moments(G0,H,H0,j,4,exp);
    a3 := a3Moments(G0,H,H0,j,4,exp);
    m101 := vExp(G0,H,H0,j,[1,0,1],exp);
    m030 := vExp(G0,H,H0,j,[0,3,0],exp);
    m202 := vExp(G0,H,H0,j,[2,0,2],exp);
    z1 := sprint_rat(z1Exp(H,H0,j));
    z2 := [sprint_rat(r):r in z2Exp(H,H0,j)];
    z2 := Sprintf("%o\\hspace{-8pt}&\\hspace{-8pt}%o\\hspace{-8pt}&\\hspace{-8pt}%o\\hspace{-8pt}&\\hspace{-8pt}%o\\hspace{-8pt}&\\hspace{-8pt}%o",z2[1],z2[2],z2[3],z2[4],z2[5]);
    z3 := sprint_rat(z3Exp(H,H0,j));
    ID := sprint_id(j eq 1 select IdentifyGroup(H) else IdentifyGroup(H0));
    return Sprintf("%o & %o & %o & %o & %o & %o & %o & %o & %o & %o & %o & %o & %o & %o\\\\\n",i,ID,m101,m030,m202,a1[2],a1[4],a2[1],a2[2],a3[2],a3[4],z1,z2,z3);
end function;

function sprint_stgroup_data(groups,i,exp)
    r:=groups[i];
    G0:=r[1];  H:=r[2]; H0:=r[3]; j:=r[4];
    a1 := a1Moments(G0,H,H0,j,4,exp);
    a2 := a2Moments(G0,H,H0,j,4,exp);
    a3 := a3Moments(G0,H,H0,j,4,exp);
    m101 := vExp(G0,H,H0,j,[1,0,1],exp);
    m030 := vExp(G0,H,H0,j,[0,3,0],exp);
    m202 := vExp(G0,H,H0,j,[2,0,2],exp);
    z1 := z1Exp(H,H0,j);
    z2 := z2Exp(H,H0,j);
    z3 := z3Exp(H,H0,j);
    return Sprintf("%4o   %5o   %5o   %5o   %5o   %5o   %5o   %5o   %6o   %5o  [%5o,%5o,%5o,%5o,%5o]  %5o", m101,m030,m202,a1[2],a1[4],a2[1],a2[2],a3[2],a3[4],z1,z2[1],z2[2],z2[3],z2[4],z2[5],z3);
end function;

function sprint_poly(f)
    return StripWhiteSpace(Sprintf("%o",f));
end function;

function sprint_factorization(N)
    q := Factorization(N);
    s := Sprintf("%o^%o", q[1][1],q[1][2]);
    for i:=2 to #q do s *:= Sprintf("*%o^%o", q[i][1], q[i][2]); end for;
    return s;
end function;

function sprint_curve(r)
    R<x,y,z>:=PolynomialRing(BaseRing(Parent(r[6])),3);
    c := StripWhiteSpace(&cat Split(Sprintf("%o",Evaluate(r[6],[x,y,z])),"*"));
    return Sprintf("&\\multicolumn{8}{l|}{$%o$}\\\\\\hline\n", c);
end function;

function sprint_monomial(c,i)
    if c eq 0 then return ""; end if;
    cstr := (c gt 0 select "\\!+\\!" else "\\!-\\!") cat IntegerToString(Abs(c));
    if i eq 0 then return cstr; end if;
    if i eq 1 then mstr := "x"; else mstr := i ge 10 select "x^{" cat IntegerToString(i) cat "}" else "x^" cat IntegerToString(i); end if;
    if c eq 1 then return "\\!+\\!" cat mstr; end if;
    if c eq -1 then return "\\!-\\!" cat mstr; end if;
    return cstr cat mstr;
end function;

function sprint_field(f)
    d := Degree(f);
    if d eq 1 then return "$\\Q$"; end if;
    if d eq 2 then 
        D := SquareFree(Integers()!Discriminant(f));
        if D eq -1 then return "$\\Q(i)$"; end if;
        if D eq -7 then return "$\\Q(a)$"; end if;
        return Sprintf("$\\Q(\\sqrt{%o})$", D);
    end if;
    if d eq 4 then
        G:=GaloisGroup(f);
        if IsAbelian(G) and Exponent(G) eq 2 then
            D := Sort([d: d in {SquareFree(Integers()!Abs(Discriminant(K[1]))): K in Subfields(SplittingField(f))|Degree(K[1]) eq 2}]);
            if D[1] eq 1 then
                if D[2] eq 7 then 
                    return Sprintf("$\\Q(a,i)$");
                else
                    return Sprintf("$\\Q(\\sqrt{%o},i)$", D[2]);
                end if;
            else
                assert 7 in D;
                if D[1] eq 7 then 
                    return Sprintf("$\\Q(a,\\sqrt{%o})$", D[2]);
                else
                    return Sprintf("$\\Q(\\sqrt{%o},a)$", D[1]);
                end if;
            end if;
        end if;
    end if;
    if d eq 8 then
        G:=GaloisGroup(f);
        if IsAbelian(G) and Exponent(G) eq 2 then
            D := Sort([d: d in {SquareFree(Integers()!Abs(Discriminant(K[1]))): K in Subfields(SplittingField(f))|Degree(K[1]) eq 2}]);
            if D[1] eq 1 then
                if D[2] eq 7 then
                    return Sprintf("$\\Q(\\sqrt{%o},\\sqrt{%o},a,i)$", D[3]);
                else if D[3] eq 7 then
                    return Sprintf("$\\Q(\\sqrt{%o},\\sqrt{%o},a,i)$", D[2]);
                else
                    return Sprintf("$\\Q(\\sqrt{%o},\\sqrt{%o},i)$", D[2], D[3]);
                end if; end if;
            else
                Ds := Sort([d: d in {SquareFree(Integers()!Discriminant(K[1])): K in Subfields(SplittingField(f))|Degree(K[1]) eq 2}]);
                assert 7 in D;
                i:= 1; while D[i] mod 7 eq 0 do i+:=1; end while;
                j:= i+1; while D[j] mod 7 eq 0 do j+:=1; end while;
                if not D[i] in Ds then D[i] := -D[i]; end if;
                if not D[j] in Ds then D[j] := -D[j]; end if;
                return Sprintf("$\\Q(\\sqrt{%o},\\sqrt{%o},a)$", D[i], D[j]);
            end if;
        end if;
    end if;
    c := Coefficients(f);
    s := &cat [sprint_monomial(Integers()!c[d-i+1],d-i):i in [0..d]];
    if s[1] eq "\\" then s := Substring(s,6,#s-5); end if;
    return Sprintf("$\\Gal(%o)$",s);
end function;

function sprint_gens(gens)
    if #gens eq 1 and gens[1] eq Identity(Parent(gens[1])) then return "id"; end if;
    s := StripWhiteSpace(&cat Split(Sprintf("%o",gens),"*"));
    return  "$" cat Substring(s,2,#s-2) cat "$";
end function;

function sprint_h(h)
    if h eq Identity(Parent(h)) then return "id"; end if;
    s := StripWhiteSpace(&cat Split(Sprintf("$%o$",h),"*"));
    if s eq "$zz$" then s := "$z$"; end if;
    return s;
end function;

function sprint_id(id)
    return Sprintf("$\\langle %o,%o \\rangle$",id[1],id[2]);
end function;

function sprint_twist(twists,i,phi,exp,STdict)
    r := twists[i];
    H := IdentifyGroup(H_group(r,phi));  H0 := IdentifyGroup(H0_group(r,phi));
    st :=stGroups(r,phi,exp,STdict);
    s := Sprintf("%o & %o & %o & %o & %o & %o & %o & %o & %o\\\\\\hline %% %o\n", i, sprint_gens(r[2]), sprint_h(r[3]), sprint_id(H), sprint_id(H0), sprint_field(r[8]), sprint_field(r[9]), st[1], st[2], r[1]);
    return s cat sprint_curve(r);
end function;

