"""Exact verification of the RSC certificate (model g: graded closure p2 <= p3 < 1/2 and typed exposures).
Usage: python3 rsc-verify.py rsc-certificate.json[.gz] [procs]
Imports rsc_verify_core.py and rsc_modelg.py from this folder.  Checks:
 (K) the poset cone K (covers R2>R1>R0, R2>O01>O00, O01>O1>O2; all e_s >= 0) is invariant under the graded
     merge (defined below) and final is K-monotone (exact);
 (Z) every z >= 0; c = min final(z) > 0;  (a) all ordered pairs z_i o z_j dominate a combination of Z (weight >= 1);
 (b) Init(m, 0, 0)/(e1 e2^m), m <= m0, and (y/2, p2/2, 0..)/(e1 e2^(m0+1)) dominate combinations of Z;
 (W) omega2, omega3 admissible: 1 <= omega_d <= 2 q_d and kappa_d' = 1 - log2(omega_d) satisfies
     kappa_2' >= rho_1, kappa_3' >= rho_1^2 + (eps3/ln 2)(1 - rho_1^2), rho_1 = (1 + (y/2) q3)^-1 (60-digit Decimal);
 (L) classification: for m <= ME, d < dmin3(m) passes test3 and d < dmin2(m) passes test2 (exact tests with
     b = 2 + 2 max(p2, p3 - p2)/y); dmin3 <= dmin2;
 (E1) failing leaves (dmin3(m) <= 1): Init_g(m, 0, c3)/(e1 e2^m), 1 <= c3 <= m - 1, dominates a combination of Z;
      if also dmin2(m) <= 1: the same for every hull-vertex count (c2 >= 1, c3) of the lattice c2 + c3 <= m - 1;
 (E2) failing internal vertices: staged products Init_g(m, c2, c3)/(e1 e2^m) o z_1 o ... o z_j with j = dmin3 - 1
      for (0, c3), 1 <= c3 <= m, and j = dmin2 - 1 for the hull-vertex counts with c2 >= 1 of the lattice c2 + c3 <= m;
 (H) the hull vertices are recomputed exactly here (modelg.expo_vertices with Fractions); every other lattice point
     is a convex combination of hull vertices, and the final exposure point is linear in Init_g while the
     verified region (conv(Z) scaled by >= 1, plus K) is convex, so the vertices suffice;
 (T) tail m > ME: q3 > e2 and (q3/e2)^(ME+1) >= e2^-(m0+1) (with omega_d >= 1), so every Init_g(m, c2, c3)/(e1 e2^m)
     dominates u_inf/(e1 e2^(m0+1)).
Init_g(m, c2, c3) = (qx y/2, qx p2/2, 0, 0, 2^-m omega2^c2 omega3^c3, 0, 0, 0), qx = q2^c2 q3^c3."""
import json, sys
from decimal import Decimal as Dd, getcontext
from fractions import Fraction as Fr
from multiprocessing import Pool
from rsc_verify_core import F, ratw, dominated, Tc, NAMES, IDX, COVERS
from rsc_modelg import expo_vertices
getcontext().prec = 60

def dec(fr): return Dd(fr.numerator) / Dd(fr.denominator)

def merge(P, Q, p2, p3):
    q2, q3, g = 1 - p2, 1 - p3, p3 / p2
    A, R0, R1, R2, O00, O01, O1, O2 = P
    a, r0, r1, r2, o00, o01, o1, o2 = Q
    kS = 2 * a + 2 * r1 + 2 * g * r2 + o00 + o01 + q2 * o1 + q3 * o2
    sR, uR = 2 * a, 2 * r2 + o00 + o01 + o1 + q2 * o2
    sO, xO, uO = a, r2, o00 / 2 + o01 + o1 + q2 * o2
    return [A * kS, R0 * uR, R1 * uR + R0 * sR, R2 * (uR + sR) + R1 * sR,
            O00 * uO, O01 * (uO + xO) + O00 * xO, O1 * (uO + xO) + (O00 + O01) * sO, O2 * (uO + xO + sO) + O1 * sO]

def final(P, q2):
    return P[0] + P[3] + P[4] / 2 + P[5] + P[6] + q2 * P[7]

