#!/usr/bin/env python3
# verify_wilson_multiplicity.py
# Ancillary file for the note "The exact multiplicity of the Wilson supercongruence in F_q[t]".
# David Niedbala Giraudin, ORCID 0009-0009-1526-1178. Technical assistance: Claude (Anthropic).
#
# Plain Python 3, no dependency, no input file, exact arithmetic only.
# Run:  python3 verify_wilson_multiplicity.py
#
# Notation of the note: q = p**s, p odd; P monic irreducible of degree d >= 2 in F_q[t];
# [i] = t**(q**i) - t;  F_d = (-1)**d * product over 1 <= i <= d-1 of [i]**(q**(d-i) - 1)  (formula (1));
# m(P) = v_P(F_d + 1);  r(P) = depth = largest r with every exponent of P congruent to 0 or 1
# modulo p**r.  Theorem: m = 1 if r = 0, m = p**r - 1 if 1 <= r <= s, m = q + 1 if r > s.
# F_9 = F_3[a]/(a**2 - a - 1); its element x + 3y (0 <= x, y <= 2) stands for x + y*a.
#
# Part 1: F_d + 1 is computed in F_q[t] modulo P**K and its P-adic valuation is read off,
#         for every monic irreducible P of small degree, and for the deep examples where this
#         is feasible, using formula (1); formula (1) itself is checked against the definition
#         (product of all nonzero polynomials of degree < d) over F_3 for d = 2, 3.
# Part 2: for all examples, the Taylor coefficients f_k of F_d(theta + u) at a root theta are
#         computed in the residue field, from F_d(theta+u) = (-1)**d prod (beta_i - u + u**(q**i))**e_i,
#         with e_i = q**(d-i) - 1 and A**(q**j - 1) = A**(q**j)/A (Frobenius on power series).
# The script stops at the first failure and ends with "ALL CHECKS PASSED".

import sys
from math import comb
import time

T0 = time.time()


def check(cond, msg):
    if not cond:
        print("FAILED: " + msg, flush=True)
        sys.exit(1)
    print("  ok  %-78s [%6.1f s]" % (msg, time.time() - T0), flush=True)


