/*
Given a group G, the corresponding group data
structure is a list G, pG, phi, psi, pGels where:
- G is a finitely presented group
- pG is a permutation group
- phi is an isomorphism G --> pG
- psi is the inverse of phi
- pGels is an enumeration of the elements of pG

We store G both as a finitely presented
group and as a permutation group because
MAGMA offers different functionalities
for the two kinds of data structures.

Throughout the code, CM types are denoted
by S, and dual CM types by R. Stilde and 
Rtilde denote the inverse images in G of
these CM types.
*/



/*
Given a group data structure and a subgroup H of G
returns representatives for the quotients H\G, also
known as a transversal of H in G.
*/
function Reps(H, Gdata)
	G, pG, phi, psi, pGels := Explode(Gdata);

	pReps := [];
	pH := phi(H);
	for g in pGels do
		if &and[ not (h*g) in pReps : h in pH ] then
			Append(~pReps, g);
		end if;
	end for;
	return [psi(g) : g in pReps];
end function;

/*
Given a subgroup H of G and a CM type S for H\G,
represented as a subset of G, computes the subset
\tilde{S} of all elements of G that map to S in
the quotient H\G.
*/
function StildeFromS(H, S, Gdata)
	G, pG, phi, psi, pGels := Explode(Gdata);

	Stilde := [];
	for s in S do
	for ph in phi(H) do
		Append(~Stilde, psi(ph*phi(s)) );
	end for;
	end for;
	return Stilde;
end function;


/*
Given a subgroup H of G, a subset S of H\G
represented as a subset of G, and a complex
conjugation i in the center of G, checks that
S corresponds to a CM type for H\G.
*/
function CheckCMType(H, S, Gdata, i)
	G, pG, phi, psi, pGels := Explode(Gdata);

	DegreeCMField := #G/#H;
	Stilde := StildeFromS(H, S, Gdata);
	pStilde := { phi(s) : s in Stilde };
	if #(pStilde meet { phi(s*i) : s in Stilde} ) ne 0 then
		return false;
	end if;
	return #S * 2 eq DegreeCMField and #({ phi(s) : s in S } join { phi(i*s) : s in S }) eq DegreeCMField;
end function;


/*
Given a subgroup H of G and a CM type S for H\G,
represented as a subset of G, computes the dual
CM type. The function returns two values: the
reflex subgroup H' and the inverse image in G
of the reflex CM type.
*/
function DualType(H, S, Gdata)
	G, pG, phi, psi, pGels := Explode(Gdata);

	Stilde := StildeFromS(H, S, Gdata);
	pStilde := [ phi(s) : s in Stilde ];
	pHdash := [g : g in pG | &and[ s*g in pStilde : s in pStilde] ];
	Hdash := [psi(h) : h in pHdash];
	Rtilde := [s^(-1) : s in Stilde];
	return Hdash, Rtilde;
end function;

/*
Given a subgroup H of G and a CM type S for H\G,
represented as a subset of G, checks if S is a
self-dual CM type.
*/
function CheckSelfDualType(H, S, Gdata)
	G, pG, phi, psi, pGels := Explode(Gdata);

	Hdash, Rtilde := DualType(H, S, Gdata);
	Hddash, Rtildetilde := DualType(Hdash, Rtilde, Gdata);
	return { psi(h) : h in phi(H) } eq { h : h in Hddash };
end function;

