"""Rate computations for the paper (Proposition 6.1, Corollary 9.3, Theorems 9.4 and G.3, Lemma 10.1).

Parts:
 (1) RS rate  e_RS(s)  = min_tau R(s,tau),  R = H(tau) - s H(tau/s) + 1 - tau/s.
 (2) Black-box lower-bound exponent for the all-levels RS hypothesis
       e_BB(s) = min( H(s), min_tau max(R(s,tau), H(tau)) ).
     Closed form at s = 1/2:  H(1/4) = 2 - (3/4) log2 3.
 (3) Black-box lower bound of Theorem 9.4 for Theorem 3.5 (spiky adversary) at the RS-theta
     budget omega* = 0.7827016:  max{H(tau): tau<=s, H(tau)-sH(tau/s) <= omega*}.
 (4) Integer-affine exponential-moment rate (deterministic two-class scheme)
       e_aff(u) = min_y  log2(1+y) + 1 - ((a+1)u-1) log2(1+y 2^a) - (1-a u) log2(1+y 2^(a+1)),
     a = floor(1/u); max over u is log2(sqrt5 - 1/2).
     Exact rational certificate of the upper bound on u in [0.15, 0.85].
 (5) Finite-n exact check of the integer-affine bound with big integers.
 (6) Finite-n evaluation of the all-levels black-box lower bound.
Run:  python rates.py
"""
import math
from fractions import Fraction as Fr
import numpy as np
from scipy.optimize import minimize_scalar

L2 = math.log2


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


def R(s, t):
    return H(t) - s * H(t / s) + 1 - t / s


def e_rs(s):
    r = minimize_scalar(lambda t: R(s, t), bounds=(1e-12, s - 1e-12), method='bounded',
                        options={'xatol': 1e-14})
    return r.fun, r.x


def e_bb(s):
    # max(R, H) is minimised where R = H (R decreasing on [0,tau*], H increasing on [0,1/2])
    # robust: dense grid then local refinement
    ts = np.linspace(1e-9, s - 1e-9, 20001)
    vals = np.array([max(R(s, t), H(t)) for t in ts])
    i = int(vals.argmin())
    lo, hi = ts[max(i - 1, 0)], ts[min(i + 1, len(ts) - 1)]
    r = minimize_scalar(lambda t: max(R(s, t), H(t)), bounds=(lo, hi), method='bounded',
                        options={'xatol': 1e-15})
    return min(r.fun, H(s)), r.x


def part12():
    print('== (1) RS and (2) black-box all-levels exponent')
    ss = np.linspace(0.15, 0.8, 651)
    rs = [(e_rs(s)[0], s) for s in ss]
    m = max(rs)
    s0 = minimize_scalar(lambda s: -e_rs(s)[0], bounds=(m[1] - 0.002, m[1] + 0.002), method='bounded',
                         options={'xatol': 1e-12}).x
    v, t0 = e_rs(s0)
    print(f'RS: max_s e_RS = {v:.10f} bits at s={s0:.7f} (tau={t0:.7f}); base {2**v:.10f}')
    bb = [(e_bb(s)[0], s) for s in ss]
    mb = max(bb)
    print(f'BB grid max: e_BB = {mb[0]:.10f} at s = {mb[1]:.4f}; base {2**mb[0]:.10f}')
    v, t = e_bb(0.5)
    print(f'BB at s=1/2: {v:.12f} (tau={t:.9f});  H(1/4) = {H(0.25):.12f}; base 4*3^(-3/4) = {4*3**-0.75:.10f}')
    print(f'   RS at s=1/2: {e_rs(0.5)[0]:.12f} = log2 sqrt3 = {0.5*L2(3):.12f}')
    # sanity: R(1/2,1/4) == H(1/4)
    print(f'   R(1/2,1/4) = {R(0.5,0.25):.12f}')
    # at the RS worst case
    s = s0
    print(f'   at RS worst s={s:.6f}: e_BB = {e_bb(s)[0]:.10f} vs e_RS = {e_rs(s)[0]:.10f}')
    # single-level gap at RS's (s,tau): log2 C(k,t)/n - lambda/n
    s, t = s0, t0
    print(f'   single-level data at RS optimum: lambda/n = {1-t/s:.6f}, log2C(k,t)/n = {s*H(t/s):.6f},'
          f' H(tau) = {H(t):.6f} -> det. black-box >= 2^({H(t):.6f} n) = {2**H(t):.6f}^n')


