"""Arb enclosure of F_4 = sup_tau f(tau), f(tau) = Pr(Y_0 > tau, min_{i<=4} Y_i <= tau) for the ideal 4-star
(Y_i = r Y_0 + s U_i, r = sqrt(3)/2 = r_44, s = 1/2, U the regular simplex vector with correlations -1/3).
See Theorem 7.1(a) and Proposition 6.9 of the paper.

With G_d(h) = Pr(max of d simplex variables <= h) (G_d(h) = 0 for h <= 0):
    G_2(h) = erf(h/sqrt 2),   G_d(h) = int_0^h d phi(x) G_{d-1}(sqrt(d/(d-2)) x) dx.          (Lemma K3)
Conditioning on Y_0 = z gives  f(tau) = Phi(-tau) - I(tau),  I(tau) = int_{z>tau} phi(z) G_4((r z - tau)/s) dz.
Writing G_4(h) = int_0^h g(x) dx with g(x) = 4 phi(x) G_3(sqrt 2 x) and exchanging the order of integration (Tonelli):
    I(tau) = Phi(-tau) G_4(x0) + int_{x0}^inf g(x) Phi(-(tau + s x)/r) dx,    x0 = max(0, (r-1) tau/s),
    I'(tau) = -phi(tau) G_4(x0) - (1/r) int_{x0}^inf g(x) phi((tau + s x)/r) dx,
    f'(tau) = -phi(tau) - I'(tau).
(The terms from differentiating the limit x0 cancel because (tau + s x0)/r = tau when x0 > 0 and G_4(0) = 0.)
Integrals are Arb integrals of entire functions over [0,1] after x = x0 + (X - x0) w; cutting the upper limit at
X = 10 lowers I by at most 4 Phi(-X), added as an interval to the enclosure (and 4 phi(0)... for I').
Upper bound for sup f: on [-2.5, 0.7] cells [t0,t1] with midpoint m, f(tau) <= f(m) + max(D_hi (t1-m), D_lo (t0-m)),
[D_lo, D_hi] an enclosure of f' on the cell (mean value theorem).
For tau >= 0.7: f <= Phi(-tau) <= Phi(-0.7); for tau <= -2.5: f <= 4 Phi(-2.5) (union bound).
Run: .venv/bin/python experiments/arb-certificates/certify_F4.py [upper bound]
"""
import sys, time
from flint import arb, acb, ctx

ctx.prec = 64
SQ2PI = (2 * arb.pi()).sqrt()
R = arb(3).sqrt() / 2
S = arb(1) / 2
XMAX = 10
TOL = arb(10) ** -13


def phi(x):
    return (-(x * x) / 2).exp() / SQ2PI


def Phi(x):                       # x: arb
    return (1 + (x / arb(2).sqrt()).erf()) / 2


def Phi_c(x):                     # x: acb
    return (1 + (x / acb(2).sqrt()).erf()) / 2


def G2(y):
    return (y / acb(2).sqrt()).erf()


def _int(f):
    return acb.integral(f, 0, 1, rel_tol=TOL, abs_tol=TOL)


def G3(y):                        # y: acb; int_0^1 y*3*phi(y v) G2(sqrt3 y v) dv
    s3 = acb(3).sqrt()
    return _int(lambda v, an: y * 3 * phi(y * v) * G2(s3 * y * v))


def G4(y):
    s2 = acb(2).sqrt()
    return _int(lambda u, an: y * 4 * phi(y * u) * G3(s2 * y * u))


def dG4(y):                       # G_4'
    return 4 * phi(y) * G3(acb(2).sqrt() * y)


def x0_of(tau, neg):
    """x0 = max(0, (r-1) tau/s); `neg` says (from exact endpoints) whether tau <= 0 on the whole cell."""
    if neg:
        return ((R - 1) * tau / S).intersection(arb(0).union(arb(10) ** 6))
    return arb(0)


def xint(tau, kernel, neg):
    """int_{x0}^{XMAX} g(x) kernel(x) dx for tau an arb (point or ball), g = 4 phi(x) G_3(sqrt 2 x)."""
    x0 = acb(x0_of(tau, neg)); L = XMAX - x0; s2 = acb(2).sqrt()
    return _int(lambda w, an: L * 4 * phi(x0 + L * w) * G3(s2 * (x0 + L * w)) * kernel(x0 + L * w)).real


