#!/usr/bin/env python3
"""p-curvature of a Fuchsian connection over F_p(t), in pure Python.

Built because a mathematician on Lemmy (PM_ME_VINTAGE_30S) asked me to spend
hours hunting a counterexample to the Grothendieck-Katz p-curvature conjecture
and to document the hunt whether or not it found anything. This file is the
instrument; sweep.py is the hunt.

  Conjecture (Grothendieck-Katz). If a linear differential system over a number
  field has zero p-curvature for almost all primes p, its solutions are
  algebraic.

A counterexample is therefore a connection whose p-curvature vanishes for every
prime we test and whose solutions are transcendental. Nobody expects one; the
useful output of a search like this is a documented list of things that looked
like candidates and were not.

WHY THIS IS COMPUTABLE WITHOUT A CAS
For a system Y' = A Y of rank n over F_p(t), the p-curvature is the F_p(t)-linear
map psi_p = A_p, where

    A_1 = A,   A_{k+1} = A_k' + A_k A.

Done naively that needs rational-function arithmetic with a gcd at every step,
and the degrees blow up. The trick that makes it cheap: write A = B/d once, with
B a polynomial matrix and d the common denominator. Then A_k = B_k / d^k, and

    A_k' = (B_k' d - k B_k d') / d^{k+1},   A_k A = B_k B / d^{k+1}

so the whole recursion collapses to a *polynomial* one with no division at all:

    B_1 = B,   B_{k+1} = B_k' d - k B_k d' + B_k B          (matrix product, in
                                                             that order)

and psi_p = 0 if and only if the polynomial matrix B_p is identically zero.
That is p steps of small-degree polynomial multiplication: milliseconds, not
minutes, and exact -- no floating point anywhere, so a zero is a real zero.

Sanity: rank 1 with A = c/t (solution t^c, algebraic). d = t, B = c, and the
recursion gives B_k = c(c-1)...(c-k+1). At k = p one factor is c - c = 0 in F_p,
so psi_p vanishes for every p. Correct, and it is a factorial-shaped identity
that would be very hard to produce by accident, which is why I use it as test 1.
"""

# ---------------------------------------------------------------- polynomials
# A polynomial is a list of coefficients mod p, lowest degree first. Trailing
# zeros are trimmed so that "is this zero" is just "is this the empty list".


def trim(a):
    while a and a[-1] == 0:
        a.pop()
    return a


def padd(a, b, p):
    n = max(len(a), len(b))
    return trim([((a[i] if i < len(a) else 0) + (b[i] if i < len(b) else 0)) % p
                 for i in range(n)])


def pscal(a, c, p):
    c %= p
    if c == 0:
        return []
    return trim([(x * c) % p for x in a])


def pmul(a, b, p):
    if not a or not b:
        return []
    out = [0] * (len(a) + len(b) - 1)
    for i, x in enumerate(a):
        if x:
            for j, y in enumerate(b):
                if y:
                    out[i + j] = (out[i + j] + x * y) % p
    return trim(out)


def pder(a, p):
    return trim([(i * a[i]) % p for i in range(1, len(a))])


def pfromroots(roots, p):
    """monic polynomial with the given roots, e.g. t(t-1)(t-lam)"""
    out = [1]
    for r in roots:
        out = pmul(out, [(-r) % p, 1], p)
    return out


# ------------------------------------------------------------------- matrices
# An n x n matrix of polynomials, as a list of rows.


def mzero(n):
    return [[[] for _ in range(n)] for _ in range(n)]


def mder(M, p):
    return [[pder(e, p) for e in row] for row in M]


def mscal(M, c, p):
    return [[pscal(e, c, p) for e in row] for row in M]


def madd(M, N, p):
    return [[padd(a, b, p) for a, b in zip(r, s)] for r, s in zip(M, N)]


def mmul(M, N, p):
    n = len(M)
    out = mzero(n)
    for i in range(n):
        for k in range(n):
            if M[i][k]:
                for j in range(n):
                    if N[k][j]:
                        out[i][j] = padd(out[i][j], pmul(M[i][k], N[k][j], p), p)
    return out


def mmulpoly(M, a, p):
    return [[pmul(e, a, p) for e in row] for row in M]


def miszero(M):
    return all(not e for row in M for e in row)


def mstr(M):
    def s(e):
        if not e:
            return "0"
        parts = []
        for i, c in enumerate(e):
            if not c:
                continue
            parts.append(str(c) if i == 0 else (f"{c}t^{i}" if i > 1 else f"{c}t"))
        return "+".join(parts)
    return "[" + "; ".join(", ".join(s(e) for e in row) for row in M) + "]"


# --------------------------------------------------------------- p-curvature


def pcurvature(B, d, p):
    """Numerator of psi_p for the system Y' = (B/d) Y over F_p(t).

    Returns the polynomial matrix B_p. psi_p == 0 exactly when it is zero.
    """
    dd = pder(d, p)
    Bk = [row[:] for row in B]
    for k in range(1, p):
        # B_{k+1} = B_k' d - k B_k d' + B_k B
        t1 = mmulpoly(mder(Bk, p), d, p)
        t2 = mscal(mmulpoly(Bk, dd, p), -k, p)
        t3 = mmul(Bk, B, p)
        Bk = madd(madd(t1, t2, p), t3, p)
    return Bk


def vanishes(B, d, p):
    return miszero(pcurvature(B, d, p))


# -------------------------------------------------------- building the system


def companion_from_second_order(q, p1, d, p):
    """y'' = (p1/d) y' + (q/d) y  ->  B, d with A = B/d for Y = (y, y')^T.

    Y' = A Y with A = [[0, 1], [q/d, p1/d]]; over the common denominator d that
    is B = [[0, d], [q, p1]].
    """
    return [[[], d[:]], [q[:], p1[:]]]