def check_cone(p2, p3):
    E = [[Fr(int(i == j)) for j in range(8)] for i in range(8)]
    gens = [[E[IDX[b]][d] - E[IDX[w]][d] for d in range(8)] for b, w in COVERS] + E
    if not all(final(g, 1 - p2) >= 0 for g in gens): return False
    for g in gens:
        for e in E:
            if any(x < 0 for x in Tc(merge(g, e, p2, p3))): return False
            if any(x < 0 for x in Tc(merge(e, g, p2, p3))): return False
    return True

def test3_holds(m, d, y, p2, p3):
    yD, p3D = dec(y), dec(p3); q3 = 1 - p3D
    eps = -(q3.ln()); dlt = (2 * q3).ln(); beta = dec(2 + 2 * max(p2, p3 - p2) / y)
    return Dd(2) ** (-m) * beta ** (-d) >= (1 + p3D / yD) * eps * yD / dlt * (1 + Dd('1e-30'))

def test2_holds(m, d, y, p2, p3):
    yD, p2D = dec(y), dec(p2); q2 = 1 - p2D
    eps = -(q2.ln()); beta = dec(2 + 2 * max(p2, p3 - p2) / y)
    return Dd(2) ** (-m) * beta ** (-d) >= eps * yD / Dd(2).ln() * (1 + Dd('1e-30'))

def omegas_ok(om2, om3, y, p2, p3):
    if not (Fr(1) <= om2 <= 2 * (1 - p2) and Fr(1) <= om3 <= 2 * (1 - p3)): return False
    yD, p3D = dec(y), dec(p3); q3 = 1 - p3D
    eps3 = -(q3.ln()); ln2 = Dd(2).ln()
    k2 = 1 - dec(om2).ln() / ln2; k3 = 1 - dec(om3).ln() / ln2
    r1 = 1 / (1 + yD / 2 * q3); r2 = r1 * r1
    mg = Dd('1e-30')
    return k2 >= r1 + mg and k3 >= r2 + eps3 / ln2 * (1 - r2) + mg

_G = {}
def _init_worker(Z, p2, p3): _G.update(Z=Z, p2=p2, p3=p3)

def _check_rows(args):
    rows, combos = args
    Z, p2, p3 = _G['Z'], _G['p2'], _G['p3']
    for (i, j), comb in zip(rows, combos):
        if not dominated(merge(Z[i], Z[j], p2, p3), ratw(comb), Z): return (i, j)
    return None

def _check_expo(args):
    """one exposure stage, rows (a, b) with a in prev, b in Z: x_a o z_b dominates a combination of target (None: Z)."""
    prev, rows, combos, target = args
    Z, p2, p3 = _G['Z'], _G['p2'], _G['p3']
    tgt = Z if target is None else target
    for (a, b), comb in zip(rows, combos):
        if not dominated(merge(prev[a], Z[b], p2, p3), ratw(comb), tgt): return (a, b)
    return None

