"""Rigorous (Arb ball arithmetic) upper bound for the loss from restricting F_psi = psi(S) prod g(t_i)
(0 <= psi <= 1) to the region R = R_k(u0, kappa) = {sum t <= 1, (3/2) v + sum_{t_i >= v} t_i <= kappa for
u0 < v <= max t}: first inequality of Lemma 4.3 of the paper (used for m = 2).
For t' = (t_1..t_{k-1}) with s' = sum t', Psi(s') = int_0^T psi(s'+t) g(t) dt and Psi_R(t') is the same integral
over the fibre {t : (t', t) in R}, an initial segment of [0, T] since R is a down-set.  Then
    Q(F_psi 1_R) >= Q(F_psi) - (k/m2) E[Psi(S')^2 - Psi_R(X')^2] / E[psi(S_k)^2]           (Lemma 4.2)
and, with cells [v_l, v_{l+1}] (v_0 = u0, v_{l+1} = 1.02 v_l, v_L = T), A_l = kappa - (3/2) v_{l+1},
    E[Psi^2 - Psi_R^2] <= sum_l G(T)^2 P(N_l >= n^a_l + 1, R_l <= b^a_l)
                          + sum_l 2 G(T) (G(T) - G(v_l)) P(N_l >= n^b_l + 1, R_l <= b^b_l),
    n^a_l = floor(A_l/T), n^b_l = floor((A_l - T)/T),
    b^a_l = min(1 - A_l, 1 - (n^a_l+1) v_l),  b^b_l = min(1 - A_l, 1 - (n^b_l+2) v_l),
where N_l = #{i < k : X_i >= v_l} and R_l = sum of the X_i < v_l (case (a): t' alone violates the region
condition and the whole fibre is lost; case (b): the violation needs t_k >= v >= v_l).  The probabilities are
bounded as in rc_loss_rig.py ((4.2) of the paper); floors are rounded down and eta = 1e-4 is added to the b's.
The assertion T < kappa/(1 + 3/2) = 2 kappa/5 is not needed for validity; it guarantees A_l - T > 0."""
import sys, json, math
import numpy as np
from flint import arb, ctx
from rc_loss_rig import mgf_trunc_arb, mgf_trunc_f, kl_tail
ctx.prec = 128

def chernoff_lower(m, b_l, c, b, wl, q, m2, cf, bf, m2f, lams):
    """Upper bound for P(sum of m iid (X | X < wl) < b_l), Arb; returns arb in [0,1]."""
    if float(b_l.upper()) <= 0:
        return arb(0)
    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()
    return logPS.exp() if float(logPS.upper()) < 0 else arb(1)

def rc_loss_rig2(k, theta, beta, delta, omega, Eh2_lower, eta=1e-4, ufrac=0.5, ratio=1.02, cv=1.5):
    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)
    s = arb(cv)
    assert float(((s + 1) * T).upper()) < float(kappa.lower()), "need T < kappa/(s+1)"
    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))
    G = lambda w: (1 + b * w / c).log() / b
    GT = G(T)
    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
    Pa_tot = arb(0); Lb_tot = arb(0)
    for l in range(len(grid) - 1):
        wl, wn = grid[l], grid[l + 1]
        q = cdf(wl); p = 1 - q
        pf = float(p.mid())
        meanN = K * pf; sdN = math.sqrt(K * pf * (1 - pf))
        nmax = min(int(meanN + 14 * sdN + 20), K - 1)
        m = K - nmax
        tailNmax = kl_tail(K, nmax + 1, p)
        # case (a)
        A_a = kappa - s * wn
        n0a = int(math.floor(np.nextafter(float((A_a / T).lower()), -np.inf)))   # float() rounds to nearest: step down
        b_a = arb(min(float((1 - A_a + arb(eta)).upper()), float((1 - (n0a + 1) * wl + arb(eta)).upper())))
        Pa = kl_tail(K, n0a + 1, p) * chernoff_lower(m, b_a, c, b, wl, q, m2, cf, bf, m2f, lams) + tailNmax
        # case (b)
        A_b = kappa - s * wn - T
        n0b = int(math.floor(np.nextafter(float((A_b / T).lower()), -np.inf)))
        b_b = arb(min(float((1 - kappa + s * wn + arb(eta)).upper()), float((1 - (n0b + 2) * wl + arb(eta)).upper())))
        Pb = kl_tail(K, n0b + 1, p) * chernoff_lower(m, b_b, c, b, wl, q, m2, cf, bf, m2f, lams) + tailNmax
        Pa_tot += Pa
        Lb_tot += 2 * GT * (GT - G(wl)) * Pb
    Eloss = GT ** 2 * Pa_tot + Lb_tot
    loss = kk * Eloss / (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,
                Pa_upper=float(Pa_tot.upper()), Lb_upper=float(Lb_tot.upper()), loss_upper=float(np.nextafter(float(loss.upper()), np.inf)))

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