#!/usr/bin/env python3
# verify_carleman_padic.py
# Ancillary file for the note "The p-adic valuations of the coefficients in Yang's
# expansion of (1+1/x)**x".
# David Niedbala Giraudin, ORCID 0009-0009-1526-1178. Technical assistance: Claude (Anthropic).
#
# Plain Python 3, no dependency, no input file, exact rational arithmetic only.
# Run:  python3 verify_carleman_padic.py      (about fifteen seconds)
# Notation of the note: p odd prime, c = p - 2,
#   (1+1/x)**x = e (1 - sum_n b_n u**n),  u = 1/(x+1),
#   1 - B(u) = exp(-sum u**k/(k(k+1))),  T(u) = exp(-sum u**k/(k+1)) = e (1-u)**(1/u),
#   S(u) = 1/(1 - B(u)) = T(u)/(1-u).
# The script stops at the first failure and ends with "ALL CHECKS PASSED".

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

T0 = time.time()
N = 250          # number of coefficients computed
NDP = 100        # range of the composition dynamic programmes
PRIMES_DP = (3, 5, 7, 11, 13)
PRIMES_ALL = (3, 5, 7, 11, 13, 17, 19, 23, 29, 31)


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


def vp_int(n, p):
    v = 0
    while n % p == 0:
        n //= p
        v += 1
    return v


def vp(x, p):
    return vp_int(x.numerator, p) - vp_int(x.denominator, p)


def sp(n, p):
    s = 0
    while n:
        s += n % p
        n //= p
    return s


def vfact(n, p):
    s, q = 0, p
    while q <= n:
        s += n // q
        q *= p
    return s


def exp_series(c, n):
    """coefficients 0..n of exp(sum_{k>=1} c[k] u**k), by E' = L' E."""
    E = [Fr(1)] + [Fr(0)] * n
    for r in range(1, n + 1):
        E[r] = sum(k * c[k] * E[r - k] for k in range(1, r + 1)) / r
    return E


def inverse_series(a, n):
    inv = [Fr(1) / a[0]] + [Fr(0)] * n
    for r in range(1, n + 1):
        inv[r] = -sum(a[k] * inv[r - k] for k in range(1, r + 1)) / a[0]
    return inv


print("p-adic valuations of Yang's coefficients b_n: verification", flush=True)

# ----------------------------------------------------------------- the three series
W = exp_series([Fr(0)] + [Fr(-1, k * (k + 1)) for k in range(1, N + 1)], N)   # 1 - B
b = [Fr(0)] + [-W[n] for n in range(1, N + 1)]
T = exp_series([Fr(0)] + [Fr(-1, k + 1) for k in range(1, N + 1)], N)
S = [sum(T[: r + 1]) for r in range(N + 1)]

# Independent route 1: the recurrence b_1 = 1/2,
# b_n = (1/n)(1/(n+1) - sum over k = 0..n-2 of b_{n-k-1}/(k+2)), obtained by differentiating.
br = [Fr(0), Fr(1, 2)]
for n in range(2, N + 1):
    br.append((Fr(1, n + 1) - sum(br[n - k - 1] / (k + 2) for k in range(n - 1))) / n)
check(br == b, "b_n from exp(-sum u**k/(k(k+1))) equals the classical recurrence, n <= %d" % N)

# Independent route 2: 1 - B(u) = (1-u) * A(u), A(u) = (1-u)**(-1/u)/e, whose numerators and
# denominators are OEIS A055505 / A055535 (Comtet).
A_num = [1, 1, 11, 7, 2447, 959, 238043, 67223, 559440199, 123377159, 29128857391,
         5267725147, 9447595434410813, 1447646915836493, 225037938358318573,
         29911565062525361, 3651003047854884043877, 38950782815463986767]
A_den = [1, 2, 24, 16, 5760, 2304, 580608, 165888, 1393459200, 309657600, 73574645760,
         13377208320, 24103053950976000, 3708162146304000, 578473294823424000,
         77129772643123200, 9440684171518279680000, 100969884187361280000]
