#!/usr/bin/env python3
"""Independent checker for the integer operator corpus (extensions/integer-operators.md).

Python standard library only; it shares no code with wai-ops. Written from the
extension's text (and, for the source forms it lowers, integer-payloads), it:

  - recomputes every vector file under ops-conformance/vectors/, stage/ and graphs/:
    each graph vector's `load` (WHG1 parse and load steps 1-8, section 10.5), its
    `profiles` (section 14.1, from the bounds pass and verdict of sections 8.1-8.5),
    and each decode's `expect` (the run of section 4, every graph operator of
    section 5), and each stage vector's `expect` (section 5E); it writes each file
    again canonically (sections 8.5 and 14.1) and requires the bytes of the
    committed file;
  - re-lowers every parameter set certificates/index.json lists, by the canonical
    lowerings of section 13, and requires each lowered graph's bytes, each set's
    digest (SPEC 3.1), and that every lowered graph loads; and requires the graph
    of each graphs/lower_*.jsonl vector to be the role of the set its id names, as
    re-lowered;
  - is a certifier (section 1): from each set's lowered graphs it computes the
    bounds certificate (the pass of sections 8.1-8.3, the classes and parameter
    conditions of section 7, the flags and reasons of section 9.2, and the canonical
    writer of section 8.5), and requires the bytes of the committed certificate,
    for every set of index.json and the two falsification sets of section 14.3,
    each of which must fail the profile it claims. The sets of external.json are
    certified too when their files are given (--external NAME=DIR); otherwise their
    certificates are checked for what needs no set. Either way, an entry's verdict
    must meet its `profile`, and each graph its `role_profiles` names must meet its
    own profile (section 11).

    python3 -I tools/wai_ops_check.py ops-conformance             # check
    python3 -I tools/wai_ops_check.py ops-conformance --write DIR  # also write the files
    python3 -I tools/wai_ops_check.py ops-conformance --external jpegai/v1=DIR ...

Exit status 1 on any difference or failure.

    python3 -I tools/wai_ops_check.py certify --capability CAP --lowering ID \\
        [--c N] --member NAME=PATH ... [--profile P] [--role-profile ROLE=P ...] \\
        [--out FILE] [--check FILE]
    python3 -I tools/wai_ops_check.py certify --whg1 ROLE=PATH ... [--profile P] ...

certifies one set, or WHG1 graphs on their own, writing the certificate to FILE or
to stdout. --role-profile (repeatable) holds the graph of one role to a profile of
its own, as a capability's `role_profiles` does (section 11): that graph meets it
when its inputs' and scalars' declared ranges meet the profile's tensor width and
every node of it fits the profile (section 8.5). Exit status: 0 the verdict meets
the profile (acc64 by default) and each named graph meets its own; 1 one does not,
and the certificate is still written; 2 a graph is malformed or refused at load;
3 --check finds a difference, and the first differing path is printed with both
values; 4 the lowering refuses the set; 5 usage, or a --role-profile role that
names no graph.

The lowering jpegai-int/1 (section 13.6) reads eleven members, norm.json and
weights/{hs_y,hs_uv,hyp_y,hyp_uv,mcm0_y,mcm1_y,mcm2_y,mcm3_y,synth_y,synth_uv}.json,
and writes six roles: hs_y and hs_uv (13.6.1), rec_y, rec_uv, syn_y and syn_uv
(13.6.2). Any other member is part of the set's digest and is not read.
"""
import hashlib
import json
import math
import os
import struct
import sys

CAP = 1 << 26
I64_MIN, I64_MAX = -(1 << 63), (1 << 63) - 1
I32_MIN, I32_MAX = -(1 << 31), (1 << 31) - 1


class Refuse(Exception):
    """A refusal with a code of section 4 or 5E, and, for a stage operator, `at`."""

    def __init__(self, code, at=None):
        super().__init__(code)
        self.code = code
        self.at = at


# --------------------------------------------------------------- section 2


def R(p, s, mode):
    """R_floor, R_half_up, R_half_even (section 2)."""
    if s == 0:
        return p
    if mode == 0:
        return p >> s
    if mode == 1:
        return (p + (1 << (s - 1))) >> s
    q = p >> s
    r = p - (q << s)
    h = 1 << (s - 1)
    if r > h or (r == h and q % 2 == 1):
        return q + 1
    return q


def exp_neg_q16(x):
    """exp_neg_q16(x), x <= 0, in the order section 2 writes it."""
    u = -x
    n = u // 45426
    g = 45426 * n
    r = x + g
    t = r + 88670
    t2 = t * t
    q = t2 >> 16
    c = 23495 * q
    y = (c >> 16) + 22544
    return 0 if n >= 63 else y >> n


def rsqrt_q16(v):
    r = math.isqrt(v)
    return 0 if r == 0 else (1 << 32) // r


def trunc_div(a, b):
    q = abs(a) // abs(b)
    return q if (a >= 0) == (b > 0) else -q


# --------------------------------------------------------------- WHG1 (section 10)

NAMES = {
    0x01: "conv", 0x02: "rescale", 0x03: "chscale", 0x04: "lut16", 0x05: "add",
    0x06: "concat", 0x07: "chunk_add", 0x08: "split", 0x09: "pixel_shuffle",
    0x0A: "pixel_unshuffle", 0x0B: "clamp", 0x0C: "relu", 0x0D: "leaky",
    0x0E: "pad_rep", 0x0F: "crop_to", 0x10: "mul", 0x11: "conv_transpose",
    0x12: "igdn", 0x13: "abs", 0x14: "crop", 0x15: "subpos_split",
    0x16: "subpos_merge", 0x17: "fill", 0x18: "layernorm_c", 0x19: "l2norm",
    0x1A: "softmax", 0x1B: "sigmoid", 0x1C: "elu", 0x1D: "gram", 0x1E: "apply",
}
CODE = {v: k for k, v in NAMES.items()}
TWO_IN = {"add", "crop_to", "mul", "gram", "apply"}


def arity(name):
    if name in TWO_IN:
        return (2, 1)
    if name == "concat":
        return (None, 1)
    if name == "subpos_split":
        return (1, 4)
    if name == "subpos_merge":
        return (4, 1)
    return (1, 1)


class Reader:
    def __init__(self, b, pos=0):
        self.b, self.p = b, pos

    def take(self, n):
        if n < 0 or self.p + n > len(self.b):
            raise Refuse("model_mismatch")
        s = self.b[self.p:self.p + n]
        self.p += n
        return s

    def u(self, n):
        return int.from_bytes(self.take(n), "little")

    def i(self, n):
        return int.from_bytes(self.take(n), "little", signed=True)

    def ints(self, count, width):
        raw = self.take(count * width)
        return [int.from_bytes(raw[k:k + width], "little", signed=True) for k in range(0, len(raw), width)]

    def zeros(self, n):
        if any(self.take(n)):
            raise Refuse("model_mismatch")


def bad():
    raise Refuse("model_mismatch")


def read_rq(rd, C):
    """The RQ block (section 10.4) over C channels."""
    rq = {}
    rq["round"], rq["bias"], rq["m_kind"], rq["s_kind"] = rd.u(1), rd.u(1), rd.u(1), rd.u(1)
    rq["m_bytes"], rq["b_bytes"] = rd.u(1), rd.u(1)
    if rd.u(2) != 0:
        bad()
    if rq["round"] > 2 or rq["bias"] > 2 or rq["m_kind"] > 2 or rq["s_kind"] not in (1, 2):
        bad()
    if rq["m_kind"] == 0:
        if rq["m_bytes"] != 0:
            bad()
    elif rq["m_bytes"] not in (4, 8):
        bad()
    if rq["bias"] == 0:
        if rq["b_bytes"] != 0:
            bad()
    elif rq["b_bytes"] not in (4, 8):
        bad()
    nm = {0: 0, 1: 1, 2: C}[rq["m_kind"]]
    rq["mult"] = rd.ints(nm, rq["m_bytes"]) if nm else []
    ns = 1 if rq["s_kind"] == 1 else C
    rq["shift"] = list(rd.take(ns))
    if any(s > 127 for s in rq["shift"]):
        bad()
    rq["biasv"] = rd.ints(C, rq["b_bytes"]) if rq["bias"] else []
    return rq


def rq_m(rq, o):
    return None if rq["m_kind"] == 0 else rq["mult"][0 if rq["m_kind"] == 1 else o]


def rq_s(rq, o):
    return rq["shift"][0 if rq["s_kind"] == 1 else o]


def rq_apply(rq, o, v):
    """rq_o(v), section 5.1."""
    t = v + (rq["biasv"][o] if rq["bias"] == 1 else 0)
    m = rq_m(rq, o)
    p = t * (1 if m is None else m)
    q = R(p, rq_s(rq, o), rq["round"])
    return q + (rq["biasv"][o] if rq["bias"] == 2 else 0)


def read_operand(rd, n_scalars):
    kind = rd.u(1)
    rd.zeros(3)
    v = rd.u(4)
    if kind == 0:
        return ("const", v)
    if kind == 1 and v < n_scalars:
        return ("scalar", v)
    bad()


