////////////////////////////////////////////////////////////
// Chabauty + Mordell-Weil Sieve                          //
//                                                        //
// Michael Stoll, 2007                                    //
////////////////////////////////////////////////////////////

// Main function:
//   Chabauty(ptJ, ptC)
// where 
// * ptC is a rational point on C, a genus 2 curve over Q
// * ptJ is a rational point on J, the Jacobian of C,
//   such that ptJ generates the Mordell-Weil group J(Q) modulo torsion.
// The function returns three values:
// 1. the set of rational points on C,
// 2. the set of primes p such that information mod p was used,
// 3. the sequence of primes q that were used to build B,
//    where B is a number such that
//    + C(Q) --> J(Q) --> J(Q)/B J(Q) is injective  and
//    + for the cosets in J(Q)/ B J(Q) not in the image of the map above,
//      the information from the set of primes p (second return value)
//      proves that they are not in the image.

// Example:
//
// P<x> := PolynomialRing(Rationals());
// C := HyperellipticCurve(P![1,8,22,22,5,6,1]);
// ptC := C![0,1];
// ptJ := C![0,1] - C![0,-1]; // known to generate J(Q)
// SetVerbose("Chabauty", 1); // to see something...
// Chabauty(ptJ, ptC);


/***********************************************************
 disclog.m
 
 General purpose discrete log for groups of smooth order
 
 M. Stoll, started 2005-06-07
 converted to package file 2006-03-05
 ***********************************************************/

// given: Abelian group G, bijective map G -> X, X some structure
// #G smooth (so that for groups of order p^f|#G, lookup is feasible)

// Exported intrinsics: 
// MakeDL(G, m)
//     m : G -> X is the map (Map), * and + defined on X

function MakeLookup1(G, m)
  return pmap<Codomain(m) -> G| [<m(g), g> : g in G]>;
end function;

