#!/usr/bin/env python3
# verify_Ep_congruence.py
# Ancillary file for the note "A congruence modulo p**2 for the p-th coefficient of e(1-z)**((1-z)/z)".
# David Niedbala Giraudin, ORCID 0009-0009-1526-1178. Technical assistance: Claude (Anthropic).
#
# Plain Python 3, no dependency, exact rational arithmetic (fractions).  Run:
#     python3 verify_Ep_congruence.py
# Notation of the note:
#   G(z)   = exp(sum_{k>=1} z**k/(k(k+1))) = sum_n E_n z**n,
#   Phi(u) = e (1+u)**(-1/u) = exp(sum_{n>=1} h_n u**n),  h_n = (-1)**(n+1)/(n+1),  phi_k = [u**k] Phi.
# E_n is computed by the recurrence (c), phi_k from the exponential.  For every odd prime p in the
# tested range the script checks the identities, the intermediate congruences of the proof and the
# theorem E_p = -p/2 (mod p**2).  It stops at the first failure and ends with "ALL CHECKS PASSED".

import sys
import time
from fractions import Fraction as Fr
from math import comb

T0 = time.time()


def check(cond, msg):
    if not cond:
        print("FAIL  " + msg, flush=True)
        sys.exit(1)


def expseries(c, N):
    """coefficients 0..N of exp(sum_{n>=1} c[n] x**n), via a' = (sum n c_n x**(n-1)) a"""
    a = [Fr(0)] * (N + 1)
    a[0] = Fr(1)
    for n in range(N):
        a[n + 1] = sum((k + 1) * c[k + 1] * a[n - k] for k in range(n + 1)) / (n + 1)
    return a


def vp(x, p):
    if x == 0:
        return 10 ** 9
    v, n, d = 0, x.numerator, x.denominator
    while n % p == 0:
        n //= p
        v += 1
    while d % p == 0:
        d //= p
        v -= 1
    return v


def red(x, M):
    """image of a p-integral rational in Z/MZ (M a power of p)"""
    return x.numerator * pow(x.denominator, -1, M) % M


def is_prime(n):
    return n > 1 and all(n % d for d in range(2, int(n ** 0.5) + 1))


PRIMES = [p for p in range(3, 200) if is_prime(p)] + [251, 307, 401]
print("Congruence E_p = -p/2 (mod p**2): verification", flush=True)
print("primes: all odd primes < 200, and 251, 307, 401 (%d primes)" % len(PRIMES), flush=True)