def part3():
    print('== (3) Theorem 9.4 black-box lower bound at the RS-theta budget')
    om = 0.7827016
    for s in [0.3, 0.35, 0.4, 0.4176, 0.45, 0.5, 0.55, 0.6]:
        ts = np.linspace(1e-6, s - 1e-6, 200001)
        best = 0.0
        for t in ts[::50]:
            if H(t) - s * H(t / s) <= om - 1e-9:
                best = max(best, H(t))
        lb = min(best, H(s))
        print(f's={s:.4f}: feasible-level max H(tau) = {best:.5f}, H(s) = {H(s):.5f};'
              f' det. black-box >= {2**lb:.4f}^n  vs randomised 2^omega* = {2**om:.5f}^n')


def ell(y, u):
    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 e_aff(u):
    r = minimize_scalar(lambda ly: ell(2.0 ** ly, u), bounds=(-30, 8), method='bounded',
                        options={'xatol': 1e-13})
    return r.fun


def part4():
    print('== (4) integer-affine exponential-moment rate')
    cstar = math.sqrt(5) - 0.5
    us = np.linspace(0.15, 0.85, 7001)
    vals = [(e_aff(u), u) for u in us]
    m = max(vals)
    print(f'grid max e_aff = {m[0]:.12f} at u={m[1]:.5f}: base {2**m[0]:.12f};  sqrt5-1/2 = {cstar:.12f}')
    # check the closed-form saddle point
    y0 = (1 + math.sqrt(5)) / 8
    print(f'  y0=(1+sqrt5)/8: ell(y0,1/2) base {2**ell(y0,0.5):.12f}, ell(y0,1/3) base {2**ell(y0,1/3+1e-15):.12f}')
    # exact rational certificate on u in [0.15,0.85]: on each a-interval ell_y(u) is affine in u,
    # so ell_y(u) <= max over the two endpoints; endpoint inequality ell_y(u_e) <= log2 c is the
    # rational inequality  (2(1+y))^q <= c^q * (1+y 2^a)^{q((a+1)u-1)} * (1+y 2^{a+1})^{q(1-a u)}
    # with q the common denominator of the two (rational) exponents.
    c = Fr(17360, 10000)  # 1.7360 < sqrt5 - 1/2
    intervals = [(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))]
    allok = True
    # rational y per interval (chosen near the numerical minimiser of the endpoint maximum)
    ylist = {1: Fr(1, 2), 3: Fr(5, 16), 4: Fr(1, 4), 5: Fr(27, 128), 6: Fr(29, 160)}
    for a, ul, ur in intervals:
        f = lambda ly: max(ell(2.0 ** ly, float(ul) + 1e-15), ell(2.0 ** ly, float(ur) - 1e-15))
        ly = minimize_scalar(f, bounds=(-30, 8), method='bounded', options={'xatol': 1e-12}).x
        y = ylist[a]
        print(f'    (numerical minimiser of the endpoint maximum: y = {2**ly:.6f}; used y = {y})')
        oks = []
        for ue in (ul, ur):
            e1 = (a + 1) * ue - 1
            e2 = 1 - a * ue
            assert e1 >= 0 and e2 >= 0
            q = math.lcm(e1.denominator, e2.denominator)
            lhs = (2 * (1 + y)) ** q
            rhs = c ** q * (1 + y * 2 ** a) ** int(e1 * q) * (1 + y * 2 ** (a + 1)) ** int(e2 * q)
            oks.append(lhs <= rhs)
            # float value for display
        allok &= all(oks)
        print(f'  a={a}: u in [{float(ul):.4f},{float(ur):.4f}], y0={float(y):.6f}: endpoint bases '
              f'{2**ell(float(y), float(ul)+1e-15):.6f}, {2**ell(float(y), float(ur)-1e-15):.6f};'
              f' exact check vs 1.7360: {oks}')
    print('  interval a=2 [1/3,1/2]: y0=(1+sqrt5)/8 gives exactly sqrt5-1/2 at both endpoints (algebraic identity).')
    print('  exact certificate passed:', allok)
    print(f'  enumeration ranges: H(0.15) base {2**H(0.15):.4f}, H(0.2) base {2**H(0.2):.4f}')