function MakeDLp1(G, m, p)
  // G a p-group
  if #G le 25 then
    return MakeLookup1(G, m);
  end if;
  invs := Invariants(G);
  // printf "MakeDLp: Invariants(G) = %o\n", invs;
  pp := ExactQuotient(invs[#invs], p);
  if pp eq 1 then
    return MakeLookup1(G, m);
  end if;
  // printf "MakeDLp: pp = %o\n", pp;
  h := hom<G -> G | [pp*G.i : i in [1..#invs]]>;
  G1 := Image(h);
  // printf "MakeDLp: Invariants(Image(h)) = %o\n", Invariants(G1);
  m1 := map<G1 -> Codomain(m) | x:->m(x)>;
  f1 := MakeLookup1(G1, m1);
  G2 := Kernel(h);
  // printf "MakeDLp: Invariants(Kernel(h)) = %o\n", Invariants(G2);
  m2 := map<G2 -> Codomain(m) | x:->m(x)>;
  f2 := MakeDLp1(G2, m2, p);
  return pmap<Codomain(m) -> G |
               x :-> f2(x - m(a)) + a where a := f1(pp*x) @@ h>;
end function;

intrinsic MakeDL(G::GrpAb, m::Map) -> Map
{ Given bijection  m : G -> X, where X has group structure, compute the
  inverse of m. Assumes that #G is smooth.}
  n := #Invariants(G);
  f := Factorization(#G);
  cofs := [&*[Integers()|f[i,1]^f[i,2] : i in [1..#f] | i ne j] : j in [1..#f]];
  _, refs := XGCD(cofs);
  assert &+[Integers()|refs[i]*cofs[i] : i in [1..#f]] eq 1;
  DLs := [**];
  for i := 1 to #f do
    p := f[i,1];
    hp := hom<G -> G | [cofs[i]*G.j : j in [1..n]]>;
    Gp := Image(hp);
    mp := map<Gp -> Codomain(m) | x:->m(x)>;
    DLp := MakeDLp1(Gp, mp, p);
    Append(~DLs, DLp);
  end for;
  return pmap<Codomain(m) -> G
               | x :-> &+[G|refs[i]*G!(DLs[i](cofs[i]*x)) : i in [1..#f]]>;
end intrinsic;

/************************************************************
 * chabauty.m                                               *
 *                                                          *
 * Michael Stoll, started 2006-07-19                        *
 *                                                          *
 * Routines for Chabauty's method on genus 2 curves         *
 ************************************************************/

declare verbose Chabauty, 3;

// Find the reduction mod p of the differential over Q_p that kills
// the subgroup generated by a point on the Jacobian
intrinsic KillingDifferentialModp(pt::JacHypPt, p::RngIntElt) -> Pt
{ Given a point pt of infinite order on the Jacobian of a genus 2 curve
  over the rationals of the form y^2 = f(x) and a prime p of good reduction,
  this returns a point in P^1(F_p) which is the image of the two zeros
  of omega mod p, where omega is a Q_p-differential on the curve that kills
  the group generated by pt. }
  require Order(pt) eq 0: "Point pt must have infinite order.";
  J := Parent(pt);
  require Dimension(J) eq 2: "Point pt must be on a genus 2 Jacobian.";
  C := Curve(J);
  Jp := BaseChange(J, GF(p));
  n := Order(Jp!pt);
  prec := 1;
  K := KummerSurface(J);
  repeat
    prec +:= 2;
    K1 := BaseChange(K, pAdicField(p, prec));
    ptK := n*K1!K!pt;
  until Min([Valuation(ptK[i]) : i in [1..3]]) eq prec - 1;
  coords := [GF(p) | ExactQuotient(Integers()!ptK[i], p^(prec-1))
                   : i in [1..3]];
  assert coords[2]^2 eq 4*coords[1]*coords[3];
  Pr1 := ProjectiveSpace(GF(p), 1);
  return coords[1] eq 0 select Pr1![1,0] else Pr1![coords[2]/2,coords[1]];
end intrinsic;

function GroupInfo(J, bound, bas, ptC, GIlb, PrimeBound, SmoothBound)
// Construct a list of "interesting" primes and information on
// C(F_p) and J(F_p)
// An entry is [* p, imbas, imC, flag *],
// where p is the prime,
//       imbas is the image of the given generators of J(Q) in G ~ J(F_p),
//       imC is the image of C(F_p) in G,
//       flag is true iff points in C(F_p) are separated in the image group.
// GIlb is a lower bound on the primes.
  bp := Seqset(BadPrimes(J));
  res := [Parent([**]) | ];
  lb := Max(3, IsEven(GIlb) select GIlb+1 else GIlb);
  for p in Type(bound) eq RngIntElt select [lb..bound by 2] else bound do
    if IsPrime(p) and p notin bp then
      Jp := BaseChange(J, GF(p));
      Cp := BaseChange(Curve(J), Bang(Rationals(), GF(p)));
      pts := Points(Cp);
      oG := #Jp;
      if Max(PrimeDivisors(oG)) lt SmoothBound
      then
        G, m := AbelianGroup(Jp);
        vprintf Chabauty, 1: " GroupInfo: p = %o...\n", p;
        I := Invariants(G);
        vprintf Chabauty, 1:"   #C(F_p) = %o, Invariants(G) = %o\n", #pts, I;
        fI := Factorization(I[#I]);
        vprintf Chabauty, 2: "   Exponent = %o\n", fI;
        imbas := [(Jp!b) @@ m : b in bas];
        orders := [Order(g) : g in imbas];
        fL := Factorization(LCM(orders));
        vprintf Chabauty, 3: "   Exponent of image = %o\n", fL;
        inj := func<pt | Jp![Cp!pt, bpt]> where bpt := Cp!ptC;
        DL := MakeDL(G, m);
        imC := {DL(inj(pt)) : pt in pts};
        Gsub := sub<G | imbas>;
        imC := {Gsub | pt : pt in imC | pt in Gsub};
        imbas := ChangeUniverse(imbas, Gsub);
        e1 := &*[Integers()| a[1]^a[2] : a in fL | a[1] lt PrimeBound];
        if e1 lt LCM(orders) then
          // kick out primes larger than PrimeBound
          Gq, qmap := quo<Gsub | [e1*g : g in Generators(Gsub)]>;
          imbas := [qmap(b) : b in imbas];
          imCold := imC;
          imC := {qmap(c) : c in imC};
          if #imC eq #Gq then
            vprintf Chabauty, 1:
                    "   Relevant part does not give information ==> discard.\n";
          else
            Append(~res, [* p, imbas, imC, #imC eq #imCold *]);
          end if;
        else // e1 = LCM(orders)
          Append(~res, [* p, imbas, imC, true *]);
        end if; // e1 lt LCM(orders)
      end if; // Max(PrimeDivisors(oG)) lt SmoothBound
    end if; // IsPrime(p) and p notin bp
  end for; // p
  return res;
end function;

function CheckKillingDifferential(f, ptJ, entry)
// Select the entries in GI (entry[4] = true), such that
// the Q_p-differential killing the group generated by ptJ does not vanish
// on any point of the curve mod p
  if not entry[4] then return entry; end if;
  p := entry[1];
  cl := KillingDifferentialModp(ptJ, p);
  fp := PolynomialRing(GF(p))!f;
  if cl[2] eq 0 select not IsSquare(Coefficient(fp, 6))
                else not IsSquare(Evaluate(fp, cl[1]/cl[2]))
  then
    vprintf Chabauty, 2:
            "Found prime p = %o with killing differential not vanishing on C(F_p).\n", p;
    return entry;
  else
    return [* p, entry[2], entry[3], false *];
  end if;
end function;

function GetMinimalDivisors(set)
  res := {};
  while not IsEmpty(set) do
    a := Min(set);
    set := {s : s in set | not IsDivisibleBy(s, a)};
    Include(~res, a);
  end while;
  return res;
end function;

function FindMinimalBs(GI)
  function bestB(entry)
    G := Universe(entry[2]);
    for d in Divisors(Exponent(G)) do
      qG, qmap := quo<G | [d*g : g in OrderedGenerators(G)]>;
      if #{qmap(c) : c in entry[3]} eq #entry[3] then
        vprintf Chabauty, 2: "Best B for p = %o is %o\n", entry[1], d;
        return d;
      end if;
    end for;
    // never reached
  end function;
  return GetMinimalDivisors({bestB(entry) : entry in GI | entry[4]});
end function;

function dryRun1(GI, Bound, eps, tors, Bs)
  vprintf Chabauty, 1: "dryRun: Bound = %o, eps = %o\n", Bound, eps;
  function frac(B, entry)
    I := Invariants(Universe(entry[2]));
    Is := [Integers() | GCD(i,B) : i in I];
    n := &*Is;
    if Is eq I then return 1.0*#entry[3]/n; end if;
    cs := {};
    for cc in entry[3] do
      Include(~cs, [c[i] mod Is[i] : i in [1..#Is]] where c := Eltseq(cc));
      if #cs eq n then return 1.0; end if;
    end for;
    return 1.0*#cs/n;
  end function;
  red := func<B | B*tors*&*[frac(B, e) : e in GI]>;
  prl := [2] cat [p : p in [3..Bound by 2] | IsPrime(p)];
  max := &*[p^Max(&cat[[Valuation(i,p) : i in Invariants(Universe(e[2]))]
                         : e in GI])
             : p in prl];
  vprintf Chabauty, 2: "        max = %o\n", max;
  pl := [p : p in prl | IsDivisibleBy(max, p)];
  // optimze according to max red(B_i)
  // if B_0 = 1, B_{i+1} = B_i*p_i
  agendas := [[<1, 1, 1.0>]]; // entries <p, B, s>
  agendav :=[1.0]; 
  reached := {1}; // lists the values of B that have been seen so far
  while true do
    v, pos := Min(agendav); // try to extend cheapest path so far
    l := agendas[pos];
    vprintf Chabauty, 1: "  dryRun: l = %o,\n   v = %o\n",
                   [l[i,1] : i in [2..#l]], v;
    Remove(~agendas, pos);
    Remove(~agendav, pos);
    B := l[#l,2];
    s := l[#l,3];
    if B eq max or (s lt eps and exists{b : b in Bs | IsDivisibleBy(B,b)}) then 
      break;  // reached success criterion
    end if;
    // list possible extensions
    new := [<p, Bp, red(Bp)> : p in pl 
                             | Bp notin reached and IsDivisibleBy(max, Bp)
                               where Bp := B*p];
    // and update agenda
    agendas cat:= [Append(l, n) : n in new];
    agendav cat:= [n[3] : n in new];
    reached join:= {n[2] : n in new};
  end while;
  vprintf Chabauty, 1: "dryRun:\n Multipliers: %o\n",
                [l[i,1] : i in [2..#l]];
  vprintf Chabauty, 1: " Predicted sizes: %o\n", [Round(l[i,3]) : i in [2..#l]];
  return [<l[i,1], l[i,3]> : i in [2..#l]], s lt eps;
end function;

// Given a set of candidates <a, t> mod B, return all pairs mod B*p lifting
// them and compatible with GI. Return "primes" used in the process as second
// value.
function LiftInformation(S, GI, B, p)
  // S = {<a, t>, ...} , a : integer, t : element of abstract torsion subgroup
  // preparation: find relevant entries in GI
  vpB := Valuation(B, p);
  Bp := B*p;
  tests := [e : e in GI | Valuation(Exponent(Universe(e[2])), p) gt vpB];
  tests := [[* e[1], [q(b) : b in e[2]], {q(c) : c in e[3]} *]
              where _, q := quo< G | [Bp*g : g in Generators(G)] >
              where G := Universe(e[2])
             : e in tests];
  tests := [e : e in tests | #e[3] lt #Universe(e[3])];
  Sort(~tests, func<a,b | #a[3]/#Universe(a[3]) - #b[3]/#Universe(b[3])>);
  if IsVerbose("Chabauty", 2) then
    printf "LiftInformation: B = %o, p = %o, #S = %o\n", B, p, #S;
    if IsEmpty(tests) then
      printf "  no relevant entries.\n";
    else
      printf "  relevant entries:\n";
      for e in tests do
        printf "    %o : #G = %o, #imC = %o\n", e[1], #Universe(e[3]), #e[3];
      end for;
    end if;
  end if;
  res := {};
  primes := {};
  for tup in S do
    cands := {<a + i*B, t> : i in [0..p-1]}
               where a := tup[1] where t := tup[2];
    size := #cands;
    for e in tests do
      cands := {c : c in cands
                  | c[1]*e[2][1]
                     + &+[Universe(e[3])|ts[i]*e[2][i+1] : i in [1..#ts]]
                     in e[3]
                     where ts := Eltseq(c[2])};
      if #cands lt size then
        size := #cands; Include(~primes, e[1]);
      end if;
    end for;
    res join:= cands;
  end for;
  if IsVerbose("Chabauty", 2) then
    if not IsEmpty(tests) then
      printf "  entries used: %o\n", primes;
    end if;
    printf "  size of conditions mod %o is %o\n", Bp, #res;
  end if;
  return res, primes;
end function;

// Compute a heuristic value indicating whether we have collected
// enough information.
function CheckInvariant(GI, tors)
  rep := 10;
  function frac(B, entry)
    I := Invariants(Universe(entry[2]));
    Is := [Integers() | GCD(i,B) : i in I];
    n := &*Is;
    if Is eq I then return 1.0*#entry[3]/n; end if;
    cs := {};
    for cc in entry[3] do
      Include(~cs, [c[i] mod Is[i] : i in [1..#Is]] where c := Eltseq(cc));
      if #cs eq n then return 1.0; end if;
    end for;
    return 1.0*#cs/n;
  end function;
  //
  Gs := [Universe(e[2]) : e in GI];
  primes := &join[Seqset(PrimeDivisors(#G)) : G in Gs];
  // for each prime, find largest exponents in invariants of product
  Bs := [1 : j in [1..rep]];
  for p in primes do
    vals := &join[{* Valuation(i, p) : i in Invariants(G) *} : G in Gs];
    m := Max(vals);
    Exclude(~vals, m);
    for j := 1 to rep do
      Bs[j] *:= p^m;
      if m ne 0 then
        if IsEmpty(vals) then
          m := 0;
        else
          m := Max(vals);
          Exclude(~vals, m);
        end if;
      end if;
    end for;
  end for;
  // compute estimated number of candidates for B
  rs := [1.0*B*tors : B in Bs];
  for e in GI do 
    rs := [rs[j]*frac(Bs[j], e) : j in [1..rep]]; 
  end for;
  return Min(rs);
end function;

function CheckPoint(gen, l, t, ptC, ht)
// Checks if l*gen + t (t a torsion point) gives
// a point on J that is of the form pt - ptC.
// Returns a flag, and (if successful) the point.
// ht is a bound for the logarithmic naive height of the point on J
  C := Parent(ptC);
  prec := 0.0;
  p := 101;
  disc := Integers()!Discriminant(Curve(Parent(ptC)));
  ht1 := ht + 2*Log(Max([Abs(ptC[1]), Abs(ptC[3])])) + Log(3);
  J := Parent(gen);
  while prec lt ht1 do
    if not IsDivisibleBy(disc, p) then
      Jp := BaseChange(J, GF(p));
      ptCp := Curve(Jp)!ptC;
      ptJp := l*(Jp!gen) + Jp!t;
      if not ( ptJp eq Jp!0 or
             ( ptCp[3] eq 0
               select (Degree(ptJp[1]) lt 2 
                        and Coefficient(ptJp[2], 3) eq -ptCp[2])
               else (Evaluate(ptJp[1], x) eq 0
                      and Evaluate(ptJp[2], x) eq -ptCp[2]/ptCp[3]^3
                     where x := ptCp[1]/ptCp[3]) ) )
      then
        vprintf Chabauty, 3: "  CheckPoint: failed at p = %o\n", p;
        return false, _;
      end if;
    end if;
    prec +:= Log(p);
    p := NextPrime(p);
  end while;
  // Now there should be a point; find it!
  ptJ := l*gen + t;
  assert ptJ eq J! 0 or
         ( ptC[3] eq 0
               select (Degree(ptJ[1]) lt 2 
                        and Coefficient(ptJ[2], 3) eq -ptC[2])
               else (Evaluate(ptJ[1], x) eq 0
                      and Evaluate(ptJ[2], x) eq -ptC[2]/ptC[3]^3
                     where x := ptC[1]/ptC[3]) );
  if ptJ eq J!0 then return true, ptC; end if;
  if ptC[3] eq 0 then
    if Degree(ptJ[1]) lt 1 then
      // other point also at infinity
      return true, C![1, Coefficient(ptJ[2], 3), 0];
    else
      return true, C![x, Evaluate(ptJ[2], x), 1]
                   where x := -Coefficient(ptJ[1], 0);
    end if;
  else
    if Degree(ptJ[1]) lt 2 then
      // other point at infinity
      return true, C![1, Coefficient(ptJ[2], 3), 0];
    else
      return true, C![x, Evaluate(ptJ[2], x), 1]
                   where x := -Coefficient(ptJ[1], 1) - ptC[1]/ptC[3];
    end if;
  end if;
end function;

intrinsic Chabauty(ptJ::JacHypPt, ptC::PtHyp : 
                   PrimeBound := 30, SmoothBound := 50, Epsilon := 1e-2)
                    -> SetEnum, SetEnum, SeqEnum
{Use Chabauty's method, combined with a Mordell-Weil sieve in order to find
 all rational points on a genus 2 curve C with Jacobian of rank 1.
 ptJ is a generator of the Mordell-Weil group modulo torsion,
 ptC is a known rational point on the curve C.}
  J := Parent(ptJ);
  C := Curve(Parent(ptC));
  require Curve(J) eq C: "ptJ must be on the Jacobian of the curve ptC is on.";
  require Genus(C) eq 2: "The parent curve of ptC must have genus 2.";
  f, h := HyperellipticPolynomials(C);
  require h eq 0: "Curve must be of the form y^2 = f(x).";
  require Order(ptJ) eq 0: "ptJ must be a point of infinite order.";
  // find torsion subgroup and hence generators of J(Q)
  T, mT := TorsionSubgroup(J);
  vprintf Chabauty, 1: " Torsion subgroup has invariants %o\n", Invariants(T);
  bas := [ptJ] cat [mT(t) : t in OrderedGenerators(T)];
  ht := Height(ptJ);
  hc := HeightConstant(J);
  // compute GroupInfo for sufficiently many primes
  oldbound := 2;
  GI := [];
  b := 10;
  est := 1.0; 
  Bs := {};
  while true do
    vprintf Chabauty, 1: " Computing GroupInfo...\n";
    while est ge Epsilon or IsEmpty(Bs) do
      b *:= 2;
      vprintf Chabauty, 2: "   for bound = %o...\n", b;
      nGI := GroupInfo(J, b, bas, ptC, oldbound, PrimeBound, SmoothBound);
      nGI := [CheckKillingDifferential(f, ptJ, entry) : entry in nGI];
      Bs := GetMinimalDivisors(Bs join FindMinimalBs(nGI));
      vprintf Chabauty, 1: "  set of minimal Bs found so far:\n    %o\n", Bs;
      oldbound := b;
      GI cat:= nGI;
      est := CheckInvariant(GI, #T);
      vprintf Chabauty, 2: "   check value is %o.\n", est;
    end while;
    run, flag := dryRun1(GI, PrimeBound, 5*Epsilon, #T, Bs);
    if not flag then
      // expected size goal not reached: need more information
      vprintf Chabauty, 1:
              "No good run found  ==>  reduce Epsilon and start over.\n";
      Epsilon *:= 0.1;
      continue;
    end if;
    result := {C|};
    primes := {};
    B := 1;
    cands := {<0, t> : t in T};
    for j := 1 to #run do
      p := run[j,1];
      cands, pr := LiftInformation(cands, GI, B, p);
      primes join:= pr;
      B *:= p;
      vprintf Chabauty, 1: "    B = %o, size = %o (predicted: %o)\n",
                           B, #cands, Round(run[j,2]);
      if IsEmpty(cands) then
        vprintf Chabauty, 1: "  No more points!\n";
        vprintf Chabauty, 1: "Final list of points: %o\n", result;
        return result, primes, [run[i,1] : i in [1..j]];
      end if; 
      // otherwise check if rational points are separated
      if exists{b : b in Bs | IsDivisibleBy(B, b)} then
        // for each candidate, try to identify a rational point
        if IsVerbose("Chabauty", 3) then
          printf "   current list:\n   ";
          for c in cands do printf " %o", c; end for;
          printf "\n";
        end if;
        oldcands := cands;
        cands := {};
        while not IsEmpty(oldcands) do
          vprintf Chabauty, 2: "  %o candidates remaining.\n", #oldcands;
          c := Rep(oldcands);
          vprintf Chabauty, 2: "   Checking candidate %o...\n", c;
          l := 2*c[1] le B select c[1] else c[1]-B;
          flag, pt := CheckPoint(bas[1], l, mT(c[2]), ptC, l^2*ht+hc);
          if flag then
            vprintf Chabauty, 1: "  Point %o on C found!\n", pt;
            Include(~result, pt);
            // remove all suitable candidates
            dl := {d : d in Bs | IsDivisibleBy(B, d)};
            for d in dl do
              Td := sub<T | [d*t : t in OrderedGenerators(T)]>;
              oldcands := {oc : oc in oldcands
                              | not (IsDivisibleBy(oc[1]-c[1], d)
                                      and oc[2]-c[2] in Td)};
              cands :=    {oc : oc in cands
                              | not (IsDivisibleBy(oc[1]-c[1], d)
                                      and oc[2]-c[2] in Td)};
            end for;
          else
            // move to cands
            Exclude(~oldcands, c);
            Include(~cands, c);
          end if; // flag
        end while; // not IsEmpty(oldcands)
        if IsEmpty(cands) then
          vprintf Chabauty, 1: "  No more points!\n";
          vprintf Chabauty, 1: "Final list of points: %o\n", result;
          return result, primes, [run[i,1] : i in [1..j]];
        end if;
      end if; // exists{b : b in Bs | IsDivisibleBy(B, b)}
    end for; // j := 1 to #run
    // run used up, but some classes left undecided
    vprintf Chabauty, 1:
            "Some classes unresolved ==> reduce Epsilon and start over.\n";
    Epsilon *:= 0.1;
  end while;
  // never reached
end intrinsic;