# ------------------------------------------------------------------ finite fields
class GF:
    def __init__(self, p, s):
        self.p, self.s, self.q = p, s, p ** s
        q = self.q
        if s == 1:
            self.mul_t = [[(x * y) % p for y in range(q)] for x in range(q)]
        elif (p, s) == (3, 2):
            def mul9(x, y):
                x0, x1, y0, y1 = x % 3, x // 3, y % 3, y // 3
                # (x0 + x1 a)(y0 + y1 a), a**2 = a + 1
                c0 = x0 * y0 + x1 * y1
                c1 = x0 * y1 + x1 * y0 + x1 * y1
                return (c0 % 3) + 3 * (c1 % 3)
            self.mul_t = [[mul9(x, y) for y in range(q)] for x in range(q)]
        else:
            raise ValueError
        self.add_t = [[self._add(x, y) for y in range(q)] for x in range(q)]
        self.neg_t = [self._neg(x) for x in range(q)]
        self.inv_t = [0] + [next(y for y in range(1, q) if self.mul_t[x][y] == 1)
                            for x in range(1, q)]

    def _add(self, x, y):
        if self.s == 1:
            return (x + y) % self.p
        return ((x % 3 + y % 3) % 3) + 3 * ((x // 3 + y // 3) % 3)

    def _neg(self, x):
        if self.s == 1:
            return (-x) % self.p
        return ((-(x % 3)) % 3) + 3 * ((-(x // 3)) % 3)


# polynomials over GF: lists of ints, index = degree, no trailing zeros
def trim(a):
    while a and a[-1] == 0:
        a.pop()
    return a


def padd(F, a, b):
    n = max(len(a), len(b))
    r = [F.add_t[a[i] if i < len(a) else 0][b[i] if i < len(b) else 0] for i in range(n)]
    return trim(r)


def pneg(F, a):
    return [F.neg_t[x] for x in a]


def psub(F, a, b):
    return padd(F, a, pneg(F, b))


def pmul(F, a, b):
    if not a or not b:
        return []
    r = [0] * (len(a) + len(b) - 1)
    mt, at = F.mul_t, F.add_t
    for i, x in enumerate(a):
        if x == 0:
            continue
        row = mt[x]
        for j, y in enumerate(b):
            if y:
                r[i + j] = at[r[i + j]][row[y]]
    return trim(r)


def pdivmod(F, a, m):
    a = a[:]
    dm = len(m) - 1
    inv_lead = F.inv_t[m[-1]]
    qt = [0] * max(len(a) - dm, 0)
    while len(a) - 1 >= dm and a:
        c = F.mul_t[a[-1]][inv_lead]
        sh = len(a) - 1 - dm
        qt[sh] = c
        for j, y in enumerate(m):
            if y:
                a[sh + j] = F.add_t[a[sh + j]][F.neg_t[F.mul_t[c][y]]]
        trim(a)
    return trim(qt), a


def pmod(F, a, m):
    return pdivmod(F, a, m)[1]


def ppowmod(F, a, e, m):
    r, a = [1], pmod(F, a, m)
    while e:
        if e & 1:
            r = pmod(F, pmul(F, r, a), m)
        a = pmod(F, pmul(F, a, a), m)
        e >>= 1
    return r


def pgcd(F, a, b):
    while b:
        a, b = b, pmod(F, a, b)
    return a


def is_irreducible(F, P):
    d, q = len(P) - 1, F.q
    t = [0, 1]
    if ppowmod(F, t, q ** d, P) != pmod(F, t, P):
        return False
    for l in {l for l in range(2, d + 1) if d % l == 0 and all(l % k for k in range(2, l))}:
        g = pgcd(F, P, psub(F, ppowmod(F, t, q ** (d // l), P), t))
        if len(g) > 1:
            return False
    return True


def depth(P, p):
    exps = [n for n, c in enumerate(P) if c]
    r = 0
    while all(n % p ** (r + 1) in (0, 1) for n in exps):
        r += 1
    return r


def m_formula(r, p, s):
    q = p ** s
    return 1 if r == 0 else (p ** r - 1 if r <= s else q + 1)


def m_direct(F, P, K):
    """v_P(F_d + 1), computed modulo P**K from the definition of F_d (returns K if >= K)."""
    d, q = len(P) - 1, F.q
    mod = [1]
    for _ in range(K):
        mod = pmul(F, mod, P)
    t = [0, 1]
    Fd = [1] if d % 2 == 0 else [F.neg_t[1]]
    for i in range(1, d):
        bi = psub(F, ppowmod(F, t, q ** i, mod), t)
        Fd = pmod(F, pmul(F, Fd, ppowmod(F, bi, q ** (d - i) - 1, mod)), mod)
    X = padd(F, Fd, [1])
    k = 0
    while k < K and X:
        Q, R = pdivmod(F, X, P)
        if R:
            break
        X, k = Q, k + 1
    return k if X else K


def v_numerator_H(F, P, K):
    """v_P of the numerator of 1/[1] + ... + 1/[d-1], i.e. of sum_i prod_{j != i} [j], mod P**K."""
    d, q = len(P) - 1, F.q
    mod = [1]
    for _ in range(K):
        mod = pmul(F, mod, P)
    t = [0, 1]
    br = [psub(F, ppowmod(F, t, q ** i, mod), t) for i in range(1, d)]
    N = []
    for i in range(d - 1):
        term = [1]
        for j in range(d - 1):
            if j != i:
                term = pmod(F, pmul(F, term, br[j]), mod)
        N = padd(F, N, term)
    k = 0
    while k < K and N:
        Q, R = pdivmod(F, N, P)
        if R:
            break
        N, k = Q, k + 1
    return k if N else K


def monic_polys(F, d):
    q = F.q
    for n in range(q ** d):
        c, x = [], n
        for _ in range(d):
            c.append(x % q)
            x //= q
        yield c + [1]


print("Exact multiplicity of the Wilson supercongruence in F_q[t]: verification", flush=True)
FIELDS = {3: GF(3, 1), 5: GF(5, 1), 9: GF(3, 2)}

# sanity of F_9: a**2 = a + 1 and every nonzero element has an inverse
F9 = FIELDS[9]
check(F9.mul_t[3][3] == 4
      and all(F9.mul_t[x][F9.inv_t[x]] == 1 for x in range(1, 9)),
      "F_9 = F_3[a]/(a**2 - a - 1): field tables consistent")

# Wilson theorem and the product formula for F_d, against the definition (product of all
# nonzero polynomials of degree < d), over F_3 for d = 2, 3
F3 = FIELDS[3]
for d in (2, 3):
    prod = [1]
    for n in range(1, 3 ** d):
        c, x = [], n
        for _ in range(d):
            c.append(x % 3)
            x //= 3
        prod = pmul(F3, prod, trim(c))
    t = [0, 1]
    # explicit product formula
    form = [1] if d % 2 == 0 else [2]
    for i in range(1, d):
        bi = psub(F3, [0] * (3 ** i) + [1], t)
        for _ in range(3 ** (d - i) - 1):
            form = pmul(F3, form, bi)
    check(prod == form, "F_%d = product of the nonzero polynomials of degree < %d over F_3" % (d, d))

# ------------------------------------------- Part 1a: all irreducible P of small degree
RANGES = ((3, 1, range(2, 8)), (5, 1, range(2, 5)), (3, 2, range(2, 4)))
for p, s, degs in RANGES:
    F = FIELDS[p ** s]
    for d in degs:
        count, hist = 0, {}
        for P in monic_polys(F, d):
            if P[0] == 0 or not is_irreducible(F, P):
                continue
            r = depth(P, p)
            mf = m_formula(r, p, s)
            md = m_direct(F, P, mf + 1)
            if md != mf:
                check(False, "q=%d P=%s: m_direct=%d, formula %d" % (p ** s, P, md, mf))
            if v_numerator_H(F, P, mf + 1) != mf - 1:
                check(False, "q=%d P=%s: Corollary (multiplicity in the numerator of H)" % (p ** s, P))
            count += 1
            hist[md] = hist.get(md, 0) + 1
        check(count > 0, "q=%d, d=%d: all %d irreducible P: m_direct = formula = 1 + v_P(num H); values %s"
              % (p ** s, d, count, dict(sorted(hist.items()))))

# ------------------------------------------------------------------ the deep examples
# (field size q, exponents -> coefficients, expected depth, expected m, direct computation?)
EX = [
    (3, {3: 1, 1: 2, 0: 2}, 1, 2, True),                      # t**3 - t - 1
    (3, {18: 1, 10: 2, 0: 1}, 2, 4, True),                    # t**18 - t**10 + 1      r > s
    (5, {5: 1, 1: 4, 0: 4}, 1, 4, True),                      # t**5 - t - 1
    (5, {51: 1, 50: 3, 26: 4, 0: 1}, 2, 6, False),            # t**51 - 2t**50 - t**26 + 1   r > s
    (9, {3: 1, 1: 8, 0: 8}, 1, 2, True),                      # t**3 + (2a+2)t + (2a+2)
    (9, {10: 1, 9: 4, 1: 1, 0: 1}, 2, 8, True),               # t**10 + (a+1)t**9 + t + 1   r = s
    (9, {54: 1, 28: 6, 27: 3, 1: 3, 0: 1}, 3, 10, False),     # t**54 + 2a t**28 + a t**27 + a t + 1
]


def mkpoly(dct):
    P = [0] * (max(dct) + 1)
    for e, c in dct.items():
        P[e] = c
    return P


for q, dct, r_exp, m_exp, direct in EX:
    F = FIELDS[q]
    P = mkpoly(dct)
    p, s = F.p, F.s
    ok = is_irreducible(F, P) and depth(P, p) == r_exp and m_formula(r_exp, p, s) == m_exp
    msg = "q=%d, d=%d: irreducible, depth %d, formula m=%d" % (q, len(P) - 1, r_exp, m_exp)
    if direct:
        ok = ok and m_direct(F, P, m_exp + 1) == m_exp and v_numerator_H(F, P, m_exp + 1) == m_exp - 1
        msg += ", direct v_P(F_d+1) = %d, v_P(num H) = %d" % (m_exp, m_exp - 1)
    check(ok, msg)


def sinv(F, P, A, K):
    """inverse of the power series A (A[0] invertible mod P)."""
    red = lambda a: pmod(F, a, P)
    d = len(P) - 1
    inv0 = ppowmod(F, A[0], F.q ** d - 2, P)
    R = [inv0] + [[] for _ in range(K - 1)]
    for k in range(1, K):
        acc = []
        for j in range(1, k + 1):
            if A[j] and R[k - j]:
                acc = padd(F, acc, pmul(F, A[j], R[k - j]))
        R[k] = red(pneg(F, pmul(F, red(acc), inv0)))
    return R


# --------------------------------------- Part 2: Taylor coefficients at a root theta
def taylor_profile(F, P, K):
    """Coefficients f_0..f_{K-1} of F_d(theta+u), theta = t mod P, as residues mod P;
    also beta_1 and the Hasse derivatives P**[j](theta)."""
    d, q, p = len(P) - 1, F.q, F.p
    t = [0, 1]
    red = lambda a: pmod(F, a, P)
    beta = {i: psub(F, ppowmod(F, t, q ** i, P), t) for i in range(1, d)}

    def smul(A, B):
        C = [[] for _ in range(K)]
        for i in range(K):
            if not A[i]:
                continue
            for j in range(K - i):
                if B[j]:
                    C[i + j] = red(padd(F, C[i + j], pmul(F, A[i], B[j])))
        return C

    def spow_frob(A, j):
        """A**(q**j - 1) = A**(q**j) / A; A**(q**j) has the coefficient c_k**(q**j) at u**(k q**j)."""
        Aq = [[] for _ in range(K)]
        for k in range(K):
            if A[k] and k * q ** j < K:
                Aq[k * q ** j] = ppowmod(F, A[k], q ** j, P)
        return smul(Aq, sinv(F, P, A, K))

    S = [[1]] + [[] for _ in range(K - 1)]
    for i in range(1, d):
        B = [beta[i], [F.neg_t[1]]] + [[] for _ in range(K - 2)]
        if q ** i < K:
            B[q ** i] = padd(F, B[q ** i], [1])
        S = smul(S, spow_frob(B, d - i))
    if d % 2:
        S = [pneg(F, c) for c in S]
    # Hasse derivatives of P at theta: coefficients of P(X + theta) = sum_j P**[j](theta) X**j
    from math import comb
    hasse = []
    for j in range(K + 2):
        acc = []
        for n, c in enumerate(P):
            if c and n >= j and comb(n, j) % p:
                acc = padd(F, acc, pmul(F, [F.mul_t[c][comb(n, j) % p]],
                                        ppowmod(F, t, n - j, P)))
        hasse.append(red(acc))
    return S, beta, hasse


for q, dct, r_exp, m_exp, direct in EX:
    F = FIELDS[q]
    P = mkpoly(dct)
    d, p, s = len(P) - 1, F.p, F.s
    K = q + 2
    S, beta, hasse = taylor_profile(F, P, K)
    red = lambda a: pmod(F, a, P)
    m_taylor = next(k for k in range(1, K) if S[k])
    ok = (red(padd(F, S[0], [1])) == []) and m_taylor == m_exp
    # f_k = -h_k(beta**-1) for k <= q-1, with sum_k h_k u**k = P'(theta) / sum_{j>=1} P**[j](theta) u**(j-1)
    Hs = sinv(F, P, [hasse[j + 1] for j in range(K)], K)
    Hs = [red(pmul(F, hasse[1], c)) for c in Hs]
    ok = ok and all(S[k] == red(pneg(F, Hs[k])) for k in range(q))
    # vanishing pattern of the Hasse derivatives: P**[j](theta) = 0 for 2 <= j <= p**r - 1
    lim = min(p ** r_exp - 1, K + 1)
    ok = ok and hasse[1] != [] and all(hasse[j] == [] for j in range(2, lim + 1))
    if p ** r_exp <= K + 1:
        ok = ok and hasse[p ** r_exp] != []
    # P**[p**r] is a nonzero polynomial of degree < d (hence nonzero at theta): some exponent n
    # with a nonzero coefficient has binomial(n, p**r) prime to p
    ok = ok and any(c and comb(n, p ** r_exp) % p for n, c in enumerate(P))
    # beta_{d-1}**q = -beta_1
    ok = ok and ppowmod(F, beta[d - 1], q, P) == red(pneg(F, beta[1]))
    if r_exp > s:
        b1inv = ppowmod(F, beta[1], q ** d - 2, P)
        ok = ok and S[q] == [] and S[q + 1] == red(pmul(F, b1inv, b1inv))
    check(ok, "q=%d, d=%d: Taylor coefficients: first nonzero at k=%d; f_k = -h_k for k<q; "
          "Hasse pattern%s" % (q, d, m_taylor, "; f_q = 0, f_{q+1} = beta_1**-2" if r_exp > s else ""))

print("ALL CHECKS PASSED", flush=True)