"""Rigorous (Arb ball arithmetic) upper bound for the loss from restricting the trial function
F_psi = psi(t_1+...+t_k) prod g(t_i) (0 <= psi <= 1) to the region R = R_k(u0, kappa) of the paper
(second inequality of Lemma 4.3, used for m >= 3):
    Q(F_psi 1_R) >= Q(F_psi) - k G(T)^2 sum_l P(N_l >= n^b_l + 1, R_l <= 1 - A_l) / (m2 E[psi(S_k)^2]),
    A_l = kappa - (3/2) v_{l+1},  n^b_l = floor((A_l - T)/T),
and each probability is bounded as in (4.2):
    P(N_l >= n+1, R_l <= b') <= P(N_l >= n+1) P(Z_1+...+Z_{K-nmax} <= b') + P(N_l > nmax),   K = k-1,
with the binomial tail P(Bin(K,p) >= n) <= exp(-K D(n/K || p)) (Hoeffding) and the Chernoff bound
P(Z_1+...+Z_M <= b') <= exp(lam b') E[e^{-lam Z}]^M, E[e^{-lam Z}] in closed form via E_1.
The argument Eh2_lower is a rigorous lower bound for E[psi(S_k)^2] (from hcert.py).  cv is the 3/2 of the
region condition; eta = 1e-4 is the slack added to 1 - A_l and the offset in kappa."""
import sys, json, math
import numpy as np
from flint import arb, ctx
from scipy.special import exp1
ctx.prec = 128

def E1(z):
    return z.expint(arb(1))

def mgf_trunc_arb(lam, c, b, w):
    """int_0^w e^{-lam t} (c+bt)^{-2} dt, closed form, ball arithmetic."""
    a = lam / b
    def prim(u):
        return -(-(lam) * (u - c) / b).exp() / u + a * (a * c).exp() * E1(a * u)
    return (prim(c + b * w) - prim(c)) / b

def mgf_trunc_f(lam, c, b, w):
    a = lam / b
    def prim(u):
        return -np.exp(-a * (u - c)) / u + a * np.exp(a * c) * exp1(a * u)
    return (prim(c + b * w) - prim(c)) / b

def kl_tail(K, n, p):
    """Upper bound for P(Bin(K,p) >= n), ball arithmetic (returns arb)."""
    if n <= 0:
        return arb(1)
    x = arb(n) / K
    if float(x.lower()) <= float(p.upper()):
        return arb(1)
    if n > K:
        return arb(0)
    if n == K:
        return p ** K
    D = x * (x / p).log() + (1 - x) * ((1 - x) / (1 - p)).log()
    return (-(K * D)).exp()

def rc_loss_rig(k, theta, beta, delta, omega, Eh2_lower, eta=1e-4, ufrac=0.5, ratio=1.02, cv=1.5):
    """cv = r/2: 1.5 for triply densely divisible moduli (r=3, f=g=t^{3/2}), 2.0 for r=4 (f=g=t^2)."""
    kk = arb(k); L = kk.log(); c = arb(theta) / L; T = arb(beta) / L; b = kk - 1
    m2 = (1 / b) * (1 / c - 1 / (c + b * T))
    rho = arb(1) / 4 + arb(omega)
    u0 = arb(ufrac) * arb(delta) / rho
    kappa = (arb(1) / 4 + arb(delta) / 2) / rho - arb(eta)
    cf = float(c.mid()); bf = float(b.mid()); m2f = float(m2.mid()); Tf = float(T.mid())
    cdf = lambda w: (1 / (m2 * b)) * (1 / c - 1 / (c + b * w))
    grid = [u0]
    while float(grid[-1].mid()) < Tf:
        nxt = grid[-1] * arb(ratio)
        grid.append(T if float(nxt.mid()) >= Tf else nxt)
    lams = np.exp(np.linspace(0, math.log(1e7), 600))
    K = k - 1
    total = arb(0); worst = None
    for l in range(len(grid) - 1):
        wl, wn = grid[l], grid[l + 1]
        a_l = kappa - arb(cv) * wn - T
        b_l = (1 - kappa + arb(eta)) + arb(cv) * wn
        q = cdf(wl); p = 1 - q
        # violation needs > a_l/T coordinates >= w_l among the k-1 others, i.e. N >= floor(a_l/T) + 1.  A smaller
        # n0 only enlarges the event, so round down: lower end of a_l/T, then one ulp down (float() rounds to nearest)
        n0 = int(math.floor(np.nextafter(float((a_l / T).lower()), -np.inf))) if float(a_l.lower()) > 0 else -1
        pf = float(p.mid())
        meanN = K * pf; sdN = math.sqrt(K * pf * (1 - pf))
        nmax = min(int(meanN + 14 * sdN + 20), K - 1)
        tailN = kl_tail(K, n0 + 1, p)
        tailNmax = kl_tail(K, nmax + 1, p)
        m = K - nmax
        # choose lambda by float search, then evaluate rigorously
        wf = float(wl.mid()); qf = float(q.mid()); blf = float(b_l.mid())
        Ef = np.array([mgf_trunc_f(lam, cf, bf, wf) for lam in lams]) / (m2f * qf)
        vals = lams * blf + m * np.log(Ef)
        lam = float(lams[int(np.argmin(vals))])
        Ee = mgf_trunc_arb(arb(lam), c, b, wl) / (m2 * q)
        logPS = arb(lam) * b_l + m * Ee.log()
        PS = logPS.exp() if float(logPS.upper()) < 0 else arb(1)
        Pl = tailN * PS + tailNmax
        total += Pl
        if worst is None or float(Pl.upper()) > float(worst[0].upper()):
            worst = (Pl, float(wl.mid()), n0, float(tailN.upper()), float(logPS.upper()))
    GT = (1 + b * T / c).log() / b
    loss = kk * GT ** 2 * total / (m2 * arb(Eh2_lower))
    return dict(k=k, theta=theta, beta=beta, delta=delta, omega=omega, cv=cv, kappa=float(kappa.mid()),
                cells=len(grid) - 1, Pbad_upper=float(total.upper()), loss_upper=float(loss.upper()),
                worst_cell=dict(P=float(worst[0].upper()), w=worst[1], n0=worst[2], tailN=worst[3], logPS=worst[4]))

if __name__ == "__main__":
    if len(sys.argv) < 8:
        sys.exit("usage: rc_loss_rig.py k theta beta delta omega Eh2_lower cv   (cv = 1.5 in the paper)")
    k, th, be, d, om, Pl = int(sys.argv[1]), float(sys.argv[2]), float(sys.argv[3]), float(sys.argv[4]), float(sys.argv[5]), float(sys.argv[6])
    cv = float(sys.argv[7])
    print(json.dumps(rc_loss_rig(k, th, be, d, om, Pl, cv=cv)), flush=True)