/*
Given a subgroup H of G and a dual CM type \tilde{R},
represented as the inverse image in G of the subset
R of H' \ G, computes the matrix of N^* \phi^*
as described before Lemma 4.3 in the companion paper.
*/
function MatrixNStarPhiStar(H, Rtilde, Gdata)
	G, pG, phi, psi, pGels := Explode(Gdata);

	iRows := Reps(H, Gdata);
	M := ZeroMatrix(Integers(), #G, #iRows);
	for i in [1..#iRows] do
		g := iRows[i];
		for r in Rtilde do
			j := Min( [ k : k in [1..#pGels] | pGels[k] eq phi(r)*phi(g)  ] );
			M[j][i] := 1;
		end for;
	end for;
	return M;
end function;


/*
Given a subgroup H of G and a CM type S for H\G,
represented as a subset of G, computes the matrix
of N^*\phi^*, where \phi^* is the reflex norm.
This is simply a wrapper for the previous function.
*/
function MatrixFromCMType(H, S, Gdata)
	Hdash, Rtilde := DualType(H, S, Gdata);
	return MatrixNStarPhiStar(H, Rtilde, Gdata);
end function;

/*
Given two CM types S1, S2 for the H\G, returns true
if they differ by composition by an automorphism of
the field fixed by H
*/
function AreEssentiallyEquivalent(H, S1, S2, Gdata)
	G, pG, phi, psi, pGels := Explode(Gdata);
	FieldAutos := [ psi(g) : g in Normaliser(phi(G),phi(H)) ] ;
	return &or [ { phi(s * auto) : s in StildeFromS(H, S1, Gdata) } eq { phi(s) : s in StildeFromS(H, S2, Gdata) } : auto in FieldAutos ] ;
end function;

/*
Given a subgroup H of G and a complex conjugation
iota, the next function lists all the simple CM
types for H\G up to automorphisms of the fixed
field of H that are induced by elements of G.
The return value is a list of subsets of G.
*/
function ListSimpleCMTypes(H, Gdata, iota)
	G, pG, phi, psi, pGels := Explode(Gdata);

	CMTypes := [];
	Representatives := Transversal(G,H);

	FieldAutos := [ psi(g) : g in Normaliser(phi(G),phi(H)) ] ;


	for Sidx in Subsets({1..#Representatives},Integers()!(#Representatives/2)) do	// we loop over all subsets of H\G of the correct size
		S := [ Representatives[j] : j in Sidx ] ;
		if CheckCMType(H, S, Gdata, iota) then			// Check that the subset represents a CM type
			if CheckSelfDualType(H, S, Gdata) then		// Check for simplicity
				if &and[ not AreEssentiallyEquivalent(H, S, cmt, Gdata) : cmt in CMTypes ] then	// Avoid duplicates, up to essential equivalence
					Append(~CMTypes, S);
				end if;
			end if;
		end if;
	end for;
	return CMTypes;
end function;

/*
Given a subgroup H of G and two simple 
CM types S_1, S_2 of H\G, represented 
as subsets of G, computes the dimensions
of the Mumford-Tate groups of two simple
abelian varieties A_1, A_2 with these CM 
types, as well as the dimension of the
Mumford-Tate group of A_1 x A_2 and the
order of the kernels of the maps
MT(A_1 x A_2) --> MT(A_i)
when these are isogenies.
*/
function ComputeData(H1, H2, S1, S2, Gdata : Verbose := false)
	G, pG, phi, psi, pGels := Explode(Gdata);
	dimension := Integers()!(#G/(2*#H1));
	if Verbose then
		"Dimension of abelian varieties", dimension;
	end if;

	M1 := MatrixFromCMType(H1, S1, Gdata);
	M2 := MatrixFromCMType(H2, S2, Gdata);
	M := HorizontalJoin(M1, M2);
	K1 := Kernel(Transpose(M1));
	K2 := Kernel(Transpose(M2));
	KTot := Kernel(Transpose(M));
	d1 := #Rows(Transpose(M1)) - Dimension(K1);
	d2 := #Rows(Transpose(M2)) - Dimension(K2);
	d := #Rows(Transpose(M)) - Dimension(KTot);
	if Verbose then
		"Dimension MT(A1):", d1;
		"Dimension MT(A2):", d2;
		"Dimension MT(A1 x A2):", d;
	end if;

	if d gt Min(d1,d2) then
		if Verbose then
			"One of the projections is not an isogeny";
		end if;
		return false, d1, d2, d;
	else
		M := ZeroMatrix(Integers(), Dimension(KTot) + Integers()!(Degree(KTot)/2), Degree(KTot) );
		
		for i in [1..#Basis(KTot)] do
		for j in [1..Degree(KTot)] do
			M[i][j] := Basis(KTot)[i][j];
		end for;
		end for;

		M2 := M;
		for k in [1..Integers()!(Degree(KTot)/2)] do
			M[k + Dimension(KTot)][k] := 1;
		end for;
		for k in [1..Integers()!(Degree(KTot)/2)] do
			M2[k + Dimension(KTot)][k + 2*dimension] := 1;
		end for;

		V := RSpace(Integers(), #Rows(Transpose(M)));
		W := sub< V | [ [M[i][j] : j in [1..#Rows(Transpose(M))]] : i in [1..#Rows(M)] ] >;
		W2 := sub< V | [ [M2[i][j] : j in [1..#Rows(Transpose(M2))]] : i in [1..#Rows(M2)] ] >;
	
		if Verbose then
			"Order of kernel of first projection", #(V/W);
			"Order of kernel of second projection", #(V/W2);
		end if;
		return true, d1, d2, d, #(V/W), #(V/W2);
	end if;
end function;


