"""Independent verifier re-computation of the rate claims of Sections 9 and 10 of the paper.

V1  lem:det-cert: exact rational endpoint check, one y per interval, both endpoints.
V2  numerical sanity check, not a verification: dense u-grid with a bounded (local) scalar
    minimisation in log2 y at each grid point, so each value estimates min_y ell_y(u) from above;
    compared with log2(sqrt5-1/2) on [0.15, 0.85] up to a stated tolerance; location of the max.
    The bound is attained with equality on [1/3, 1/2], so the check is equality-tight.
V3  lem:det-time at finite n: max over all guesses (k-, W-) of alg:dr-integer (fixed guess range:
    1 <= k- <= s, 0 <= W- <= a k-, W+ >= (a+1) k+ if k+ >= 1, W+ = 0 if k+ = 0) of the best (t-, t+)
    bound, compared with n log2(sqrt5 - 1/2).
V4  balanced-profile finite-n value at n = 1000 (claimed 1.735878).
V5  black-box numbers: H(1/4); 1.8107 example; thm:dr-bb-moments at sigma = 0.4176 (worst sigma of
    Algorithm RS) and at sigma = 0.4172629 (where the certified bound of Algorithm RS-theta is attained).
"""
import math, sys
from fractions import Fraction as Fr
import numpy as np
from scipy.special import gammaln
from scipy.optimize import minimize_scalar

L2 = math.log2
LN2 = math.log(2)
C = math.sqrt(5) - 0.5
LC = L2(C)


def H(x):
    return 0.0 if x <= 0 or x >= 1 else -x * L2(x) - (1 - x) * L2(1 - x)


def ell(y, u, a=None):
    if a is None:
        a = math.floor(1 / u)
    return L2(1 + y) + 1 - ((a + 1) * u - 1) * L2(1 + y * 2 ** a) - (1 - a * u) * L2(1 + y * 2 ** (a + 1))


def lb(n, k):  # log2 binom
    if k < 0 or k > n:
        return -math.inf
    return (gammaln(n + 1) - gammaln(k + 1) - gammaln(n - k + 1)) / LN2


# ---------------- V1 ----------------
def v1():
    print('V1  exact endpoint check, bound c = 1.7360 (rational)')
    c = Fr(17360, 10000)
    assert float(c) < C
    ys = [Fr(1, 2), Fr(5, 16), Fr(1, 4), Fr(27, 128), Fr(29, 160)]
    ivs = {1: (Fr(1, 2), Fr(85, 100)), 3: (Fr(1, 4), Fr(1, 3)), 4: (Fr(1, 5), Fr(1, 4)),
           5: (Fr(1, 6), Fr(1, 5)), 6: (Fr(15, 100), Fr(1, 6))}

    def ok(y, u, a):
        c1 = (a + 1) * u - 1
        c2 = 1 - a * u
        assert c1 >= 0 and c2 >= 0
        D = math.lcm(c1.denominator, c2.denominator)
        lhs = (2 * (1 + y)) ** D
        rhs = c ** D * (1 + y * 2 ** a) ** int(c1 * D) * (1 + y * 2 ** (a + 1)) ** int(c2 * D)
        return lhs <= rhs

    allok = True
    for a, (lo, hi) in ivs.items():
        good = [y for y in ys if ok(y, lo, a) and ok(y, hi, a)]
        print(f'   a={a} [{float(lo):.4f},{float(hi):.4f}]: y valid at both endpoints: {[str(y) for y in good]}')
        allok &= bool(good)
    # a = 2 algebraic identity, numerically
    y0 = (1 + math.sqrt(5)) / 8
    print(f'   a=2: 2^ell(y0) at u=1/3: {2**ell(y0, 1/3, 2):.15f}, at u=1/2: {2**ell(y0, 1/2, 2):.15f}, sqrt5-1/2={C:.15f}')
    # coverage of [0.15,0.85] by the intervals plus [1/3,1/2]
    print('   all intervals certified:', allok)


# ---------------- V2 ----------------
V2_TOL = 1e-11  # stated tolerance on ell (bits per vertex) for floating-point and minimiser error


def v2():
    print('V2  dense u-grid, bounded scalar minimisation over log2 y in [-30, 8] (xatol 1e-12) per grid point')
    worst = (-1, None, None)
    for u in np.linspace(0.15, 0.85, 3501):
        r = minimize_scalar(lambda ly: ell(2.0 ** ly, u), bounds=(-30, 8), method='bounded',
                            options={'xatol': 1e-12})
        if r.fun > worst[0]:
            worst = (r.fun, u, 2.0 ** r.x)
    print(f'   max_u min_y ell = {worst[0]:.13f} (bound {LC:.13f}) at u={worst[1]:.4f}, y={worst[2]:.5f}; base {2**worst[0]:.9f}')
    print(f'   difference to bound = {worst[0] - LC:.2e}; no grid value exceeds the bound by more than {V2_TOL:.0e}:',
          worst[0] <= LC + V2_TOL)
    print('   The bound is attained with equality on [1/3, 1/2] (a=2, y0 = (1+sqrt5)/8), so this check is equality-tight')
    print('   and cannot detect an excess below the tolerance. An earlier version minimised over a coarse y-grid and')
    print('   printed False (0.7958234394 against 0.7958234391): an artefact of that grid, not of the bound.')


