#!/usr/bin/env python3
"""R(x, A) beyond the band by a rigorous floating-point sweep (paper "Analytic q-inequalities and a proof of Bergeron's conjecture", Section 9.7).

For x >= 3 and 2x(x-1) <= A <= LAM x^2 the script checks, for every breakpoint A (a point of xZ or (x-1)Z) and
every degree D < k <= floor(A+/2), where A+ is the next breakpoint, that

    [q^k] N_A < [q^k] P_A,        P_A = (x-1) q^x [(x-r) G(x,D-1) + r G(x,D)],
                                  N_A = sum_{e=D+1}^{E+1} w_e q^e G(x-2,e)          (Proposition 8.23),

with A = xD + r = (x-1)E + s, w_{D+1} = (x-1)(x-r), w_e = x(x-1) for D+2 <= e <= E, w_{E+1} = x s.  In degrees
k <= D, N_A vanishes.  Since every coefficient of D_A = P_A - N_A is affine in A between consecutive breakpoints,
this proves R(x, A) for 2x(x-1) <= A <= LAM x^2 (the last interval is closed with the first breakpoint >= LAM x^2).

Arithmetic.  G(m, n) = [m+n choose m]_q is computed for m <= x by the positive recurrence
    G(m, n+1) = G(m, n) + q^{n+1} G(m-1, n+1)
in IEEE double precision, truncated at the largest degree needed.  Every operation adds two nonnegative numbers, so
each computed coefficient equals the true one times (1 + delta) with |delta| <= (1+u)^{m+n} - 1, u = 2^-53
(induction on m + n; zeros are exact).  P_A and N_A add at most (E-D+4) further roundings of nonnegative terms, and
products with integers below 2^20 one more each.  The script checks N < RATIO_MAX * P with RATIO_MAX = 0.999, far
above the accumulated relative error (reported), and no value comes near overflow (checked).  For large x all values
are scaled by one power of two, 2^-S, which is exact and keeps them in range.

Usage:  python3 row_sweep.py XMIN XMAX LAM      e.g.  python3 row_sweep.py 22 30 16
        python3 row_sweep.py 22-40:64 41-60:32 61-200:16      (blocks with their own LAM)
"""
import sys, time, math
import numpy as np

U = 2.0 ** -53
RATIO_MAX = 0.999


def breakpoints(x, Alo, Ahi):
    """sorted breakpoints A in [Alo, Ahi'] where Ahi' is the first breakpoint >= Ahi."""
    pts = set()
    a = (Alo + x - 1) // x * x
    while a <= Ahi + x:
        pts.add(a); a += x
    b = (Alo + x - 2) // (x - 1) * (x - 1)
    while b <= Ahi + x:
        pts.add(b); b += x - 1
    pts = sorted(p for p in pts if p >= Alo)
    out = []
    for p in pts:
        out.append(p)
        if p >= Ahi:
            break
    return out