def verify(cert, procs=14):
    assert cert.get('model') == 'g'
    y, p2, p3 = F(cert['y']), F(cert['p2']), F(cert['p3'])
    om2, om3, e1, e2 = F(cert['omega2']), F(cert['omega3']), F(cert['e1']), F(cert['e2'])
    q2, q3 = 1 - p2, 1 - p3
    assert 0 < p2 <= p3 < Fr(1, 2) and y > 0 and e1 > 0 and 0 < e2 < q3 < 1
    if not omegas_ok(om2, om3, y, p2, p3): return False, "omega"
    assert check_cone(p2, p3), "cone"
    Z = [[F(v) for v in z] for z in cert['Z']]
    assert all(all(v >= 0 for v in z) for z in Z)
    c = min(final(z, q2) for z in Z); assert c > 0
    n = len(Z)
    keys = [(i, j) for i in range(n) for j in range(n)]
    tasks = []
    for s in range(0, len(keys), 2000):
        ks = keys[s:s + 2000]
        tasks.append((ks, [cert['combos'][f"{i},{j}"] for i, j in ks]))
    pool = Pool(procs, initializer=_init_worker, initargs=(Z, p2, p3))
    for r in pool.imap_unordered(_check_rows, tasks):
        if r is not None: pool.terminate(); return False, f"pair {r}"
    m0 = cert['m0']
    def initv(m, c2=0, c3=0):
        s = q2 ** c2 * q3 ** c3
        return [s * y / 2, s * p2 / 2, Fr(0), Fr(0), Fr(1, 2 ** m) * om2 ** c2 * om3 ** c3, Fr(0), Fr(0), Fr(0)]
    for m in range(m0 + 1):
        P = [x / (e1 * e2 ** m) for x in initv(m)]
        if not dominated(P, ratw(cert['inits'][str(m)]), Z): return False, f"init {m}"
    uinf = [x / (e1 * e2 ** (m0 + 1)) for x in [y / 2, p2 / 2, 0, 0, 0, 0, 0, 0]]
    if not dominated(uinf, ratw(cert['inits']['inf']), Z): return False, "init inf"
    ME = cert['ME']
    dmin3 = {int(k): v for k, v in cert['dmin3'].items()}; dmin2 = {int(k): v for k, v in cert['dmin2'].items()}
    for m in range(1, ME + 1):
        assert dmin3[m] <= dmin2[m], "dmin order"
        for d in range(1, dmin3[m]):
            if not test3_holds(m, d, y, p2, p3): return False, f"classification test3 m={m} d={d}"
        for d in range(1, dmin2[m]):
            if not test2_holds(m, d, y, p2, p3): return False, f"classification test2 m={m} d={d}"
    # (E1) leaves
    for m in range(2, ME + 1):
        if dmin3[m] <= 1:
            V3, Vm = expo_vertices(m - 1, q2, q3, om2, om3, exact=True)
            for c2, c3 in V3 + (Vm if dmin2[m] <= 1 else []):
                P = [x / (e1 * e2 ** m) for x in initv(m, c2, c3)]
                key = f"{m},{c2},{c3}"
                if key not in cert['leafexp'] or not dominated(P, ratw(cert['leafexp'][key]), Z):
                    return False, f"leaf exposure m={m} c=({c2},{c3})"
    # (E2) internal
    done = {}
    for ex in cert['internal']:
        m, c2, c3, j = ex['m'], ex['c2'], ex['c3'], ex['j']
        Y0 = [[x / (e1 * e2 ** m) for x in initv(m, c2, c3)]]
        stages = [Y0] + [[[F(v) for v in z] for z in Y] for Y in ex['stages']]
        assert len(stages) == j
        for i in range(1, j + 1):
            prev = stages[i - 1]; target = stages[i] if i < j else None
            cmb = ex['combos'][i - 1]
            rows = [(a, b) for a in range(len(prev)) for b in range(n)]
            etasks = [(prev, rows[s:s + 1000], [cmb[f"{a},{b}"] for a, b in rows[s:s + 1000]], target) for s in range(0, len(rows), 1000)]
            for r in pool.imap_unordered(_check_expo, etasks):
                if r is not None: pool.terminate(); return False, f"internal exposure m={m} c=({c2},{c3}) stage {i} {r}"
        done[(m, c2, c3)] = min(j, done.get((m, c2, c3), j))
    for m in range(1, ME + 1):
        V3, Vm = expo_vertices(m, q2, q3, om2, om3, exact=True)
        j3, j2 = max(1, dmin3[m] - 1), max(1, dmin2[m] - 1)
        for c2, c3 in V3:
            if done.get((m, c2, c3), 10 ** 9) > j3: return False, f"internal exposure missing m={m} c=(0,{c3}) j<={j3}"
        for c2, c3 in Vm:
            if done.get((m, c2, c3), 10 ** 9) > j2: return False, f"internal exposure missing m={m} c=({c2},{c3}) j<={j2}"
    pool.close(); pool.join()
    assert om2 >= 1 and om3 >= 1 and (q3 / e2) ** (ME + 1) >= 1 / e2 ** (m0 + 1), "tail"
    return True, dict(y=y, p2=p2, p3=p3, omega2=om2, omega3=om3, e1=e1, e2=e2, c=c, npts=n)

if __name__ == '__main__':
    import gzip, time
    args = [a for a in sys.argv[1:] if not a.isdigit()]
    procs = int([a for a in sys.argv[1:] if a.isdigit()][0]) if any(a.isdigit() for a in sys.argv[1:]) else 4
    for f in args:
        t = time.time()
        with (gzip.open(f, 'rt') if f.endswith('.gz') else open(f)) as fh:
            cert = json.load(fh)
        ok, info = verify(cert, procs=procs)
        print(f, "VERIFIED" if ok else "FAILED", info if not ok else {k: (float(v) if not isinstance(v, int) else v) for k, v in info.items()}, "time %.0f s" % (time.time() - t))
