#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
verify_margin.py  --  rigorous verification of the Cesaro margin M(N) from the primes.

    M(N) = R(N) - C(N),   C(N) = 2 sum_{n<N} Lambda(n) n^{-1/2} log(N/n),
    R(N) = 8 sqrt(N) - c0 log N + K0 - tau(log N),
    c0 = log(pi) + gamma + pi/2 + 3 log 2,  K0 = -16 + pi^2/2 + 4G,
    tau(L) = 2 sum_{k>=1} exp(-(2k+1/2)L)/(2k+1/2)^2.

Between consecutive prime powers q < q' (a = log q, b = log q'), with
A = sum_{n<=q} Lambda(n) n^{-1/2} and B = sum_{n<=q} Lambda(n) n^{-1/2} log n constant,
    m(L) = M(e^L) = r(L) - 2 A L + 2 B,   m'' = r'' > 0   (N >= 2),
so m is convex on [a, b]:  min m >= max( m(a) - max(0,-m'(a))(b-a),  m(b) - max(0,m'(b))(b-a) ),
                           max m  = max( m(a), m(b) ).

Arithmetic: IEEE double precision with an explicit a-priori bound for every rounding error
(sqrt, +, -, *, / correctly rounded; numpy's log checked at start-up to be within 2 ulp).
Sums of segment totals are kept exactly (math.fsum + fractions.Fraction).

The work is split in segments of length SEG, processed in parallel; each segment is saved
to the folder  margin_ckpt/  so the run can be stopped and resumed (just run it again).

Segments already computed are reused by any later run with the same SEG (e.g. a 1e10 test
run is part of the 1e12 run).

Usage (Windows / Linux / Mac, Python 3.8+ with numpy; in PyCharm just edit SETTINGS and Run):
    python verify_margin.py                 # N up to 1e12, all cores but one
    python verify_margin.py --X 1e11        # smaller run (about 10x faster)
    python verify_margin.py --X 1e12 --workers 6 --seg 5e7

Output: margin_result.json and a summary on screen.
Pedro Caceres, 2026 -- companion to "The Cesaro margin of the primes".
"""
import argparse, json, math, os, sys, time
from fractions import Fraction
from decimal import Decimal, getcontext
import numpy as np

# ======================= SETTINGS (edit here if you run it from PyCharm without arguments) ==========
X_MAX = 1e12        # upper end of the verification (try 1e10 first: a few minutes)
WORKERS = 0         # 0 = all cores but one
SEG_LEN = 5e7       # segment length; do not change between runs that share margin_ckpt
# ====================================================================================================

EPS = 2.0 ** -52            # bound for the relative rounding error of one operation (generous: u = 2^-53)
OV = 4096                   # overlap to find the next prime power (prime gaps < 1550 below 1.8e19)
NSUB = 200                  # blocks for the cumulative sums
HB = 2e-5                   # width in L of the sub-blocks kept in the checkpoint

# ---------------------------------------------------------------- constants to 40 digits
getcontext().prec = 50
_pi = Decimal("3.141592653589793238462643383279502884197169")
_g = Decimal("0.5772156649015328606065120900824024310422")
_G = Decimal("0.9159655941772190150546035149323841107741")
_l2 = Decimal("0.6931471805599453094172321214581765680755")
_lpi = Decimal("1.144729885849400174143427351353058711647")
C0 = float(_lpi + _g + _pi / 2 + 3 * _l2)          # 5.3721834...
K0 = float(-16 + _pi * _pi / 2 + 4 * _G)            # -7.4013354...


def small_primes(n):
    s = np.ones(n + 1, dtype=bool); s[:2] = False
    for p in range(2, int(n ** 0.5) + 1):
        if s[p]: s[p * p::p] = False
    return np.nonzero(s)[0].astype(np.int64)


_SP = None
_PP = None   # prime powers p^k, k>=2, up to X+SEG+OV:  (q, log p, k)


def init_worker(X, seg):
    global _SP, _PP
    lim = int(math.isqrt(int(X + seg + OV))) + 2
    _SP = small_primes(lim)
    qs, ps, ks = [], [], []
    top = X + seg + OV
    for p in _SP.tolist():
        q, k = p * p, 2
        while q <= top:
            qs.append(q); ps.append(p); ks.append(k); q *= p; k += 1
    o = np.argsort(qs)
    _PP = (np.array(qs, dtype=np.int64)[o], np.array(ps, dtype=np.int64)[o], np.array(ks, dtype=np.int64)[o])


def sieve(lo, hi):
    """primes in [lo, hi), lo >= 2"""
    out = []
    if lo <= 2 < hi: out.append(np.array([2], dtype=np.int64))
    lo1 = max(lo, 3) | 1
    if lo1 >= hi: return np.concatenate(out) if out else np.zeros(0, np.int64)
    size = (hi - lo1 + 1) // 2
    arr = np.ones(size, dtype=bool)
    sp = _SP[1:]
    sp = sp[sp * sp < hi]
    st = np.maximum(sp * sp, ((lo1 + sp - 1) // sp) * sp)
    st = np.where(st % 2 == 0, st + sp, st)
    j0 = (st - lo1) // 2
    for p, j in zip(sp.tolist(), j0.tolist()):
        if j < size: arr[j::p] = False
    pr = lo1 + 2 * np.nonzero(arr)[0].astype(np.int64)
    pr = pr[pr >= 2]
    out.append(pr)
    return np.concatenate(out)


def r_and_dr(sq, L):
    """r(L) and r'(L) with sq = sqrt(e^L) computed directly from the integer."""
    tau = np.zeros_like(L); dtau = np.zeros_like(L)
    for k in range(1, 9):
        c = 2 * k + 0.5
        e = np.exp(-c * L)
        tau += 2 * e / (c * c); dtau += 2 * e / c
    return 8 * sq - C0 * L + K0 - tau, 4 * sq - C0 + dtau


def process(job):
    """One segment [lo, hi). Returns sums and sub-block bounds computed with A^ = 2 sqrt(lo) - c0/2,
    B^ = 2 sqrt(lo)(log lo - 2) in place of the unknown A0, B0 (corrected exactly afterwards)."""
    idx, lo, hi, ckdir = job
    fn = os.path.join(ckdir, f"seg_{idx:06d}.npz")
    if os.path.exists(fn):
        return idx
    pr = sieve(lo, hi + OV)
    qq, pp, kk = _PP
    sel = (qq >= lo) & (qq < hi + OV)
    q = np.concatenate([pr, qq[sel]]); p = np.concatenate([pr, pp[sel]]); k = np.concatenate([np.ones(pr.size, np.int64), kk[sel]])
    o = np.argsort(q, kind="stable"); q, p, k = q[o], p[o], k[o]
    nin = int(np.searchsorted(q, hi))           # prime powers in [lo, hi)
    assert nin < q.size, "no prime power in the overlap"
    qf = q[:nin + 1].astype(np.float64)         # exact (< 2^53)
    logp = np.log(p[:nin].astype(np.float64))
    sq = np.sqrt(qf)                            # correctly rounded
    ta = logp / sq[:nin]                        # Lambda(n) n^-1/2      rel. err <= 4 EPS
    tb = ta * (k[:nin] * logp)                  # ... * log n           rel. err <= 6 EPS
    Ahat = 2 * math.sqrt(lo) - C0 / 2 if lo > 2 else 0.0
    Bhat = 2 * math.sqrt(lo) * (math.log(lo) - 2) if lo > 2 else 0.0
    # block-wise cumulative sums: exact block offsets (fsum), cumsum inside blocks
    nb = max(1, min(NSUB, nin))
    edges = np.linspace(0, nin, nb + 1).astype(np.int64)
    A = np.empty(nin); B = np.empty(nin)
    offA = Fraction(0); offB = Fraction(0); cerrA = 0.0; cerrB = 0.0
    for s, e in zip(edges[:-1].tolist(), edges[1:].tolist()):
        if e == s: continue
        ca = np.cumsum(ta[s:e]); cb = np.cumsum(tb[s:e])
        A[s:e] = (Ahat + float(offA)) + ca
        B[s:e] = (Bhat + float(offB)) + cb
        n = e - s
        cerrA = max(cerrA, n * EPS * float(ca[-1])); cerrB = max(cerrB, n * EPS * float(cb[-1]))
        offA += Fraction(math.fsum(ta[s:e])); offB += Fraction(math.fsum(tb[s:e]))
    SA = math.fsum(ta); SB = math.fsum(tb)
    # local error bound for A_i, B_i (term errors, cumsum inside a block, offset roundings, additions)
    eA = 4 * EPS * SA + cerrA + nb * EPS * SA + 4 * EPS * (abs(Ahat) + SA)
    eB = 6 * EPS * SB + cerrB + nb * EPS * SB + 4 * EPS * (abs(Bhat) + SB)
    a = np.log(qf[:nin]); b = np.log(qf[1:nin + 1])
    ra, dra = r_and_dr(sq[:nin], a); rb, drb = r_and_dr(sq[1:nin + 1], b)
    ma = ra - 2 * A * a + 2 * B; mb = rb - 2 * A * b + 2 * B
    dma = dra - 2 * A; dmb = drb - 2 * A
    w = b - a
    lb = np.maximum(ma - np.maximum(0, -dma) * w, mb - np.maximum(0, dmb) * w)
    ub = np.maximum(ma, mb)
    # rounding error of the evaluation (operands up to |8 sq|, |c0 L|, |2 A L|, |2B|), error of log
    # (<= 2 ulp, independent of sqrt), tau truncation (< 1e-30 for N >= 2^?; bounded by first omitted term)
    mag = 8 * sq[1:nin + 1] + C0 * b + 10 + 2 * np.abs(A) * b + 2 * np.abs(B)
    err = 10 * EPS * mag + (2 * np.abs(A) + C0 + 1) * 4 * EPS * b + 2 * b * eA + 2 * eB \
        + (np.abs(dma) + np.abs(dmb)) * (w * 1e-12 + 8 * EPS * b) + 2 * np.exp(-18.5 * a) / 18.5 ** 2 / (1 - np.exp(-2 * a))
    lb = lb - err; ub = ub + err
    # sub-blocks: consecutive intervals grouped in bins of width HB in L (so the linear correction loses <= 2|dA| HB/2)
    gid = np.floor((a - a[0]) / HB).astype(np.int64)
    starts = np.concatenate([[0], np.nonzero(np.diff(gid))[0] + 1])
    ends = np.concatenate([starts[1:], [nin]])
    i1 = starts + np.array([int(np.argmin(lb[s:e])) for s, e in zip(starts.tolist(), ends.tolist())], dtype=np.int64)
    i2 = starts + np.array([int(np.argmax(ub[s:e])) for s, e in zip(starts.tolist(), ends.tolist())], dtype=np.int64)
    aL = a[starts]; bL = b[ends - 1]
    rows = np.stack([lb[i1], ub[i2], (aL + bL) / 2, (bL - aL) / 2 * (1 + 1e-12) + 4 * EPS * bL, qf[i1], qf[i2]], axis=1)
    np.savez(fn + ".tmp.npz", rows=rows, SA=SA, SB=SB, Ahat=Ahat, Bhat=Bhat, n=nin,
             q0=qf[0], qnext=qf[nin], ulpA=math.ulp(SA), ulpB=math.ulp(SB))
    os.replace(fn + ".tmp.npz", fn)
    return idx


def check_log():
    """numpy log within 2 ulp of the correctly rounded value, on 200k random integers + powers of 2 neighbours."""
    getcontext().prec = 40
    rng = np.random.default_rng(1)
    xs = np.concatenate([rng.integers(2, 10 ** 13, 200000), 2 ** np.arange(1, 44), 2 ** np.arange(1, 44) - 1, 2 ** np.arange(1, 44) + 1]).astype(np.int64)
    lg = np.log(xs.astype(np.float64))
    worst = 0.0
    for x, v in zip(xs.tolist()[:20000] + xs.tolist()[-129:], lg.tolist()[:20000] + lg.tolist()[-129:]):
        ex = Decimal(x).ln()
        u = math.ulp(v)
        worst = max(worst, float(abs(Decimal(v) - ex) / Decimal(u)))
    return worst


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--X", type=float, default=X_MAX, help="upper end")
    ap.add_argument("--seg", type=float, default=SEG_LEN, help="segment length")
    ap.add_argument("--workers", type=int, default=WORKERS or max(1, (os.cpu_count() or 2) - 1))
    ap.add_argument("--ckpt", default="margin_ckpt")
    ap.add_argument("--report", type=float, nargs="*", default=[1e3, 1e6, 1e8, 1e10, 1e11])
    args = ap.parse_args()
    X = int(args.X); SEG = int(args.seg)
    assert X < 2 ** 52 and SEG % 2 == 0
    print(f"Cesaro margin verification up to N = {X:.3e}, segments of {SEG:.1e}, {args.workers} workers")
    w = check_log()
    print(f"numpy log: worst error on test set = {w:.3f} ulp (the bound used is 2 ulp)")
    if w > 2: sys.exit("numpy log less accurate than assumed; aborting")
    os.makedirs(args.ckpt, exist_ok=True)
    meta = os.path.join(args.ckpt, "meta.json")
    if os.path.exists(meta):
        m = json.load(open(meta))
        if m["SEG"] != SEG: sys.exit(f"checkpoint uses seg={m['SEG']}; use --seg {m['SEG']} or delete {args.ckpt}")
    json.dump({"SEG": SEG}, open(meta, "w"))
    nseg = (X - 2) // SEG + 1
    jobs = [(i, 2 + i * SEG, 2 + (i + 1) * SEG, args.ckpt) for i in range(nseg)]
    todo = [j for j in jobs if not os.path.exists(os.path.join(args.ckpt, f"seg_{j[0]:06d}.npz"))]
    print(f"{nseg} segments, {nseg - len(todo)} already done")
    t0 = time.time()
    if todo:
        if args.workers > 1:
            from multiprocessing import Pool
            with Pool(args.workers, initializer=init_worker, initargs=(X, SEG)) as pool:
                for c, _ in enumerate(pool.imap_unordered(process, todo), 1):
                    el = time.time() - t0
                    if c % 10 == 0 or c == len(todo):
                        print(f"  {c}/{len(todo)} segments, {el / 60:.1f} min, ETA {el / c * (len(todo) - c) / 60:.1f} min", flush=True)
        else:
            init_worker(X, SEG)
            for c, j in enumerate(todo, 1):
                process(j)
                el = time.time() - t0
                if c % 10 == 0 or c == len(todo):
                    print(f"  {c}/{len(todo)} segments, {el / 60:.1f} min, ETA {el / c * (len(todo) - c) / 60:.1f} min", flush=True)
    # ------------------------------------------------ assembly: exact prefix sums and corrections
    print("assembling ...")
    A0 = Fraction(0); B0 = Fraction(0); eApre = 0.0; eBpre = 0.0
    rep = sorted(r for r in args.report if r < X) or [2.0]
    best = {r: (math.inf, None) for r in rep}; worst_ub = {r: (-math.inf, None) for r in rep}
    nterms = 0; last_next = None
    for i in range(nseg):
        d = np.load(os.path.join(args.ckpt, f"seg_{i:06d}.npz"))
        if last_next is not None: assert float(d["q0"]) == last_next, "gap between segments"
        last_next = float(d["qnext"])
        dA = float(A0 - Fraction(float(d["Ahat"]))); dB = float(B0 - Fraction(float(d["Bhat"])))
        rows = d["rows"]; ok = np.isfinite(rows[:, 0])
        rows = rows[ok]
        Lc, hw = rows[:, 2], rows[:, 3]
        corr_err = 2 * (Lc + hw) * eApre + 2 * eBpre + 8 * EPS * (abs(dA) * (Lc + hw) + abs(dB)) + 4 * EPS * (np.abs(rows[:, 0]) + 1)
        lbf = rows[:, 0] + 2 * dB - 2 * dA * Lc - 2 * abs(dA) * hw - corr_err
        ubf = rows[:, 1] + 2 * dB - 2 * dA * Lc + 2 * abs(dA) * hw + corr_err
        for r in rep:
            s = (Lc + hw) > math.log(r)     # sub-blocks reaching beyond r
            if s.any():
                j = int(np.argmin(np.where(s, lbf, np.inf))); j2 = int(np.argmax(np.where(s, ubf, -np.inf)))
                if lbf[j] < best[r][0]: best[r] = (float(lbf[j]), float(rows[j, 4]))
                if ubf[j2] > worst_ub[r][0]: worst_ub[r] = (float(ubf[j2]), float(rows[j2, 5]))
        A0 += Fraction(float(d["SA"])); B0 += Fraction(float(d["SB"]))
        eApre += 4 * EPS * float(d["SA"]) + float(d["ulpA"]); eBpre += 6 * EPS * float(d["SB"]) + float(d["ulpB"])
        nterms += int(d["n"])
    alpha = float(A0) - 2 * math.sqrt(last_next) + C0 / 2
    res = {"X": X, "covered_up_to": last_next, "prime_powers": nterms, "A(X)": float(A0), "B(X)": float(B0),
           "error_bound_A": eApre, "error_bound_B": eBpre,
           "least_lower_bound": {f"N>={r:g}": {"bound": best[r][0], "near_N": best[r][1]} for r in rep},
           "largest_upper_bound": {f"N>={r:g}": {"bound": worst_ub[r][0], "near_N": worst_ub[r][1]} for r in rep},
           "minutes": (time.time() - t0) / 60}
    json.dump(res, open("margin_result.json", "w"), indent=1)
    print(f"\nprime powers up to {last_next:.6e}: {nterms}   A = {float(A0):.10f}  (error <= {eApre:.1e})")
    for r in rep:
        print(f"  N in [{r:g}, {X:g}]:  {best[r][0]:.6f} <= M(N) <= {worst_ub[r][0]:.6f}"
              f"   (min near N={best[r][1]:.0f}, max near N={worst_ub[r][1]:.0f})")
    ok = all(best[r][0] > 0 for r in rep) and all(worst_ub[r][0] < 0.18484 for r in rep)
    print("\nRESULT:", "0 < M(N) < 0.18484 on the whole range" if ok else "the sandwich is NOT certified on part of the range -- see margin_result.json")
    print("Written margin_result.json")


if __name__ == "__main__":
    main()