def check_x(x, lam, verbose=False):
    """returns dict with the worst ratio N/P over all checked (A, k), the number of breakpoints, and error bounds."""
    t0 = time.time()
    Alo = 2 * x * (x - 1)
    Ahi = int(math.ceil(lam * x * x))
    bps = breakpoints(x, Alo, Ahi)
    # degrees needed: k <= floor(A+/2) for every breakpoint A but the last
    Kmax = bps[-1] // 2 + 1
    # the largest index n needed: E + 1 at the last-but-one breakpoint (the last one only closes an interval),
    # but we also evaluate the last breakpoint itself (for its left interval) up to degree floor(A_last/2)
    need = []
    for i, A in enumerate(bps):
        D, r = divmod(A, x)
        E, s = divmod(A, x - 1)
        K = bps[i + 1] // 2 if i + 1 < len(bps) else A // 2
        need.append((A, D, r, E, s, K))
    Nmax = max(E for (_, _, _, E, _, _) in need) + 1
    W = max(E - D for (_, D, _, E, _, _) in need) + 4         # window of stored G(x, .), G(x-2, .)
    # arrays G[m] = 2^-S G(m, n) for the current n, degrees 0..Kmax.  The scale 2^-S (exact in binary floating point)
    # keeps every value below 2^990 (G(m, n) <= binom(x + Nmax, x), times weights at most x^2 (W + 2)); every nonzero
    # value is at least 2^-S >= 2^-900, far from underflow.  S = 0 whenever that bound is below 2^990.
    lg2 = (math.lgamma(x + Nmax + 1) - math.lgamma(x + 1) - math.lgamma(Nmax + 1)) / math.log(2) \
        + math.log2(x * x * (W + 2))
    S = max(0, int(math.ceil(lg2 - 990)))
    if S > 900:
        raise RuntimeError("scale out of range")
    G = np.zeros((x + 1, Kmax + 1))
    G[:, 0] = 2.0 ** -S
    winX = {}     # n -> copy of G(x, n)
    winM2 = {}    # n -> copy of G(x-2, n)
    worst = 0.0; worst_at = None
    idx = 0
    n = 0
    maxval = 0.0
    while idx < len(need):
        # store the current column in the windows
        winX[n] = G[x].copy(); winM2[n] = G[x - 2].copy()
        for key in [k for k in winX if k < n - W]:
            del winX[key]; del winM2[key]
        # process every breakpoint whose largest index E + 1 equals n
        while idx < len(need) and need[idx][3] + 1 == n:
            A, D, r, E, s, K = need[idx]
            if K > D:
                ks = np.arange(D + 1, K + 1)
                P = (x - 1) * ((x - r) * winX[D - 1][ks - x] + r * winX[D][ks - x])
                Nn = np.zeros(len(ks))
                for e in range(D + 1, E + 2):
                    w = (x - 1) * (x - r) if e == D + 1 else (x * (x - 1) if e <= E else x * s)
                    if w == 0:
                        continue
                    sh = ks - e
                    ok = sh >= 0
                    Nn[ok] += w * winM2[e][sh[ok]]
                if not np.all(P > 0):
                    raise RuntimeError("P vanishes at x=%d A=%d" % (x, A))
                ratio = Nn / P
                j = int(np.argmax(ratio))
                if ratio[j] > worst:
                    worst = float(ratio[j]); worst_at = (A, int(ks[j]), round(A / x / x, 3))
                maxval = max(maxval, float(P.max()), float(Nn.max()))
            idx += 1
        if idx >= len(need):
            break
        # step n -> n + 1
        n += 1
        if n <= Kmax:
            for m in range(1, x + 1):
                G[m, n:] += G[m - 1, :Kmax + 1 - n]
    # relative error of every stored coefficient: (1+u)^(x+Nmax) - 1 <= 1.01 (x+Nmax) u, plus the combinations in P, N
    err = 1.01 * (x + Nmax + W + 8) * U
    require_small = err < 1e-9
    if not require_small:
        raise RuntimeError("error bound too large")
    ok = worst < RATIO_MAX and maxval < 1e300
    return dict(x=x, lam=lam, nbp=len(bps), Amax=bps[-1], worst=worst, worst_at=worst_at, relerr=err,
                maxval=maxval, scale=S, ok=ok, time=time.time() - t0)


def main():
    """python3 row_sweep.py XMIN XMAX LAM, or blocks: python3 row_sweep.py 22-40:64 41-60:32 61-100:16"""
    if ':' in sys.argv[1]:
        jobs = []
        for spec in sys.argv[1:]:
            rng, lam = spec.split(':')
            a, b = rng.split('-')
            jobs += [(x, float(lam)) for x in range(int(a), int(b) + 1)]
    else:
        xmin = int(sys.argv[1]); xmax = int(sys.argv[2]); lam = float(sys.argv[3])
        jobs = [(x, lam) for x in range(xmin, xmax + 1)]
    allok = True
    t0 = time.time()
    print("R(x, A) for 2x(x-1) <= A <= LAM x^2: N_A < %.3f P_A in degrees D < k <= A+/2 at every breakpoint; %s"
          % (RATIO_MAX, " ".join(sys.argv[1:])), flush=True)
    for (x, lam) in jobs:
        r = check_x(x, lam)
        allok &= r['ok']
        print("x=%3d: %5d breakpoints up to A=%d; max N/P = %.4f at (A, k, lam) = %s; rel. error <= %.1e; max %.1e%s; %s  %.0fs"
              % (x, r['nbp'], r['Amax'], r['worst'], r['worst_at'], r['relerr'], r['maxval'],
                 (" (scaled by 2^-%d)" % r['scale']) if r['scale'] else "", "ok" if r['ok'] else "FAILED", r['time']),
              flush=True)
    print("ALL VERIFIED" if allok else "NOT VERIFIED", "(%.0fs)" % (time.time() - t0))
    return 0 if allok else 1


if __name__ == "__main__":
    sys.exit(main())