"""Rigorous rate certificate for component-count costs with c = 2 (lem:cc-cert), and the decimal
part (B3) of the c = 3 identities, in ball arithmetic.  Replaces the 50-digit decimal evaluation of
component_count_rate_certificate.py (whose exact-rational parts B1, B2 are repeated here).

Notation as in the paper: M_z = [[1, z], [1, 2z]], spr(M_z) its spectral radius, and
    phi_{2,sigma}(y) = log2(1 + y) - sigma log2(spr(M_y) / 2)          (c = 2, so z = cy/2 = y).

Arithmetic.  All parameters are exact rationals (python-flint fmpq).  Logarithms are evaluated in arb
ball arithmetic (each operation returns a ball provably containing the exact real result, i.e.
interval arithmetic with outward rounding); a check passes only if the upper endpoint of the
enclosure is below the bound.  The spectral radius is bracketed by exact rationals with the
Collatz-Wielandt inequalities for the positive matrix M_z: for every positive vector v,
    min_i (M_z v)_i / v_i  <=  spr(M_z)  <=  max_i (M_z v)_i / v_i ,
applied with v = (z, l - 1) for a rational l slightly below, and one slightly above, spr(M_z).  The
lower bracket is used in phi (phi decreases in spr), the upper one in the prefactor.  As a cross-check,
the closed form spr(M_z) = (1 + 2z + sqrt(1 + 4z^2))/2 is enclosed in arb and must lie in the bracket.

Checks (Part A, omega = 0.71755, sigma_0 = 0.7175, sigma_1 = 0.81):
  (A0) omega < log2(1.6444);
  (A1) sigma_0 * log2(2) = sigma_0 <= omega (exact);
  (A2) H(sigma_1) < omega and sigma_1 >= 1/2;
  (A3) 925 consecutive intervals [a_j, b_j] of width 1e-4 cover [sigma_0, sigma_1] (exact), each with
       a rational y_j > 0 such that max(phi_{2,a_j}(y_j), phi_{2,b_j}(y_j)) < omega, where phi is
       evaluated with the lower Collatz-Wielandt bound for spr(M_{y_j});
  (A4) upper Collatz-Wielandt bound of spr(M_{y_j}) / min(1, y_j) <= 16 (exact rationals).
The y_j are chosen in floating point (as in component_count_rate_certificate.py) and rounded to six
significant digits; only the rigorous checks matter.
Part B (c = 3):
  (B1) exact: lam(6/5) = 3 is the larger root of x^2 - (1 + 12/5) x + 6/5, and 1 + y* = 9/5;
  (B2) exact: 5^45 > 3^64;
  (B3) arb: H(26/45) > log2(9/5).
Usage: python3 component_count_rate_interval.py   (paths relative to this script; a few seconds)
Requires python-flint (tested with 0.9.0).
"""
import math
import sys
from fractions import Fraction
from flint import arb, fmpq, ctx

ctx.prec = 256
LN2 = arb(2).log()


def lg(x):
    return arb(x).log() / LN2


def H(p):
    p = arb(p)
    return -(p * p.log() + (1 - p) * (1 - p).log()) / LN2


def spr_closed(z):
    z = arb(z)
    return (1 + 2 * z + (1 + 4 * z * z).sqrt()) / 2


def cw_ratios(z, l):
    """Collatz-Wielandt ratios of M_z at v = (z, l - 1), exact; requires z > 0 and l > 1."""
    v0, v1 = z, l - 1
    r0 = (v0 + z * v1) / v0
    r1 = (v0 + 2 * z * v1) / v1
    return min(r0, r1), max(r0, r1)


def spr_bracket(z):
    """Exact rationals lo <= spr(M_z) <= hi via Collatz-Wielandt."""
    zf = float(z.p) / float(z.q)
    s = (1 + 2 * zf + math.sqrt(1 + 4 * zf * zf)) / 2
    lo = hi = None
    for eps in (1e-12, 1e-9, 1e-6):
        for l in (Fraction(s - eps * s), Fraction(s + eps * s), Fraction(s)):
            l = fmpq(l.numerator, l.denominator)
            if l <= 1:
                continue
            a, b = cw_ratios(z, l)
            lo = a if lo is None else max(lo, a)
            hi = b if hi is None else min(hi, b)
        if lo is not None and hi is not None and hi - lo < fmpq(1, 10 ** 8):
            break
    return lo, hi


def phif(rho, y):
    s = (1 + 2 * y + math.sqrt(1 + 4 * y * y)) / 2
    return math.log2(1 + y) - rho * math.log2(s / 2)