mod_p3 = []
for p in PRIMES:
    N = p + 1
    inner = [Fr(0)] + [Fr(1, k * (k + 1)) for k in range(1, N + 1)]
    E = expseries(inner, N)
    h = [Fr(0)] + [Fr((-1) ** (n + 1), n + 1) for n in range(1, N + 1)]
    phi = expseries(h, N)

    # Lemma 1 (a), (b): exact coefficientwise identities up to z**(p+1)
    check(all(E[n] == sum((-1) ** k * phi[k] for k in range(n + 1)) for n in range(N + 1)),
          "p=%d: (a) E_n = sum_{k<=n} (-1)**k phi_k" % p)
    check(all(E[n] == sum(comb(n - 1, k - 1) * phi[k] for k in range(1, n + 1)) for n in range(1, N + 1)),
          "p=%d: (b) E_n = sum_k binom(n-1,k-1) phi_k" % p)
    check(all(phi[k] == (-1) ** k * (E[k] - (E[k - 1] if k else 0)) for k in range(N + 1)),
          "p=%d: (-1)**k phi_k = E_k - E_{k-1}" % p)

    # Lemma 2 (p-integrality and polar parts)
    psi = expseries([h[n] if n <= p - 2 else Fr(0) for n in range(N + 1)], N)
    check(all(psi[k] == phi[k] for k in range(p - 1)), "p=%d: psi_k = phi_k for k <= p-2" % p)
    check(phi[p - 1] == psi[p - 1] - Fr(1, p) and vp(psi[p - 1], p) >= 0,
          "p=%d: phi_{p-1} = psi_{p-1} - 1/p, psi_{p-1} p-integral" % p)
    check(vp(phi[p] + Fr(1, p), p) >= 0, "p=%d: phi_p + 1/p p-integral" % p)
    check(all(vp(h[n], p) >= 0 for n in range(1, p - 1)) and all(vp(psi[k], p) >= 0 for k in range(p)),
          "p=%d: h_n (n <= p-2) and psi_k (k <= p-1) p-integral" % p)
    check(all(vp(E[n], p) >= 0 for n in range(p - 1)) and all(vp(phi[k], p) >= 0 for k in range(p - 1)),
          "p=%d: E_n, phi_n p-integral for n <= p-2" % p)

    # Step A
    P2 = p * p
    H = [Fr(0)] * p
    for k in range(1, p):
        H[k] = H[k - 1] + Fr(1, k)
    check(all(red(comb(p - 1, j) - (-1) ** j * (1 - p * H[j]), P2) == 0 for j in range(p)),
          "p=%d: binom(p-1,j) = (-1)**j (1 - p H_j) mod p**2" % p)
    Hstar = sum((-1) ** (k - 1) * H[k - 1] * phi[k] for k in range(1, p - 1))
    check(E[p] == E[p - 2] + phi[p - 1] - phi[p], "p=%d: E_p = E_{p-2} + phi_{p-1} - phi_p" % p)
    first = sum((-1) ** (k - 1) * phi[k] for k in range(1, p - 1)) - p * Hstar + (p - 1) * phi[p - 1] + phi[p]
    check(red(E[p] - first, P2) == 0, "p=%d: first display of Step A mod p**2" % p)
    check(red(2 * E[p] - p * (psi[p - 1] - Hstar), P2) == 0, "p=%d: 2E_p = p(psi_{p-1} - H*) mod p**2" % p)
    check(red(psi[p - 1] - Hstar + 1, p) == 0, "p=%d: (2) psi_{p-1} - H* = -1 mod p" % p)
    # Step B
    rhsB = sum((-1) ** k * (Fr(1, k) + 1) * phi[k] for k in range(1, p - 1))
    check((p - 1) * psi[p - 1] == sum(j * h[j] * phi[p - 1 - j] for j in range(1, p - 1)),
          "p=%d: (p-1) psi_{p-1} = sum j h_j phi_{p-1-j} (exact)" % p)
    check(red(psi[p - 1] - rhsB, p) == 0, "p=%d: psi_{p-1} = sum (-1)**k (1+1/k) phi_k mod p" % p)
    lhsC = sum((-1) ** k * H[k] * phi[k] for k in range(1, p - 1))
    check(red(psi[p - 1] - Hstar - (E[p - 2] - 1 + lhsC), p) == 0, "p=%d: second display of Step B mod p" % p)
    # Step C
    check(lhsC == H[p - 2] * E[p - 2] - 1 - sum(E[k] / (k + 1) for k in range(1, p - 2)),
          "p=%d: summation by parts (exact)" % p)
    check(red(H[p - 2] - 1, p) == 0, "p=%d: H_{p-2} = 1 mod p" % p)
    check(red(lhsC + E[p - 2], p) == 0, "p=%d: (3) sum (-1)**k H_k phi_k = -E_{p-2} mod p" % p)
    check(red(2 * E[p - 2] - 1 - sum(E[j] / (j + 1) for j in range(1, p - 2)), p) == 0,
          "p=%d: 2E_{p-2} = 1 + sum_{1<=j<=p-3} E_j/(j+1) mod p" % p)
    # Theorem
    check(vp(E[p], p) == 1 and red(E[p] / p + Fr(1, 2), p) == 0, "p=%d: E_p = -p/2 mod p**2" % p)
    if red(E[p] + Fr(p, 2), p ** 3) == 0:
        mod_p3.append(p)
    if p in (3, 5, 7, 101, 199, 251, 307, 401):
        print("  ok  p=%d: all steps and E_p = -p/2 mod p**2   [%6.1f s]" % (p, time.time() - T0), flush=True)

print("  all %d primes passed" % len(PRIMES), flush=True)
print("  primes in the range with E_p = -p/2 even modulo p**3: %s" % mod_p3, flush=True)
print("ALL CHECKS PASSED", flush=True)