#!/usr/bin/env python3
"""Export the RSC certificate to Lean data modules for the kernel checkers of RSCCheck.lean.

Usage:
    python formal/scripts/export_rsc_cert.py <rsc-certificate.json.gz> <outdir> [--only PREFIX]
                                             [--selfcheck] [--procs N]

<outdir> is normally formal/MaxLeafFormal/RSCData.  The generated files are data, not proof: Lean
re-checks every condition in the kernel (RSCCheck.lean), so a wrong export can only make a check
fail.  The exporter aborts (non-zero exit, message on stderr) on any inconsistency.

Output is deterministic (byte-identical across runs).  --only PREFIX regenerates only the files whose
name starts with PREFIX (and then writes no MANIFEST.sha256).  --selfcheck re-parses the generated
files and re-verifies every exported condition in exact integer arithmetic, mirroring RSCCheck.lean.
Standard library only.
"""
import argparse
import gzip
import hashlib
import json
import os
import re
import sys
from fractions import Fraction as Fr
from math import lcm

GZ_SHA = "9465bd0f7f74dc3a6376724645bfbde97b417035583c89ba787e347813177651"
JSON_SHA = "188c5b7206fe3ba487513caa859ae68777efdae15d184f94092dd8904a55b48b"
NS = "MaxLeafFormal.RSCData"
ZQ = NS + ".Z"
NZ = 1024            # |Z|
BLK = 8              # rows per block
CH1_GROUP = 8        # chain-1 entries per module
PY_EXPECTED_PARAMS = dict(
    y=Fr(31, 100), p2=Fr(7, 100), p3=Fr(77, 1000), omega2=Fr(272657, 250000),
    omega3=Fr(288677, 250000), e1=Fr(145729958, 189014575), e2=Fr(145729958, 189014575))


class ExportError(Exception):
    pass


def fail(msg):
    raise ExportError(msg)


# ----------------------------------------------------------------------------------------------
# paper functions and exact helpers

def paper_d3(nu):
    return 3 if nu == 1 else (2 if nu in (2, 3) else 1)


def paper_d2(nu):
    return 4 if nu == 1 else (3 if nu == 2 else (2 if nu == 3 else 1))


def jof(d):
    return max(1, d - 1)


def j_target(nu, c2):
    return jof(paper_d3(nu)) if c2 == 0 else jof(paper_d2(nu))


def ratw(lst, den=10 ** 12):
    # identical to rsc_verify_core.ratw
    return [(k, Fr(w).limit_denominator(den)) for k, w in lst if w > 0]


def frac(pair):
    return Fr(pair[0], pair[1])