def part5(n=300):
    print(f'== (5) finite-n exact check, n = {n}')
    cstar = math.sqrt(5) - 0.5
    worst = (-1, None)
    worstA = (-1, None)
    for s in range(math.ceil(0.15 * n) + 1, math.floor(0.8 * n)):
        W = n - 2
        a = W // s
        ka = (a + 1) * s - W
        kb = W - a * s
        assert ka >= 0 and kb >= 0 and ka + kb == s
        # f_t = coefficient of y^t in (1+2^a y)^ka (1+2^{a+1} y)^kb
        poly = [1]
        for (mult, wt) in ((ka, a), (kb, a + 1)):
            for _ in range(mult):
                new = poly + [0]
                for i in range(len(poly)):
                    new[i + 1] += poly[i] * 2 ** wt
                poly = new
        best = None
        bestA = None
        for t in range(1, s + 1):
            val = math.comb(n, t) * 2 ** (n - 1) / poly[t] if False else None
            # use logs of exact integers
            lv = math.log2(math.comb(n, t)) + (n - 1) - math.log2(poly[t])
            if best is None or lv < best:
                best = lv
            # Theorem 4.1 applied to the RS hypothesis, with lambda = n-1-t(n-2)/s
            la = math.log2(math.comb(n, t)) - math.log2(math.comb(s, t)) + (n - 1 - t * (n - 2) / s)
            if bestA is None or la < bestA:
                bestA = la
        if best > worst[0]:
            worst = (best, s)
        if bestA > worstA[0]:
            worstA = (bestA, s)
    print(f'  integer-affine: max_s min_t log2 cost / n = {worst[0]/n:.6f} at s={worst[1]} (s/n={worst[1]/n:.4f});'
          f' base {2**(worst[0]/n):.6f}; target sqrt5-1/2 = {cstar:.6f}')
    print(f'  Theorem 4.1 applied to the RS hypothesis:  max_s min_t log2 cost / n = {worstA[0]/n:.6f} at s={worstA[1]}; base {2**(worstA[0]/n):.6f}')




def part6(n=1000):
    """Finite-n evaluation of the all-levels black-box lower bound (Theorem G.3 of the paper)
    for RS's hypothesis lambda_0 = n, lambda_j = n-1-j(n-2)/s, against the bound of Theorem 4.1."""
    print(f'== (6) finite-n black-box bound, n = {n}')
    for s in [n // 2, round(0.4176 * n)]:
        lgk = lambda m, j: math.log2(math.comb(m, j))
        best_bb, best_A = None, None
        for j in range(0, s + 1):
            lam = n if j == 0 else n - 1 - j * (n - 2) / s
            RA = lgk(n, j) - lgk(s, j) + lam
            H_j = max(1, math.ceil(math.log2(lam + 1) + lgk(s, j)))
            Lq = lgk(n, j) + math.log2(lam) - math.log2(H_j) if lam > 0 else -1e9
            Lj = max(RA, Lq)
            best_bb = Lj if best_bb is None else min(best_bb, Lj)
            best_A = RA if best_A is None else min(best_A, RA)
        lb = min(lgk(n, s), best_bb) - math.log2(3)
        print(f'  s={s}: deterministic black-box >= 2^{lb:.2f} = {2**(lb/n):.6f}^n;'
              f'  Theorem 4.1 (randomised) min_t R_t = 2^{best_A:.2f} = {2**(best_A/n):.6f}^n (times poly)')


if __name__ == '__main__':
    part12()
    part3()
    part4()
    part5(300)
    part5(1000)
    part6(1000)
    part6(4000)