/****************************************************
 *                                                  *
 * M. Kirschmer and C. Leedham-Green,               *
 * Computing with subgroups of the modular group    *
 *                                                  *
 * Complimentary Magma code, M. Kirschmer, May 2012 *
 *                                                  *
 ****************************************************/

declare attributes GrpMat : PSLData;

RF:= recformat< 
  Index ,
   Comm : SeqEnum,
      W : SeqEnum,
     IJ : Assoc,
      T : SetEnum
>;

G:= SL(2, Integers());
S:= G ! [0,-1,1,0];
T:= G ! [1,1,0,1];
U:= S*T;

// A continued fraction algorithm that approximates from above.
// i.e., x = a_0 - 1/(a_1 - 1/a2 ...)
MyCF:= function(x)
//  if x eq 0 then return []; end if;
  res:= [];
  while true do
    t:= Ceiling(x);
    Append(~res, t);
    if x eq t then return res; end if;
    x:= 1/(t-x);
  end while;
end function;

intrinsic WordST(g :: Mtrx[RngInt]) -> GrpFPElt
{}
  require Ncols(g) eq 2 and Nrows(g) eq 2 and Determinant(g) eq 1: "Not an element of SL(2,Z).";
  F<s,t>:= FreeGroup(2);
  if g[2,1] eq 0 then
    w:= s^(g[1,1] le 0 select 2 else 0) * t^(Sign(g[1,1]) * g[1,2]);
  elif g[1,1] eq 0 then
    w:= (g[2,1] le 0 select s^-1 else s) * t^(Sign(g[2,1]) * g[2,2]);
  else
    s1:= Sign(g[1,1]); s2:= Sign(g[2,1]);
    if s1 ne s2 then
      gg:= S^s1*g;
      w:= s^s2;
    elif s1 eq -1 and s2 eq -1 then
      gg:= -g;
      w:= s^2;
    else
      gg:= g;
      w:= F ! 1;
    end if;
    C:= MyCF( gg[1,1] / gg[2,1] );
    for c in C do w:= w * t^c * s; end for;
    w :=  w * t^(1+ gg[1,2] div gg[1,1]);
  end if;
  // testing
  f:= hom< F -> SL(2,Integers()) | S, T >;
  assert f(w) eq g;
  return w;
end intrinsic;

intrinsic ReduceSU(w :: GrpFPElt) -> GrpFPElt
{}
  ww:= w; g:= SUToMtrx(w);				// debug
  e:= Eltseq(w);
  n:= #e;
  i:= 1;
  while i lt n do
    t:= Abs(e[i]);
    if t eq Abs(e[i+1]) then
      if t eq 1 or e[i] ne e[i+1] then
        Remove(~e, i);
        Remove(~e, i);
        i := Max(i-1, 1);
        n := n-2;
      else
	Remove(~e, i);
        e[i] := -e[i];
        n:= n-1;
      end if;
    else
      if e[i] eq -1 then e[i]:= 1; end if;
      i +:= 1;
    end if;
  end while;
  if n ge 1 and e[n] eq -1 then e[n]:= 1; end if;
  w:= Parent(w) ! e;
  assert SUToMtrx(w) in {g, -g};	// debug
  return w;
end intrinsic;

intrinsic SUToMtrx(g::GrpFPElt) -> Mtrx
{}
  ff:= hom< Parent(g) -> SL(2,Integers()) | S, U >;
  return ff(g);
end intrinsic;

intrinsic WordSU(g :: Mtrx[RngInt], G::GrpFP) -> GrpFPElt
{}
  h:= WordST(g);
  f:= hom< Parent(h) -> G | G.1, G.1^-1*G.2 >;
  hh:= ReduceSU(f(h));

  // testing
  assert SUToMtrx(hh) in {g, -g};

  return hh;
end intrinsic;

intrinsic WordSU(g :: Mtrx[RngInt]) -> GrpFPElt
{}
  return WordSU(g, FreeGroup(2));
end intrinsic;