# ---------------- V3 ----------------
def v3(n):
    """For each s in (0.15n, 0.8n): max over guesses of min_(t-,t+) [log2 C(n,t) - log2 C(k-,t-) - log2 C(k+,t+)
    + alpha - t- m- - t+ m+], alpha = n-1, W = n-2."""
    W = n - 2
    alpha = n - 1
    lbn = np.array([lb(n, t) for t in range(n + 1)])
    worst_over_s = (-1, None)
    for s in range(int(0.15 * n) + 1, math.ceil(0.8 * n)):
        if not (0.15 * n < s < 0.8 * n):
            continue
        a = W // s
        worst = (-math.inf, None)
        for km in range(1, s + 1):
            kp = s - km
            tm = np.arange(km + 1)[:, None]
            tp = np.arange(kp + 1)[None, :]
            lbm = np.array([lb(km, i) for i in range(km + 1)])[:, None]
            lbp = np.array([lb(kp, j) for j in range(kp + 1)])[None, :]
            T = tm + tp
            base = lbn[T] - lbm - lbp
            for Wm in range(0, a * km + 1):
                Wp = W - Wm
                if Wp < (a + 1) * kp:
                    continue
                if kp == 0 and Wp != 0:  # guess range of alg:dr-integer: W+ = 0 if k+ = 0
                    continue
                mm = Wm / km
                mp = Wp / kp if kp else 0.0
                theta = alpha - tm * mm - tp * mp
                val = base + theta
                val[0, 0] = math.inf
                v = val.min()
                if v > worst[0]:
                    worst = (v, (km, Wm))
        if worst[0] > worst_over_s[0]:
            worst_over_s = (worst[0], s, worst[1], a)
    v, s, g, a = worst_over_s
    print(f'V3  n={n}: max_s max_guess min_(t-,t+) log2 bound /n = {v/n:.6f}  (base {2**(v/n):.6f}) at s={s}, guess (k-,W-)={g}, a={a};'
          f'  log2(sqrt5-1/2) = {LC:.6f}; excess bits = {v - n*LC:.2f}')
    # also: is the balanced guess the worst?
    km_bal = s - (W - a * s)
    print(f'     balanced guess for this s would be (k-,W-)=({km_bal},{a*km_bal})')


# ---------------- V4 ----------------
def v4(n=1000):
    W = n - 2
    alpha = n - 1
    lbn = np.array([lb(n, t) for t in range(n + 1)])
    best = (-1, None)
    for s in range(int(0.15 * n) + 1, math.ceil(0.8 * n)):
        a = W // s
        kb = W - a * s
        ka = s - kb
        # log2 f_t = log2 sum_i C(ka,i) C(kb,t-i) 2^{a i + (a+1)(t-i)}
        la = np.array([lb(ka, i) + a * i for i in range(ka + 1)])
        lbb = np.array([lb(kb, j) + (a + 1) * j for j in range(kb + 1)])
        M = la[:, None] + lbb[None, :]
        lf = np.full(s + 1, -np.inf)
        for t in range(1, s + 1):
            i = np.arange(max(0, t - kb), min(ka, t) + 1)
            vals = M[i, t - i]
            mx = vals.max()
            lf[t] = mx + np.log2(np.exp2(vals - mx).sum())
        cost = lbn[1:s + 1] + alpha - lf[1:]
        v = cost.min()
        if v > best[0]:
            best = (v, s)
    print(f'V4  n={n}: balanced-profile max_s min_t = 2^({best[0]/n:.6f} n) = {2**(best[0]/n):.6f}^n at s={best[1]}')


# ---------------- V5 ----------------
def v5():
    print('V5  black-box numbers')
    print(f'   H(1/4) = {H(0.25):.10f}, 2^H = {2**H(0.25):.6f}, 4*3^(-3/4) = {4*3**-0.75:.6f}')
    s, t = 0.41759, 0.28083
    print(f'   example: lambda/n={1-t/s:.4f}, sigma H(tau/sigma)={s*H(t/s):.4f}, 2^H(tau)={2**H(t):.4f}, 2^H(sigma)={2**H(s):.4f}')
    om = L2(1.72035)
    for s in (0.4176, 0.4172629):
        ts = np.linspace(1e-6, s - 1e-6, 400001)
        feas = [tt for tt in ts[::10] if H(tt) - s * H(tt / s) <= om]
        tm = max(feas)
        print(f'   thm:dr-bb-moments at sigma={s}: largest feasible tau={tm:.4f}, H={H(tm):.4f}, 2^H={2**H(tm):.4f}; cap 2^H(sigma)={2**H(s):.4f}')


if __name__ == '__main__':
    v1(); v2(); v5(); v4(1000)
    for n in [int(x) for x in sys.argv[1:]] or [60, 120]:
        v3(n)