def best_y(rho):
    grid = [-6 + 10 * i / 2000 for i in range(2001)]
    s = min(grid, key=lambda t: phif(rho, math.exp(t)))
    step = 10 / 2000
    for _ in range(80):
        s = min((s - step, s, s + step), key=lambda t: phif(rho, math.exp(t)))
        step /= 1.5
    return "%.6g" % math.exp(s)


def dec(s):
    s = s.strip()
    if "e" in s.lower():
        m, e = s.lower().split("e")
        return dec(m) * fmpq(10) ** int(e) if int(e) >= 0 else dec(m) / fmpq(10) ** (-int(e))
    if "." in s:
        ip, fp = s.split(".")
        return fmpq(int(ip + fp), 10 ** len(fp))
    return fmpq(int(s))


def part_a():
    W = fmpq(71755, 10 ** 5)
    rho0, rho1 = fmpq(7175, 10 ** 4), fmpq(81, 100)
    a0 = arb(W) < lg(fmpq(16444, 10 ** 4))
    a1 = rho0 <= W
    h1 = H(rho1)
    a2 = (h1 < arb(W)) and rho1 >= fmpq(1, 2)
    ok = a0 and a1 and a2
    print("Part A (c = 2), omega = 0.71755")
    print(f"  (A0) log2(1.6444) in {lg(fmpq(16444, 10 ** 4)).str(15, radius=True)} > omega: {a0}")
    print(f"  (A1) sigma_0 = 0.7175 <= omega (exact): {a1}")
    print(f"  (A2) H(0.81) in {h1.str(15, radius=True)} < omega: {a2}")
    h = fmpq(1, 10 ** 4)
    a = rho0
    count, fails = 0, 0
    worst = None
    worst_pref = fmpq(0)
    worst_gap = None
    while a < rho1:
        b = a + h
        y = dec(best_y(float((a + b).p) / float((a + b).q) / 2))
        lo, hi = spr_bracket(y)
        cl = spr_closed(y)
        consistent = (arb(lo) <= cl) and (cl <= arb(hi))
        u = None
        for s in (a, b):
            v = lg(1 + y) - arb(s) * lg(lo / 2)
            if u is None or v.upper() > u.upper():
                u = v
        pref = hi / min(fmpq(1), y)
        good = y > 0 and (u < arb(W)) and pref <= 16 and consistent
        if not good:
            fails += 1
            print("FAIL", a, b, y, u, pref, consistent)
        if worst is None or u.upper() > worst[0].upper():
            worst = (u, a, b, y)
        worst_pref = max(worst_pref, pref)
        gap = hi - lo
        worst_gap = gap if worst_gap is None else max(worst_gap, gap)
        a = b
        count += 1
    cover = a == rho1
    print(f"  (A3) intervals: {count}, width 1e-4, cover [0.7175, 0.81] exactly: {cover}; failures: {fails}")
    print(f"       largest max(phi(a_j,y_j), phi(b_j,y_j)) in {worst[0].str(15, radius=True)}")
    print(f"       attained on [{worst[1]}, {worst[2]}] with y = {worst[3]}")
    print(f"       certified margin omega - max >= {(arb(W) - worst[0].upper()).str(5)}")
    print(f"       widest Collatz-Wielandt bracket for spr: {float(worst_gap):.3e}")
    print(f"  (A4) max_j (upper bound of spr(M_y_j)) / min(1, y_j) = {float(worst_pref):.6f} (<= 16): "
          f"{worst_pref <= 16}")
    return ok and cover and fails == 0


def part_b():
    c, z = fmpq(3), fmpq(6, 5)
    b1 = c * c - (1 + 2 * z) * c + z == 0 and z / c < c and 1 + 2 * z / c == fmpq(9, 5)
    b2 = 5 ** 45 > 3 ** 64
    hb, lb = H(fmpq(26, 45)), lg(fmpq(9, 5))
    b3 = hb > lb
    print("Part B (c = 3)")
    print(f"  (B1) lam(6/5) = 3 and 1 + y* = 9/5 (exact): {b1}")
    print(f"  (B2) 5^45 > 3^64 (exact): {b2}")
    print(f"  (B3) H(26/45) in {hb.str(15, radius=True)} > log2(9/5) in {lb.str(15, radius=True)}: {b3}")
    return b1 and b2 and b3


if __name__ == "__main__":
    ok = part_a()
    ok = part_b() and ok
    print("CERTIFICATE OK" if ok else "CERTIFICATE FAILED")
    sys.exit(0 if ok else 1)