def read_params(name, pb, cin, n_scalars):
    rd = Reader(pb)
    P = {}
    if name == "conv":
        for k in ("cout", "kh", "kw", "sh", "sw", "ph", "pw", "groups"):
            P[k] = rd.u(4)
        P["w_bytes"] = rd.u(1)
        rd.zeros(3)
        if min(P["cout"], P["kh"], P["kw"], P["sh"], P["sw"], P["groups"]) == 0 or P["w_bytes"] not in (1, 2, 4, 8):
            bad()
        if cin % P["groups"] or P["cout"] % P["groups"]:
            bad()
        P["rq"] = read_rq(rd, P["cout"])
        P["w"] = rd.ints(P["cout"] * (cin // P["groups"]) * P["kh"] * P["kw"], P["w_bytes"])
    elif name == "conv_transpose":
        for k in ("cout", "kh", "kw", "sh", "sw", "ph", "pw", "oph", "opw"):
            P[k] = rd.u(4)
        P["w_bytes"] = rd.u(1)
        rd.zeros(3)
        if min(P["cout"], P["kh"], P["kw"], P["sh"], P["sw"]) == 0 or P["w_bytes"] not in (1, 2, 4, 8):
            bad()
        P["rq"] = read_rq(rd, P["cout"])
        P["w"] = rd.ints(cin * P["cout"] * P["kh"] * P["kw"], P["w_bytes"])
    elif name == "rescale":
        P["rq"] = read_rq(rd, cin)
    elif name == "chscale":
        P["scalar"], P["n_rows"] = rd.u(4), rd.u(4)
        P["round"], P["m_bytes"], P["shift"] = rd.u(1), rd.u(1), rd.u(1)
        rd.zeros(1)
        if P["scalar"] >= n_scalars or P["n_rows"] == 0 or P["m_bytes"] not in (4, 8) or P["shift"] > 127 or P["round"] > 2:
            bad()
        P["M"] = rd.ints(P["n_rows"] * cin, P["m_bytes"])
    elif name == "mul":
        P["rq"] = read_rq(rd, cin)
        if P["rq"]["bias"] != 0:
            bad()
    elif name == "lut16":
        P["e_in"], P["e_out"] = rd.i(1), rd.i(1)
        rd.zeros(2)
        P["T"] = rd.ints(65536, 2)
        if P["T"][0] != 0:
            bad()
    elif name in ("add", "concat", "chunk_add", "relu", "abs", "crop_to", "softmax", "sigmoid"):
        pass
    elif name == "split":
        P["start"], P["count"] = rd.u(4), rd.u(4)
        if P["count"] == 0:
            bad()
    elif name in ("pixel_shuffle", "pixel_unshuffle"):
        P["r"] = rd.u(4)
        if P["r"] == 0:
            bad()
    elif name == "clamp":
        P["lo"], P["hi"] = rd.i(8), rd.i(8)
        if P["lo"] > P["hi"]:
            bad()
    elif name == "leaky":
        P["a"], P["f"] = rd.i(8), rd.u(4)
        if P["f"] > 63:
            bad()
    elif name == "pad_rep":
        P["m"] = rd.u(4)
        if P["m"] == 0:
            bad()
    elif name == "crop":
        P["ops"] = [read_operand(rd, n_scalars) for _ in range(4)]
    elif name in ("subpos_split", "subpos_merge"):
        P["order"] = list(rd.take(4))
        if sorted(P["order"]) != [0, 1, 2, 3]:
            bad()
    elif name == "fill":
        P["kind"] = rd.u(1)
        rd.zeros(7)
        P["value"] = rd.i(8)
        if not (P["kind"] == 0 or (P["kind"] == 1 and 0 <= P["value"] < n_scalars)):
            bad()
    elif name == "igdn":
        P["gf"], P["frac"] = rd.u(4), rd.u(4)
        P["g_bytes"] = rd.u(1)
        rd.zeros(7)
        if P["gf"] > 127 or P["frac"] > 127 or P["g_bytes"] not in (4, 8):
            bad()
        P["gamma"] = rd.ints(cin * cin, P["g_bytes"])
        P["beta"] = rd.ints(cin, 8)
    elif name == "layernorm_c":
        P["eps"] = rd.u(8)
        P["g_bytes"] = rd.u(1)
        rd.zeros(7)
        if P["eps"] == 0 or P["g_bytes"] not in (4, 8):
            bad()
        P["gamma"] = rd.ints(cin, P["g_bytes"])
        P["beta"] = rd.ints(cin, 8)
    elif name == "l2norm":
        P["eps"] = rd.u(8)
        if P["eps"] == 0:
            bad()
    elif name == "elu":
        P["alpha"] = rd.i(8)
    elif name in ("gram", "apply"):
        P["groups"] = rd.u(4)
        P["rq"] = read_rq(rd, cin)
        if P["groups"] == 0 or P["rq"]["m_kind"] != 0 or P["rq"]["bias"] != 0 or P["rq"]["s_kind"] != 1:
            bad()
    if rd.p != len(pb):
        bad()
    return P


def parse_whg1(b):
    """Load steps 1-6 of section 10.5. Every refusal is model_mismatch."""
    if len(b) < 32 or b[0:4] != b"WHG1" or b[4] != 1 or b[5] != 1:
        bad()
    rd = Reader(b, 6)
    r0 = rd.u(2)
    n_t, n_o = rd.u(4), rd.u(4)
    n_i, n_out, n_s = rd.u(2), rd.u(2), rd.u(2)
    r1 = rd.u(2)
    max_elems = rd.u(4)
    r2 = rd.u(4)
    if r0 or r1 or r2:
        bad()
    if not (n_t and n_o and n_i and n_out) or not (1 <= max_elems <= CAP):
        bad()
    if 32 + 16 * n_s + 32 * n_i + 4 * n_out + 24 * n_t + 12 * n_o > len(b):
        bad()
    g = {"max_elems": max_elems, "scalars": [], "inputs": [], "outputs": [], "tensors": [], "ops": []}
    for _ in range(n_s):
        lo, hi = rd.i(8), rd.i(8)
        if lo > hi:
            bad()
        g["scalars"].append((lo, hi))
    for _ in range(n_i):
        t, mh, mw = rd.u(4), rd.u(4), rd.u(4)
        kind = rd.u(1)
        rd.zeros(3)
        src = rd.u(4)
        rd.zeros(4)
        init = rd.i(8)
        if (mh == 0) != (mw == 0) or kind > 1 or t >= n_t:
            bad()
        if kind == 0 and (src or init):
            bad()
        g["inputs"].append({"tensor": t, "max_h": mh, "max_w": mw, "kind": kind, "src": src, "init": init})
    for _ in range(n_out):
        t = rd.u(4)
        if t >= n_t:
            bad()
        g["outputs"].append(t)
    for _ in range(n_t):
        ch = rd.u(4)
        ex = rd.i(1)
        pol = rd.u(1)
        rd.zeros(2)
        lo, hi = rd.i(8), rd.i(8)
        if ch == 0 or pol > 2:
            bad()
        if pol == 0 and (lo or hi):
            bad()
        if pol != 0 and lo > hi:
            bad()
        g["tensors"].append({"c": ch, "e": ex, "policy": pol, "lo": lo, "hi": hi})
    for _ in range(n_o):
        code, ni, no = rd.u(1), rd.u(1), rd.u(1)
        rd.zeros(1)
        name = NAMES.get(code)
        if name is None:
            bad()
        ai, ao = arity(name)
        if no != ao or (ai is None and ni < 2) or (ai is not None and ni != ai):
            bad()
        ins = [rd.u(4) for _ in range(ni)]
        outs = [rd.u(4) for _ in range(no)]
        if any(t >= n_t for t in ins + outs):
            bad()
        plen = rd.u(4)
        pb = rd.take(plen)
        cin = g["tensors"][ins[0]]["c"]
        g["ops"].append({"name": name, "ins": ins, "outs": outs, "p": read_params(name, pb, cin, n_s)})
    if rd.p != len(b):
        bad()
    rules(g)
    return g


def rules(g):
    """Load step 6: the graph rules of section 10.3, state, exponents (3, 5.6)."""
    T = g["tensors"]
    defined = [False] * len(T)
    for i in g["inputs"]:
        if defined[i["tensor"]] or T[i["tensor"]]["policy"] == 0:
            bad()
        defined[i["tensor"]] = True
    if g["inputs"][0]["kind"] != 0:
        bad()
    if len({i["max_h"] != 0 for i in g["inputs"]}) != 1:
        bad()
    for op in g["ops"]:
        if not all(defined[t] for t in op["ins"]):
            bad()
        for t in op["outs"]:
            if defined[t]:
                bad()
            defined[t] = True
    if not all(defined):
        bad()
    for i in g["inputs"]:
        if i["kind"] == 1:
            if i["src"] >= len(g["outputs"]):
                bad()
            rec, out = T[i["tensor"]], T[g["outputs"][i["src"]]]
            if out["c"] != rec["c"] or not (rec["lo"] <= i["init"] <= rec["hi"]):
                bad()
    for op in g["ops"]:
        n, P = op["name"], op["p"]
        C = T[op["ins"][0]]["c"]
        cs = [T[t]["c"] for t in op["ins"]]
        want = C
        if n in ("conv", "conv_transpose"):
            want = P["cout"]
        elif n == "add":
            want = C if cs[1] == C else None
        elif n == "mul":
            want = C if cs[1] in (C, 1) else None
        elif n == "concat":
            want = sum(cs)
        elif n == "chunk_add":
            want = C // 2 if C % 2 == 0 else None
        elif n == "split":
            want = P["count"] if P["start"] + P["count"] <= C else None
        elif n == "pixel_shuffle":
            want = C // (P["r"] ** 2) if C % (P["r"] ** 2) == 0 else None
        elif n == "pixel_unshuffle":
            want = C * P["r"] ** 2
        elif n == "subpos_merge":
            want = C if all(c == C for c in cs) else None
        elif n == "fill":
            want = T[op["outs"][0]]["c"]
        elif n in ("gram", "apply"):
            want = C if cs[1] == C and C % P["groups"] == 0 else None
        elif n == "chscale":
            lo, hi = g["scalars"][P["scalar"]]
            want = C if lo >= 0 and hi <= P["n_rows"] - 1 else None
        if want is None or any(T[t]["c"] != want for t in op["outs"]):
            bad()
        if n in ("add", "concat", "chunk_add"):
            es = [T[t]["e"] for t in op["ins"] + op["outs"]]
            if not (all(e == -128 for e in es) or (all(e != -128 for e in es) and len(set(es)) == 1)):
                bad()
        if n == "lut16":
            if P["e_in"] != T[op["ins"][0]]["e"] or P["e_out"] != T[op["outs"][0]]["e"]:
                bad()


# --------------------------------------------------------------- plans (section 4, 10.5 step 7)


class Big(Exception):
    pass


class Shape(Exception):
    pass


def capped(v):
    if v > CAP:
        raise Big()
    return v


def op_sizes(g, op, hw, scal):
    """The output sizes of one operator, from its inputs' sizes. `scal` is None at load
    (scalars by their ranges) or the scalar values."""
    n, P = op["name"], op["p"]
    h, w = hw[op["ins"][0]]
    same = lambda o: hw[o] == (h, w)
    if n == "conv":
        hp, wp = capped(h + 2 * P["ph"]), capped(w + 2 * P["pw"])
        if hp < P["kh"] or wp < P["kw"]:
            raise Shape()
        return [((hp - P["kh"]) // P["sh"] + 1, (wp - P["kw"]) // P["sw"] + 1)]
    if n == "conv_transpose":
        uh = capped((h - 1) * P["sh"] + P["kh"] + P["oph"])
        uw = capped((w - 1) * P["sw"] + P["kw"] + P["opw"])
        capped(P["cout"] * uh * uw)
        if uh <= 2 * P["ph"] or uw <= 2 * P["pw"]:
            raise Shape()
        return [(uh - 2 * P["ph"], uw - 2 * P["pw"])]
    if n in ("add", "mul"):
        if not same(op["ins"][1]):
            raise Shape()
        return [(h, w)]
    if n in ("concat", "subpos_merge"):
        # subpos_merge's output sides are sizes it forms: capped before the shape rule.
        out = (capped(2 * h), capped(2 * w)) if n == "subpos_merge" else (h, w)
        if not all(same(t) for t in op["ins"]):
            raise Shape()
        return [out]
    if n == "pixel_shuffle":
        return [(capped(h * P["r"]), capped(w * P["r"]))]
    if n == "pixel_unshuffle":
        if h % P["r"] or w % P["r"]:
            raise Shape()
        return [(h // P["r"], w // P["r"])]
    if n == "pad_rep":
        m = P["m"]
        return [(capped(-(-h // m) * m), capped(-(-w // m) * m))]
    if n == "crop_to":
        rh, rw = hw[op["ins"][1]]
        if rh > h or rw > w:
            raise Shape()
        return [(rh, rw)]
    if n == "crop":
        y0, x0, ch, cw = P["ops"]

        def side(start, length, size):
            if scal is not None:
                val = lambda o: o[1] if o[0] == "const" else scal[o[1]]
                s0, ln = val(start), val(length)
                if s0 < 0 or ln < 1 or s0 + ln > size:
                    raise Shape()
                return ln
            s0 = start[1] if start[0] == "const" else 0
            ln = length[1] if length[0] == "const" else min(g["scalars"][length[1]][1], size - s0)
            if ln < 1 or s0 + ln > size:
                raise Shape()
            return ln

        return [(side(y0, ch, h), side(x0, cw, w))]
    if n == "subpos_split":
        if h % 2 or w % 2:
            raise Shape()
        return [(h // 2, w // 2)] * 4
    if n == "gram":
        if not same(op["ins"][1]):
            raise Shape()
        return [(1, g["tensors"][op["ins"][0]]["c"] // P["groups"])]
    if n == "apply":
        d = g["tensors"][op["ins"][0]]["c"] // P["groups"]
        if (h, w) != (1, d):
            raise Shape()
        return [hw[op["ins"][1]]]
    return [(h, w)]


def walk(g, hw, scal, load):
    for op in g["ops"]:
        try:
            outs = op_sizes(g, op, hw, scal)
        except Big:
            raise Refuse("too_large")
        except Shape:
            raise Refuse("model_mismatch" if load else "geometry_mismatch")
        for t, (oh, ow) in zip(op["outs"], outs):
            if g["tensors"][t]["c"] * oh * ow > g["max_elems"]:
                raise Refuse("too_large")
            hw[t] = (oh, ow)


def load(b):
    """Load steps 1-8 of section 10.5. Step 8 is the bounds pass and the verdict
    (section 8): a graph whose verdict meets neither profile is refused, so a graph
    that loads meets acc64 at least, and g["verdict"] says whether it meets acc32."""
    g = load_graph(b)
    if not g["verdict"]["acc64"]:
        bad()
    return g


def load_graph(b):
    """Load steps 1-7, then step 8's bounds pass with the refusals the pass makes
    (lut16's range, an empty `check` interval), but not the profile against the
    verdict: what a certifier needs to write a certificate whatever the verdict."""
    g = parse_whg1(b)
    if g["inputs"][0]["max_h"] == 0:
        if any(op["name"] in ("gram", "l2norm", "softmax") for op in g["ops"]):
            bad()
        g["sizes"] = None
    else:
        hw = {}
        for i in g["inputs"]:
            c = g["tensors"][i["tensor"]]["c"]
            if c * i["max_h"] * i["max_w"] > g["max_elems"]:
                raise Refuse("too_large")
            hw[i["tensor"]] = (i["max_h"], i["max_w"])
        walk(g, hw, None, True)
        g["sizes"] = hw
    g["pass"] = bounds_pass(g)
    g["verdict"] = verdict_of([("", g)])
    return g


# --------------------------------------------------------------- bounds (section 8)

SAT = 1 << 255
# Section 7.1: the width of each class under each profile, in bits.
WIDTH = {
    "acc64": {"tensor": 64, "acc": 64, "wide": 128, "product": 128},
    "acc32": {"tensor": 32, "acc": 32, "wide": 32, "product": 64},
}
# The class of each named intermediate (sections 5 and 2).
CLASS = {
    "conv.term": "acc", "conv.acc": "acc", "convT.term": "acc", "convT.acc": "acc",
    "rq.t": "acc", "rq.p": "product", "rq.q": "acc", "rq.y": "acc",
    "mul.p": "acc", "leaky.p": "acc",
    "gram.term": "acc", "gram.acc": "acc", "apply.term": "acc", "apply.acc": "acc",
    "igdn.x2": "wide", "igdn.gx2": "product", "igdn.term": "wide", "igdn.N": "wide",
    "igdn.r": "acc", "igdn.xr": "wide",
    "ln.S": "wide", "ln.T": "wide", "ln.m": "wide", "ln.v": "wide", "ln.dk": "wide", "ln.ng": "wide",
    "ln.d": "acc", "rsq.num": "acc", "ln.k": "acc", "ln.n": "acc",
    "l2.T": "wide", "l2.vk": "wide", "l2.k": "acc",
    "sm.diff": "acc", "sm.sum": "acc", "sm.num": "acc", "sig.num": "acc", "elu.p": "wide",
    "exp.negx": "acc", "exp.nl": "acc", "exp.t2": "acc", "exp.c": "acc",
}


def hull(ivs):
    return (min(v[0] for v in ivs), max(v[1] for v in ivs))


def isqrt_ceil(n):
    r = math.isqrt(n)
    return r if r * r == n else r + 1


class Node:
    """One node of the bounds pass. Every +, -, x and floor division (a shift
    included) goes through sat (section 8.1), and the node is saturated when sat
    changes any result. min, max, |.|, negation and isqrt cannot leave
    [-2^255, 2^255] when their operands lie in it. `rec` holds each named
    intermediate's interval, hulled over channels."""

    def __init__(self):
        self.saturated = False
        self.rec = {}

    def sat(self, v):
        if v > SAT:
            self.saturated = True
            return SAT
        if v < -SAT:
            self.saturated = True
            return -SAT
        return v

    def add(self, a, b):
        return self.sat(a + b)

    def sub(self, a, b):
        return self.sat(a - b)

    def mul(self, a, b):
        return self.sat(a * b)

    def div(self, a, b):
        return self.sat(a // b)

    def shr(self, a, s):
        return self.sat(a >> s)

    def note(self, key, iv):
        old = self.rec.get(key)
        self.rec[key] = iv if old is None else (min(old[0], iv[0]), max(old[1], iv[1]))
        return iv

    def term(self, w, x):
        """The term interval of a weight w over x (section 8.2)."""
        a, b = self.mul(w, x[0]), self.mul(w, x[1])
        return (min(a, b), max(a, b))

    def corner(self, x, y):
        c = [self.mul(x[0], y[0]), self.mul(x[0], y[1]), self.mul(x[1], y[0]), self.mul(x[1], y[1])]
        return (min(c), max(c))

    def A(self, terms, b):
        """The partial-sum interval A(K, b) of section 8.2."""
        lo, hi = min(0, b), max(0, b)
        for t in terms:
            lo = self.add(lo, min(0, t[0]))
            hi = self.add(hi, max(0, t[1]))
        return (lo, hi)

    def R(self, p, s, mode):
        """R_floor, R_half_up, R_half_even (section 2), as written there."""
        if s == 0:
            return p
        if mode == 0:
            return self.shr(p, s)
        if mode == 1:
            return self.shr(self.add(p, 1 << (s - 1)), s)
        q = self.shr(p, s)
        r = self.sub(p, self.mul(q, 1 << s))
        h = 1 << (s - 1)
        return self.add(q, 1) if r > h or (r == h and q % 2 == 1) else q

    def rq(self, t, m, s, mode, post):
        """The rq.p, rq.q and rq.y rules of section 8.3 for one channel."""
        self.note("rq.t", t)
        P = self.term(m, t)
        a = (1 << (s - 1)) if mode != 0 and s >= 1 else 0
        self.note("rq.p", (P[0], self.add(P[1], a)))
        q = (self.R(P[0], s, mode), self.R(P[1], s, mode))
        self.note("rq.q", q)
        return self.note("rq.y", (self.add(q[0], post), self.add(q[1], post)))

    def exp(self, u):
        """The exp_neg_q16 rule of section 8.3, given the interval of u = -x."""
        self.note("exp.negx", u)
        self.note("exp.nl", (self.mul(45426, self.div(u[0], 45426)), self.mul(45426, self.div(u[1], 45426))))
        self.note("exp.t2", (self.mul(43245, 43245), self.mul(88670, 88670)))
        self.note("exp.c", (0, self.mul(23495, self.div(self.mul(88670, 88670), 1 << 16))))


def bpre(rq, o):
    return rq["biasv"][o] if rq["bias"] == 1 else 0


def bpost(rq, o):
    return rq["biasv"][o] if rq["bias"] == 2 else 0


def rq_mult(rq, o):
    m = rq_m(rq, o)
    return 1 if m is None else m


def node_bounds(nb, g, op, iv):
    """The producer results of one operator (section 8.3), recording its named
    intermediates in nb. Refuses (model_mismatch) when lut16's input range is not
    within [-32767, 32767] (section 5.6)."""
    T = g["tensors"]
    n, P = op["name"], op["p"]
    x = iv[op["ins"][0]]
    l, h = x
    C = T[op["ins"][0]]["c"]
    if n == "conv":
        rq, cg, kk = P["rq"], C // P["groups"], P["kh"] * P["kw"]
        outs = []
        for o in range(P["cout"]):
            taps = [nb.term(P["w"][(o * cg + ci) * kk + k], x) for ci in range(cg) for k in range(kk)]
            nb.note("conv.term", hull(taps))
            acc = nb.note("conv.acc", nb.A(taps, bpre(rq, o)))
            outs.append(nb.rq(acc, rq_mult(rq, o), rq_s(rq, o), rq["round"], bpost(rq, o)))
        return [hull(outs)]
    if n == "conv_transpose":
        rq, cout, kh, kw, sh, sw = P["rq"], P["cout"], P["kh"], P["kw"], P["sh"], P["sw"]
        outs = []
        for o in range(cout):
            every, by_res = [], {}
            for ci in range(C):
                for i in range(kh):
                    for j in range(kw):
                        t = nb.term(P["w"][((ci * cout + o) * kh + i) * kw + j], x)
                        every.append(t)
                        by_res.setdefault((i % sh, j % sw), []).append(t)
            nb.note("convT.term", hull(every))
            b = bpre(rq, o)
            accs = [nb.A(by_res[r], b) for r in sorted(by_res)]
            if len(by_res) < sh * sw:
                # A residue that no tap has: A over no terms.
                accs.append(nb.A([], b))
            acc = nb.note("convT.acc", hull(accs))
            outs.append(nb.rq(acc, rq_mult(rq, o), rq_s(rq, o), rq["round"], bpost(rq, o)))
        return [hull(outs)]
    if n == "rescale":
        rq = P["rq"]
        t = lambda c: (nb.add(l, bpre(rq, c)), nb.add(h, bpre(rq, c)))
        return [hull([nb.rq(t(c), rq_mult(rq, c), rq_s(rq, c), rq["round"], bpost(rq, c)) for c in range(C)])]
    if n == "chscale":
        lo_r, hi_r = g["scalars"][P["scalar"]]
        s, mode = P["shift"], P["round"]
        a = (1 << (s - 1)) if mode != 0 and s >= 1 else 0
        outs = []
        for c in range(C):
            ps, qs = [], []
            for r in range(lo_r, hi_r + 1):
                Pm = nb.term(P["M"][r * C + c], x)
                ps.append((Pm[0], nb.add(Pm[1], a)))
                qs.append((nb.R(Pm[0], s, mode), nb.R(Pm[1], s, mode)))
            nb.note("rq.t", x)
            nb.note("rq.p", hull(ps))
            q = nb.note("rq.q", hull(qs))
            outs.append(nb.note("rq.y", q))
        return [hull(outs)]
    if n == "mul":
        rq = P["rq"]
        mp = nb.note("mul.p", nb.corner(x, iv[op["ins"][1]]))
        return [hull([nb.rq(mp, rq_mult(rq, c), rq_s(rq, c), rq["round"], 0) for c in range(C)])]
    if n in ("gram", "apply"):
        rq = P["rq"]
        t = nb.note(n + ".term", nb.corner(x, iv[op["ins"][1]]))
        if n == "gram":
            L = max(g["sizes"][u][0] * g["sizes"][u][1] for u in op["ins"])
        else:
            L = C // P["groups"]
        acc = nb.note(n + ".acc", (nb.mul(L, min(0, t[0])), nb.mul(L, max(0, t[1]))))
        return [nb.rq(acc, 1, rq_s(rq, 0), rq["round"], 0)]
    if n == "add":
        b = iv[op["ins"][1]]
        return [(nb.add(l, b[0]), nb.add(h, b[1]))]
    if n == "chunk_add":
        return [(nb.mul(2, l), nb.mul(2, h))]
    if n == "clamp":
        return [(min(max(l, P["lo"]), P["hi"]), min(max(h, P["lo"]), P["hi"]))]
    if n == "relu":
        return [(max(l, 0), max(h, 0))]
    if n == "abs":
        return [(l, h) if l >= 0 else ((-h, -l) if h <= 0 else (0, max(-l, h)))]
    if n == "leaky":
        parts = []
        if l < 0:
            p = nb.note("leaky.p", nb.corner((l, min(h, -1)), (P["a"], P["a"])))
            parts.append((nb.shr(p[0], P["f"]), nb.shr(p[1], P["f"])))
        else:
            nb.note("leaky.p", (0, 0))
        if h >= 0:
            parts.append((max(l, 0), h))
        return [hull(parts)]
    if n == "lut16":
        if l < -32767 or h > 32767:
            bad()
        vals = P["T"][l + 32768:h + 32769]
        return [(min(vals), max(vals))]
    if n == "igdn":
        l2, h2 = nb.mul(l, l), nb.mul(h, h)
        x2 = nb.note("igdn.x2", (0, max(l2, h2)) if l <= 0 <= h else (min(l2, h2), max(l2, h2)))
        outs = []
        for i in range(C):
            terms = []
            for j in range(C):
                gm = P["gamma"][i * C + j]
                gx = nb.note("igdn.gx2", nb.corner((gm, gm), x2))
                terms.append(nb.note("igdn.term", (nb.shr(gx[0], P["gf"]), nb.shr(gx[1], P["gf"]))))
            nb.note("igdn.N", nb.A(terms, P["beta"][i]))
            # The final norm, a left fold over j from beta_i; not recorded.
            f_lo = f_hi = P["beta"][i]
            for t in terms:
                f_lo, f_hi = nb.add(f_lo, t[0]), nb.add(f_hi, t[1])
            r = nb.note("igdn.r", (math.isqrt(max(0, f_lo)), math.isqrt(max(0, f_hi))))
            xr = nb.note("igdn.xr", nb.corner(x, r))
            outs.append((nb.shr(xr[0], P["frac"]), nb.shr(xr[1], P["frac"])))
        return [hull(outs)]
    if n == "layernorm_c":
        e = P["eps"]
        r0 = math.isqrt(e)
        nb.note("ln.S", (nb.mul(C, min(0, l)), nb.mul(C, max(0, h))))
        half = C // 2
        nb.note("ln.m", (nb.sub(nb.mul(C, l), half), nb.add(nb.mul(C, h), half)))
        d = nb.sub(h, l)
        nb.note("ln.d", (nb.sub(l, h), d))
        dd = nb.mul(d, d)
        nb.note("ln.T", (0, nb.mul(C, dd)))
        nb.note("ln.v", (e, nb.add(dd, e)))
        nb.note("rsq.num", (1 << 32, 1 << 32))
        k = nb.div(1 << 32, r0)
        nb.note("ln.k", (0, k))
        ic = isqrt_ceil(C)
        big = nb.mul(nb.mul(ic, 1 << 32), nb.add(r0, 1))
        D = min(nb.mul(d, k), -nb.div(-big, r0))
        nb.note("ln.dk", (-D, D))
        nmax = nb.add(nb.div(nb.mul(nb.mul(ic, 1 << 16), nb.add(r0, 1)), r0), 1)
        N = min(nb.add(nb.shr(nb.mul(d, k), 16), 1), nmax)
        nb.note("ln.n", (-N, N))
        outs = []
        for c in range(C):
            gm = P["gamma"][c]
            ng = nb.note("ln.ng", nb.corner((-N, N), (gm, gm)))
            outs.append((nb.add(nb.shr(ng[0], 16), P["beta"][c]), nb.add(nb.shr(ng[1], 16), P["beta"][c])))
        return [hull(outs)]
    if n == "l2norm":
        e = P["eps"]
        hh, ww = g["sizes"][op["ins"][0]]
        X = max(abs(l), abs(h))
        nb.note("l2.T", (0, nb.add(nb.mul(nb.mul(hh, ww), nb.mul(X, X)), e)))
        nb.note("rsq.num", (1 << 32, 1 << 32))
        nb.note("l2.k", (0, nb.div(1 << 32, math.isqrt(e))))
        nb.note("l2.vk", (-(1 << 32), 1 << 32))
        return [(-(1 << 16), 1 << 16)]
    if n == "softmax":
        nb.note("sm.diff", (nb.sub(l, h), 0))
        nb.exp((0, nb.sub(h, l)))
        wm = g["sizes"][op["ins"][0]][1]
        nb.note("sm.sum", (0, nb.mul(wm, 65553)))
        nb.note("sm.num", (0, nb.add(nb.mul(65553, 1 << 16), nb.div(nb.mul(wm, 65553), 2))))
        return [(0, 1 << 16)]
    if n == "sigmoid":
        nb.exp((l, h) if l >= 0 else ((-h, -l) if h <= 0 else (0, max(-l, h))))
        nb.note("sig.num", (0, nb.mul(65553, 1 << 16)))
        return [(0, 1 << 16)]
    if n == "elu":
        nb.exp((max(0, -h), max(0, -l)))
        al = P["alpha"]
        p = nb.note("elu.p", nb.corner((al, al), (-(1 << 16), 17)))
        parts = []
        if h > 0:
            parts.append((max(l, 1), h))
        if l <= 0:
            parts.append((nb.shr(p[0], 16), nb.shr(p[1], 16)))
        return [hull(parts)]
    if n in ("concat", "subpos_merge"):
        return [hull([iv[t] for t in op["ins"]])]
    if n == "subpos_split":
        return [x] * 4
    if n == "fill":
        return [(P["value"], P["value"]) if P["kind"] == 0 else g["scalars"][P["value"]]]
    # split, pixel_shuffle, pixel_unshuffle, pad_rep, crop_to (its first input), crop.
    return [x]


PROFILES = ("acc32", "acc64")
POLICY = ("derived", "clamp", "check")
_RQ = ["rq.t", "rq.p", "rq.q", "rq.y"]
# Table 8.5-A: the intermediates each operator records, in definition order. Every
# other operator records none.
RECORDS = {
    "conv": ["conv.term", "conv.acc"] + _RQ,
    "conv_transpose": ["convT.term", "convT.acc"] + _RQ,
    "rescale": _RQ,
    "chscale": _RQ,
    "mul": ["mul.p"] + _RQ,
    "gram": ["gram.term", "gram.acc"] + _RQ,
    "apply": ["apply.term", "apply.acc"] + _RQ,
    "leaky": ["leaky.p"],
    "igdn": ["igdn.x2", "igdn.gx2", "igdn.term", "igdn.N", "igdn.r", "igdn.xr"],
    "layernorm_c": ["ln.S", "ln.m", "ln.d", "ln.T", "ln.v", "rsq.num", "ln.k", "ln.dk", "ln.n", "ln.ng"],
    "l2norm": ["l2.T", "rsq.num", "l2.k", "l2.vk"],
    "softmax": ["sm.diff", "exp.negx", "exp.nl", "exp.t2", "exp.c", "sm.sum", "sm.num"],
    "sigmoid": ["exp.negx", "exp.nl", "exp.t2", "exp.c", "sig.num"],
    "elu": ["exp.negx", "exp.nl", "exp.t2", "exp.c", "elu.p"],
}
# Section 8.5, k_c: the operators that record it, and the term interval of section 6
# it is taken over.
KC_TERM = {"conv": "conv.term", "conv_transpose": "convT.term", "igdn": "igdn.term", "gram": "gram.term", "apply": "apply.term"}
# Section 8.5, mac: the operators that record it; w_bits for the first three only.
MAC = ("conv", "conv_transpose", "chscale", "mul", "gram", "apply")
# Section 5's table: the operators marked "no" in the wai-ops/1-npu column.
NPU_NO = ("igdn", "layernorm_c", "l2norm", "softmax", "sigmoid", "elu", "gram", "apply")


def bits(iv):
    """bits([lo, hi]) of section 2: the smallest N >= 1 such that both fit iN."""
    return max((v if v >= 0 else ~v).bit_length() + 1 for v in iv)


def acc32_params(op):
    """The parameter conditions of section 7.2 the node fails, as their keys, in that
    table's order."""
    n, P = op["name"], op["p"]
    f32 = lambda v: I32_MIN <= v <= I32_MAX
    keys = []
    rq = P.get("rq")
    if rq is not None:
        if not all(map(f32, rq["mult"])):
            keys.append("rq.mult")
        if not all(s <= 63 for s in rq["shift"]):
            keys.append("rq.shift")
        if rq["bias"] == 2 and not all(map(f32, rq["biasv"])):
            keys.append("rq.post_bias")
    if n in ("conv", "conv_transpose") and not all(map(f32, P["w"])):
        keys.append("weights")
    if n == "chscale":
        if not all(map(f32, P["M"])):
            keys.append("chscale.m")
        if P["shift"] > 63:
            keys.append("chscale.s")
    if n == "igdn":
        if not all(map(f32, P["gamma"])):
            keys.append("igdn.gamma")
        if P["gf"] > 63 or P["frac"] > 63:
            keys.append("igdn.shift")
    if n == "leaky" and not f32(P["a"]):
        keys.append("leaky.a")
    return keys


def node_npu(g, op):
    """The node's flags of section 9.2, in ascending order, and its reasons of section
    8.5 other than "acc32"."""
    n, P = op["name"], op["p"]
    flags, reasons = set(), set()
    if n == "lut16":
        flags.add("lut16")
    if n in ("add", "chunk_add"):
        state = {i["tensor"] for i in g["inputs"] if i["kind"] == 1}
        if any(t in state for t in op["ins"]):
            flags.add("state_add")
    rq = None
    if n == "chscale":
        flags.add("chscale")
        # A chscale always has a multiplier, M[r][c], and no bias.
        rq = (P["round"], True, P["M"], [P["shift"]], False)
    elif "rq" in P:
        q = P["rq"]
        rq = (q["round"], q["m_kind"] != 0, q["mult"], q["shift"], q["bias"] == 2)
    if rq is not None:
        rnd, has_mult, mults, shifts, post = rq
        if rnd == 2:
            flags.add("rq_half_even")
        if rnd == 0 and has_mult:
            flags.add("rq_floor_mult")
        if post:
            flags.add("rq_post_bias")
        if any(m < 0 for m in mults):
            flags.add("rq_neg_mult")
        if any(not 0 <= s <= 31 for s in shifts):
            flags.add("rq_shift_range")
    if n in NPU_NO:
        reasons.add("op:" + n)
    if n == "conv":
        cin, G = g["tensors"][op["ins"][0]]["c"], P["groups"]
        if (G != 1 and G != cin) or (G == cin and G > 1 and P["cout"] != cin):
            reasons.add("conv_groups")
    if n == "crop" and any(o[0] == "scalar" for o in P["ops"]):
        reasons.add("crop_scalar")
    return sorted(flags), reasons


def width_failures(nd, p):
    """The node's recorded intermediates and producer results (`out.<k>`) that do not
    meet profile p's width for their class (section 7.1), as (bits, key)."""
    W = WIDTH[p]
    f = [(bits(r), k) for k, r in nd["rec"].items() if bits(r) > W[CLASS[k]]]
    f += [(bits(o["result"]), "out.%d" % j) for j, o in enumerate(nd["outs"]) if bits(o["result"]) > W["tensor"]]
    return f


def bounds_pass(g):
    """The bounds pass of sections 8.1-8.3 over a graph that passed load steps 1-7,
    in graph order. Refuses (model_mismatch) where the pass refuses the graph at
    load: a `check` interval that is empty (section 8.1), or lut16's input range
    (section 5.6). Returns one record per node: its recorded intermediates, whether
    it saturated, its input tensors' intervals, each output's producer result and
    interval after its policy, its acc32_params, its npu flags and reasons, and the
    profiles it fits (section 8.5)."""
    T = g["tensors"]
    iv = {i["tensor"]: (T[i["tensor"]]["lo"], T[i["tensor"]]["hi"]) for i in g["inputs"]}
    nodes = []
    for k, op in enumerate(g["ops"]):
        nb = Node()
        ins = [iv[t] for t in op["ins"]]
        results = node_bounds(nb, g, op, iv)
        if sorted(nb.rec) != sorted(RECORDS.get(op["name"], [])):
            raise AssertionError("%s records %s" % (op["name"], sorted(nb.rec)))
        outs = []
        for t, r in zip(op["outs"], results):
            rec = T[t]
            if rec["policy"] == 0:
                ti = r
            elif rec["policy"] == 1:
                ti = (min(max(r[0], rec["lo"]), rec["hi"]), min(max(r[1], rec["lo"]), rec["hi"]))
            else:
                ti = (max(r[0], rec["lo"]), min(r[1], rec["hi"]))
                if ti[0] > ti[1]:
                    bad()
            iv[t] = ti
            outs.append({"tensor": t, "iv": ti, "result": r, "policy": rec["policy"]})
        flags, reasons = node_npu(g, op)
        nd = {"index": k, "op": op, "rec": nb.rec, "saturated": nb.saturated, "ins": ins, "outs": outs,
              "acc32_params": acc32_params(op), "flags": flags, "reasons": reasons}
        nd["fits"] = {}
        for p in PROFILES:
            fit = not nd["saturated"] and not width_failures(nd, p)
            if p == "acc32":
                fit = fit and not nd["acc32_params"] and all(o["policy"] != 2 for o in outs)
            nd["fits"][p] = fit
        nodes.append(nd)
    return nodes


def node_failure(nd, p, name):
    """The failure of section 8.5 for a node that does not fit profile p."""
    f = {"graph": name, "index": nd["index"], "where": "node"}
    if nd["saturated"]:
        f["key"] = "saturated"
        return f
    w = width_failures(nd, p)
    if w:
        f["bits"], f["key"] = min(w, key=lambda e: (-e[0], e[1].encode()))
    else:
        f["key"] = nd["acc32_params"][0] if nd["acc32_params"] else "policy.check"
    return f


def verdict_of(graphs):
    """The verdict of section 8.5 over [(name, g)], each g carrying its pass."""
    v = {}
    for p in PROFILES:
        tw = WIDTH[p]["tensor"]
        fail = None
        for name, g in graphs:
            T = g["tensors"]
            for k, i in enumerate(g["inputs"]):
                r = (T[i["tensor"]]["lo"], T[i["tensor"]]["hi"])
                if bits(r) > tw:
                    fail = {"bits": bits(r), "graph": name, "index": k, "key": "range", "where": "input"}
                    break
            if fail is None:
                for k, r in enumerate(g["scalars"]):
                    if bits(r) > tw:
                        fail = {"bits": bits(r), "graph": name, "index": k, "key": "range", "where": "scalar"}
                        break
            if fail is None:
                for nd in g["pass"]:
                    if not nd["fits"][p]:
                        fail = node_failure(nd, p, name)
                        break
            if fail is not None:
                break
        v[p] = fail is None
        if fail is not None:
            v[p + "_first_failure"] = fail
    reasons, flags = set(), set()
    if not v["acc32"]:
        reasons.add("acc32")
    for _, g in graphs:
        for nd in g["pass"]:
            reasons |= nd["reasons"]
            flags |= set(nd["flags"])
    v["npu_reasons"] = sorted(reasons, key=lambda s: s.encode())
    v["npu"] = not reasons
    v["npu_flags"] = sorted(flags, key=lambda s: s.encode())
    return v


# --------------------------------------------------------------- the operators (section 5)


class T:
    """A tensor [c, h, w], elements in CHW order."""

    def __init__(self, c, h, w, d):
        self.c, self.h, self.w, self.d = c, h, w, d

    def at(self, c, y, x):
        return self.d[(c * self.h + y) * self.w + x]


def ev_conv(x, P, oh, ow):
    cin, H, W = x.c, x.h, x.w
    G = P["groups"]
    cg, og = cin // G, P["cout"] // G
    kh, kw = P["kh"], P["kw"]
    d = []
    for o in range(P["cout"]):
        base = (o // og) * cg
        for y in range(oh):
            for xx in range(ow):
                acc = 0
                for ci in range(cg):
                    for i in range(kh):
                        u = y * P["sh"] + i - P["ph"]
                        if not 0 <= u < H:
                            continue
                        for j in range(kw):
                            v = xx * P["sw"] + j - P["pw"]
                            if 0 <= v < W:
                                acc += P["w"][((o * cg + ci) * kh + i) * kw + j] * x.at(base + ci, u, v)
                d.append(rq_apply(P["rq"], o, acc))
    return T(P["cout"], oh, ow, d)


def ev_conv_t(x, P, oh, ow):
    """The gather form of section 5.3."""
    cin, H, W = x.c, x.h, x.w
    kh, kw, sh, sw, cout = P["kh"], P["kw"], P["sh"], P["sw"], P["cout"]
    d = []
    for o in range(cout):
        for y in range(oh):
            for xx in range(ow):
                acc = 0
                for i in range(kh):
                    ty = y + P["ph"] - i
                    if ty < 0 or ty % sh or ty // sh >= H:
                        continue
                    for j in range(kw):
                        tx = xx + P["pw"] - j
                        if tx < 0 or tx % sw or tx // sw >= W:
                            continue
                        for ci in range(cin):
                            acc += P["w"][((ci * cout + o) * kh + i) * kw + j] * x.at(ci, ty // sh, tx // sw)
                d.append(rq_apply(P["rq"], o, acc))
    return T(cout, oh, ow, d)


def per_channel(x, f):
    plane = x.h * x.w
    return T(x.c, x.h, x.w, [f(k // plane, v) for k, v in enumerate(x.d)])


def ev(g, op, vals, hw, scal):
    n, P = op["name"], op["p"]
    x = vals[op["ins"][0]]
    oh, ow = hw[op["outs"][0]]
    oc = [g["tensors"][t]["c"] for t in op["outs"]]
    if n == "conv":
        return [ev_conv(x, P, oh, ow)]
    if n == "conv_transpose":
        return [ev_conv_t(x, P, oh, ow)]
    if n == "rescale":
        return [per_channel(x, lambda c, v: rq_apply(P["rq"], c, v))]
    if n == "chscale":
        r = scal[P["scalar"]]
        return [per_channel(x, lambda c, v: R(v * P["M"][r * x.c + c], P["shift"], P["round"]))]
    if n == "mul":
        b = vals[op["ins"][1]]
        plane = x.h * x.w
        d = []
        for c in range(x.c):
            cb = c if b.c == x.c else 0
            for k in range(plane):
                d.append(rq_apply(P["rq"], c, x.d[c * plane + k] * b.d[cb * plane + k]))
        return [T(x.c, x.h, x.w, d)]
    if n == "lut16":
        return [T(x.c, x.h, x.w, [P["T"][v + 32768] if -32768 <= v <= 32767 else None for v in x.d])]
    if n == "add":
        b = vals[op["ins"][1]]
        return [T(x.c, x.h, x.w, [p + q for p, q in zip(x.d, b.d)])]
    if n == "concat":
        d = []
        for t in op["ins"]:
            d += vals[t].d
        return [T(oc[0], x.h, x.w, d)]
    if n == "chunk_add":
        half = (x.c // 2) * x.h * x.w
        return [T(x.c // 2, x.h, x.w, [x.d[k] + x.d[k + half] for k in range(half)])]
    if n == "split":
        plane = x.h * x.w
        return [T(P["count"], x.h, x.w, x.d[P["start"] * plane:(P["start"] + P["count"]) * plane])]
    if n == "pixel_shuffle":
        r = P["r"]
        co = x.c // (r * r)
        d = [0] * (co * oh * ow)
        for c in range(co):
            for i in range(r):
                for j in range(r):
                    for hh in range(x.h):
                        for ww in range(x.w):
                            d[(c * oh + hh * r + i) * ow + ww * r + j] = x.at(c * r * r + i * r + j, hh, ww)
        return [T(co, oh, ow, d)]
    if n == "pixel_unshuffle":
        r = P["r"]
        d = [0] * (x.c * r * r * oh * ow)
        for c in range(x.c):
            for i in range(r):
                for j in range(r):
                    for hh in range(oh):
                        for ww in range(ow):
                            d[((c * r * r + i * r + j) * oh + hh) * ow + ww] = x.at(c, hh * r + i, ww * r + j)
        return [T(x.c * r * r, oh, ow, d)]
    if n == "clamp":
        return [T(x.c, x.h, x.w, [min(max(v, P["lo"]), P["hi"]) for v in x.d])]
    if n == "relu":
        return [T(x.c, x.h, x.w, [max(v, 0) for v in x.d])]
    if n == "leaky":
        return [T(x.c, x.h, x.w, [v if v >= 0 else (v * P["a"]) >> P["f"] for v in x.d])]
    if n == "abs":
        return [T(x.c, x.h, x.w, [abs(v) for v in x.d])]
    if n == "pad_rep":
        return [T(x.c, oh, ow, [x.at(c, min(y, x.h - 1), min(xx, x.w - 1)) for c in range(x.c) for y in range(oh) for xx in range(ow)])]
    if n in ("crop_to", "crop"):
        y0 = x0 = 0
        if n == "crop":
            val = lambda o: o[1] if o[0] == "const" else scal[o[1]]
            y0, x0 = val(P["ops"][0]), val(P["ops"][1])
        return [T(x.c, oh, ow, [x.at(c, y0 + y, x0 + xx) for c in range(x.c) for y in range(oh) for xx in range(ow)])]
    if n == "subpos_split":
        outs = []
        for k in range(4):
            a, b2 = P["order"][k] >> 1, P["order"][k] & 1
            outs.append(T(x.c, oh, ow, [x.at(c, 2 * y + a, 2 * xx + b2) for c in range(x.c) for y in range(oh) for xx in range(ow)]))
        return outs
    if n == "subpos_merge":
        d = [0] * (x.c * oh * ow)
        for k in range(4):
            a, b2 = P["order"][k] >> 1, P["order"][k] & 1
            t = vals[op["ins"][k]]
            for c in range(x.c):
                for y in range(t.h):
                    for xx in range(t.w):
                        d[(c * oh + 2 * y + a) * ow + 2 * xx + b2] = t.at(c, y, xx)
        return [T(x.c, oh, ow, d)]
    if n == "fill":
        v = P["value"] if P["kind"] == 0 else scal[P["value"]]
        return [T(oc[0], x.h, x.w, [v] * (oc[0] * x.h * x.w))]
    if n == "igdn":
        C, plane = x.c, x.h * x.w
        out = [0] * (C * plane)
        for p in range(plane):
            xs = [x.d[j * plane + p] for j in range(C)]
            for i in range(C):
                N = P["beta"][i] + sum((P["gamma"][i * C + j] * xs[j] * xs[j]) // (1 << P["gf"]) for j in range(C))
                r = math.isqrt(max(0, N))
                out[i * plane + p] = (xs[i] * r) // (1 << P["frac"])
        return [T(C, x.h, x.w, out)]
    if n == "layernorm_c":
        C, plane = x.c, x.h * x.w
        out = [0] * (C * plane)
        for p in range(plane):
            xs = [x.d[c * plane + p] for c in range(C)]
            S = sum(xs)
            sg = (S > 0) - (S < 0)
            M = S + (C // 2) * sg
            m = trunc_div(M, C)
            ds = [v - m for v in xs]
            V = sum(d * d for d in ds) // C
            k = rsqrt_q16(V + P["eps"])
            for c in range(C):
                nn = (ds[c] * k) >> 16
                out[c * plane + p] = ((nn * P["gamma"][c]) >> 16) + P["beta"][c]
        return [T(C, x.h, x.w, out)]
    if n == "l2norm":
        plane = x.h * x.w
        out = []
        for c in range(x.c):
            vs = x.d[c * plane:(c + 1) * plane]
            k = rsqrt_q16(sum(v * v for v in vs) + P["eps"])
            out += [(v * k) >> 16 for v in vs]
        return [T(x.c, x.h, x.w, out)]
    if n == "softmax":
        out = []
        for row in range(x.c * x.h):
            vs = x.d[row * x.w:(row + 1) * x.w]
            m = max(vs)
            es = [exp_neg_q16(v - m) for v in vs]
            s = sum(es)
            out += [(e * 65536 + s // 2) // s for e in es]
        return [T(x.c, x.h, x.w, out)]
    if n == "sigmoid":
        def sig(v):
            if v >= 0:
                return (1 << 32) // (65536 + exp_neg_q16(-v))
            e = exp_neg_q16(v)
            return (e * 65536) // (65536 + e)
        return [T(x.c, x.h, x.w, [sig(v) for v in x.d])]
    if n == "elu":
        return [T(x.c, x.h, x.w, [v if v > 0 else (P["alpha"] * (exp_neg_q16(v) - 65536)) >> 16 for v in x.d])]
    if n == "gram":
        b = vals[op["ins"][1]]
        G = P["groups"]
        dd = x.c // G
        plane = x.h * x.w
        out = [0] * (x.c * dd)
        for gi in range(G):
            for i in range(dd):
                for j in range(dd):
                    s = sum(x.d[(gi * dd + i) * plane + p] * b.d[(gi * dd + j) * plane + p] for p in range(plane))
                    out[(gi * dd + i) * dd + j] = R(s, P["rq"]["shift"][0], P["rq"]["round"])
        return [T(x.c, 1, dd, out)]
    if n == "apply":
        V = vals[op["ins"][1]]
        G = P["groups"]
        dd = x.c // G
        plane = V.h * V.w
        out = [0] * (V.c * plane)
        for gi in range(G):
            for i in range(dd):
                for p in range(plane):
                    s = sum(x.d[(gi * dd + i) * dd + j] * V.d[(gi * dd + j) * plane + p] for j in range(dd))
                    out[(gi * dd + i) * plane + p] = R(s, P["rq"]["shift"][0], P["rq"]["round"])
        return [T(V.c, V.h, V.w, out)]
    raise Refuse("model_mismatch")


def apply_policy(rec, t):
    """Returns False when a `check` tensor holds an element outside its range."""
    if rec["policy"] == 1:
        t.d = [min(max(v, rec["lo"]), rec["hi"]) for v in t.d]
    elif rec["policy"] == 2:
        return all(rec["lo"] <= v <= rec["hi"] for v in t.d)
    return True


def run(g, data_inputs, scalars, state):
    """One decode (section 4): the data inputs in input-table order; each state input
    from `state` (section 10.3), which the decode updates when it completes."""
    allin = []
    di = 0
    si = 0
    if len(data_inputs) != sum(1 for i in g["inputs"] if i["kind"] == 0):
        raise Refuse("model_mismatch")
    for i in g["inputs"]:
        if i["kind"] == 0:
            allin.append(data_inputs[di])
            di += 1
        else:
            if state.get("s") is None:
                x0 = data_inputs[0]
                c = g["tensors"][i["tensor"]]["c"]
                allin.append(T(c, x0.h, x0.w, [i["init"]] * (c * x0.h * x0.w)))
            else:
                allin.append(state["s"][si])
            si += 1
    hw = {}
    # Step 1: each input's channel count, then its size.
    for i, t in zip(g["inputs"], allin):
        if t.c != g["tensors"][i["tensor"]]["c"] or t.h < 1 or t.w < 1:
            raise Refuse("model_mismatch")
        if i["max_h"] and (t.h > i["max_h"] or t.w > i["max_w"]):
            raise Refuse("too_large")
        if t.c * t.h * t.w > g["max_elems"]:
            raise Refuse("too_large")
        hw[i["tensor"]] = (t.h, t.w)
    # Step 2: each scalar against its range.
    if len(scalars) != len(g["scalars"]):
        raise Refuse("model_mismatch")
    for v, (lo, hi) in zip(scalars, g["scalars"]):
        if not lo <= v <= hi:
            raise Refuse("field_range")
    # Step 3: the run-time plan.
    walk(g, hw, scalars, False)
    # Step 4: evaluation, the inputs' policies first.
    vals = {}
    for i, t in zip(g["inputs"], allin):
        t = T(t.c, t.h, t.w, list(t.d))
        if not apply_policy(g["tensors"][i["tensor"]], t):
            raise Refuse("field_range")
        vals[i["tensor"]] = t
    for op in g["ops"]:
        outs = ev(g, op, vals, hw, scalars)
        for t, ten in zip(op["outs"], outs):
            # A value outside i64 is one load step 8 excludes; should one be formed
            # all the same, the decode is refused rather than carried on.
            if any(v is None or not I64_MIN <= v <= I64_MAX for v in ten.d):
                raise Refuse("model_mismatch")
            if not apply_policy(g["tensors"][t], ten):
                raise Refuse("field_range")
            vals[t] = ten
    outs = [vals[t] for t in g["outputs"]]
    state["s"] = [outs[i["src"]] for i in g["inputs"] if i["kind"] == 1]
    return outs


# --------------------------------------------------------------- canonical JSON (8.5, 14.1)


def cw(v):
    """The canonical form: ints as decimal strings, members sorted by key bytes."""
    if v is None:
        return "null"
    if v is True:
        return "true"
    if v is False:
        return "false"
    if isinstance(v, int):
        return '"%d"' % v
    if isinstance(v, str):
        return '"' + v.replace("\\", "\\\\").replace('"', '\\"') + '"'
    if isinstance(v, list):
        return "[" + ",".join(cw(x) for x in v) + "]"
    if isinstance(v, dict):
        keys = sorted(v, key=lambda k: k.encode())
        return "{" + ",".join(cw(k) + ":" + cw(v[k]) for k in keys) + "}"
    raise TypeError(v)


def dec(s):
    """A canonical decimal integer string."""
    if not isinstance(s, str) or not s or s == "-0":
        raise ValueError("not a canonical integer: %r" % (s,))
    d = s[1:] if s[0] == "-" else s
    if not d.isdigit() or (len(d) > 1 and d[0] == "0") or not d.isascii():
        raise ValueError("not a canonical integer: %r" % (s,))
    return int(s)


def tensor_in(v):
    c, h, w = (dec(x) for x in v["shape"])
    d = [dec(x) for x in v["data"]]
    if len(d) != c * h * w:
        raise ValueError("a tensor's data does not fill its shape")
    return T(c, h, w, d)


def tensor_out(t):
    return {"data": list(t.d), "shape": [t.c, t.h, t.w]}


def graph_line(v):
    """A graph vector, recomputed: `load` (steps 1-8), `profiles` (section 14.1: both
    when the load is refused, else exactly the profiles the verdict meets) and every
    decode's `expect`."""
    b = bytes.fromhex(v["graph"])
    out = {"graph": b.hex(), "id": v["id"], "decodes": []}
    try:
        g = load(b)
        out["load"] = "ok"
        out["profiles"] = [p for p in ("acc32", "acc64") if g["verdict"][p]]
    except Refuse as e:
        out["load"] = e.code
        out["profiles"] = ["acc32", "acc64"]
        g = None
    if g is not None:
        state = {}
        for d in v["decodes"]:
            ins = [tensor_in(t) for t in d["inputs"]]
            sc = [dec(s) for s in d["scalars"]]
            try:
                res = run(g, ins, sc, state)
                exp = {"outputs": [tensor_out(t) for t in res]}
            except Refuse as e:
                exp = {"refuse": e.code}
            out["decodes"].append({"expect": exp, "inputs": [tensor_out(t) for t in ins], "scalars": sc})
    return out


# --------------------------------------------------------------- the certificate (8.5)


def ivj(iv):
    return {"bits": bits(iv), "hi": iv[1], "lo": iv[0]}


def node_json(nd):
    op = nd["op"]
    n, P = op["name"], op["p"]
    d = {
        "acc32_params": nd["acc32_params"],
        "fits": [p for p in PROFILES if nd["fits"][p]],
        "index": nd["index"],
        "intermediates": {k: ivj(nd["rec"][k]) for k in RECORDS.get(n, [])},
        "npu_flags": nd["flags"],
        "op": n,
        "out": [{"bits": bits(o["iv"]), "hi": o["iv"][1], "lo": o["iv"][0], "policy": POLICY[o["policy"]],
                 "result": ivj(o["result"]), "tensor": o["tensor"]} for o in nd["outs"]],
        "saturated": nd["saturated"],
    }
    if n in KC_TERM:
        t = nd["rec"][KC_TERM[n]]
        m = max(abs(t[0]), abs(t[1]))
        d["k_c"] = "unbounded" if m == 0 else I32_MAX // m
    if n in MAC:
        d["mac"] = {"a_bits": [bits(x) for x in nd["ins"]]}
        if n in ("conv", "conv_transpose", "chscale"):
            w = P["M"] if n == "chscale" else P["w"]
            d["mac"]["w_bits"] = bits((min(w), max(w)))
    return d


def graph_json(name, b, g):
    T = g["tensors"]
    inputs = []
    for i in g["inputs"]:
        rec = T[i["tensor"]]
        e = {"hi": rec["hi"], "kind": "state" if i["kind"] else "data", "lo": rec["lo"], "max_h": i["max_h"],
             "max_w": i["max_w"], "policy": POLICY[rec["policy"]], "tensor": i["tensor"]}
        if i["kind"]:
            e["init"], e["src"] = i["init"], i["src"]
        inputs.append(e)
    return {
        "graph_sha256": hashlib.sha256(b).hexdigest(),
        "inputs": inputs,
        "name": name,
        "nodes": [node_json(nd) for nd in g["pass"]],
        "scalars": [{"hi": hi, "index": k, "lo": lo} for k, (lo, hi) in enumerate(g["scalars"])],
    }


def canonical(v):
    """The canonical JSON of section 8.5, as bytes. Every string is printable ASCII;
    a role or capability name outside that alphabet cannot be certified."""
    def strings(x):
        if isinstance(x, str):
            yield x
        elif isinstance(x, list):
            for y in x:
                yield from strings(y)
        elif isinstance(x, dict):
            for k, y in x.items():
                yield k
                yield from strings(y)
    for s in strings(v):
        if any(not 0x20 <= ord(ch) <= 0x7E for ch in s):
            raise ValueError("a string outside printable ASCII cannot be certified: %r" % s)
    return cw(v).encode("ascii")


def certificate(capability, lowering, set_sha256, roles, role_profiles=None):
    """The bounds certificate of section 8.5 for the graphs [(role, WHG1 bytes)], in
    the order given: (its canonical bytes, its verdict). Each graph passes load steps
    1-7 and the bounds pass, or the Refuse of the step that fails propagates. With
    `role_profiles` ({role: profile}, section 11), it returns a third value, the
    failure of section 8.5 of each named graph that does not meet its profile, judged
    on that graph alone ({role: failure}; empty when every one meets it)."""
    gs = [(name, b, load_graph(b)) for name, b in roles]
    v = verdict_of([(name, g) for name, _, g in gs])
    cert = {
        "capability": capability,
        "certificate": "wai-ops/1 bounds",
        "graphs": [graph_json(name, b, g) for name, b, g in gs],
        "lowering": lowering,
        "set_sha256": set_sha256,
        "verdict": v,
    }
    if role_profiles is None:
        return canonical(cert), v
    fails = {}
    for name, p in role_profiles.items():
        # A graph meets a profile when its inputs' and scalars' declared ranges meet
        # the profile's tensor width and every node of it fits the profile: the
        # verdict of section 8.5 over that graph alone.
        rv = verdict_of([(n, g) for n, _, g in gs if n == name])
        if not rv[p]:
            fails[name] = rv[p + "_first_failure"]
    return canonical(cert), v, fails


ABSENT = "(absent)"


def json_diffs(a, b, path="", out=None, limit=None):
    """The JSON paths at which two parsed values differ, in document order, each
    with both values: [(path, a's value, b's value)]."""
    out = [] if out is None else out
    if limit is not None and len(out) >= limit:
        return out
    if isinstance(a, dict) and isinstance(b, dict):
        for k in sorted(set(a) | set(b), key=lambda s: s.encode()):
            p = path + ("." + k if k.isidentifier() else '["%s"]' % k)
            if k not in a or k not in b:
                out.append((p, a.get(k, ABSENT), b.get(k, ABSENT)))
            else:
                json_diffs(a[k], b[k], p, out, limit)
    elif isinstance(a, list) and isinstance(b, list):
        for k in range(max(len(a), len(b))):
            p = "%s[%d]" % (path, k)
            if k >= len(a) or k >= len(b):
                out.append((p, a[k] if k < len(a) else ABSENT, b[k] if k < len(b) else ABSENT))
            else:
                json_diffs(a[k], b[k], p, out, limit)
    elif a != b:
        out.append((path or "$", a, b))
    return out[:limit] if limit is not None else out


def first_difference(got, want):
    """Part of what a `--check` mismatch prints: the first differing JSON path, with
    both values."""
    try:
        a, b = json.loads(got), json.loads(want)
    except ValueError:
        return "the committed file is not JSON"
    d = json_diffs(a, b, limit=1)
    if not d:
        return "the same JSON, written differently"
    p, x, y = d[0]
    return "first difference at %s: computed %s, committed %s" % (p, cw(x) if x is not ABSENT else x, cw(y) if y is not ABSENT else y)


# --------------------------------------------------------------- stage operators (5E)


class Rans:
    """integer-payloads section 2's decoder state."""

    def __init__(self, b):
        self.b = b
        self.x = (int.from_bytes(b[4:8], "little") << 32) | int.from_bytes(b[0:4], "little")
        self.i = 2
        self.out = False

    def renorm(self):
        if self.x < (1 << 31):
            if 4 * self.i + 4 <= len(self.b):
                self.x = (self.x << 32) | int.from_bytes(self.b[4 * self.i:4 * self.i + 4], "little")
                self.i += 1
            else:
                self.out = True

    def bits(self, n):
        v = self.x & ((1 << n) - 1)
        self.x >>= n
        self.renorm()
        return v


def sat32(v):
    return max(I32_MIN, min(I32_MAX, v))


def rans64_decode(i, m):
    table_rule, escape, end = m["table"], m["escape"], m["end"]
    lo = hi = None
    if escape == "strict":
        lo, hi = dec(m["lo"]), dec(m["hi"])
        if not (I32_MIN <= lo <= hi <= I32_MAX):
            raise Refuse("model_mismatch", -1)
    tables = [([dec(c) for c in t["cdf"]], dec(t["offset"])) for t in i["tables"]]
    reads = [dec(r) for r in i["reads"]]
    for cdf, o in tables:
        n = len(cdf)
        ok_n = 3 <= n <= 4097 if table_rule == "t4097" else 2 <= n <= 65537
        if not ok_n or cdf[0] != 0 or cdf[-1] != 65536 or any(cdf[k] >= cdf[k + 1] for k in range(n - 1)) or not I32_MIN <= o <= I32_MAX:
            raise Refuse("model_mismatch", -1)
    if any(not 0 <= r < len(tables) for r in reads):
        raise Refuse("model_mismatch", -1)
    stream = bytes.fromhex(i["stream"])
    if len(stream) < 8 or len(stream) % 4:
        raise Refuse("stream_exhausted", 0)
    d = Rans(stream)
    syms = []
    for k, r in enumerate(reads):
        cdf, o = tables[r]
        n = len(cdf)
        cum = d.x & 0xFFFF
        s = max(j for j in range(n - 1) if cdf[j] <= cum)
        d.x = (cdf[s + 1] - cdf[s]) * (d.x >> 16) + cum - cdf[s]
        d.renorm()
        code = None
        if s < n - 2:
            sym = s + o
            if escape == "lenient":
                sym = sat32(sym)
            elif not lo <= sym <= hi:
                code = "field_range"
        elif escape == "lenient":
            v = d.bits(4)
            count = v
            groups = 0
            while v == 15 and groups < 8 and not d.out:
                v = d.bits(4)
                count += v
                groups += 1
            count = min(count, 8)
            raw = 0
            for j in range(count):
                raw |= d.bits(4) << (4 * j)
            value = -(raw >> 1) - 1 if raw & 1 else (raw >> 1) + (n - 2)
            sym = sat32(sat32(value) + o)
        else:
            count = d.bits(4)
            raw, last = 0, None
            if count <= 8:
                for j in range(count):
                    last = d.bits(4)
                    raw += last << (4 * j)
            value = -(raw // 2) - 1 if raw % 2 else raw // 2 + (n - 2)
            sym = value + o
            if count > 8 or (count >= 1 and last == 0):
                code = "stream_noncanonical"
            elif not lo <= sym <= hi:
                code = "field_range"
        if end == "canonical" and d.out:
            raise Refuse("stream_exhausted", k)
        if code:
            raise Refuse(code, k)
        syms.append(sym)
    if end == "canonical":
        if d.out or d.x != (1 << 31) or 4 * d.i != len(stream):
            raise Refuse("stream_noncanonical", len(reads))
    elif d.out:
        raise Refuse("stream_exhausted", len(reads))
    return {"symbols": syms}


class Win:
    """E8's window and states over memory M (8 zero bytes, then the stream)."""

    def __init__(self, stream):
        self.M = bytes(8) + stream
        self.p = len(stream)
        self.h = self.le64(self.p)
        self.pos = 0
        self.last_in = 0
        if self.h == 0:
            raise Refuse("stream_exhausted", 0)
        self.b = self.h.bit_length() - 1
        self.h -= 1 << self.b
        self.s = [0, 0]
        try:
            self.s[0] = self.read(8)
            self.s[1] = self.read(8)
            self.flush()
        except Refuse:
            raise Refuse("stream_exhausted", 0)

    def le64(self, i):
        return int.from_bytes(self.M[i:i + 8], "little")

    def read(self, k):
        if k > self.b:
            raise Refuse("stream_exhausted", self.pos)
        self.b -= k
        r = self.h >> self.b
        self.h -= r << self.b
        return r

    def flush(self):
        d = 7 ^ (self.b // 8)
        if d > self.p:
            raise Refuse("stream_exhausted", self.pos)
        self.p -= d
        self.h = self.le64(self.p) % (1 << (56 + self.b % 8))
        self.b = 56 + self.b % 8

    def step(self, which, t):
        self.s[which] = ((t >> 16) | self.read(t >> 24)) % 256

    def at(self, k):
        """The `in` or `out` step of position k runs: a refusal in it is at k."""
        self.pos = k

    def at_in(self, k):
        self.pos = k
        self.last_in = k

    def group_flush(self):
        """A flush that belongs to no single position (section 14.1, `at`): it reports
        the last position whose `in` step ran."""
        self.pos = self.last_in
        self.flush()


def wrap16(a):
    return ((a + 32768) % 65536) - 32768


def latent_tables(pmf):
    w = len(pmf[0]) if pmf else 0
    if w < 1 or any(len(r) != w or min(r) < 0 or sum(r) != 256 for r in pmf):
        raise Refuse("model_mismatch", -1)
    tabs = []
    for P in pmf:
        C = [0]
        for p in P:
            C.append(C[-1] + p)
        F = lambda k: C[k] - C[k] // 2
        G = lambda k: C[k] // 2 + 128
        sym = lambda k: k // 2 if k % 2 == 0 else -((k + 1) // 2)
        slot_sym, slot_next = [None] * 256, [None] * 256
        k = 0
        while C[k] < 256:
            for u in range(F(k), F(k + 1)):
                slot_sym[u], slot_next[u] = sym(k), P[k] + (u - F(k))
            for u in range(G(k), G(k + 1)):
                slot_sym[u], slot_next[u] = sym(k), P[k] + (F(k + 1) - F(k)) + (u - G(k))
            k += 1
        t = []
        for u in range(256):
            nb = 0
            while slot_next[u] * (1 << nb) < 256:
                nb += 1
            base = (slot_next[u] * (1 << nb)) ^ 256
            t.append(nb * (1 << 24) + base * (1 << 16) + slot_sym[u] % 65536)
        tabs.append(t)
    return tabs


def mans_sgm(i):
    pmf = [[dec(x) for x in r] for r in i["pmf"]]
    bounds = [dec(x) for x in i["bounds"]]
    spans = [([bool(b) for b in s["mask"]], [dec(x) for x in s["sigma"]]) for s in i["spans"]]
    tabs = latent_tables(pmf)
    if any(not 0 <= b <= 255 for b in bounds):
        raise Refuse("model_mismatch", -1)
    nd = min(len(tabs), len(bounds))
    if any(len(m) != len(sg) or any(not 0 <= x < nd for x in sg) for m, sg in spans):
        raise Refuse("model_mismatch", -1)
    w = Win(bytes.fromhex(i["memory"]))
    out = []
    base = 0
    for mask, sigma in spans:
        L = len(mask)
        sy = [0] * L

        def inb(which, k):
            w.at_in(base + k)
            if mask[k]:
                t = tabs[sigma[k]][w.s[which]]
                w.step(which, t)
                sy[k] = wrap16(t % 65536)

        def outb(k):
            w.at(base + k)
            if mask[k]:
                beta = bounds[sigma[k]]
                if sy[k] + beta == 0:
                    f = w.read(1)
                    v = w.read(3 if f == 1 else 16)
                    sy[k] = wrap16(v // 2 + beta) if v % 2 == 0 else wrap16(-(v // 2) - beta)
                    w.flush()

        k = 0
        while k + 3 < L:
            inb(0, k); inb(1, k + 1); inb(0, k + 2); inb(1, k + 3)
            w.group_flush()
            for q in range(4):
                outb(k + q)
            k += 4
        if L % 4 >= 2:
            inb(0, k); outb(k); inb(1, k + 1); outb(k + 1)
            k += 2
            w.group_flush()
        if L % 2 == 1:
            inb(0, k); outb(k)
            w.group_flush()
        out.append(sy)
        base += L
    return {"symbols": out}


def mans_fac(i):
    rows = [[dec(x) for x in r] for r in i["rows"]]
    lens = [dec(x) for x in i["lengths"]]
    if len(rows) != len(lens) or any(len(r) != 63 or r[0] < 0 or r[62] > 255 or any(r[k] > r[k + 1] for k in range(62)) for r in rows):
        raise Refuse("model_mismatch", -1)
    Zp = [0] * 512
    nb = 0
    for i_ in range(511, 0, -1):
        if i_ * (1 << nb) < 256:
            nb += 1
        Zp[i_] = ((i_ * (1 << nb)) ^ 256) + 256 * nb
    w = Win(bytes.fromhex(i["memory"]))
    Tz = [0] * 256
    out = []
    base = 0
    for r, L in zip(rows, lens):
        prev = 0
        for a in range(63):
            k = r[a] - prev
            for j in range(prev, r[a]):
                Tz[j] = Zp[k + j - prev] * 65536 + a
            prev = r[a]
        Tz[255] = Zp[1] * 65536 + 63
        sy = [0] * L

        def inb(which, k):
            w.at_in(base + k)
            t = Tz[w.s[which]]
            w.step(which, t)
            sy[k] = t % 256

        def outb(k):
            w.at(base + k)
            if sy[k] == 63:
                sy[k] = w.read(6)

        k = 0
        while k + 3 < L:
            inb(0, k); inb(1, k + 1); inb(0, k + 2); inb(1, k + 3)
            for q in range(4):
                outb(k + q)
            w.group_flush()
            k += 4
        if L % 4 >= 2:
            inb(0, k); outb(k); inb(1, k + 1); outb(k + 1)
            k += 2
        if L % 2 == 1:
            inb(0, k); outb(k)
        w.group_flush()
        out.append(sy)
        base += L
    return {"symbols": out}


def rdiv(n, d):
    if d < 1:
        return 0
    sg = (n > 0) - (n < 0)
    return sg * ((2 * abs(n) + d) // (2 * d))


def stage_expect(op, i, m):
    D = lambda k: dec(i[k])
    if op == "rans64_decode":
        return rans64_decode(i, m)
    if op == "dequantize":
        f, s, b = D("f"), D("sym"), D("base")
        if not 0 <= f <= 62:
            raise Refuse("field_range", 0)
        v = s * (1 << f) + b
        if not I64_MIN <= v <= I64_MAX:
            raise Refuse("field_range", 0)
        return {"value": v}
    if op == "bucket":
        sigma, e = D("sigma"), [dec(x) for x in i["edges"]]
        if m["bucket"] == "ceil":
            if not e:
                raise Refuse("model_mismatch", -1)
            s2 = max(sigma, e[0])
            return {"index": (len(e) - 1) - sum(1 for k in range(len(e) - 1) if s2 <= e[k])}
        if not e or any(e[k] > e[k + 1] for k in range(len(e) - 1)):
            raise Refuse("model_mismatch", -1)
        return {"index": sum(1 for x in e if x <= sigma)}
    if op == "coded":
        return {"coded": D("v") > D("thr")}
    if op == "scale_index":
        n, s = D("n"), D("s")
        if n < 1 or not 0 <= s <= 127:
            raise Refuse("model_mismatch", -1)
        return {"index": min(max(R(D("v"), s, 1), 0), n - 1)}
    if op == "block_stat":
        x = tensor_in(i["x"])
        b, s, lo, hi = D("b"), D("s"), D("lo"), D("hi")
        if b < 1 or not 0 <= s <= 127 or lo > hi:
            raise Refuse("model_mismatch", -1)
        d = list(x.d)
        for c in range(x.c):
            for by in range(0, x.h, b):
                for bx in range(0, x.w, b):
                    cells = [(y, xx) for y in range(by, min(by + b, x.h)) for xx in range(bx, min(bx + b, x.w))]
                    v = min(max(R(sum(x.at(c, y, xx) for y, xx in cells), s, 1), lo), hi)
                    for y, xx in cells:
                        d[(c * x.h + y) * x.w + xx] = v
        return {"out": tensor_out(T(x.c, x.h, x.w, d))}
    if op == "band_lookup":
        e = [dec(x) for x in i["edges"]]
        table = [[dec(x) for x in r] for r in i["table"]]
        row = D("row")
        if not 0 <= row < len(table) or len(table[row]) != len(e) + 1:
            raise Refuse("model_mismatch", -1)
        return {"value": table[row][sum(1 for k in e if D("v") >= k)]}
    if op == "mans_decode":
        return mans_sgm(i) if i["mode"] == "sgm" else mans_fac(i)
    if op == "residual_step":
        pred, r, q, lo, hi = D("pred"), D("r"), D("q"), D("lo"), D("hi")
        if lo >= hi or not lo <= pred <= hi or not 1 <= q <= 65535 or not I32_MIN <= r <= I32_MAX:
            raise Refuse("model_mismatch", -1)
        return {"value": min(max(pred + r * q, lo), hi)}
    if op == "mc_predict":
        f = tensor_in(i["frame"])
        y, x, c, dy, dx, pr = D("y"), D("x"), D("c"), D("dy"), D("dx"), D("precision")
        if not (0 <= c < f.c and 0 <= y < f.h and 0 <= x < f.w) or pr not in (1, 2):
            raise Refuse("model_mismatch", -1)
        p = lambda Y, X: f.at(c, min(max(Y, 0), f.h - 1), min(max(X, 0), f.w - 1))
        if pr == 1:
            return {"value": p(y + dy, x + dx)}
        iy, fy, ix, fx = y + (dy >> 1), dy & 1, x + (dx >> 1), dx & 1
        if not fy and not fx:
            v = p(iy, ix)
        elif not fy:
            v = (p(iy, ix) + p(iy, ix + 1) + 1) >> 1
        elif not fx:
            v = (p(iy, ix) + p(iy + 1, ix) + 1) >> 1
        else:
            v = (p(iy, ix) + p(iy, ix + 1) + p(iy + 1, ix) + p(iy + 1, ix + 1) + 2) >> 2
        return {"value": v}
    if op == "up2_predict":
        prev = tensor_in(i["prev"])
        d0, d1 = D("d0"), D("d1")
        hp, wp = -(-d0 // 2), -(-d1 // 2)
        if d0 < 1 or d1 < 1 or (prev.h, prev.w) != (hp, wp):
            raise Refuse("geometry_mismatch", -1)
        p = lambda c, Y, X: prev.at(c, min(Y, hp - 1), min(X, wp - 1))
        d = []
        for c in range(prev.c):
            for y in range(d0):
                for x in range(d1):
                    iy, fy, ix, fx = y >> 1, y & 1, x >> 1, x & 1
                    if not fy and not fx:
                        v = p(c, iy, ix)
                    elif not fy:
                        v = (p(c, iy, ix) + p(c, iy, ix + 1) + 1) >> 1
                    elif not fx:
                        v = (p(c, iy, ix) + p(c, iy + 1, ix) + 1) >> 1
                    else:
                        v = (p(c, iy, ix) + p(c, iy, ix + 1) + p(c, iy + 1, ix) + p(c, iy + 1, ix + 1) + 2) >> 2
                    d.append(v)
        return {"out": tensor_out(T(prev.c, d0, d1, d))}
    if op == "rate_predict":
        prev = tensor_in(i["prev"])
        num, den, d0 = D("rate_num"), D("rate_den"), D("d0")
        n = prev.w
        if prev.h != 1 or num < 1 or den < 1 or n * num // den != d0:
            raise Refuse("geometry_mismatch", -1)
        d = []
        for c in range(prev.c):
            for k in range(d0):
                s = k * den
                lo, rho = s // num, s % num
                hi = min(lo + 1, n - 1)
                a, b = prev.at(c, 0, lo), prev.at(c, 0, hi)
                d.append(rdiv(a * num + (b - a) * rho, num))
        return {"out": tensor_out(T(prev.c, 1, d0, d))}
    if op == "coding_order":
        c, h, w = (dec(x) for x in i["shape"])
        cb = None if i["order"] == "raster" else i["anchor"]
        idx = [(cc * h + y) * w + x for cc in range(c) for y in range(h) for x in range(w) if cb is None or (((y + x) % 2 == 0) == cb)]
        return {"indices": idx}
    raise ValueError("unknown stage operator %r" % op)


def stage_line(v):
    try:
        exp = stage_expect(v["op"], v["in"], v["modes"])
    except Refuse as e:
        exp = {"at": e.at, "refuse": e.code}
    return {"expect": exp, "id": v["id"], "in": v["in"], "modes": v["modes"], "op": v["op"]}


# --------------------------------------------------------------- lowerings (section 13)


class WB:
    """A WHG1 graph under construction with section 13.0's numbering."""

    def __init__(self):
        self.scalars, self.inputs, self.tensors, self.ops = [], [], [], []

    def tensor(self, c, policy=0, lo=0, hi=0):
        self.tensors.append((c, policy, lo, hi))
        return len(self.tensors) - 1

    def input(self, c, lo, hi, mh=0, mw=0):
        t = self.tensor(c, 2, lo, hi)
        self.inputs.append((t, mh, mw))
        return t

    def op(self, code, ins, out_rec, params):
        t = self.tensor(*out_rec)
        self.ops.append((code, ins, [t], params))
        return t

    def op_n(self, code, ins, out_recs, params):
        """An operator of several outputs, which take the next indices in their order."""
        ts = [self.tensor(*r) for r in out_recs]
        self.ops.append((code, ins, ts, params))
        return ts

    def bytes(self, outputs):
        b = bytearray(b"WHG1")
        b += bytes([1, 1, 0, 0])
        b += struct.pack("<IIHHHHII", len(self.tensors), len(self.ops), len(self.inputs), len(outputs), len(self.scalars), 0, CAP, 0)
        for lo, hi in self.scalars:
            b += struct.pack("<qq", lo, hi)
        for t, mh, mw in self.inputs:
            b += struct.pack("<IIIB3xI4xq", t, mh, mw, 0, 0, 0)
        for t in outputs:
            b += struct.pack("<I", t)
        for c, pol, lo, hi in self.tensors:
            b += struct.pack("<Ibb2xqq", c, -128, pol, lo, hi)
        for code, ins, outs, params in self.ops:
            b += struct.pack("<BBBx", code, len(ins), len(outs))
            for t in ins + outs:
                b += struct.pack("<I", t)
            b += struct.pack("<I", len(params)) + params
        return bytes(b)


def ints_le(v, width):
    return b"".join(x.to_bytes(width, "little", signed=True) for x in v)


def rq_block(round_, bias, mult, mkind, mbytes, shifts, skind, biasv, bbytes):
    """The RQ block of section 10.4."""
    b = struct.pack("<BBBBBBH", round_, bias, mkind, skind, mbytes, bbytes, 0)
    b += ints_le(mult, mbytes) if mkind else b""
    b += bytes(shifts)
    b += ints_le(biasv, bbytes) if bias else b""
    return b


def smallest(v, allowed):
    for w in allowed:
        lo, hi = -(1 << (8 * w - 1)), (1 << (8 * w - 1)) - 1
        if all(lo <= x <= hi for x in v):
            return w
    return 8


class LowerRefuse(Exception):
    pass


def lref():
    raise LowerRefuse()


def p_conv(cout, kh, kw, sh, sw, ph, pw, groups, w_bytes, rq, w):
    return struct.pack("<8IB3x", cout, kh, kw, sh, sw, ph, pw, groups, w_bytes) + rq + ints_le(w, w_bytes)


def p_convt(cout, kh, kw, sh, sw, ph, pw, oph, opw, w_bytes, rq, w):
    return struct.pack("<9IB3x", cout, kh, kw, sh, sw, ph, pw, oph, opw, w_bytes) + rq + ints_le(w, w_bytes)


def p_clamp(lo, hi):
    return struct.pack("<qq", lo, hi)


def int_synth_lower(js):
    """int_synth/1 (section 13.1)."""
    try:
        root = json.loads(js)
    except ValueError:
        lref()
    layers = root.get("layers") if isinstance(root, dict) else None
    if not isinstance(layers, list) or not layers:
        lref()

    def u(l, k, default=None):
        v = l.get(k, default)
        if v is None or isinstance(v, bool) or not isinstance(v, int) or not 0 <= v <= 0xFFFFFFFF:
            lref()
        return v

    def i32s(l, k):
        a = l.get(k)
        if not isinstance(a, list) or any(isinstance(x, bool) or not isinstance(x, int) or not I32_MIN <= x <= I32_MAX for x in a):
            lref()
        return a

    if not isinstance(layers[0], dict):
        lref()
    c = u(layers[0], "in_ch")
    if c == 0:
        lref()
    wb = WB()
    x = wb.input(c, -128, 127)
    for l in layers:
        if not isinstance(l, dict):
            lref()
        ic, oc, k, s, p = u(l, "in_ch"), u(l, "out_ch"), u(l, "k"), u(l, "stride"), u(l, "pad")
        op_, sh = u(l, "out_pad", 0), u(l, "shift", 0)
        relu = l.get("relu", False)
        if not isinstance(relu, bool):
            lref()
        w, b = i32s(l, "weight"), i32s(l, "bias")
        if ic != c or oc == 0 or k == 0 or s == 0 or sh > 127 or len(w) != ic * oc * k * k or len(b) != oc:
            lref()
        rq = rq_block(0, 1, [], 0, 0, [sh], 1, b, 4)
        x = wb.op(0x11, [x], (oc,), p_convt(oc, k, k, s, s, p, p, op_, op_, 4, rq, w))
        if relu:
            x = wb.op(0x0C, [x], (oc,), b"")
        c = oc
    x = wb.op(0x0B, [x], (c,), p_clamp(0, 255))
    return [("synth", wb.bytes([x]))]


def parse_net(b, magic, audio):
    """WHM1 / WAM1 (integer-payloads section 9)."""
    if b[:4] != magic:
        lref()
    rd = Reader(b, 4)
    try:
        afrac, gf, shift, n = rd.u(4), rd.u(4), rd.u(4), rd.u(4)
        layers = []
        for _ in range(n):
            tag = rd.u(1)
            if tag == 0:
                inc, out, k, s, p, op_ = (rd.u(4) for _ in range(6))
                w = rd.ints(inc * out * k * (1 if audio else k), 2)
                layers.append(("convT", inc, out, k, s, p, op_, w, rd.ints(out, 8), rd.ints(out, 8)))
            elif tag == 1 and not audio:
                inc, out, k, s, p = (rd.u(4) for _ in range(5))
                w = rd.ints(out * inc * k * k, 2)
                layers.append(("conv", inc, out, k, s, p, 0, w, rd.ints(out, 8), rd.ints(out, 8)))
            elif tag == 2:
                c = rd.u(4)
                layers.append(("igdn", c, rd.ints(c * c, 4), rd.ints(c, 8)))
            elif tag == 3 and not audio:
                layers.append(("relu",))
            elif tag == 4 and not audio:
                layers.append(("leaky", rd.i(8), rd.u(4)))
            else:
                lref()
    except Refuse:
        lref()
    if rd.p != len(b):
        lref()
    return {"afrac": afrac, "gf": gf, "shift": shift, "layers": layers, "audio": audio}


def largest_x(pred):
    """The largest X in [0, 2^63] with pred(X), pred monotone; None if pred(0) fails."""
    if not pred(0):
        return None
    lo, hi = 0, 1 << 63
    if pred(hi):
        return hi
    while hi - lo > 1:
        mid = (lo + hi) // 2
        if pred(mid):
            lo = mid
        else:
            hi = mid
    return lo


def x_l(net, layer, c):
    """X_l of section 13.2 for a layer of tag 0, 1, 2 or 4 reading c channels; the set
    is refused when P_l(0) fails."""
    I64, I128 = (1 << 63) - 1, (1 << 127) - 1
    kind = layer[0]
    if kind in ("conv", "convT"):
        k = layer[3]
        terms = c * k * (1 if net["audio"] else k)
        Wm, Mm, Bm = max(map(abs, layer[7]), default=0), max(map(abs, layer[8]), default=0), max(map(abs, layer[9]), default=0)
        p = lambda X: terms * Wm * X <= I64 and (terms * Wm * X * Mm) // (1 << net["shift"]) + 1 + Bm <= I64
    elif kind == "igdn":
        G, Bb = max(map(abs, layer[2]), default=0), max(map(abs, layer[3]), default=0)

        def p(X):
            if G * X * X > I128:
                return False
            N = Bb + c * ((G * X * X) // (1 << net["gf"]) + 1)
            if N > I128:
                return False
            r = math.isqrt(N)
            return r <= I64 and (X * r) // (1 << net["afrac"]) + 1 <= I64
    else:
        a = abs(layer[1])
        p = lambda X: X * a <= I64
    x = largest_x(p)
    if x is None:
        lref()
    return x


def x_record(net, k, c):
    """The tensor record of x_k on c channels (section 13.2, Range checks). x_k is read
    by layer k, if there is one. When that layer is of tag 0, 1, 2 or 4, x_k is
    check [-X_k, min(X_k, I64)]; that includes the network's input. A network input
    that no such layer reads is check [-2^63, I64]. Every other tensor is derived."""
    L = net["layers"]
    if k < len(L) and L[k][0] in ("convT", "conv", "igdn", "leaky"):
        x = x_l(net, L[k], c)
        return (c, 2, -x, min(x, I64_MAX))
    if k == 0:
        return (c, 2, I64_MIN, I64_MAX)
    return (c,)


def net_layers(wb, net, x):
    L = net["layers"]
    for k, l in enumerate(L):
        c = wb.tensors[x][0]
        oc = l[2] if l[0] in ("conv", "convT") else c
        rec = x_record(net, k + 1, oc)
        if l[0] in ("conv", "convT"):
            _, inc, out, kk, s, p, op_, w, mult, bias = l
            if inc != c or out == 0 or kk == 0 or s == 0:
                lref()
            rq = rq_block(0, 2, mult, 2, 8, [net["shift"]], 1, bias, 8)
            if l[0] == "conv":
                x = wb.op(0x01, [x], rec, p_conv(out, kk, kk, s, s, p, p, 1, 2, rq, w))
            elif net["audio"]:
                x = wb.op(0x11, [x], rec, p_convt(out, 1, kk, 1, s, 0, p, 0, op_, 2, rq, w))
            else:
                x = wb.op(0x11, [x], rec, p_convt(out, kk, kk, s, s, p, p, op_, op_, 2, rq, w))
        elif l[0] == "igdn":
            if l[1] != c:
                lref()
            x = wb.op(0x12, [x], rec, struct.pack("<IIB7x", net["gf"], net["afrac"], 4) + ints_le(l[2], 4) + ints_le(l[3], 8))
        elif l[0] == "relu":
            x = wb.op(0x0C, [x], rec, b"")
        else:
            if l[2] > 63:
                lref()
            x = wb.op(0x0D, [x], rec, struct.pack("<qI", l[1], l[2]))
    return x


def net_check(net):
    if net["shift"] > 127 or net["gf"] > 127 or not 1 <= net["afrac"] <= 30 or not net["layers"]:
        lref()
    for l in net["layers"]:
        if l[0] in ("conv", "convT"):
            return l[1]
        if l[0] == "igdn":
            return l[1]
    lref()


def rgb8(wb, x, afrac):
    c = wb.tensors[x][0]
    t = wb.op(0x0B, [x], (c,), p_clamp(0, 1 << afrac))
    return wb.op(0x02, [t], (c,), rq_block(1, 0, [255], 1, 4, [afrac], 1, [], 0))


def whm1_lower(b, output_step):
    net = parse_net(b, b"WHM1", False)
    c0 = net_check(net)
    wb = WB()
    r = x_record(net, 0, c0)
    x = wb.input(r[0], r[2], r[3])
    x = net_layers(wb, net, x)
    if output_step:
        x = rgb8(wb, x, net["afrac"])
    return wb.bytes([x])


def wam1_lower(b):
    net = parse_net(b, b"WAM1", True)
    c0 = net_check(net)
    wb = WB()
    r = x_record(net, 0, c0)
    x = wb.input(r[0], r[2], r[3])
    wb.scalars.append((1, 1 << 26))
    x = net_layers(wb, net, x)
    c = wb.tensors[x][0]
    lo = -((1 << 63) // 32767)
    hi = (I64_MAX - (1 << (net["afrac"] - 1))) // 32767
    crop = struct.pack("<B3xIB3xIB3xIB3xI", 0, 0, 0, 0, 0, 1, 1, 0)
    t = wb.op(0x14, [x], (c, 2, lo, hi), crop)
    t = wb.op(0x02, [t], (c,), rq_block(1, 0, [32767], 1, 4, [net["afrac"]], 1, [], 0))
    t = wb.op(0x0B, [t], (c,), p_clamp(-32768, 32767))
    return wb.bytes([t])


def wmb1_lower(b):
    if b[:4] != b"WMB1":
        lref()
    rd = Reader(b, 4)
    try:
        secs = [rd.take(rd.u(4)) for _ in range(4)]
    except Refuse:
        lref()
    if rd.p != len(b):
        lref()
    nets = [parse_net(s, b"WHM1", False) for s in secs[:3]]
    if len({n["afrac"] for n in nets}) != 1:
        lref()
    try:
        tables = json.loads(secs[3])
    except ValueError:
        lref()
    a = tables.get("afrac") if isinstance(tables, dict) else None
    if isinstance(a, bool) or not isinstance(a, int) or a != nets[2]["afrac"]:
        lref()
    return [("h_s", net_graph(nets[0], False)), ("entropy", net_graph(nets[1], False)), ("g_s", net_graph(nets[2], True))]


def net_graph(net, step):
    c0 = net_check(net)
    wb = WB()
    r = x_record(net, 0, c0)
    x = wb.input(r[0], r[2], r[3])
    x = net_layers(wb, net, x)
    if step:
        x = rgb8(wb, x, net["afrac"])
    return wb.bytes([x])


def world_lower(tb, gb, c):
    tn, gn = parse_net(tb, b"WHM1", False), parse_net(gb, b"WHM1", False)
    cin = net_check(tn)
    net_check(gn)
    if tn["afrac"] != gn["afrac"] or c < 1 or cin <= c:
        lref()
    A = cin - c
    wb = WB()
    z = wb.input(c, I64_MIN, I64_MAX)
    parts = [z]
    for _ in range(A):
        wb.scalars.append((I64_MIN, I64_MAX))
        s = len(wb.scalars) - 1
        parts.append(wb.op(0x17, [z], (1,), struct.pack("<B7xq", 1, s)))
    x0 = wb.op(0x06, parts, x_record(tn, 0, cin), b"")
    x = net_layers(wb, tn, x0)
    one = 1 << tn["afrac"]
    out = wb.op(0x0B, [x], (wb.tensors[x][0],), p_clamp(-one, one))
    return [("T", wb.bytes([out])), ("g_s", net_graph(gn, True))]


JPEGAI_MEMBERS = (
    "norm.json", "weights/hs_y.json", "weights/hs_uv.json", "weights/hyp_y.json", "weights/hyp_uv.json",
    "weights/mcm0_y.json", "weights/mcm1_y.json", "weights/mcm2_y.json", "weights/mcm3_y.json",
    "weights/synth_y.json", "weights/synth_uv.json",
)


def jpegai_lower(files):
    """jpegai-int/1 (section 13.6): the scale graphs of 13.6.1, then the graphs of
    13.6.2, in section 13.0's role order. `files` maps member names to bytes; a member
    the lowering reads that the set lacks refuses it. Other members are not read."""
    if any(m not in files for m in JPEGAI_MEMBERS):
        lref()
    roles = jpegai_scale_lower(files["norm.json"], files["weights/hs_y.json"], files["weights/hs_uv.json"])
    W = {k: q16_file(files["weights/%s.json" % k]) for k in ("hyp_y", "hyp_uv", "mcm0_y", "mcm1_y", "mcm2_y", "mcm3_y", "synth_y", "synth_uv")}
    roles.append(("rec_y", rec_graph(W["hyp_y"], [W["mcm%d_y" % i] for i in range(4)], 160)))
    roles.append(("rec_uv", rec_graph(W["hyp_uv"], None, 96)))
    roles.append(("syn_y", syn_y_graph(W["synth_y"])))
    roles.append(("syn_uv", syn_uv_graph(W["synth_uv"])))
    return roles


def q16_file(b):
    """One weights file of section 13.6.2, *The members*: {name: member} for each
    member whose value is an object of kind "conv" or "q16"; every other member is
    ignored. Refused: a file that is not a JSON object, two members of one object with
    one name (anywhere in the file), and a member of either kind that lacks one of its
    fields or holds another value."""
    def pairs(ps):
        if len({k for k, _ in ps}) != len(ps):
            lref()
        return dict(ps)

    try:
        root = json.loads(b, object_pairs_hook=pairs)
    except ValueError:
        lref()
    if not isinstance(root, dict):
        lref()
    isint = lambda v: isinstance(v, int) and not isinstance(v, bool)
    out = {}
    for name, m in root.items():
        if not isinstance(m, dict) or m.get("kind") not in ("conv", "q16"):
            continue
        if m["kind"] == "conv":
            dims, w, mult, shift = m.get("dims"), m.get("w"), m.get("mult"), m.get("shift")
            if not isinstance(dims, list) or not all(isint(d) and d >= 0 for d in dims):
                lref()
            n = 1
            for d in dims:
                n *= d
            if not isinstance(w, list) or len(w) != n or not all(isint(x) and I64_MIN <= x <= I64_MAX for x in w):
                lref()
            if not isint(mult) or not I64_MIN <= mult <= I64_MAX or not isint(shift) or not 0 <= shift <= 127:
                lref()
            out[name] = {"kind": "conv", "dims": dims, "w": w, "mult": mult, "shift": shift}
        else:
            v = m.get("v")
            if not isinstance(v, list) or not all(isint(x) and I64_MIN <= x <= I64_MAX for x in v):
                lref()
            out[name] = {"kind": "q16", "v": v}
    return out


A32 = (2, I32_MIN, I32_MAX)
Q16_RQ = rq_block(0, 0, [], 0, 0, [16], 1, [], 0)  # mul: no multiplier, shift 16, no bias, floor
ATT_RQ = rq_block(1, 0, [], 0, 0, [16], 1, [], 0)  # gram, apply: shift 16, half_up
SUBLATTICE = bytes([0, 3, 1, 2])  # (0,0), (1,1), (0,1), (1,0), each 2a + b


class Q16Net:
    """The blocks of section 13.6.2 over one graph under construction, reading the
    modules of the weights file F (which a graph may switch, as rec_y does)."""

    def __init__(self):
        self.wb = WB()
        self.F = {}

    def c(self, x):
        return self.wb.tensors[x][0]

    def weight(self, m):
        e = self.F.get(m + ".weight")
        if e is None or e["kind"] != "conv":
            lref()
        return e

    def q16(self, name):
        e = self.F.get(name)
        if e is None or e["kind"] != "q16":
            lref()
        return e["v"]

    def bias(self, m, cout):
        """m.bias, cout values, or None when the file has none."""
        e = self.F.get(m + ".bias")
        if e is None:
            return None
        if e["kind"] != "q16" or len(e["v"]) != cout:
            lref()
        return e["v"]

    def rq(self, e, shift, bias):
        """C's RQ: floor; one multiplier; one shift; a post bias, or none."""
        mult = e["mult"]
        if bias is None:
            return rq_block(0, 0, [mult], 1, smallest([mult], (4, 8)), [shift], 1, [], 0)
        return rq_block(0, 2, [mult], 1, smallest([mult], (4, 8)), [shift], 1, bias, smallest(bias, (4, 8)))

    def wbytes(self, w):
        if any(not I64_MIN <= v <= I64_MAX for v in w):
            lref()
        return smallest(w, (1, 2, 4, 8))

    def conv(self, x, m, k, s, p, g, shift=None):
        """C(m; k, s, p, g); g = "dw" is groups = c, with cout = c."""
        c = self.c(x)
        e = self.weight(m)
        dims = e["dims"]
        dw = g == "dw"
        if dw:
            g = c
        if len(dims) != 4 or c % g:
            lref()
        cout = dims[0]
        if cout == 0 or cout % g or dims != [cout, c // g, k, k] or (dw and cout != c):
            lref()
        rq = self.rq(e, e["shift"] if shift is None else shift, self.bias(m, cout))
        w = e["w"]
        return self.wb.op(0x01, [x], (cout,) + A32, p_conv(cout, k, k, s, s, p, p, g, self.wbytes(w), rq, w))

    def convt(self, x, m, k, op_):
        """T(m; k, op): stride 2, padding 1, output padding op, with C's RQ."""
        c = self.c(x)
        e = self.weight(m)
        dims = e["dims"]
        if len(dims) != 4 or dims[0] != c or dims[1] == 0 or dims[2:] != [k, k]:
            lref()
        cout = dims[1]
        rq = self.rq(e, e["shift"], self.bias(m, cout))
        w = e["w"]
        return self.wb.op(0x11, [x], (cout,) + A32, p_convt(cout, k, k, 2, 2, 1, 1, op_, op_, self.wbytes(w), rq, w))

    def same(self, *xs):
        """The inputs of each mul and each add have one channel count."""
        if len({self.c(x) for x in xs}) != 1:
            lref()
        return self.c(xs[0])

    def add(self, a, b):
        return self.wb.op(0x05, [a, b], (self.same(a, b),) + A32, b"")

    def mul(self, a, b):
        return self.wb.op(0x10, [a, b], (self.same(a, b),) + A32, Q16_RQ)

    def unary(self, code, x, params=b""):
        return self.wb.op(code, [x], (self.c(x),), params)

    def relu6(self, x):
        return self.unary(0x0B, x, p_clamp(0, 6 << 16))

    def split(self, x, start, count):
        if start + count > self.c(x):
            lref()
        return self.wb.op(0x08, [x], (count,), struct.pack("<II", start, count))

    def concat(self, xs):
        return self.wb.op(0x06, list(xs), (sum(self.c(x) for x in xs),), b"")

    def shuffle2(self, x):
        if self.c(x) % 4:
            lref()
        return self.wb.op(0x09, [x], (self.c(x) // 4,), struct.pack("<I", 2))

    def merge(self, xs):
        return self.wb.op(0x16, list(xs), (self.same(*xs),), SUBLATTICE)

    def R(self, x, m):
        t = self.conv(x, m + ".conv1", 3, 1, 1, 1)
        t = self.unary(0x0C, t)
        t = self.conv(t, m + ".conv2", 3, 1, 1, 1)
        return self.add(x, t)

    def G(self, x, m):
        """The gate x·(f(x) + 1)."""
        t = self.relu6(x)
        t = self.conv(t, m + ".conv", 3, 1, 1, 4)
        t = self.conv(t, m + ".conv2", 1, 1, 0, 1)
        c = self.c(t)
        t = self.wb.op(0x02, [t], (c,) + A32, rq_block(0, 2, [], 0, 0, [0], 1, [1 << 16] * c, 4))
        return self.mul(x, t)

    def CAB(self, x):
        t = self.R(x, "CAB.residual_trunk.0")
        t = self.R(t, "CAB.residual_trunk.1")
        s = self.conv(x, "CAB.subscale", 3, 2, 1, 1)
        s = self.R(s, "CAB.residual_mask2.0")
        s = self.R(s, "CAB.residual_mask3.0")
        s = self.convt(s, "CAB.upscale", 3, 1)
        s = self.unary(0x1B, s)
        return self.add(x, self.mul(t, s))

    def D(self, x):
        """The denormalisation, clamp [0, 255·2^16]."""
        c = self.c(x)
        rq = rq_block(0, 2, [128], 1, 4, [0], 1, [127 * (1 << 16) + (1 << 15)] * c, 4)
        return self.wb.op(0x02, [x], (c, 1, 0, 255 << 16), rq)

    def layernorm(self, x, m):
        c = self.c(x)
        gm, bt = self.q16(m + ".weight"), self.q16(m + ".bias")
        if len(gm) != c or len(bt) != c:
            lref()
        gb = smallest(gm, (4, 8))
        return self.wb.op(0x18, [x], (c,) + A32, struct.pack("<QB7x", 42949, gb) + ints_le(gm, gb) + ints_le(bt, 8))

    def TAB(self, x, m):
        c = self.c(x)
        if c % 4:
            lref()
        d = c // 4
        sa = self.weight(m + ".prep_data.conv1")["shift"]
        sf = self.weight(m + ".ffn.project_in")["shift"]
        t = self.layernorm(x, m + ".prep_data.norm1")
        t = self.conv(t, m + ".prep_data.conv1", 1, 1, 0, 1, sa)
        if self.c(t) != 3 * c:
            lref()
        t = self.conv(t, m + ".prep_data.conv2", 3, 1, 1, "dw", sa)
        q, k, v = self.split(t, 0, c), self.split(t, c, c), self.split(t, 2 * c, c)
        q = self.unary(0x19, q, struct.pack("<Q", 42949))
        k = self.unary(0x19, k, struct.pack("<Q", 42949))
        a = self.wb.op(0x1D, [q, k], (c,) + A32, struct.pack("<I", 4) + ATT_RQ)
        temp = self.q16(m + ".attn.temperature")
        if len(temp) != 4:
            lref()
        mult = [temp[i // d] for i in range(c)]
        a = self.wb.op(0x02, [a], (c,) + A32, rq_block(1, 0, mult, 2, smallest(mult, (4, 8)), [16], 1, [], 0))
        a = self.unary(0x1A, a)
        a = self.wb.op(0x1E, [a, v], (c,) + A32, struct.pack("<I", 4) + ATT_RQ)
        a = self.conv(a, m + ".attn.project_out", 1, 1, 0, 1, sa)
        x2 = self.add(x, a)
        t = self.layernorm(x2, m + ".ffn.norm1")
        t = self.conv(t, m + ".ffn.project_in", 1, 1, 0, 1, sf)
        if self.c(t) % 2:
            lref()
        h = self.c(t) // 2
        t = self.conv(t, m + ".ffn.dwconv", 3, 1, 1, "dw", sf)
        a, b = self.split(t, 0, h), self.split(t, h, h)
        u = self.mul(self.unary(0x1C, a, struct.pack("<q", 1 << 16)), b)
        u = self.conv(u, m + ".ffn.project_out", 1, 1, 0, 1, sf)
        return self.add(x2, u)


def rec_graph(hyp, mcms, C):
    """rec_y (mcms: the four context stages' files) or rec_uv (mcms None), 13.6.2."""
    n = Q16Net()
    n.F = hyp
    z = n.wb.input(C, -31, 32, 4, 4)
    r = n.wb.input(C, I32_MIN, I32_MAX, 16, 16)
    t = n.wb.op(0x02, [z], (C,), rq_block(0, 0, [1 << 16], 1, 4, [0], 1, [], 0))
    t = n.conv(t, "conv1", 1, 1, 0, 1)
    t = n.convt(t, "conv2", 4, 0)
    t = n.relu6(t)
    t = n.conv(t, "conv3", 3, 1, 1, 1)
    t = n.relu6(t)
    psi = n.conv(t, "conv4", 3, 1, 1, 1)
    if n.c(psi) != 4 * C:
        lref()
    if mcms is None:
        out = n.add(r, n.merge([n.split(psi, i * C, C) for i in range(4)]))
        return n.wb.bytes([out])
    rs = n.wb.op_n(0x15, [r], [(C,)] * 4, SUBLATTICE)
    ys = []
    for i in range(4):
        n.F = mcms[i]
        p = n.split(psi, i * C, C)
        if i == 0:
            f = p
        else:
            ctx = ys[0] if i == 1 else n.concat(ys)
            f = n.conv(ctx, "conv.0", 1, 1, 0, 1)
            f = n.conv(f, "conv.1", 3, 1, 1, 5)
            f = n.concat([f, p])
        f = n.unary(0x0C, n.conv(f, "fusion_pred_net.conv1", 1, 1, 0, 1))
        f = n.unary(0x0C, n.conv(f, "fusion_pred_net.conv2", 1, 1, 0, 1))
        mu = n.conv(f, "fusion_pred_net.conv3", 1, 1, 0, 1)
        if n.c(mu) != C:
            lref()
        ys.append(n.add(rs[i], mu))
    return n.wb.bytes([n.merge(ys)])


def syn_y_graph(F):
    n = Q16Net()
    n.F = F
    x = n.wb.input(160, I32_MIN, I32_MAX, 16, 16)
    x = n.R(x, "first_stage.0.0")
    x = n.convt(x, "first_stage.1", 3, 1)
    x = n.G(x, "first_stage.2")
    x = n.convt(x, "conv2_t", 3, 1)
    x = n.CAB(x)
    x = n.G(x, "iact2")
    x = n.conv(x, "conv3_t", 1, 1, 0, 1)
    x = n.shuffle2(x)
    x = n.conv(x, "TAM.ds_conv", 3, 2, 1, 1)
    x = n.TAB(x, "TAM.TABs.0")
    x = n.TAB(x, "TAM.TABs.1")
    x = n.convt(x, "TAM.us_conv", 3, 1)
    x = n.G(x, "iact3")
    x = n.convt(x, "conv4_t", 3, 1)
    if n.c(x) != 1:
        lref()
    return n.wb.bytes([n.D(x)])


def syn_uv_graph(F):
    n = Q16Net()
    n.F = F
    y = n.wb.input(160, I32_MIN, I32_MAX, 16, 16)
    uv = n.wb.input(96, I32_MIN, I32_MAX, 16, 16)
    x = n.concat([y, uv])
    x = n.conv(x, "first_stage.conv1", 3, 1, 1, 1)
    x = n.concat([x, uv])
    x = n.convt(x, "conv2_t", 3, 1)
    x = n.CAB(x)
    x = n.G(x, "iact2")
    x = n.conv(x, "conv3_t", 1, 1, 0, 1)
    x = n.shuffle2(x)
    x = n.TAB(x, "TAM.TABs.0")
    x = n.TAB(x, "TAM.TABs.1")
    x = n.G(x, "iact3")
    x = n.convt(x, "conv4_t", 3, 1)
    if n.c(x) != 8:
        lref()
    x = n.shuffle2(x)
    return n.wb.bytes([n.D(x)])


def jpegai_scale_lower(norm_b, hy_b, huv_b):
    """jpegai-int/1, section 13.6.1: the scale graphs hs_y and hs_uv."""
    try:
        norm, hy, huv = json.loads(norm_b), json.loads(hy_b), json.loads(huv_b)
    except ValueError:
        lref()
    isint = lambda v: isinstance(v, int) and not isinstance(v, bool)
    ints = lambda v: v if isinstance(v, list) and all(isint(x) for x in v) else lref()
    if not isinstance(norm, dict):
        lref()
    pmf = norm.get("pmf")
    if not isinstance(pmf, list) or len(pmf) < 32:
        lref()
    widths = set()
    for row in pmf:
        r = ints(row)
        if not r or min(r) < 0 or sum(r) != 256:
            lref()
        widths.add(len(r))
    if len(widths) != 1:
        lref()
    bounds = ints(norm.get("bounds"))
    if len(bounds) < 32 or any(not 0 <= v <= 255 for v in bounds):
        lref()
    if not isint(norm.get("per_channel")) or norm["per_channel"] != 16:
        lref()
    for k, rows in (("cdfs_y", 160), ("cdfs_uv", 96)):
        v = ints(norm.get(k))
        if len(v) != rows * 63 or any(not 0 <= x <= 255 for x in v) or any(v[r * 63 + j] > v[r * 63 + j + 1] for r in range(rows) for j in range(62)):
            lref()
    gy, guv = ints(norm.get("gain_y")), ints(norm.get("gain_uv"))
    if len(gy) != 160 or len(guv) != 96 or any(not I64_MIN <= v <= I64_MAX for v in gy + guv):
        lref()

    def member(o, k, dims):
        m = o.get(k)
        if not isinstance(m, dict) or m.get("dims") != dims:
            lref()
        v = ints(m.get("v"))
        n = 1
        for d in dims:
            n *= d
        if len(v) != n or any(not I64_MIN <= x <= I64_MAX for x in v):
            lref()
        return v

    def plane(o, cz, gains):
        if not isinstance(o, dict):
            lref()

        def d0(k):
            m = o.get(k)
            dims = m.get("dims") if isinstance(m, dict) else None
            if not isinstance(dims, list) or not dims or not isint(dims[0]) or not 1 <= dims[0] <= 0xFFFFFFFF:
                lref()
            return dims[0]

        c1, c2 = d0("conv1.weight"), d0("depthwise.weight")
        c3 = 16 * cz
        wb = WB()
        x = wb.input(cz, -31, 32, 4, 4)
        blocks = (("conv1", c1, cz, 1, 0, "onnx::Pow_85"), ("depthwise", c2, c1, 3, 1, "onnx::Pow_91"), ("pointwise", c3, c2, 1, 0, "onnx::Pow_97"))
        for bi, (name, cout, cin, k, pad, sh) in enumerate(blocks):
            w = member(o, name + ".weight", [cout, cin, k, k])
            bias = member(o, name + ".bias", [cout])
            shifts = member(o, sh, [1, cout, 1, 1])
            if any(not 0 <= s <= 127 for s in shifts):
                lref()
            x = wb.op(0x0B, [x], (wb.tensors[x][0],), p_clamp(-128, 127))
            rq = rq_block(0, 1, [], 0, 0, shifts, 2, bias, smallest(bias, (4, 8)))
            wbytes = smallest(w, (1, 2, 4, 8))
            x = wb.op(0x01, [x], (cout,), p_conv(cout, k, k, 1, 1, pad, pad, 1, wbytes, rq, w))
            if bi < 2:
                x = wb.op(0x0C, [x], (cout,), b"")
        x = wb.op(0x09, [x], (cz,), struct.pack("<I", 4))
        x = wb.op(0x13, [x], (cz,), b"")
        x = wb.op(0x0B, [x], (cz,), p_clamp(0, 3967))
        x = wb.op(0x02, [x], (cz,), rq_block(0, 1, [], 0, 0, [0], 1, gains, smallest(gains, (4, 8))))
        return wb.bytes([x])

    return [("hs_y", plane(hy, 160, gy)), ("hs_uv", plane(huv, 96, guv))]


def lower_set(entry, files):
    """The set's lowered graphs, [(role, WHG1 bytes)] in section 13.0's role order.
    Raises LowerRefuse when the lowering refuses the set."""
    low, cap = entry["lowering"], entry["capability"]
    if low == "int_synth/1":
        return int_synth_lower(files["int_synth.json"])
    if low == "whm1/1" and cap == "wai.splat.int_codec":
        return [("g_s", whm1_lower(files["g_s.whm"], True))]
    if low == "whm1/1":
        return [("h_s", whm1_lower(files["h_s.whm"], False)), ("g_s", whm1_lower(files["g_s.whm"], True))]
    if low == "wam1/1":
        return [("g_s", wam1_lower(files["model.wam"]))]
    if low == "wmb1/1":
        return wmb1_lower(files["model.wmb"])
    if low == "world/1":
        return world_lower(files["T.whm"], files["g_s.whm"], dec(entry["c"]))
    if low == "jpegai-int/1":
        return jpegai_lower(files)
    raise ValueError("no lowering %r" % low)


def cert_path(root, entry):
    """Section 14.3: certificates/<capability>/<set digest>.json."""
    return os.path.join(root, "certificates", entry["capability"], entry["set_sha256"] + ".json")


def check_set_certificate(root, entry, roles, write):
    """Recomputes a set's certificate from its lowered graphs and requires the bytes
    of the committed one. Returns the verdict."""
    got, v = certificate(entry["capability"], entry["lowering"], entry["set_sha256"], roles)
    path = cert_path(root, entry)
    if write:
        p = os.path.join(write, os.path.relpath(path, root))
        os.makedirs(os.path.dirname(p), exist_ok=True)
        open(p, "wb").write(got)
    want = open(path, "rb").read()
    if got != want:
        raise ValueError("the certificate differs: %s" % first_difference(got, want))
    return v


def set_digest(files):
    if len(files) == 1:
        return hashlib.sha256(next(iter(files.values()))).hexdigest()
    listing = "".join("%s  %s\n" % (hashlib.sha256(files[n]).hexdigest(), n) for n in sorted(files, key=lambda s: s.encode()))
    return hashlib.sha256(listing.encode()).hexdigest()


# --------------------------------------------------------------- main


def check(root, write, external=None):
    wai = os.path.dirname(os.path.abspath(root))
    failures = 0
    nfiles = nvec = 0
    for sub, fn in (("vectors", graph_line), ("stage", stage_line), ("graphs", graph_line)):
        d = os.path.join(root, sub)
        for name in sorted(os.listdir(d)):
            if not name.endswith(".jsonl"):
                continue
            raw = open(os.path.join(d, name), "rb").read()
            out = []
            try:
                for line in raw.decode("ascii").split("\n")[:-1]:
                    out.append(cw(fn(json.loads(line))) + "\n")
            except Exception as e:
                failures += 1
                print("FAIL %s/%s: %s: %s" % (sub, name, type(e).__name__, e))
                continue
            got = "".join(out).encode()
            nfiles += 1
            nvec += len(out)
            if write:
                os.makedirs(os.path.join(write, sub), exist_ok=True)
                open(os.path.join(write, sub, name), "wb").write(got)
            if got != raw:
                failures += 1
                a, b = got.split(b"\n"), raw.split(b"\n")
                k = next((k for k in range(min(len(a), len(b))) if a[k] != b[k]), min(len(a), len(b)))
                print("DIFF %s/%s: line %d differs" % (sub, name, k + 1))
    print("rewrote %d vector files, %d vectors" % (nfiles, nvec))
    nsets = ngraphs = ncerts = 0
    lowered = {}
    # The sets index.json lists, then the falsification sets of section 14.3. Each is
    # re-lowered, its lowered graphs compared, and its certificate recomputed from
    # them and compared. An index set's graphs each load (steps 1-8), so its verdict
    # meets acc64; a falsification set's verdict fails the profile it claims.
    listed = [(e, None) for e in json.load(open(os.path.join(root, "certificates", "index.json")))["sets"]]
    listed += [(e, e["claim"]) for e in json.load(open(os.path.join(root, "certificates", "falsification.json")))["sets"]]
    for e, claim in listed:
        what = "falsification set" if claim else "set"
        try:
            files = {m: open(os.path.join(wai, p), "rb").read() for m, p in e["members"].items()}
            if set_digest(files) != e["set_sha256"]:
                raise ValueError("the set's digest differs")
            try:
                roles = lower_set(e, files)
            except LowerRefuse:
                raise ValueError("the lowering refused the set")
            if len(roles) != len(e["roles"]):
                raise ValueError("the roles differ")
            for (role, b), path in zip(roles, e["roles"]):
                if not path.endswith("/%s/%s.whg1" % (e["set_sha256"], role)):
                    raise ValueError("role %s is at %s" % (role, path))
                if write:
                    p = os.path.join(write, os.path.relpath(os.path.join(wai, path), root))
                    os.makedirs(os.path.dirname(p), exist_ok=True)
                    open(p, "wb").write(b)
                if open(os.path.join(wai, path), "rb").read() != b:
                    raise ValueError("%s: the lowered graph differs" % role)
                if e.get("synthetic") and not claim:
                    lowered[(FORM_FILE[e["lowering"]], e["name"], role)] = b
                ngraphs += 1
            nsets += 1
            v = check_set_certificate(root, e, roles, write)
            ncerts += 1
            if claim is None and not v["acc64"]:
                raise ValueError("a lowered graph does not load: the verdict fails acc64")
            if claim is not None and v[claim]:
                raise ValueError("its verdict meets %s, the profile it must fail" % claim)
        except Refuse as ex:
            failures += 1
            print("FAIL %s %s (%s): a lowered graph is refused at load: %s" % (what, e.get("name"), e.get("capability"), ex.code))
        except Exception as ex:
            failures += 1
            print("FAIL %s %s (%s): %s" % (what, e.get("name"), e.get("capability"), ex))
    print("re-lowered %d sets, %d graphs; recomputed %d certificates byte for byte" % (nsets, ngraphs, ncerts))
    failures += check_lower_vectors(root, lowered)
    failures += check_external(root, external or {}, write)
    return failures


def graph_meets(gj, p):
    """Section 11, `role_profiles`, read from a certificate's graph: its inputs' and
    scalars' declared ranges meet profile p's tensor width, and every node fits p."""
    tw = WIDTH[p]["tensor"]
    ranges = [(dec(r["lo"]), dec(r["hi"])) for r in gj["inputs"] + gj["scalars"]]
    return all(bits(r) <= tw for r in ranges) and all(p in nd["fits"] for nd in gj["nodes"])


def check_external(root, external, write):
    """The sets of external.json, whose certificates are committed although the sets
    are not. With the set's files (`external` maps an entry's name to the directory
    that holds them), the certificate is recomputed and compared. Without them, only
    what needs no set: the certificate is canonical, names its set, its capability
    and its lowering, its verdict meets the profile the entry gives, and each graph
    the entry's `role_profiles` names meets its own profile (section 11)."""
    failures = n_full = n_part = 0
    for e in json.load(open(os.path.join(root, "certificates", "external.json")))["sets"]:
        try:
            want = open(cert_path(root, e), "rb").read()
            d = e["name"] in external
            if d:
                files = {m: open(os.path.join(external[e["name"]], f), "rb").read() for m, f in e["members"].items()}
                if set_digest(files) != e["set_sha256"]:
                    raise ValueError("the set's digest differs")
                try:
                    roles = lower_set(e, files)
                except LowerRefuse:
                    raise ValueError("the lowering refused the set")
                check_set_certificate(root, e, roles, write)
            c = json.loads(want)
            if canonical(c) != want:
                raise ValueError("the certificate is not canonical")
            if (c["set_sha256"], c["capability"], c["lowering"]) != (e["set_sha256"], e["capability"], e["lowering"]):
                raise ValueError("the certificate names another set")
            if c["verdict"][e["profile"]] is not True:
                raise ValueError("its verdict does not meet %s" % e["profile"])
            for role, rp in sorted(e.get("role_profiles", {}).items()):
                gj = [g for g in c["graphs"] if g["name"] == role]
                if len(gj) != 1 or not graph_meets(gj[0], rp):
                    raise ValueError("its graph %s does not meet %s" % (role, rp))
            if d:
                n_full += 1
            else:
                n_part += 1
        except Refuse as ex:
            failures += 1
            print("FAIL external set %s: a lowered graph is refused at load: %s" % (e["name"], ex.code))
        except Exception as ex:
            failures += 1
            print("FAIL external set %s: %s" % (e["name"], ex))
    print("external sets: %d certificates recomputed byte for byte, %d checked without their set" % (n_full, n_part))
    return failures


FORM_FILE = {
    "int_synth/1": "lower_int_synth", "whm1/1": "lower_whm1", "wam1/1": "lower_wam1",
    "wmb1/1": "lower_wmb1", "world/1": "lower_world", "jpegai-int/1": "lower_jpegai",
}


def check_lower_vectors(root, lowered):
    """Each vector of graphs/lower_<form>.jsonl has the id <file>/<set>/<role>/...: its
    graph must be that role of that synthetic set as re-lowered here, and every role
    of every synthetic set must have one. `lowered` maps (file, set, role) to bytes."""
    failures = 0
    seen = set()
    d = os.path.join(root, "graphs")
    for name in sorted(os.listdir(d)):
        if not (name.startswith("lower_") and name.endswith(".jsonl")):
            continue
        f = name[:-len(".jsonl")]
        for k, line in enumerate(open(os.path.join(d, name), "rb").read().decode("ascii").split("\n")[:-1]):
            v = json.loads(line)
            parts = v["id"].split("/")
            key = tuple(parts[:3])
            if len(parts) < 4 or parts[0] != f or key not in lowered:
                failures += 1
                print("FAIL graphs/%s line %d: %s names no lowered role" % (name, k + 1, v["id"]))
            elif bytes.fromhex(v["graph"]) != lowered[key]:
                failures += 1
                print("FAIL graphs/%s line %d: %s is not the role as lowered" % (name, k + 1, v["id"]))
            seen.add(key)
    for key in sorted(set(lowered) - seen):
        failures += 1
        print("FAIL %s: no lowering vector for this role" % "/".join(key))
    print("matched %d lowered roles to their lowering vectors" % len(seen & set(lowered)))
    return failures


def certify_main(args):
    """`certify`: one set, or WHG1 graphs on their own, to a certificate (section 8.5),
    with the exit status of the bounds tool: 0 the verdict meets the profile, and each
    graph a --role-profile names meets its own (section 11, `role_profiles`); 1 one of
    them does not (the certificate is still written); 2 a graph is malformed or
    refused at load; 3 `--check` finds a difference; 4 the lowering refuses the set;
    5 usage, a --role-profile role among them that names no graph."""
    opts = {"--capability": None, "--lowering": None, "--c": None, "--profile": "acc64", "--out": None, "--check": None}
    members, graphs, role_profiles = [], [], {}
    k = 0
    while k < len(args):
        a = args[k]
        if k + 1 >= len(args) or (a not in opts and a not in ("--member", "--whg1", "--role-profile")):
            print(CERTIFY_USAGE, file=sys.stderr)
            return 5
        v = args[k + 1]
        if a in ("--member", "--whg1", "--role-profile"):
            if "=" not in v:
                print(CERTIFY_USAGE, file=sys.stderr)
                return 5
            name, val = v.split("=", 1)
            if a == "--role-profile":
                if name in role_profiles or val not in PROFILES:
                    print(CERTIFY_USAGE, file=sys.stderr)
                    return 5
                role_profiles[name] = val
            else:
                (members if a == "--member" else graphs).append((name, val))
        else:
            opts[a] = v
        k += 2
    if bool(members) == bool(graphs) or opts["--profile"] not in PROFILES or (members and not (opts["--capability"] and opts["--lowering"])):
        print(CERTIFY_USAGE, file=sys.stderr)
        return 5
    try:
        if graphs:
            roles = [(name, open(path, "rb").read()) for name, path in graphs]
            capability, lowering, digest = opts["--capability"], "whg1", None
        else:
            files = {name: open(path, "rb").read() for name, path in members}
            capability, lowering, digest = opts["--capability"], opts["--lowering"], set_digest(files)
            entry = {"capability": capability, "lowering": lowering, "c": opts["--c"], "set_sha256": digest}
            try:
                roles = lower_set(entry, files)
            except (LowerRefuse, KeyError, ValueError) as ex:
                print("the lowering refused the set%s" % (": %s" % ex if str(ex) else ""), file=sys.stderr)
                return 4
    except OSError as ex:
        print(ex, file=sys.stderr)
        return 5
    unknown = sorted(set(role_profiles) - {name for name, _ in roles})
    if unknown:
        print("--role-profile names no graph of the set: %s" % ", ".join(unknown), file=sys.stderr)
        print(CERTIFY_USAGE, file=sys.stderr)
        return 5
    try:
        got, v, role_fails = certificate(capability, lowering, digest, roles, role_profiles)
    except Refuse as ex:
        print("a graph is refused at load: %s" % ex.code, file=sys.stderr)
        return 2
    if opts["--out"]:
        open(opts["--out"], "wb").write(got)
    else:
        sys.stdout.write(got.decode("ascii"))
        sys.stdout.flush()
    p = opts["--profile"]
    if opts["--check"]:
        want = open(opts["--check"], "rb").read()
        if got != want:
            print("--check: %s" % first_difference(got, want), file=sys.stderr)
            return 3
    status = 0
    if not v[p]:
        print("the verdict does not meet %s: %s" % (p, cw(v[p + "_first_failure"])), file=sys.stderr)
        status = 1
    for name, rp in role_profiles.items():
        if name in role_fails:
            print("graph %s does not meet %s: %s" % (name, rp, cw(role_fails[name])), file=sys.stderr)
            status = 1
        else:
            print("graph %s meets %s" % (name, rp), file=sys.stderr)
    if status:
        return status
    print("the verdict meets %s: %s" % (p, cw(v)), file=sys.stderr)
    return 0


CERTIFY_USAGE = """usage:
  wai_ops_check.py certify --capability CAP --lowering ID [--c N] --member NAME=PATH ...
                           [--profile acc32|acc64] [--role-profile ROLE=acc32|acc64 ...]
                           [--out FILE] [--check FILE]
  wai_ops_check.py certify --whg1 ROLE=PATH ... [--capability CAP]
                           [--profile acc32|acc64] [--role-profile ROLE=acc32|acc64 ...]
                           [--out FILE] [--check FILE]"""


def main(argv):
    if len(argv) >= 2 and argv[1] == "certify":
        return certify_main(argv[2:])
    if len(argv) < 2 or len(argv) % 2:
        print(__doc__)
        return 2
    write, external = None, {}
    for k in range(2, len(argv), 2):
        a, v = argv[k], argv[k + 1]
        if a == "--write":
            write = v
        elif a == "--external" and "=" in v:
            name, d = v.split("=", 1)
            external[name] = d
        else:
            print(__doc__)
            return 2
    failures = check(argv[1], write, external)
    if failures:
        print("%d failures" % failures)
        return 1
    print("ops-conformance: every vector file, every lowered graph and every certificate rewritten byte for byte")
    return 0


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