// given g in PSL(2,Z)'. Write g in the two free generators
// F.1 = SUSU^-1 and F.2 = SU^-1SU
// g is assumed to be a reduced word in S and U.
intrinsic CommutatorWord( g::GrpFPElt, F::GrpFP ) -> GrpFPElt
{}
  w:= F ! 1;
  e:= Eltseq(g);
  n:= #e;
  i:= 1;
  while i lt n do
    case e[i]:
      when  1:		// S
        if e[i+1] eq 2 then	// SUS
	  w *:= F.1;
	  if e[i+3] eq -2 then i +:= 4; else i +:= 3; e[i]:= -2; end if;
        else			//SU^-1S
	  w *:= F.2;
  	  if e[i+3] eq  2 then i +:= 4; else i +:= 3; e[i]:=  2; end if;
        end if;
      when  2:		// US
        w *:= F.1^-1;
	if e[i+2] eq -2 then i +:= 4; else i +:= 1; e[i+1]:= -2; end if;
      when -2:		// U^-1S
        w *:= F.2^-1;
	if e[i+2] eq  2 then i +:= 4; else i +:= 1; e[i+1]:=  2; end if;
    end case;
  end while;
  assert i eq n+1;
  return w;
end intrinsic;

// Returns the pair <i,j> in [0,1] x [-1,0,1] such that 3*i+2*j=idx mod 6
function IndexToPair(idx)
  case idx: 
    when 0: return 0, 0;
    when 1: return 1,-1;
    when 2: return 0, 1;
    when 3: return 1, 0;
    when 4: return 0,-1;
    when 5: return 1, 1;
    else error "made a boo-boo";
  end case;
end function;

procedure Preprocessing(H)
  if assigned H`PSLData then return; end if;

  m:= Ngens(H);
  FH:= SLPGroup(m);
  F2:= FreeGroup(2);
  Gens:= [ F2 | WordSU(H.i, F2): i in [1..m] ];
  R:= Integers(6);
  act:= [ R | 3*Weight(x, F2.1) + 2*Weight(x, F2.2) : x in Gens ];

  // Do the preprocessing.
  touched:= {};
  IJ:= AssociativeArray(R);
  Stab:= [ FH | 1 ];
  StabSU:= [ F2 | 1 ];

  while #touched ne 6 do
    // Get the smallest entry of R not reached so far.
    ok:= exists(oo){ R ! i: i in [0..5] | (R ! i) notin touched }; assert ok;

    // The standard orbit enumeration.
    first:= oo eq 0;
    Orbit:= {@ oo @};
    Path:= [ FH | 1 ];
    PathSU:= [ F2 | 1 ];
    i:= 1;
    while i le #Orbit do
      o:= Orbit[i];
      for j in [1..m] do
        x:= o + act[j];
        idx:= Index(Orbit, x);
	if idx eq 0 then
	  Include(~Orbit, x);
	  Append(~Path, Path[i] * FH.j );
	  Append(~PathSU, ReduceSU(PathSU[i] * Gens[j]));
	elif first then
	  // Lets see if we already have this word.
	  p:= ReduceSU(PathSU[i] * Gens[j] * PathSU[idx]^-1);
	  if p ne One(F2) and p notin StabSU and ReduceSU(p^-1) notin StabSU then
	    Append(~StabSU, p);
	    Append(~Stab, Path[i] * FH.j * Path[idx]^-1);
	  end if;
	end if;
      end for;
      i +:= 1;
    end while;

    // update the IJ list.
    for i in [1..#Orbit] do
      IJ[ Orbit[i] ] := < oo,  Path[i], PathSU[i] >;
      Include(~touched, Orbit[i]);
    end for;
  end while;

  // Now write StabSU as words in x, y:
  C:= [ F2 | CommutatorWord(c, F2) : c in StabSU ];
  Nielsen(~C, ~Stab);
  assert #Stab eq #C;
  H`PSLData:= rec< RF | Comm:= Stab, W:= C, IJ:= IJ >;
end procedure;


mymin := func< x | x lt xinv select x else xinv where xinv:= x^-1 >;
clearrow:= procedure(~M, i)
  for ii in [1..Ncols(M)] do M[ii,i]:= 0; M[i,ii]:= 0; end for; M[i,i]:= 1;
end procedure;


// Perform Nielsen reduction on N and do the same operations to C.
intrinsic Nielsen(~W::[], ~C::[])
{}
  F:= Universe(W); One:= F ! 1;