A = inverse_series(T, len(A_num) - 1)
check(all(A[n] == Fr(A_num[n], A_den[n]) for n in range(len(A_num))),
      "1/T(u) = (1-u)**(-1/u)/e has the OEIS coefficients A055505/A055535, n <= 17")
check(all(b[n] == A[n - 1] - A[n] for n in range(1, len(A_num))),
      "b_n = a_{n-1} - a_n with a_n = A055505(n)/A055535(n), n <= 17")
Winv = inverse_series(W, N)
check(all(S[r] == Winv[r] for r in range(N + 1)),
      "S(u) = 1/(1 - B(u)) = T(u)/(1-u), r <= %d" % N)
check([b[n] for n in range(1, 7)] == [Fr(1, 2), Fr(1, 24), Fr(1, 48), Fr(73, 5760),
                                      Fr(11, 1280), Fr(3625, 580608)],
      "b_1..b_6 = 1/2, 1/24, 1/48, 73/5760, 11/1280, 3625/580608")


# ----------------------------------------------------- the knapsack by brute force
def weight_b(k, p):
    return vp_int(k, p) + vp_int(k + 1, p)


def weight_t(k, p):
    return vp_int(k + 1, p)


def dp_scores(n_max, p, weight):
    """best[m][t] = max sum of weights over compositions of t into m parts, and the number
    of compositions attaining it.  Returns M(n) and the number of optimal compositions,
    for n = 1..n_max, the score being sum of weights + v_p(m!)."""
    NEG = None
    best = [[NEG] * (n_max + 1) for _ in range(n_max + 1)]
    cnt = [[0] * (n_max + 1) for _ in range(n_max + 1)]
    best[0][0], cnt[0][0] = 0, 1
    wt = [0] + [weight(k, p) for k in range(1, n_max + 1)]
    for m in range(1, n_max + 1):
        for t in range(m, n_max + 1):
            bv, bc = NEG, 0
            for k in range(1, t - m + 2):
                pv = best[m - 1][t - k]
                if pv is None:
                    continue
                s = pv + wt[k]
                if bv is None or s > bv:
                    bv, bc = s, cnt[m - 1][t - k]
                elif s == bv:
                    bc += cnt[m - 1][t - k]
            best[m][t], cnt[m][t] = bv, bc
    M, C = [None], [None]
    for n in range(1, n_max + 1):
        sc = [(best[m][n] + vfact(m, p), cnt[m][n]) for m in range(1, n + 1)
              if best[m][n] is not None]
        top = max(s for s, _ in sc)
        M.append(top)
        C.append(sum(k for s, k in sc if s == top))
    return M, C