def f_enclosure(tau, neg=None):
    neg = tau.upper() <= 0 if neg is None else neg
    x0 = x0_of(tau, neg)
    I = Phi(-tau) * G4(acb(x0)).real + xint(tau, lambda x: Phi_c(-(tau + S * x) / R), neg)
    I = I + arb(0).union(4 * Phi(-arb(XMAX)))
    return Phi(-tau) - I


def fprime_enclosure(tau, neg):
    x0 = x0_of(tau, neg)
    J = xint(tau, lambda x: phi((tau + S * x) / R), neg) + arb(0).union(4 * phi(arb(0)) * Phi(-arb(XMAX)))
    return -phi(tau) + phi(tau) * G4(acb(x0)).real + J / R


def cell_bound(t0, t1, neg):
    """Upper bound for sup of f on the cell [t0, t1] (arb points, not straddling 0)."""
    m = (t0 + t1) / 2
    fm = f_enclosure(m, neg)
    D = fprime_enclosure(t0.union(t1), neg)
    up = fm.upper() + max(D.upper() * (t1 - m), D.lower() * (t0 - m))
    return up


def dyadic(fr):
    return arb(fr.numerator) / fr.denominator


def _work(cell):
    a, b = cell
    import numpy as np
    u = cell_bound(dyadic(a), dyadic(b), b <= 0)
    return float(np.nextafter(float(u.mid()) + float(u.rad()), np.inf))


def main(ub_str, lb_tau="-0.16367"):
    import multiprocessing as mp, json
    from fractions import Fraction
    import numpy as np
    UB = arb(ub_str); UBf = float(np.nextafter(float(ub_str), -np.inf))
    print("tail bounds: Phi(-0.7) =", Phi(arb(-7) / 10).str(10), "; 4 Phi(-2.5) =", (4 * Phi(arb(-5) / 2)).str(10))
    lower = f_enclosure(arb(lb_tau))
    print("lower bound: f(%s) = %s" % (lb_tau, lower.str(20, radius=True)))
    assert Phi(arb(-7) / 10) < UB and 4 * Phi(arb(-5) / 2) < UB
    w = Fraction(1, 16)
    pending = []
    x = Fraction(0)
    while x < Fraction(7, 10):
        pending.append((x, min(x + w, Fraction(7, 10)))); x += w
    x = Fraction(-5, 2)
    while x < 0:
        pending.append((x, x + w)); x += w
    done = []; t_start = time.time()
    pool = mp.get_context("fork").Pool(8)
    while pending:
        ups = pool.map(_work, pending, chunksize=1)
        nxt = []
        for (a, b), up in zip(pending, ups):
            if up < UBf:
                done.append((a, b, up))
            elif b - a < Fraction(1, 2 ** 24):
                print("FAILED cell", float(a), float(b), up); return
            else:
                mid = (a + b) / 2
                nxt += [(a, mid), (mid, b)]
        pending = nxt
        print("%d certified, %d to refine, %.0fs" % (len(done), len(pending), time.time() - t_start), flush=True)
    done.sort()
    assert done[0][0] == Fraction(-5, 2) and done[-1][1] == Fraction(7, 10)
    assert all(done[i][1] == done[i + 1][0] for i in range(len(done) - 1))
    top = max(done, key=lambda c: c[2])
    print("certified cells:", len(done), "tiling [-2.5, 0.7]; sup f <=", UB.str(12))
    print("largest cell bound:", top[2], "on", float(top[0]), float(top[1]), "; smallest cell", float(min(c[1] - c[0] for c in done)))
    print("ENCLOSURE:", lower.lower().str(15), "<= F_4 <=", UB.str(15))
    json.dump([[str(a), str(b), up] for a, b, up in done], open(__file__.replace("certify_F4.py", "F4-cells.json"), "w"))


if __name__ == "__main__":
    main(*(sys.argv[1:] or ["0.24577197475"]))