//  W:= Setseq({ F | mymin(w) : w in W | w ne One });
  for i:= #W to 1 by -1 do
    w:= W[i];
    if w eq One then Remove(~W, i); Remove(~C, i);
    else
      winv:= w^-1;
      if winv lt w then C[i] := C[i]^-1; W[i]:= winv; end if;
    end if;
  end for;

  n:= #W;
  M:= MatrixRing(GF(2), n) ! 1;

  repeat
  i:= 1;
  while i le n do
   // j:= 1;
    j:= i+1;
    while j le n do
      if M[i,j] eq 0 then
        m:= #W[i] ge #W[j] select i else j;
        for k,l in [-1, 1] do
          x:= W[i]^k*W[j]^l;
          if #x eq 0 then
            Remove(~W, m); Remove(~C, m); RemoveRow(~M, m); RemoveColumn(~M, m); n -:= 1;
            i:= 0; j:= n; break k;
          elif #x lt #W[m] then
	    clearrow(~M, m);
            xinv:= x^-1;
	    if x lt xinv then
	      W[m]:= x; C[m] := C[i]^k*C[j]^l;
	    else
              W[m]:= xinv; C[m]:= C[j]^-l * C[i]^-k;
	    end if;
            i:= 0; j:= n; break k; 
          end if;
        end for;
        if i ne 0 then M[i,j]:= 1; end if;
      end if;
      j +:= 1;
    end while;
    i +:= 1;
  end while;

  done:= true;
  I:= [ i: i in [1..#W] | IsEven(#W[i]) ];
  W2:= W[I]; C2:= C[I];
  ParallelSort(~W2, ~C2);
  for i in [1..#W2] do
    nn:= #W2[i];
    nn2:= nn div 2;
    qinv:= F ! Eltseq(W2[i])[nn2+1..nn];
    q:= qinv^-1;
    for j in [1..#W] do
      m:= #W[j];
      if m le nn2 or W[j] eq W2[i] then continue; end if;
      if m-nn2 eq #(qinv * W[j]) then
	//W[j]:= mymin(W2[i] * W[j]);
        x:= W2[i] * W[j]; xinv:= x^-1;
	if x lt xinv then
	  W[j]:= x; C[j]:= C2[i] * C[j];
	else
	  W[j]:= xinv; C[j]:= C[j]^-1 * C2[i]^-1;
	end if;
        clearrow(~M, j);
        done:= false;
	break i;
      elif m-nn2 eq #(W[j] * q) then
//        W[j]:= mymin(W[j] * W2[i]^-1);
        x:= W[j] * W2[i]^-1; xinv:= x^-1;
	if x le xinv then
	  W[j]:= x; C[j]:= C[j] * C2[i]^-1;
	else
	  W[j]:= xinv; C[j]:= C2[i] * C[j]^-1;
	end if;
	clearrow(~M, j);
        done:= false;
	break i;
      end if;
    end for;
  end for;
  until done;
end intrinsic;

/*
function IsNReduced(W)
  F:= Universe(W);
  W:= Set(W);
  W join:= {w^-1 : w in W };
  if exists(t){<u,v> : u,v in W | l ne 0 and l lt #u where l:= #(u*v) } then return false, t; end if;

  if exists(t){<u,v,w> : u,v,w in W | #(u*v) * #(v*w) ne 0 and #(u*v*w) le #u - #v + #w} then return false, t; end if;

  return true;
end function;
*/

intrinsic IndexInPSL2Z(H::GrpMat[RngInt]) -> RngIntElt
{}
  require Degree(H) eq 2 and forall{i: i in [1..Ngens(H)] | Determinant(H.i) eq 1} : "Not a subgroup of PSL(2,Z)";
  Preprocessing(H);

  if not assigned H`PSLData`Index then
    W:= H`PSLData`W;
    IJ:= H`PSLData`IJ;
    k:= #[ k: k in Keys(IJ) | IJ[k][1] eq 0 ];

    // compute coset table of some G that contains H as a free factor.
    F:= Universe(W);
    T:= {F | 1};
    for w in W do
      e:= Eltseq(w);
      n:= #e div 2;
      T join:= {F | e[1..i] : i in [1..n]};
      e:= Eltseq(w^-1);
      if IsEven(#e) then n:= n-1; end if;
      T join:= {F | e[1..i] : i in [1..n]};
    end for;

    H`PSLData`Index:= #W-1 eq #T select (6 div k) * #T else Infinity();
    H`PSLData`T:= T;
  end if;

  return H`PSLData`Index;
end intrinsic;


HtoMatrix:= func< H, h | f(h) where f:= hom< Parent(h) -> G | [H.i : i in [1..Ngens(H)]]>  >;
CtoMatrix:= func< h | f(h) where f:= hom< Parent(h) -> G |  S*U*S*U^-1,  S*U^-1*S*U>  >;


intrinsic Membership(H::GrpMat[RngInt], x::GrpMatElt[RngInt] : CosetRep:= false, Word:= false) -> BoolElt, ., GrpSLPElt
{}
  require Degree(H) eq 2 and forall{i: i in [1..Ngens(H)] | Determinant(H.i) eq 1} : "Not a subgroup of PSL(2,Z)";
  require Ncols(x) eq 2 and Determinant(x) eq 1 : "The matrix does not lie in PSL(2,Z).";

  Preprocessing(H);

  W:= H`PSLData`W;
  IJ:= H`PSLData`IJ;

  F:= Universe(W);
  w:= WordSU(x, F);
  R:= Integers(6);
  o:= R ! (3*Weight(w, F.1) + 2*Weight(w, F.2));
  i, h, hSU:= Explode(IJ[o]);
  if i ne 0 then
    if not CosetRep then return false, _, _; end if;
    s,u:= IndexToPair(i);
    _, rep:= Membership(H, x * U^-u * S^-s : CosetRep );
    return false, < rep[1], i >, _;
  end if;

  w:= ReduceSU(hSU^-1 * w);
  w:= CommutatorWord(w, F);

  // This should be done in a tree but I don't have one...
  C:= H`PSLData`Comm;
  word:= h;
  n:= #w;
  cont:= true;
  while n ge 0 and cont do
    cont:= false;
    for i in [1..#W] do
      xx:= W[i]^-1 * w;
      if #xx lt n then if Word then word:= word * C[i];    end if; w:= xx; cont:= true; break; end if;
      xx:= W[i] * w;
      if #xx le n then if Word then word:= word * C[i]^-1; end if; w:= xx; cont:= true; break; end if;
    end for;
    n:= #w;
  end while;

  // If H has finite index, then w must be in the list of coset representatives.
//  assert Type(Hidx) eq Infty or w in L;

  if Word and n eq 0 then
    assert HtoMatrix(H, word) in {x, -x};
    return true, < w, i >, word;
  end if;
  return n eq 0, < w, i >, _;
end intrinsic;

intrinsic CosetToMatrix(Coset::.) -> Mtrx
{}
  i,j:= IndexToPair(Coset[2]);
  return CtoMatrix(Coset[1]) * S^i*U^j;
end intrinsic;

intrinsic MyCosetReps(H::GrpMat[RngInt]) -> []
{}
  idx:= IndexInPSL2Z(H);
  require Type(idx) ne Infty : "The group must have finite index in PSL(2, Z).";
  W:= H`PSLData`W;
  L:= Setseq(H`PSLData`T);
  LMat:= [ CtoMatrix(l) : l in L ];
  IJ:= H`PSLData`IJ;
  Orbs:= [ k: k in Keys(IJ) | IJ[k][1] eq k ];
  OrbMat:= [];
  for o in Orbs do
    i,j:= IndexToPair(o);
    Append(~OrbMat, S^i*U^j);  
  end for;
  assert idx eq #L * #Orbs;

  return [ l * o: l in LMat, o in OrbMat ], [ <l, o> : l in L, o in Orbs ] ;
end intrinsic;

intrinsic MyIsCongruence(H::GrpMat[RngInt]) -> BoolElt
{}
  idx:= IndexInPSL2Z(H);
  if Type(idx) eq Infty then return false; end if;
//  if idx le 6 then return true, idx eq 1 select 1 else 2; end if;
  C, Sy:= MyCosetReps(H);
  t:= []; for c in C do _, s:= Membership(H, c*T : CosetRep, Word:= false); Append(~t, Index(Sy, s)); end for;
  R:= (T*S)^-1 * S;
  r:= []; for c in C do _, s:= Membership(H, c*R : CosetRep, Word:= false); Append(~r, Index(Sy, s)); end for;
  Sn:= Sym( #C );
  t:= Sn ! t; r:= Sn ! r;

  // Now the fun starts.
  n:= Order(t);
  if IsOdd(n) then return IsId((r^2 * t^-InverseMod(2, n))^3); end if;
  v:= Valuation(n, 2); m:= n div 2^v;
  if m eq 1 then
    s:= t^20 * r^InverseMod(5, n) * t^-4 *r^-1;
    x:= t*r^-1*t;
    return s*x*s eq x and r*s eq r^25 * s and IsId( (s*r^5*t*r^-1*t)^3 );
  end if;
  e:= 2^v; c:= CRT([0,1], [e, m]); d:= CRT([1,0], [e,m]);
  a:= t^c; b:= r^c; l:= t^d; rr:= r^d;
  s:= l^20 * rr^InverseMod(5,e) * l^-4 * rr^-1;
  x:= (a*b^-1*a)^2; y:= l*rr^-1*l;
  return IsId((a,rr)) and IsId(x^2) and x eq (b^-1*a)^3 and x eq (b^2 * a^-InverseMod(2, m))^3
         and s*y*s eq y and rr*s eq s*rr^25 and y^2 eq (s*rr^5*l*rr^-1*l)^3;
end intrinsic;


intrinsic EmbedInSL2Z( G::GrpMat ) -> BoolElt, GrpMatElt
{}
  require Degree(G) eq 2 : "Degree of group must be 2";
  R:= BaseRing(G);
  require IsField(R) and Characteristic(R) eq 0 : "Base ring must be a field of characteristic 0";

  if exists{ g : g in Generators(G) | not IsCoercible(Integers(), Trace(g)) } or
     exists{ g : g in Generators(G) | Determinant(g) ne 1 } then return false, _; end if;

  error if forall(t){ <i,j> : i in [1..j-1], j in [1..Ngens(G)] | G.i*G.j eq G.j*G.i }, "Not implemented - group is abelian.";
 
  Gen4:= [G ! 1, G.t[1], G.t[2], G.t[1] * G.t[2] ];
  V:= Matrix([ Eltseq(x): x in Gen4 ]);
  W:= Matrix([ Eltseq(x): x in Generators(G) ]);

  ok, X:= IsConsistent(V, W);
  ok:= ok and CanChangeRing(X, Rationals());
  if not ok then return false, _; end if;

  W:= Matrix([ Eltseq(Gen4[2]*x): x in Gen4 ]) ;
  ok, X:= IsConsistent(V, W); assert ok;
  X2:= ChangeRing(X, Rationals());
  W:= Matrix([ Eltseq(Gen4[3]*x): x in Gen4 ]) ;
  ok, X:= IsConsistent(V, W); assert ok;
  X3:= ChangeRing(X, Rationals());

  A:= sub< MatrixAlgebra(Rationals(), 4) | X2, X3 >;
  ok, Q, phi:= IsQuaternionAlgebra(A);
  assert ok;
  ok, M, tau:= IsMatrixRing(Q:Isomorphism);
  assert ok;

  B:= Basis(M);
  X:= Matrix([ &cat [ Eltseq(Matrix(G.t[i]) * b - b * (A.i @ phi @ tau)) : i in [1..2] ] : b in B ]);
  KX:= Kernel(X);
  assert Dimension(KX) ne 0;
  T:= Matrix(2, Eltseq(KX.1));
  assert IsInvertible(T);
  T:= GL(2, R) ! T;
  H:= ChangeRing(G^T, Rationals());

  try
    _, TT:= IntegralGroup(H);
  catch e
    return false, _;
  end try;
  return true, T * (GL(2, R) ! TT);
end intrinsic;

/* example
K:= QuadraticField(-1);
r:= Matrix(K , 2, [K.1, 2+K.1*3, 1, K.1-7 ]);
G:= ChangeRing(SL(2, Integers()), K)^GL(2,K) ! r;
ok, T:= EmbedInSL2Z(G);
G^T;
*/

test:= function(n)
  G:= SL(2, Integers());
  for i in [1..n] do
    for k in [1..5] do
    <i,k>;
      case k:
        when 1: C:= CongruenceSubgroup(i);
	when 2: C:= Gamma0(i);
	when 3: C:= Gamma1(i);
	when 4: C:= GammaUpper0(i);
	when 5: C:= GammaUpper1(i);
      end case;
      H:= sub< G | [ Matrix(g) : g in Generators(C) ] >;
      assert Index(C) eq IndexInPSL2Z(H);
    end for;
  end for;
  return true;
end function;

function testmem(H)
  F:= FreeGroup(Ngens(H));
  for i in [1..100] do
    x:= Random(F, 2, 10);
    m:= HtoMatrix(H, x);
    ok, t, w:= Membership(H, m : Word);
    assert ok;

    a:= Random(-1000, 1000);
    b:= Random(-1000, 1000);
    g, r, s:= XGCD(a,b);

    m:= G ! [a div g, b div g, -s, r];
    ok, t := Membership(H, m : CosetRep, Word);
    if not ok then
      ok, t, w:= Membership(H, m*CosetToMatrix(t)^-1 : Word);
      assert ok;
    end if;
  end for;

  return true;
end function;



/*
function getex(i)
  F<a,b>:= FreeGroup(2);
  assert i ge 2;
  W:= [a*b, a^2*b^2];
  for j in [2..i-1] do
    x:= IsOdd(j) select b else a^-1;
    Append(~W, a*W[j]*x);
  end for;
  return W;
end function;
 */