def power_solutions(alpha, beta, p):
    """The order-2 operator whose solutions are t^alpha and t^beta.

    Derived from the Wronskian determinant, not looked up:
      y'' = ((a+b-1)/t) y' - (ab/t^2) y
    Both solutions are algebraic whenever a, b are rational, so a correct
    p-curvature routine MUST return zero here. This is test 3.
    """
    d = [0, 0, 1]                                   # t^2
    p1 = pscal([0, 1], (alpha + beta - 1) % p, p)   # ((a+b-1)/t) = ((a+b-1)t)/t^2
    q = [(-alpha * beta) % p]                       # -ab/t^2
    return companion_from_second_order(q, p1, d, p)


def fuchsian_4pt(exps, lam, p):
    """An order-2 Fuchsian operator with singularities at 0, 1, lam, infinity.

    y'' + P y' + Q y = 0 with
        P = e0/t + e1/(t-1) + el/(t-lam)
        Q = (c0 t + c1) / (t (t-1) (t-lam))        [accessory parameters]
    exps = (e0, e1, el, c0, c1). Four singular points in the finite plane plus
    infinity puts this outside the rigid (hypergeometric) case, which is the
    cheapest way I know to leave the territory where the conjecture is a theorem.
    """
    e0, e1, el, c0, c1 = [x % p for x in exps]
    d = pfromroots([0, 1, lam], p)                       # t(t-1)(t-lam)
    # P over the common denominator d: e0 (t-1)(t-lam) + e1 t(t-lam) + el t(t-1)
    Pn = padd(padd(pscal(pfromroots([1, lam], p), e0, p),
                   pscal(pfromroots([0, lam], p), e1, p), p),
              pscal(pfromroots([0, 1], p), el, p), p)
    q = pscal(padd([c1], [0, c0], p), -1, p)             # y'' = -P y' - Q y
    p1 = pscal(Pn, -1, p)
    return companion_from_second_order(q, p1, d, p)


# ------------------------------------------------------------------ self-test
# Every claim in the sweep rests on this routine being right, so it gets checked
# against cases whose answers are known from analysis, not from my code.


def selftest():
    ok = True

    def check(name, got, want):
        nonlocal ok
        good = got == want
        ok = ok and good
        print(("  pass  " if good else "  FAIL  ") + f"{name}: got {got}, want {want}")

    print("rank 1, A = c/t  (solution t^c, algebraic)  -> expect vanishing")
    for p in (5, 7, 11, 13):
        for c in (1, 2, 3, p - 1):
            check(f"p={p} c={c}", vanishes([[[c]]], [0, 1], p), True)

    print("rank 1, A = 1  (solution e^t, transcendental) -> expect NON-vanishing")
    for p in (5, 7, 11, 13, 17):
        check(f"p={p}", vanishes([[[1]]], [1], p), False)

    print("rank 1, A = c  (solution e^(ct)) -> vanishing iff c == 0")
    for p in (7, 11):
        check(f"p={p} c=3", vanishes([[[3]]], [1], p), False)
        check(f"p={p} c=0", vanishes([[[0]]], [1], p), True)

    print("rank 2, solutions t^a, t^b (both algebraic) -> expect vanishing")
    for p in (7, 11, 13, 17):
        for (a, b) in ((1, 2), (2, 5), (3, 1)):
            check(f"p={p} a={a} b={b}", vanishes(power_solutions(a, b, p), [0, 0, 1], p), True)

    print("rank 2, solutions 1, e^t (y'' = y') -> expect NON-vanishing")
    for p in (7, 11, 13):
        B = [[[], [1]], [[], [1]]]     # A = [[0,1],[0,1]], d = 1
        check(f"p={p}", vanishes(B, [1], p), False)

    print("block-diagonal rank 2: vanishes iff both blocks do")
    for p in (7, 11):
        # diag(c/t, 1) over common denominator t: B = [[c,0],[0,t]], d = t
        B = [[[3], []], [[], [0, 1]]]
        check(f"p={p} diag(3/t, 1)", vanishes(B, [0, 1], p), False)
        B = [[[3], []], [[], [2]]]      # diag(3/t, 2/t): both algebraic
        check(f"p={p} diag(3/t, 2/t)", vanishes(B, [0, 1], p), True)

    # Gauge invariance is the strongest check available that needs no outside
    # knowledge: psi_p is covariant under Y = P Z, so *whether* it vanishes must
    # not depend on the gauge. If a bug produced spurious zeros it would almost
    # certainly break here.
    print("gauge invariance: A -> P^-1 (A P - P') must preserve vanishing")
    for p in (7, 11, 13):
        for (a, b) in ((1, 2), (2, 3)):
            B = power_solutions(a, b, p)
            d = [0, 0, 1]
            base = vanishes(B, d, p)
            # P = [[1, t], [0, 1]], det 1, so P^-1 = [[1,-t],[0,1]]: no new
            # denominator is introduced and the gauged system is still B'/d.
            P = [[[1], [0, 1]], [[], [1]]]
            Pinv = [[[1], [0, p - 1]], [[], [1]]]
            AP = mmul(B, P, p)                     # (B/d) P
            Pp = mmulpoly(mder(P, p), d, p)        # P' d, to share the denominator
            G = mmul(Pinv, madd(AP, mscal(Pp, -1, p), p), p)
            check(f"p={p} a={a} b={b} gauge", vanishes(G, d, p), base)

    print("\nSELFTEST " + ("OK" if ok else "FAILED"))
    return ok


if __name__ == "__main__":
    import sys
    raise SystemExit(0 if selftest() else 1)