"""Certified bounds on the closed-form rate of Algorithm RS.

F*(sigma) = H(sigma) + (1 - sigma) * log2(1 - 2^(-1/sigma))  (0 < sigma < 1),
the minimum of the one-sample-size function with beta = 1 and rho = 1/sigma
(equation (eq:co-Fstar) of the paper with p = 2^(-1/sigma) < sigma).

Checks, in ball arithmetic with outward rounding (Arb through python-flint, 256 bits):
  (1) F*(sigma) <= log2(1.744835)  for sigma in [0.2, 0.8]    (Proposition prop:rs-rate(c));
      outside [0.2, 0.8] the bound F*(sigma) <= H(sigma) <= H(0.2) < log2(1.744835) is elementary;
  (2) F*(sigma) <= 0.7807           for sigma in [0.15, 0.31]  (Proposition prop:rst-rate, mean-bound range);
  (3) F*(sigma) <= 0.7802           for sigma in [0.54, 0.80]  (same).
Each range is covered adaptively by closed intervals with exact rational endpoints; an interval is
accepted when the enclosure of F* on it lies below the enclosure of the target, and bisected
otherwise.  Every interval endpoint is an exact rational, so the covering is exact.
"""
import sys
from fractions import Fraction
from flint import arb, fmpq, ctx

ctx.prec = 256


def F_star(lo, hi):
    """Enclosure of F* on [lo, hi] (exact rationals)."""
    # a ball containing the whole interval [lo, hi]
    s = arb(fmpq(lo.numerator, lo.denominator)).union(arb(fmpq(hi.numerator, hi.denominator)))
    ln2 = arb(2).log()
    H = -(s * s.log() + (1 - s) * (1 - s).log()) / ln2
    p = (-(1 / s) * ln2).exp()           # 2^(-1/sigma)
    return H + (1 - s) * (1 - p).log() / ln2


def certify(lo, hi, target, name):
    """Cover [lo, hi] by intervals on which F* <= target is certified."""
    target_ball = arb(fmpq(target.numerator, target.denominator)) if isinstance(target, Fraction) else target
    stack = [(Fraction(lo), Fraction(hi))]
    count = 0
    worst = None
    min_len = None
    while stack:
        a, b = stack.pop()
        enc = F_star(a, b)
        if enc < target_ball:          # the whole enclosure lies below the whole target enclosure
            count += 1
            up = enc.upper()
            if worst is None or up > worst:
                worst = up
            if min_len is None or b - a < min_len:
                min_len = b - a
            continue
        if b - a < Fraction(1, 10**9):
            print(f"FAIL on {name}: interval [{a},{b}] not certified; enclosure {enc}")
            sys.exit(1)
        m = Fraction(a + b, 2)
        stack.append((a, m))
        stack.append((m, b))
    print(f"{name}: certified F* <= {target_ball} on [{lo},{hi}] with {count} intervals; "
          f"largest certified upper bound {worst}; shortest interval {float(min_len):.3e}")


if __name__ == "__main__":
    log_c = arb(fmpq(1744835, 1000000)).log() / arb(2).log()
    print("log2(1.744835) =", log_c)
    certify(Fraction(1, 5), Fraction(4, 5), log_c, "(1) RS constant")
    certify(Fraction(15, 100), Fraction(31, 100), Fraction(7807, 10000), "(2) RS-theta lower range")
    certify(Fraction(54, 100), Fraction(80, 100), Fraction(7802, 10000), "(3) RS-theta upper range")
    # elementary bound outside [0.2, 0.8]: H(0.2) < log2(1.744835)
    s = arb(fmpq(1, 5)); ln2 = arb(2).log()
    H02 = -(s * s.log() + (1 - s) * (1 - s).log()) / ln2
    print("H(0.2) =", H02, " < log2(1.744835):", H02 < log_c)
    if not (H02 < log_c):
        print("CERTIFICATE FAILED")
        sys.exit(1)
    print("CERTIFICATE OK")
    # numerical maximum (not part of the proof): golden-section search
    phi = (5 ** 0.5 - 1) / 2
    def f(x):
        import math
        return -(x*math.log2(x)+(1-x)*math.log2(1-x)) + (1-x)*math.log2(1-2**(-1/x))
    a, b = 0.3, 0.5
    for _ in range(200):
        c = b - phi*(b-a); d = a + phi*(b-a)
        if f(c) > f(d): b = d
        else: a = c
    print(f"numerical maximum (floating point, not a proof): F*({(a+b)/2:.7f}) = {f((a+b)/2):.10f}")