def enc_pt(pt):
    """A point (8 stored [num, den] pairs) as 8 numerators over the common denominator, then D."""
    fs = [frac(p) for p in pt]
    if len(fs) != 8:
        fail("point with %d coordinates" % len(fs))
    D = 1
    for f in fs:
        D = lcm(D, f.denominator)
    return [f.numerator * (D // f.denominator) for f in fs] + [D]


def enc_comb(lst, bound, what):
    """[count, k1, a1, b1, ...] with w = a/b from ratw; indices < bound; weight sum >= 1."""
    ws = ratw(lst)
    s = Fr(0)
    out = [len(ws)]
    for k, w in ws:
        if not (isinstance(k, int) and 0 <= k < bound):
            fail("%s: index %r out of range (< %d)" % (what, k, bound))
        s += w
        out += [k, w.numerator, w.denominator]
    if s < 1:
        fail("%s: weight sum %s < 1" % (what, s))
    return out


def flat_combs(cmb, a, bound, nb, what):
    """The 1024 flat combinations of row a (keys "a,b", b = 0..1023), concatenated."""
    out = []
    for b in range(NZ):
        key = "%d,%d" % (a, b)
        if key not in cmb:
            fail("%s: missing combination %s" % (what, key))
        out += enc_comb(cmb[key], bound, "%s %s" % (what, key))
    return out


def sha256_bytes(b):
    return hashlib.sha256(b).hexdigest()


# ----------------------------------------------------------------------------------------------
# Lean text

class Ctx:
    gz_sha = GZ_SHA
    json_sha = JSON_SHA
    exporter_sha = ""
    method = "m2"   # m1: `decide +kernel`; m2: `native_decide` (Serge, 2026-10-29: "use native_decide")


def lean_file(imports, ns, body):
    head = [
        "/-",
        "Generated by formal/scripts/export_rsc_cert.py -- do not edit.",
        "Data, not proof: Lean re-checks everything.",
        "input rsc-certificate.json.gz sha256: " + Ctx.gz_sha,
        "input rsc-certificate.json sha256:    " + Ctx.json_sha,
        "exporter sha256:                      " + Ctx.exporter_sha,
        "method:                               " + Ctx.method,
        "-/",
    ]
    head += ["import " + i for i in imports]
    head += ["", "set_option maxRecDepth 1000000", "set_option maxHeartbeats 0", "",
             "namespace " + ns, "", "noncomputable section", "", "open MaxLeafFormal.RSCCheck", ""]
    tail = ["", "end", "end " + ns, ""]
    return "\n".join(head) + "\n".join(body) + "\n".join(tail)


def num_rows(rows, indent="  "):
    """JSON-compatible multi-line list of flat number lists (closing bracket at column 0)."""
    if not rows:
        return "[]"
    return "[\n" + ",\n".join(indent + "[" + ", ".join(map(str, r)) + "]" for r in rows) + "\n]"


def def_rows(name, ty, rows):
    return ["def %s : %s := %s" % (name, ty, num_rows(rows))]


def thm(name, stmt):
    """A shard theorem: a checker call `= true`, decided by the kernel (M1) or by compiled code (M2)."""
    tac = "native_decide" if Ctx.method == "m2" else "decide +kernel"
    return "theorem %s : %s := by %s" % (name, stmt, tac)


def thm_kernel(name, stmt):
    """A trivial size side condition, always decided by the kernel."""
    return "theorem %s : %s := by decide +kernel" % (name, stmt)


def stmt_Z_ok():
    return thm("Z_ok", "checkZ Z = true")


def stmt_c1(r):
    return thm("ok", "checkBlock %s.length (chunks32 %s) %s ((%s.drop (8 * %d)).take 8) "
                     "(rows.map (decodeCombs 1024)) = true" % (ZQ, ZQ, ZQ, ZQ, r))


def stmt_c2():
    return thm("ok", "checkC2 %s.length (chunks32 %s) combos = true" % (ZQ, ZQ))


def stmt_leaf():
    return thm("ok", "checkLeaves %s.length (chunks32 %s) entries = true" % (ZQ, ZQ))


def stmt_ch1():
    return thm("ok", "checkChain1 %s.length (chunks32 %s) %s entries = true" % (ZQ, ZQ, ZQ))


def st_mod(k):
    return "%s.St_%d_pts" % (NS, k)


def x_expr(k, i, local=False):
    K = "" if local else st_mod(k) + "."
    return "(([initI %skey.1 %skey.2.1 %skey.2.2] :: %sstages).getD %d [])" % (K, K, K, K, i - 1)


def yt_expr(k, i, local=False):
    K = "" if local else st_mod(k) + "."
    return "((%sstages ++ [%s]).getD %d [])" % (K, ZQ, i - 1)


def stmt_st(k, i, r):
    X, Yt = x_expr(k, i), yt_expr(k, i)
    return thm("ok", "checkBlock %s.length (chunks32 %s) %s ((%s.drop (8 * %d)).take 8) "
                     "(rows.map (decodeCombs 1024)) = true" % (Yt, Yt, ZQ, X, r))


def stmt_nb(k, i, nb):
    return thm_kernel("nb_%d" % i, "decide ((%s.length + 7) / 8 ≤ %d) = true" % (x_expr(k, i, True), nb))


def stmt_lat5():
    return thm("lat5_ok", "checkLattice5 chainTable lat5 = true")


def stmt_lat4():
    return thm("lat4_ok", "checkLattice4 leafTable lat4 = true")


# ----------------------------------------------------------------------------------------------
# lattice witnesses

_uw_cache = {}


def uw(c2, c3):
    r = _uw_cache.get((c2, c3))
    if r is None:
        r = (Fr(93, 100) ** c2 * Fr(923, 1000) ** c3,
             Fr(272657, 250000) ** c2 * Fr(288677, 250000) ** c3)
        _uw_cache[(c2, c3)] = r
    return r


def lattice5_points():
    return [(nu, c2, c3) for nu in range(1, 37) for c2 in range(0, nu + 1)
            for c3 in range(0, nu - c2 + 1) if c2 + c3 >= 1]


def lattice4_points():
    return [(nu, c2, c3) for nu in range(4, 37) for c2 in range(0, nu)
            for c3 in range(0, nu - c2) if c2 + c3 >= 1]


def cross(o, a, b):
    return (a[0] - o[0]) * (b[1] - o[1]) - (a[1] - o[1]) * (b[0] - o[0])


def hull_vertices(pts):
    """Strict convex hull vertices (monotone chain) of a list of (x, y, tag)."""
    P = sorted(set((x, y, t) for x, y, t in pts), key=lambda p: (p[0], p[1]))
    if len(P) <= 2:
        return P
    lo = []
    for p in P:
        while len(lo) >= 2 and cross(lo[-2], lo[-1], p) <= 0:
            lo.pop()
        lo.append(p)
    up = []
    for p in reversed(P):
        while len(up) >= 2 and cross(up[-2], up[-1], p) <= 0:
            up.pop()
        up.append(p)
    return lo[:-1] + up[:-1]


def fan_witness(target, cands):
    """target = (c2, c3); cands = list of (tag, (c2, c3)).  Returns (comps, mode) or None.
    comps = [(tag or None for (0,0), Fraction lambda)] with sum lambda = 1, sum lambda*(u, w) = target."""
    tu, tw = uw(*target)
    # coordinates relative to P(0,0) = (1, 1): dx = 1 - u > 0, dy = w - 1 > 0
    tdx, tdy = 1 - tu, tw - 1
    if not (tdx > 0 and tdy > 0):
        fail("degenerate target %r" % (target,))
    groups = {}
    for tag, c in cands:
        u, w = uw(*c)
        dx, dy = 1 - u, w - 1
        if not (dx > 0 and dy > 0):
            fail("degenerate candidate %r" % (c,))
        s = dy / dx
        if s not in groups or dx > groups[s][0]:
            groups[s] = (dx, dy, tag)
    return _fan(tdx, tdy, groups, "fan") or _fan_hull(tdx, tdy, groups)


def _fan(tdx, tdy, groups, mode):
    if not groups:
        return None
    st = tdy / tdx
    slopes = sorted(groups)
    if st in groups:
        dx, dy, tag = groups[st]
        a = tdx / dx
        if a <= 1:
            return ([(None, 1 - a), (tag, a)], mode)
        return None
    for i in range(len(slopes) - 1):
        s1, s2 = slopes[i], slopes[i + 1]
        if s1 < st < s2:
            ax, ay, ta = groups[s1]
            bx, by, tb = groups[s2]
            det = ax * by - bx * ay
            al = (tdx * by - bx * tdy) / det
            be = (ax * tdy - ay * tdx) / det
            if al >= 0 and be >= 0 and al + be <= 1:
                return ([(None, 1 - al - be), (ta, al), (tb, be)], mode)
            return None
    return None


def _fan_hull(tdx, tdy, groups):
    pts = [(Fr(0), Fr(0), None)] + [(dx, dy, tag) for dx, dy, tag in groups.values()]
    hv = [p for p in hull_vertices(pts) if p[2] is not None]
    g2 = {}
    for dx, dy, tag in hv:
        g2[dy / dx] = (dx, dy, tag)
    return _fan(tdx, tdy, g2, "hull")


def verify_witness(target, comps, resolve):
    """Exact check; comps = [(tag or None, lambda)]; resolve(tag) -> (c2, c3)."""
    sl = Fr(0)
    su = Fr(0)
    sw = Fr(0)
    for tag, lam in comps:
        if not lam > 0:
            fail("non-positive weight in witness for %r" % (target,))
        c = (0, 0) if tag is None else resolve(tag)
        u, w = uw(*c)
        sl += lam
        su += lam * u
        sw += lam * w
    tu, tw = uw(*target)
    if not (sl == 1 and su == tu and sw == tw):
        fail("witness for %r does not verify" % (target,))


def build_lattice(CT, LT, stats):
    """CT = chain table [((nu, c2, c3), j)]; LT = leaf table [(nu, c2, c3)].  Returns (lat5, lat4)."""
    by5 = {}
    for idx, (key, j) in enumerate(CT):
        by5.setdefault(key[0], []).append((idx, (key[1], key[2]), j))
    by4 = {}
    for idx, key in enumerate(LT):
        by4.setdefault(key[0], []).append((idx, (key[1], key[2])))
    maxdig = [0]

    def entry(m, comps):
        row = [1, len(comps)]
        for kind_idx, lam in comps:
            kind, idx = kind_idx
            row += [kind, idx, lam.numerator, lam.denominator]
            maxdig[0] = max(maxdig[0], len(str(lam.numerator)), len(str(lam.denominator)))
        return row

    out = {}
    for which, pts in (("5", lattice5_points()), ("4", lattice4_points())):
        rows = []
        st = dict(direct=0, combo=0, fan=0, hull=0, maxdig=0, maxcomps=0)
        for pt in pts:
            nu, c2, c3 = pt
            if which == "5":
                J = j_target(nu, c2)
                cl = [(idx, c) for idx, c, j in by5.get(nu, []) if j <= J]
            else:
                cl = list(by4.get(nu, []))
            direct = [idx for idx, c in cl if c == (c2, c3)]
            if direct:
                rows.append([0, min(direct)])
                st["direct"] += 1
                continue
            # one candidate per count pair (the lowest table index)
            seen = {}
            for idx, c in cl:
                if c not in seen:
                    seen[c] = idx
            cands = [(idx, c) for c, idx in sorted(seen.items())]
            res = fan_witness((c2, c3), cands)
            if res is None:
                fail("lattice%s: no convex-combination witness for %r" % (which, pt))
            comps, mode = res
            tab = {idx: c for idx, c in cands}
            verify_witness((c2, c3), comps, lambda t: tab[t])
            crow = [((0, 0), lam) if tag is None else ((1, tag), lam) for tag, lam in comps if lam > 0]
            rows.append(entry(len(crow), crow))
            st["combo"] += 1
            st[mode] += 1
            st["maxcomps"] = max(st["maxcomps"], len(crow))
        st["maxdig"] = maxdig[0]
        stats["lattice" + which] = st
        out[which] = rows
        maxdig[0] = 0
    return out["5"], out["4"]


# ----------------------------------------------------------------------------------------------
# plan and generation

class Plan:
    pass


def check_input_and_params(cert):
    if cert.get("model") != "g":
        fail("model is %r, expected 'g'" % (cert.get("model"),))
    for k, v in PY_EXPECTED_PARAMS.items():
        if frac(cert[k]) != v:
            fail("parameter %s = %s, expected %s" % (k, frac(cert[k]), v))
    if cert["m0"] != 24 or cert["ME"] != 36:
        fail("m0/ME = %r/%r, expected 24/36" % (cert["m0"], cert["ME"]))
    mism = []
    for name, paper, d in (("dmin3", paper_d3, cert["dmin3"]), ("dmin2", paper_d2, cert["dmin2"])):
        for nu in range(1, 37):
            f = d[str(nu)]
            p = paper(nu)
            if f != p:
                mism.append((name, nu, f, p))
            if jof(f) != jof(p) or (f <= 1) != (p <= 1):
                fail("%s(%d) = %d disagrees with the paper's %d after j and the test d<=1" % (name, nu, f, p))
    if sorted(map(int, cert["dmin3"])) != list(range(1, 37)) or sorted(map(int, cert["dmin2"])) != list(range(1, 37)):
        fail("dmin3/dmin2 do not have exactly the keys 1..36")
    return mism


def build_plan(cert):
    P = Plan()
    P.cert = cert
    P.Zraw = [enc_pt(z) for z in cert["Z"]]
    if len(P.Zraw) != NZ:
        fail("|Z| = %d" % len(P.Zraw))
    Zf = [tuple(frac(p) for p in z) for z in cert["Z"]]
    for a, b in zip(Zf, Zf[1:]):
        if not a < b:
            fail("Z is not strictly increasing lexicographically")
    for z in Zf:
        if any(c < 0 for c in z):
            fail("negative coordinate in Z")
    internal = cert["internal"]
    P.ch1 = []     # entries with j == 1, file order
    P.st = []      # entries with j >= 2, file order
    for e in internal:
        j = e["j"]
        if len(e["combos"]) != j or len(e["stages"]) != j - 1:
            fail("internal entry %r inconsistent (j=%d, %d stages, %d combos)" %
                 ((e["m"], e["c2"], e["c3"]), j, len(e["stages"]), len(e["combos"])))
        if j == 1:
            P.ch1.append(e)
        elif j >= 2:
            P.st.append(e)
        else:
            fail("internal entry with j = %r" % j)
    P.leaf_keys = sorted(tuple(int(t) for t in k.split(",")) for k in cert["leafexp"])
    for k in P.leaf_keys:
        if len(k) != 3:
            fail("bad leaf key")
    P.n_ch1_groups = (len(P.ch1) + CH1_GROUP - 1) // CH1_GROUP
    # tables
    P.CT = [((e["m"], e["c2"], e["c3"]), 1) for e in P.ch1] + \
           [((e["m"], e["c2"], e["c3"]), e["j"]) for e in P.st]
    P.LT = list(P.leaf_keys)
    # stage sizes
    P.st_info = []   # per St entry: (key, j, [|X_i| for i in 1..j])
    for e in P.st:
        j = e["j"]
        sizes = [1] + [len(Y) for Y in e["stages"]]
        P.st_info.append(((e["m"], e["c2"], e["c3"]), j, sizes))
    return P


def nbl(n):
    return (n + BLK - 1) // BLK


def file_list(P):
    """[(name, family, producer)] in generation order."""
    cert = P.cert
    files = []

    def add(name, fam, fn):
        files.append((name, fam, fn))

    # Z
    def gen_Z():
        body = ["/-- The 1024 points of `Z` (8 numerators and the common denominator), original order. -/"]
        body += def_rows("Zraw", "List (List ℕ)", P.Zraw)
        body += ["", "def Z : List IPt := Zraw.map decodePt", "", stmt_Z_ok()]
        return lean_file(["MaxLeafFormal.RSCCheck"], NS, body), NZ
    add("Z.lean", "Z", gen_Z)

    # C1
    def gen_c1(r):
        def f():
            rows = []
            for i in range(BLK * r, BLK * r + BLK):
                row = []
                for j in range(NZ):
                    key = "%d,%d" % (i, j)
                    if key not in cert["combos"]:
                        fail("C1: missing combination %s" % key)
                    row += enc_comb(cert["combos"][key], NZ, "C1 " + key)
                rows.append(row)
            body = ["/-- Rows %d..%d of C1: for each j the combination of `Z` for the pair (i, j). -/" %
                    (BLK * r, BLK * r + BLK - 1)]
            body += def_rows("rows", "List (List ℕ)", rows)
            body += ["", stmt_c1(r)]
            return lean_file(["MaxLeafFormal.RSCData.Z"], "%s.C1_%03d" % (NS, r), body), BLK * NZ
        return f
    for r in range(NZ // BLK):
        add("C1_%03d.lean" % r, "C1", gen_c1(r))

    # C2
    def gen_c2():
        rows = []
        for m in list(range(25)) + ["inf"]:
            rows.append(enc_comb(cert["inits"][str(m)], NZ, "C2 init %s" % m))
        body = ["/-- Single combinations for `Init(0), ..., Init(24)` and `Init(inf)`. -/"]
        body += def_rows("raw", "List (List ℕ)", rows)
        body += ["", "def combos : List (List (ℕ × ℕ × ℕ)) := raw.map fun l => (decodeCombs 1 l).headD []",
                 "", stmt_c2()]
        return lean_file(["MaxLeafFormal.RSCData.Z"], NS + ".C2", body), 26
    add("C2.lean", "C2", gen_c2)

    # Leaf
    def gen_leaf():
        rows = []
        for (m, c2, c3) in P.leaf_keys:
            rows.append([m, c2, c3] + enc_comb(cert["leafexp"]["%d,%d,%d" % (m, c2, c3)], NZ,
                                              "leaf %d,%d,%d" % (m, c2, c3)))
        body = ["/-- Leaf entries `[m, c2, c3] ++ combination`, sorted by `(m, c2, c3)`. -/"]
        body += def_rows("raw", "List (List ℕ)", rows)
        body += ["", "def entries : List ((ℕ × ℕ × ℕ) × List (ℕ × ℕ × ℕ)) := raw.map decodeLeaf",
                 "", stmt_leaf()]
        return lean_file(["MaxLeafFormal.RSCData.Z"], NS + ".Leaf", body), len(rows)
    add("Leaf.lean", "Leaf", gen_leaf)

    # Ch1
    def gen_ch1(g):
        def f():
            rows = []
            for e in P.ch1[CH1_GROUP * g: CH1_GROUP * (g + 1)]:
                key = (e["m"], e["c2"], e["c3"])
                rows.append(list(key) + flat_combs(e["combos"][0], 0, NZ, NZ, "chain1 %r" % (key,)))
                if len(e["combos"][0]) != NZ:
                    fail("chain1 %r: %d combinations" % (key, len(e["combos"][0])))
            body = ["/-- Chains of length one: `[m, c2, c3]` followed by the 1024 combinations of `Z`. -/"]
            body += def_rows("raw", "List (List ℕ)", rows)
            body += ["", "def entries : List ((ℕ × ℕ × ℕ) × List (List (ℕ × ℕ × ℕ))) := raw.map (decodeChain1 1024)",
                     "", stmt_ch1()]
            return lean_file(["MaxLeafFormal.RSCData.Z"], "%s.Ch1_%03d" % (NS, g), body), len(rows) * NZ
        return f
    for g in range(P.n_ch1_groups):
        add("Ch1_%03d.lean" % g, "Ch1", gen_ch1(g))

    # St
    for k, e in enumerate(P.st):
        key, j, sizes = P.st_info[k]

        def gen_pts(k=k, e=e, key=key, j=j, sizes=sizes):
            def f():
                stages = [[enc_pt(p) for p in Y] for Y in e["stages"]]
                lines = ["/-- Longer chain %d: the exposure point `(m, c2, c3)` and its stages `Y_1, ..., Y_%d`. -/" %
                         (k, j - 1)]
                lines.append("def key : ℕ × ℕ × ℕ := (%d, %d, %d)" % key)
                lines.append("")
                st_txt = "[\n" + ",\n".join(
                    "  [\n" + ",\n".join("    [" + ", ".join(map(str, p)) + "]" for p in Y) + "\n  ]"
                    for Y in stages) + "\n]"
                lines.append("def stagesRaw : List (List (List ℕ)) := " + st_txt)
                lines.append("")
                lines.append("def stages : List (List IPt) := stagesRaw.map (·.map decodePt)")
                lines.append("")
                lines.append(thm("pts_ok", "stages.all okPts = true"))
                for i in range(1, j + 1):
                    lines.append("")
                    lines.append(stmt_nb(k, i, nbl(sizes[i - 1])))
                return lean_file(["MaxLeafFormal.RSCData.Z"], "%s.St_%d_pts" % (NS, k), lines), 0
            return f
        add("St_%d_pts.lean" % k, "St", gen_pts())

        for i in range(1, j + 1):
            nx = sizes[i - 1]
            ny = NZ if i == j else sizes[i]
            for r in range(nbl(nx)):
                def gen_blk(k=k, e=e, key=key, i=i, r=r, nx=nx, ny=ny):
                    def f():
                        cmb = e["combos"][i - 1]
                        if len(cmb) != nx * NZ:
                            fail("St_%d stage %d: %d combinations, expected %d" % (k, i, len(cmb), nx * NZ))
                        rows = [flat_combs(cmb, a, ny, ny, "St_%d stage %d row %d" % (k, i, a))
                                for a in range(BLK * r, min(BLK * r + BLK, nx))]
                        body = ["/-- Stage %d of the longer chain %d, rows %d..%d of `X`. -/" %
                                (i, k, BLK * r, min(BLK * r + BLK, nx) - 1)]
                        body += def_rows("rows", "List (List ℕ)", rows)
                        body += ["", stmt_st(k, i, r)]
                        return lean_file([st_mod(k)], "%s.St_%d_%d_%03d" % (NS, k, i, r), body), len(rows) * NZ
                    return f
                add("St_%d_%d_%03d.lean" % (k, i, r), "St", gen_blk())

    # Tables
    def gen_tables():
        mods = ["Ch1_%03d" % g for g in range(P.n_ch1_groups)]
        imports = ["%s.%s" % (NS, m) for m in mods] + [st_mod(k) for k in range(len(P.st))] + [NS + ".Leaf"]
        terms = ["(Ch1_%03d.entries.map fun e => (e.1, 1))" % g for g in range(P.n_ch1_groups)]
        terms += ["[(St_%d_pts.key, St_%d_pts.stages.length + 1)]" % (k, k) for k in range(len(P.st))]
        expr = terms[-1]
        for t in reversed(terms[:-1]):
            expr = "%s ++ (%s)" % (t, expr)
        body = ["/-- Certified chains `((nu, c2, c3), j)`: chains of length one, then the longer chains. -/",
                "def chainTable : List ((ℕ × ℕ × ℕ) × ℕ) :=",
                "  " + expr, "",
                "/-- Certified leaves. -/",
                "def leafTable : List ((ℕ × ℕ × ℕ) × Unit) := Leaf.entries.map fun e => (e.1, ())"]
        return lean_file(imports, NS, body), 0
    add("Tables.lean", "Tables", gen_tables)

    # Lattice
    def gen_lattice():
        stats = {}
        lat5, lat4 = build_lattice(P.CT, P.LT, stats)
        P.lattice_stats = stats
        body = ["/-- Witnesses for (C5): `[0, k]` or `[1, m, kind, idx, a, b, ...]`. -/"]
        body += def_rows("lat5", "List (List ℕ)", lat5)
        body += ["", "/-- Witnesses for (C4). -/"]
        body += def_rows("lat4", "List (List ℕ)", lat4)
        body += ["", stmt_lat5(), "", stmt_lat4()]
        return lean_file([NS + ".Tables"], NS, body), len(lat5) + len(lat4)
    add("Lattice.lean", "Lattice", gen_lattice)

    # All
    def gen_all():
        mods = ["Z"] + ["C1_%03d" % r for r in range(NZ // BLK)] + ["C2", "Leaf"]
        mods += ["Ch1_%03d" % g for g in range(P.n_ch1_groups)]
        for k in range(len(P.st)):
            key, j, sizes = P.st_info[k]
            mods.append("St_%d_pts" % k)
            for i in range(1, j + 1):
                for r in range(nbl(sizes[i - 1])):
                    mods.append("St_%d_%d_%03d" % (k, i, r))
        mods += ["Tables", "Lattice"]
        imports = [NS + "." + m for m in mods]
        # `List.Forall P [a₀, …, aₖ]` unfolds to `P a₀ ∧ … ∧ P aₖ` (no trailing `True`), and to `P a₀`
        # for a single element, so the proof is a flat anonymous-constructor tuple, bare for one element.
        def forall_tuple(items):
            return items[0] if len(items) == 1 else "⟨%s⟩" % ", ".join(items)
        c1 = forall_tuple(["⟨_, C1_%03d.ok⟩" % r for r in range(NZ // BLK)])
        terms = ["chain1_table Z_ok Ch1_%03d.ok" % g for g in range(P.n_ch1_groups)]
        for k in range(len(P.st)):
            key, j, sizes = P.st_info[k]
            facts = []
            for i in range(1, j + 1):
                nb = nbl(sizes[i - 1])
                blocks = forall_tuple(["⟨_, St_%d_%d_%03d.ok⟩" % (k, i, r) for r in range(nb)])
                facts.append("stage_of_blocks Z_ok (stage_X_ok St_%d_pts.pts_ok St_%d_pts.key %d) "
                             "(stage_Yt_ok Z_ok St_%d_pts.pts_ok %d) %d St_%d_pts.nb_%d %s" %
                             (k, k, i - 1, k, i - 1, nb, k, i, blocks))
            terms.append("chain_stages_table Z_ok St_%d_pts.key St_%d_pts.stages St_%d_pts.pts_ok %s" %
                         (k, k, k, forall_tuple(facts) if len(facts) == 1 else
                          "⟨%s⟩" % ",\n      ".join(facts)))
        expr = terms[-1]
        for t in reversed(terms[:-1]):
            expr = "forall_mem_append' (%s)\n    (%s)" % (t, expr)
        body = [
            "theorem c1_fact : ∀ z ∈ finsetOf Z, ∀ z' ∈ finsetOf Z, merge p₂ p₃ z z' ∈ PZ (finsetOf Z) :=",
            "  c1_of_blocks Z_ok %s" % c1, "",
            "theorem chainTable_ok : ∀ e ∈ chainTable, ChainFact Z e :=",
            "  " + expr, "",
            "theorem leafTable_ok : ∀ e ∈ leafTable, LeafFact Z e := leaf_table Z_ok Leaf.ok"]
        txt = lean_file(imports, NS, body)
        txt = txt.replace("open MaxLeafFormal.RSCCheck\n",
                          "open MaxLeafFormal.RSCCheck\nopen MaxLeafFormal.RSCCert MaxLeafFormal.RSCParams\n"
                          "open MaxLeafFormal.Antinorm8 (PZ)\n", 1)
        return txt, 0
    add("All.lean", "All", gen_all)
    return files


def count_conditions(P):
    n = {"C1": NZ * NZ, "C2": 26, "Leaf": len(P.leaf_keys), "Ch1": len(P.ch1) * NZ}
    st = 0
    for key, j, sizes in P.st_info:
        st += sum(sizes) * NZ
    n["St"] = st
    n["Lattice5"] = len(lattice5_points())
    n["Lattice4"] = len(lattice4_points())
    return n


def generate(cert, outdir, only, exporter_sha):
    P = build_plan(cert)
    files = file_list(P)
    os.makedirs(outdir, exist_ok=True)
    produced = {}
    for name, fam, fn in files:
        if only is not None and not name.startswith(only):
            continue
        # Lattice needs only table metadata, the others read the certificate lazily
        txt, ncond = fn()
        data = txt.encode("utf-8")
        with open(os.path.join(outdir, name), "w", encoding="utf-8", newline="\n") as fh:
            fh.write(txt)
        produced[name] = (fam, len(data), sha256_bytes(data))
    return P, files, produced


def write_manifest(outdir, produced, exporter_sha):
    lines = ["# input-gz sha256 %s" % GZ_SHA, "# input-json sha256 %s" % JSON_SHA,
             "# exporter sha256 %s" % exporter_sha,
             "# python %d.%d.%d" % sys.version_info[:3]]
    for name in sorted(produced):
        lines.append("%s  %s" % (produced[name][2], name))
    with open(os.path.join(outdir, "MANIFEST.sha256"), "w", encoding="utf-8", newline="\n") as fh:
        fh.write("\n".join(lines) + "\n")


# ----------------------------------------------------------------------------------------------
# self-check: re-parse the generated files and mirror RSCCheck.lean in exact integer arithmetic

E_N, E_D = 145729958, 189014575


def merge_n(P, Q):
    P0, P1, P2, P3, P4, P5, P6, P7 = P
    Q0, _, Q2, Q3, Q4, Q5, Q6, Q7 = Q
    S = 2000 * Q0 + 2000 * Q2 + 2200 * Q3 + 1000 * Q4 + 1000 * Q5 + 930 * Q6 + 923 * Q7
    Cs = 2000 * Q0
    Cu = 2000 * Q3 + 1000 * Q4 + 1000 * Q5 + 1000 * Q6 + 930 * Q7
    Os = 1000 * Q0
    Ox = 1000 * Q3
    Ou = 500 * Q4 + 1000 * Q5 + 1000 * Q6 + 930 * Q7
    return [P0 * S, P1 * Cu, P2 * Cu + P1 * Cs, P3 * (Cu + Cs) + P2 * Cs, P4 * Ou,
            P5 * (Ou + Ox) + P4 * Ox, P6 * (Ou + Ox) + (P4 + P5) * Os, P7 * (Ou + Ox + Os) + P6 * Os]


def merge_i(P, Q):
    return (merge_n(P[0], Q[0]), 1000 * P[1] * Q[1])


def init_i(nu, c2, c3):
    s = E_D ** (1 + nu)
    uN = 93 ** c2 * 923 ** c3
    uD = 100 ** c2 * 1000 ** c3
    wN = 272657 ** c2 * 288677 ** c3
    wD = 2 ** nu * 250000 ** (c2 + c3)
    D = E_N ** (1 + nu) * 200 * uD * wD
    return ([s * uN * 31 * wD, s * uN * 7 * wD, 0, 0, s * wN * 200 * uD, 0, 0, 0], D)


INIT_INF = ([E_D ** 26 * 31, E_D ** 26 * 7, 0, 0, 0, 0, 0, 0], E_N ** 26 * 200)
ZERO_I = ([0] * 8, 1)


def in_k(d):
    d0, d1, d2, d3, d4, d5, d6, d7 = d
    return 0 <= d0 and 0 <= d3 + min(0, d2 + min(0, d1)) + min(0, d5 + min(0, d4) + min(0, d6 + min(0, d7)))


def dom(n, Y, P, comb):
    for k, a, b in comb:
        if not (k < n and 0 < b):
            return False
    V = [0] * 8
    T, s, t = 1, 0, 1
    for k, a, b in reversed(comb):
        y = Y[k] if k < len(Y) else ZERO_I
        by = b * y[1]
        V = [v * by + a * nn * T for v, nn in zip(V, y[0])]
        T *= by
        s = s * b + a * t
        t *= b
    if t > s:
        return False
    return in_k([m * T - P[1] * v for m, v in zip(P[0], V)])


def check_block(n, Y, Zs, xs, combos):
    if len(xs) != len(combos):
        return False
    for x, cx in zip(xs, combos):
        if len(cx) != len(Zs):
            return False
        for z, c in zip(Zs, cx):
            if not dom(n, Y, merge_i(x, z), c):
                return False
    return True


def take_comb(c, l):
    out = []
    p = 0
    for _ in range(c):
        if p + 3 > len(l):
            return [], []
        out.append((l[p], l[p + 1], l[p + 2]))
        p += 3
    return out, l[p:]


def decode_combs(m, l):
    out = []
    for _ in range(m):
        if not l:
            break
        c = l[0]
        comb, l = take_comb(c, l[1:])
        out.append(comb)
    return out


def decode_pt(l):
    return (list(l[:8]), l[8]) if len(l) == 9 else ZERO_I


def ok_pts(Y):
    return all(len(P[0]) == 8 and 0 < P[1] for P in Y)


def fin_n(N):
    return 2000 * N[0] + 2000 * N[3] + 1000 * N[4] + 2000 * N[5] + 2000 * N[6] + 1860 * N[7]


def lt_lex(D, D2, a, b):
    for x, y in zip(a, b):
        if x * D2 < y * D:
            return True
        if x * D2 != y * D:
            return False
    return False


def check_z(Z):
    return (len(Z) == 1024 and ok_pts(Z) and all(lt_lex(P[1], Q[1], P[0], Q[0]) for P, Q in zip(Z, Z[1:]))
            and all(6060 * 2000 * z[1] < 10000 * fin_n(z[0]) for z in Z))


def uF(c2, c3):
    return (93 ** c2 * 923 ** c3, 100 ** c2 * 1000 ** c3)


def wF(c2, c3):
    return (272657 ** c2 * 288677 ** c3, 250000 ** (c2 + c3))


def comb_witness(c2, c3, comps):
    """comps = [((c2', c3'), a, b)]."""
    if not all(0 < e[2] for e in comps):
        return False
    S = (0, 1)
    for e in reversed(comps):
        S = (S[0] * e[2] + e[1] * S[1], S[1] * e[2])
    if S[0] != S[1]:
        return False
    for F, tgt in ((uF, uF(c2, c3)), (wF, wF(c2, c3))):
        S = (0, 1)
        for e in reversed(comps):
            fk = F(*e[0])
            S = (S[0] * (e[2] * fk[1]) + e[1] * fk[0] * S[1], S[1] * (e[2] * fk[1]))
        if S[0] * tgt[1] != tgt[0] * S[1]:
            return False
    return True


def take_comp(m, l):
    out = []
    for q in range(m):
        if 4 * q + 4 > len(l):
            return out, False
        out.append(tuple(l[4 * q: 4 * q + 4]))
    return out, len(l) == 4 * m


def lat_ok(T, five, pt, dat):
    """five: table entries (key, j); otherwise entries are keys only."""
    def key(k):
        return T[k][0] if five else T[k]
    if len(dat) == 2 and dat[0] == 0:
        k = dat[1]
        return k < len(T) and key(k) == pt and (not five or T[k][1] <= j_target(pt[0], pt[1]))
    if len(dat) >= 2 and dat[0] == 1:
        m = dat[1]
        comps, exact = take_comp(m, dat[2:])
        if not exact:
            return False
        out = []
        for kind, idx, a, b in comps:
            if kind == 0:
                out.append(((0, 0), a, b))
            elif kind == 1:
                if not (idx < len(T) and key(idx)[0] == pt[0]):
                    return False
                if five and not T[idx][1] <= j_target(pt[0], pt[1]):
                    return False
                out.append(((key(idx)[1], key(idx)[2]), a, b))
            else:
                return False
        return comb_witness(pt[1], pt[2], out)
    return False


# --- parsing the generated files

DEF_RE = re.compile(r"^def (\w+) : [^\n]*:= (\[.*?\n\]|\[\])$", re.S | re.M)


def read_file(outdir, name):
    with open(os.path.join(outdir, name), "r", encoding="utf-8", newline="") as fh:
        return fh.read()


def parse_defs(txt):
    return {m.group(1): json.loads(m.group(2)) for m in DEF_RE.finditer(txt)}


_zcache = {}


def load_Z(outdir):
    if outdir not in _zcache:
        d = parse_defs(read_file(outdir, "Z.lean"))
        _zcache[outdir] = [decode_pt(p) for p in d["Zraw"]]
    return _zcache[outdir]


def need_line(txt, line, what):
    if line not in txt.split("\n"):
        raise ExportError("%s: expected line not found: %s" % (what, line[:200]))


def parse_key(txt):
    m = re.search(r"^def key : ℕ × ℕ × ℕ := \((\d+), (\d+), (\d+)\)$", txt, re.M)
    return tuple(int(x) for x in m.groups())


def load_pts(outdir, k):
    txt = read_file(outdir, "St_%d_pts.lean" % k)
    d = parse_defs(txt)
    return parse_key(txt), [[decode_pt(p) for p in Y] for Y in d["stagesRaw"]], txt


def chain_keys(outdir):
    """The chain table entries parsed from Tables.lean, Ch1 and St pts files."""
    tab = read_file(outdir, "Tables.lean")
    T = []
    for m in re.finditer(r"\(Ch1_(\d+)\.entries\.map|\[\(St_(\d+)_pts\.key", tab):
        if m.group(1) is not None:
            d = parse_defs(read_file(outdir, "Ch1_%s.lean" % m.group(1)))
            T += [((r[0], r[1], r[2]), 1) for r in d["raw"]]
        else:
            key, stages, _ = load_pts(outdir, int(m.group(2)))
            T.append((key, len(stages) + 1))
    return T


def leaf_keys(outdir):
    d = parse_defs(read_file(outdir, "Leaf.lean"))
    return [(r[0], r[1], r[2]) for r in d["raw"]]


def selfcheck_file(args):
    outdir, name = args
    txt = read_file(outdir, name)
    base = name[:-5]
    try:
        if base == "Z":
            need_line(txt, stmt_Z_ok(), name)
            Z = load_Z(outdir)
            if not check_z(Z):
                raise ExportError("checkZ failed")
            return name, "Z", NZ, None
        Z = load_Z(outdir)
        n = len(Z)
        m = re.fullmatch(r"C1_(\d+)", base)
        if m:
            r = int(m.group(1))
            need_line(txt, stmt_c1(r), name)
            rows = parse_defs(txt)["rows"]
            combos = [decode_combs(1024, row) for row in rows]
            if not check_block(n, Z, Z, Z[8 * r: 8 * r + 8], combos):
                raise ExportError("C1 block %d failed" % r)
            return name, "C1", 8 * 1024, None
        if base == "C2":
            need_line(txt, stmt_c2(), name)
            raw = parse_defs(txt)["raw"]
            combos = [(decode_combs(1, l) or [[]])[0] for l in raw]
            ok = len(combos) == 26 and all(dom(n, Z, init_i(i, 0, 0), combos[i]) for i in range(25)) \
                and dom(n, Z, INIT_INF, combos[25])
            if not ok:
                raise ExportError("C2 failed")
            return name, "C2", 26, None
        if base == "Leaf":
            need_line(txt, stmt_leaf(), name)
            raw = parse_defs(txt)["raw"]
            cnt = 0
            keys = [tuple(r[:3]) for r in raw]
            if keys != sorted(keys):
                raise ExportError("Leaf not sorted")
            for r in raw:
                comb = (decode_combs(1, r[3:]) or [[]])[0]
                if not dom(n, Z, init_i(r[0], r[1], r[2]), comb):
                    raise ExportError("leaf %r failed" % (tuple(r[:3]),))
                cnt += 1
            return name, "Leaf", cnt, None
        m = re.fullmatch(r"Ch1_(\d+)", base)
        if m:
            need_line(txt, stmt_ch1(), name)
            raw = parse_defs(txt)["raw"]
            for r in raw:
                combs = decode_combs(1024, r[3:])
                if not check_block(n, Z, Z, [init_i(r[0], r[1], r[2])], [combs]):
                    raise ExportError("chain1 %r failed" % (tuple(r[:3]),))
            return name, "Ch1", len(raw) * 1024, None
        m = re.fullmatch(r"St_(\d+)_pts", base)
        if m:
            k = int(m.group(1))
            key, stages, _ = load_pts(outdir, k)
            if not all(ok_pts(Y) for Y in stages):
                raise ExportError("St pts not ok")
            need_line(txt, thm("pts_ok", "stages.all okPts = true"), name)
            sizes = [1] + [len(Y) for Y in stages]
            for i in range(1, len(sizes) + 1):
                need_line(txt, stmt_nb(k, i, nbl(sizes[i - 1])), name)
            return name, "St", 0, None
        m = re.fullmatch(r"St_(\d+)_(\d+)_(\d+)", base)
        if m:
            k, i, r = (int(x) for x in m.groups())
            key, stages, _ = load_pts(outdir, k)
            need_line(txt, stmt_st(k, i, r), name)
            Xs = [[init_i(*key)]] + stages
            Yt = (stages + [Z])[i - 1]
            X = Xs[i - 1]
            rows = parse_defs(txt)["rows"]
            combos = [decode_combs(1024, row) for row in rows]
            if not check_block(len(Yt), Yt, Z, X[8 * r: 8 * r + 8], combos):
                raise ExportError("stage block failed")
            return name, "St", len(rows) * 1024, None
        if base == "Lattice":
            need_line(txt, stmt_lat5(), name)
            need_line(txt, stmt_lat4(), name)
            d = parse_defs(txt)
            CT = chain_keys(outdir)
            LT = leaf_keys(outdir)
            l5, l4 = lattice5_points(), lattice4_points()
            if len(l5) != len(d["lat5"]) or len(l4) != len(d["lat4"]):
                raise ExportError("lattice length")
            for pt, dat in zip(l5, d["lat5"]):
                if not lat_ok(CT, True, pt, dat):
                    raise ExportError("lattice5 %r failed" % (pt,))
            for pt, dat in zip(l4, d["lat4"]):
                if not lat_ok(LT, False, pt, dat):
                    raise ExportError("lattice4 %r failed" % (pt,))
            return name, "Lattice", len(l5) + len(l4), None
        if base in ("Tables", "All"):
            return name, base, 0, None
        raise ExportError("unknown file")
    except Exception as ex:   # reported by the caller
        return name, "?", 0, "%s: %r" % (name, ex)


def run_selfcheck(outdir, names, procs):
    import multiprocessing
    jobs = [(outdir, nm) for nm in names]
    # heavy files first
    jobs.sort(key=lambda j: (0 if j[1].startswith(("Lattice", "St_", "Ch1_")) else 1, j[1]))
    totals = {}
    errors = []
    with multiprocessing.Pool(procs) as pool:
        for k, (name, fam, cnt, err) in enumerate(pool.imap_unordered(selfcheck_file, jobs, chunksize=1)):
            if err:
                errors.append(err)
            else:
                totals[fam] = totals.get(fam, 0) + cnt
            if (k + 1) % 50 == 0:
                print("  selfcheck: %d / %d files" % (k + 1, len(jobs)), flush=True)
    return totals, errors


# ----------------------------------------------------------------------------------------------

def main():
    try:
        sys.stdout.reconfigure(encoding="utf-8")
    except Exception:
        pass
    ap = argparse.ArgumentParser()
    ap.add_argument("gz")
    ap.add_argument("outdir")
    ap.add_argument("--only")
    ap.add_argument("--selfcheck", action="store_true")
    ap.add_argument("--procs", type=int, default=6)
    ap.add_argument("--method", choices=["m1", "m2"], default="m2")
    a = ap.parse_args()
    with open(os.path.abspath(__file__), "rb") as fh:
        exporter_sha = sha256_bytes(fh.read())
    Ctx.exporter_sha = exporter_sha
    Ctx.method = a.method
    with open(a.gz, "rb") as fh:
        gzb = fh.read()
    if sha256_bytes(gzb) != GZ_SHA:
        fail("SHA-256 of the gz file is %s, expected %s" % (sha256_bytes(gzb), GZ_SHA))
    raw = gzip.decompress(gzb)
    if sha256_bytes(raw) != JSON_SHA:
        fail("SHA-256 of the decompressed JSON is %s, expected %s" % (sha256_bytes(raw), JSON_SHA))
    cert = json.loads(raw)
    del raw, gzb
    mism = check_input_and_params(cert)
    print("exporter sha256 %s ; python %s" % (exporter_sha, sys.version.split()[0]))
    print("input checks passed; dmin raw-value mismatches vs the paper's d (same j and same '<= 1'): %d" % len(mism))
    if mism:
        print("  " + ", ".join("%s(%d)=%d vs %d" % t for t in mism[:6]) + (" ..." if len(mism) > 6 else ""))
    outdir = os.path.abspath(a.outdir)
    P, files, produced = generate(cert, outdir, a.only, exporter_sha)
    cond = count_conditions(P)
    if a.only is None:
        write_manifest(outdir, produced, exporter_sha)
    fams = {}
    for name, (fam, size, h) in produced.items():
        f = fams.setdefault(fam, [0, 0])
        f[0] += 1
        f[1] += size
    print("files generated: %d, total bytes: %d" % (len(produced), sum(v[1] for v in produced.values())))
    for fam in sorted(fams):
        print("  %-8s files %5d  bytes %12d" % (fam, fams[fam][0], fams[fam][1]))
    print("conditions per family (full certificate):")
    for k in sorted(cond):
        print("  %-9s %d" % (k, cond[k]))
    print("chainTable entries: %d (chain-1: %d, longer chains: %d); leafTable entries: %d" %
          (len(P.CT), len(P.ch1), len(P.st), len(P.LT)))
    print("Ch1 modules: %d; longer chains: %s" % (P.n_ch1_groups, [(k[0], k[1], k[2], j) for k, j, s in P.st_info]))
    if hasattr(P, "lattice_stats"):
        for w in ("5", "4"):
            s = P.lattice_stats["lattice" + w]
            print("lattice%s: points %d, certified directly %d, combinations %d (plain fan %d, hull fan %d), "
                  "max components %d, max digits of a weight numerator/denominator %d" %
                  (w, s["direct"] + s["combo"], s["direct"], s["combo"], s["fan"], s["hull"],
                   s["maxcomps"], s["maxdig"]))
    del cert, P
    if a.selfcheck:
        names = sorted(n for n in produced)
        print("selfcheck on %d files with %d processes" % (len(names), a.procs), flush=True)
        totals, errors = run_selfcheck(outdir, names, a.procs)
        print("selfcheck conditions verified per family: " + ", ".join("%s %d" % kv for kv in sorted(totals.items())))
        if errors:
            for e in errors[:20]:
                print("SELFCHECK FAILURE: " + e, file=sys.stderr)
            fail("selfcheck failed in %d files" % len(errors))
        print("selfcheck passed")


if __name__ == "__main__":
    try:
        main()
    except ExportError as ex:
        print("EXPORT ERROR: %s" % ex, file=sys.stderr)
        sys.exit(1)