#!/usr/bin/env python3
"""Sweep for Grothendieck-Katz counterexample candidates, and document the misses.

A candidate is a connection over Q whose p-curvature vanishes for every prime we
test. Almost all of those are connections with algebraic solutions, i.e. not
counterexamples at all -- so the job is to enumerate the vanishing set and then
account for every member of it. What has value here is the accounting, not a
hoped-for hit.

Three families, in increasing order of how much they tell us:

  A. y'' = ((a+b-1)/t) y' - (ab/t^2) y, solutions t^a, t^b, a b rational.
     Solutions are algebraic by construction, so the whole family must vanish.
     If any member fails, my instrument is broken and nothing below means
     anything. This is a control, not a search.

  B. Gauss hypergeometric 2F1(a,b;c). Three singular points, rigid. The
     algebraic cases are classical (Schwarz). I report the vanishing set as
     data; I deliberately do NOT hardcode the Schwarz list from memory and
     claim agreement with it, because a remembered list is exactly the kind of
     thing I would get subtly wrong and then "confirm".

  C. Order-2 Fuchsian with four finite singular points 0, 1, lam plus infinity,
     with two accessory parameters. Not rigid, so the local data no longer
     determines the connection, and this is the cheapest place I know to be
     outside the settled cases. This is the actual search.

usage: sweep.py [family] ...
"""
import sys
from itertools import product

from pcurv import (vanishes, power_solutions, fuchsian_4pt, mstr,
                   companion_from_second_order, pscal, padd, pfromroots)

PRIMES = [7, 11, 13, 17, 19, 23, 29, 31]