def M_formula(n, p):
    """Proposition 6: M(n) = max_{0 <= j <= n/(p-1)} [ j + v_p((n - j(p-2))!) ]."""
    return max(j + vfact(n - j * (p - 2), p) for j in range(n // (p - 1) + 1))


def in_U(m, p):
    """Criterion m in U_p (Theorem 7): i + s_p(m + i(p-2)) > s_p(m) for 1 <= i < s_p(m)."""
    c = p - 2
    return all(i + sp(m + i * c, p) > sp(m, p) for i in range(1, sp(m, p)))


def unique_predicted(n, p):
    return n == 1 or (n % (p - 1) == 0 and in_U(n // (p - 1), p))


for p in PRIMES_DP:
    for name, wfun in (("b", weight_b), ("T", weight_t)):
        M, C = dp_scores(NDP, p, wfun)
        check(all(M[n] == M_formula(n, p) for n in range(1, NDP + 1)),
              "p=%d, %s: brute-force M(n) = formula of Proposition 6, n <= %d" % (p, name, NDP))
        check(all((C[n] == 1) == unique_predicted(n, p) for n in range(1, NDP + 1)),
              "p=%d, %s: unique optimum <=> criterion of Theorem 7, n <= %d" % (p, name, NDP))
        seq = b if name == "b" else T
        check(all(vp(seq[n], p) >= -M[n] for n in range(1, NDP + 1)),
              "p=%d, %s: v_p >= -M(n) (Lemma 4), n <= %d" % (p, name, NDP))

# Remark 8: strict inequality v_3(b_n) > -M(n) at exactly 23 indices n <= 80, the first being 5
strict = [n for n in range(1, 81) if vp(b[n], 3) > -M_formula(n, 3)]
check(len(strict) == 23 and strict[0] == 5,
      "p=3: v_3(b_n) > -M(n) at exactly 23 indices n <= 80, the smallest being 5")

# criterion with all i >= 1 equals criterion with i < s_p(m)
for p in PRIMES_DP:
    c = p - 2
    check(all(in_U(m, p) == all(i + sp(m + i * c, p) > sp(m, p) for i in range(1, m + 1))
              for m in range(1, 3000)),
          "p=%d: the conditions for i < s_p(m) and for 1 <= i <= m agree, m < 3000" % p)

# ------------------------------------------------------------- exact valuations (Theorem 1)
for p in PRIMES_ALL:
    good = [n for n in range(2, N + 1) if unique_predicted(n, p)]
    ok = True
    for n in good:
        m = n // (p - 1)
        target = -(m + (m - sp(m, p)) // (p - 1))
        if vp(b[n], p) != target or vp(T[n], p) != target or target != -M_formula(n, p):
            ok = False
    check(ok and len(good) > 0,
          "p=%d: v_p(b_n) = v_p(T_n) = -(m + v_p(m!)) on N_p (%d indices <= %d)"
          % (p, len(good), N))

# ------------------------------------------------------------- the indices p**K - 1 (Theorem 3)
for p in PRIMES_ALL:
    K = 1
    while p ** K - 1 <= N:
        R, RK = p ** K - 1, (p ** K - 1) // (p - 1)
        D = (p ** (K + 1) - p - K * (p - 1)) // (p - 1) ** 2
        assert (p - 1) ** 2 * D == p ** (K + 1) - p - K * (p - 1)
        assert D == sum((K - j) * p ** j for j in range(K)) == RK + vfact(RK, p)
        for name, seq in (("b", b), ("T", T), ("S", S)):
            vals = [vp(seq[n], p) for n in range(1, R + 1)]
            if not (vals[-1] == -D and all(v > -D for v in vals[:-1])):
                check(False, "p=%d K=%d %s: minimality" % (p, K, name))
        check(in_U(RK, p) and M_formula(R - 1, p) == D - 1 if R > 1 else True,
              "p=%d K=%d: R_K in U_p, M(p**K-2) = D_K - 1, b, T, S minimal at p**K-1 only"
              % (p, K))
        K += 1

# check of M(p**K - 2) <= D_K - 1 for larger K (formula only)
for p in (3, 5, 7, 11):
    for K in range(1, 8 if p == 3 else 5):
        RK = (p ** K - 1) // (p - 1)
        D = RK + vfact(RK, p)
        assert in_U(RK, p)
        if p ** K - 2 >= 1:
            assert M_formula(p ** K - 2, p) <= D - 1
check(True, "M(p**K-2) <= D_K - 1 and R_K in U_p for p=3 (K<=7), p=5,7,11 (K<=4)")

# ------------------------------------------- density of U_p (Lemma 10 and Theorem 2)
for p in PRIMES_ALL:
    c, X = p - 2, 60000
    Lam = lambda y: y + c * sp(y, p)
    cnt, ok_delta = 0, True
    for m in range(1, X + 1):
        if in_U(m, p):
            cnt += 1
            # jump: Delta = min_{i>=1} [i + s_p(m+ic) - s_p(m)] must equal p-1
            delta = min(i + sp(m + i * c, p) - sp(m, p) for i in range(1, sp(m, p) + p + 1))
            if delta != p - 1:
                ok_delta = False
    lo = Fr(X, p - 1) - Fr(c * (c + 1), 2 * (p - 1))
    E = cnt - Fr(X, p - 1)
    hi = c * (1 + math.log(X + c) / math.log(p))
    check(ok_delta and lo <= cnt and float(E) <= hi,
          "p=%d: Delta = p-1 on U_p, and %d = #U_p in [1,%d] = X/(p-1) %+.2f within bounds"
          % (p, cnt, X, float(E)))

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