def reduced(pairs):
    """Distinct rationals in lowest terms. Written after the control family
    reported four failures that were really -1/2 and -2/4 being the same number:
    an unreduced duplicate is excluded at every prime, tests zero primes, and
    then looks like a failure. A test suite that can be fooled by 2/4 is not
    yet a test suite."""
    from math import gcd
    out = []
    for (n, d) in pairs:
        if d == 0:
            continue
        g = gcd(abs(n), d) or 1
        r = (n // g, d // g)
        if r not in out:
            out.append(r)
    return out


def rat(num, den, p):
    """num/den in F_p, or None if the reduction is not defined at p."""
    if den % p == 0:
        return None
    return (num % p) * pow(den % p, p - 2, p) % p


def survives(build, params, primes=PRIMES):
    """True if p-curvature vanishes at every prime where the params reduce.

    Cheapest prime first: nearly everything dies at p=7 and never costs more.
    """
    tested = 0
    for p in primes:
        got = build(params, p)
        if got is None:      # bad reduction at this p; skip, don't count
            continue
        B, d = got
        if not vanishes(B, d, p):
            return False, tested
        tested += 1
    return tested > 0, tested


# ----------------------------------------------------------------- family A
def build_A(params, p):
    (an, ad), (bn, bd) = params
    a, b = rat(an, ad, p), rat(bn, bd, p)
    if a is None or b is None:
        return None
    # Bad reduction: if a == b in F_p the two solutions t^a, t^b collide, the
    # solution space drops to dimension 1, and vanishing is not expected. This
    # is the "almost all p" of the conjecture statement, and it is not
    # decoration -- 88 of 420 control operators tripped it before I excluded it.
    if a == b:
        return None
    return power_solutions(a, b, p), [0, 0, 1]


def family_A():
    print("== family A (control): solutions t^a, t^b, must ALL vanish")
    rats = reduced([(n, d) for d in (1, 2, 3, 4) for n in range(-3, 4)])
    bad = []
    total = 0
    for a, b in product(rats, rats):
        if a == b:
            continue
        total += 1
        ok, _ = survives((lambda pr, p: build_A(pr, p)), (a, b))
        if not ok:
            bad.append((a, b))
    print(f"   {total} operators, failures: {len(bad)}"
          + (f"  FIRST {bad[:5]}" if bad else "  <- control holds"))
    return not bad


# ----------------------------------------------------------------- family B
def build_B(params, p):
    """t(1-t) y'' + (c - (a+b+1)t) y' - a b y = 0, as Y' = (B/d) Y.

    y'' = [((a+b+1)t - c)/(t(1-t))] y' + [ab/(t(1-t))] y
    Common denominator d = t(1-t) = -t(t-1).
    """
    (an, ad), (bn, bd), (cn, cd) = params
    a, b, c = rat(an, ad, p), rat(bn, bd, p), rat(cn, cd, p)
    if a is None or b is None or c is None:
        return None
    d = pscal(pfromroots([0, 1], p), -1, p)              # t - t^2
    p1 = padd([(-c) % p], [0, (a + b + 1) % p], p)       # (a+b+1)t - c
    q = [(a * b) % p]
    return companion_from_second_order(q, p1, d, p), d


def family_B():
    print("== family B: Gauss hypergeometric 2F1(a,b;c), 3 singular points")
    rats = reduced([(n, d) for d in (1, 2, 3, 4, 6) for n in range(-3, 5)])
    hits = []
    total = 0
    for a, b, c in product(rats, rats, rats):
        total += 1
        ok, ntested = survives(build_B, (a, b, c))
        if ok:
            hits.append((a, b, c, ntested))
    print(f"   {total} parameter triples tested over primes {PRIMES}")
    print(f"   vanishing for every prime tested: {len(hits)}")
    for h in hits[:40]:
        (an, ad), (bn, bd), (cn, cd), nt = h
        print(f"     a={an}/{ad} b={bn}/{bd} c={cn}/{cd}   ({nt} primes)")
    if len(hits) > 40:
        print(f"     ... and {len(hits) - 40} more")
    import json
    json.dump(hits, open('/tmp/hitsB.json', 'w'))
    classify_B(hits)
    return hits


def classify_B(hits):
    """Account for the vanishing hypergeometric cases.

    Decidable, no remembered lists: the hypergeometric equation is REDUCIBLE
    exactly when one of a, b, c-a, c-b is an integer, and a reducible equation
    has a lower-order factor whose solutions are elementary -- those are not
    counterexample candidates and I expect the bulk to sit there. What is left
    is irreducible with vanishing p-curvature, i.e. finite projective
    monodromy, i.e. classical Schwarz territory. I print the exponent
    differences (1-c, c-a-b, b-a) so that someone holding the Schwarz list can
    check my arithmetic against it. I am not asserting the comparison myself."""
    from fractions import Fraction as F
    red, irred = [], []
    for (a, b, c, nt) in hits:
        A, B_, C = F(*a), F(*b), F(*c)
        if any(x.denominator == 1 for x in (A, B_, C - A, C - B_)):
            red.append((a, b, c))
        else:
            irred.append((A, B_, C, nt))
    print(f"   reducible (one of a, b, c-a, c-b in Z): {len(red)} of {len(hits)}")
    print(f"   irreducible with vanishing p-curvature: {len(irred)}")
    seen = {}
    for (A, B_, C, nt) in irred:
        lmn = (1 - C, C - A - B_, B_ - A)
        key = tuple(sorted(abs(x) for x in lmn))
        seen[key] = seen.get(key, 0) + 1
        if seen[key] == 1:
            print(f"     a={A} b={B_} c={C}  exponent diffs (lam,mu,nu)=({lmn[0]},{lmn[1]},{lmn[2]})")
    print(f"   distinct exponent-difference triples among the irreducible: {len(seen)}")


# ----------------------------------------------------------------- family C
def build_C(params, p):
    exps, (lamn, lamd) = params
    lam = rat(lamn, lamd, p)
    if lam is None or lam % p in (0, 1):     # lam must be a fourth point
        return None
    vals = []
    for (n, d) in exps:
        v = rat(n, d, p)
        if v is None:
            return None
        vals.append(v)
    d = pfromroots([0, 1, lam], p)
    return fuchsian_4pt(vals, lam, p), d


def family_C(limit=None):
    print("== family C (the search): 4 finite singular points 0, 1, lam, plus infinity")
    # Exponent-difference-like parameters at the three finite points, taken from
    # the rationals where finite local monodromy lives, plus two accessory
    # parameters that the local data does NOT determine. lam = 2 and lam = 1/2
    # (a non-integral cross-ratio) both get a run.
    small = reduced([(n, d) for d in (1, 2, 3) for n in range(-2, 3)])
    acc = [(n, 1) for n in range(-3, 4)]
    hits, total = [], 0
    for lam in [(2, 1), (1, 2), (3, 1), (-1, 1)]:
        for e0, e1, el in product(small, small, small):
            for c0, c1 in product(acc, acc):
                total += 1
                if limit and total > limit:
                    break
                ok, nt = survives(build_C, ((e0, e1, el, c0, c1), lam))
                if ok:
                    hits.append(((e0, e1, el, c0, c1), lam, nt))
            if limit and total > limit:
                break
        if limit and total > limit:
            break
    print(f"   {total} operators tested over primes {PRIMES}")
    print(f"   vanishing for every prime tested: {len(hits)}")
    for (exps, lam, nt) in hits[:60]:
        e = " ".join(f"{n}/{d}" for (n, d) in exps)
        print(f"     lam={lam[0]}/{lam[1]}  [e0 e1 el c0 c1] = {e}   ({nt} primes)")
    if len(hits) > 60:
        print(f"     ... and {len(hits) - 60} more")
    import json
    json.dump(hits, open('/tmp/hitsC.json', 'w'))
    return hits


def classify_C(hits):
    """Account for every survivor. A survivor is only interesting if it is not
    reducible and not visibly algebraic."""
    print("== accounting for family C survivors")
    if not hits:
        print("   nothing to account for: the vanishing set is empty")
        return
    trivial = [h for h in hits if h[0][3] == (0, 1) and h[0][4] == (0, 1)]
    print(f"   accessory parameters both zero (Q == 0, so y'' = -P y' and the")
    print(f"   system is reducible with the constant solution): {len(trivial)}")
    rest = [h for h in hits if h not in trivial]
    print(f"   remaining: {len(rest)}")
    for (exps, lam, nt) in rest[:40]:
        e = " ".join(f"{n}/{d}" for (n, d) in exps)
        print(f"     lam={lam[0]}/{lam[1]}  [{e}]  ({nt} primes)  <- needs a reason")


if __name__ == "__main__":
    which = sys.argv[1] if len(sys.argv) > 1 else "all"
    if which in ("A", "all"):
        if not family_A():
            print("CONTROL FAILED -- stopping, nothing below would be trustworthy")
            raise SystemExit(1)
    if which in ("B", "all"):
        family_B()
    if which in ("C", "all"):
        lim = int(sys.argv[2]) if len(sys.argv) > 2 else None
        classify_C(family_C(lim))