#!/usr/bin/env python3
"""Second implementation of extensions/splat-exact.md (`wai.splat.exact`), over the corpus
`live-conformance/splat/`.

Written from the extension's text. It shares no code with wai-rs: it reads each glTF asset as §3
has it read, parses each `WAI1` envelope, decodes each `WIS1` payload with the second
implementations the repository already holds for the steps that come before rendering (the
payload parser and rANS decoder of int-payload-conformance/int_payload_verify.py, and the
`whm1/1` lowering and interpreter of tools/wai_ops_check.py), renders every view with its own
integer renderer written from §5, verifies every record with WAI's strict Ed25519 rule
(claim-conformance/_ed25519.py) and every link hash with BLAKE3 written from its specification
(claim-conformance/_blake3.py), and reads §10's codes from the extension.

    python3 -I splat_exact_verify.py <live-conformance/splat> [--png DIR]

It gives every case of the corpus the outcome, the reason and the place the corpus states, and
checks that every reason of §10 refuses a case. It then falsifies itself three ways: it mutates
copies of the corpus and requires every copy to fail; it drops each of its own checks in turn and
requires the corpus to catch the weakened verifier; and it runs its renderer with each of a set of
changed arithmetic rules and requires the corpus's frames to catch every one. Exit status 0 only if
all of that holds. --png writes each committed frame as a PNG. Python 3.8 or later; standard
library only.
"""

import copy
import hashlib
import importlib.util
import json
import math
import os
import re
import struct
import sys
import unicodedata
import zlib
from fractions import Fraction

HERE = os.path.dirname(os.path.abspath(__file__))
WAI = os.path.join(HERE, "..")
sys.path.insert(0, os.path.join(WAI, "claim-conformance"))
import _ed25519 as ED  # noqa: E402
from _blake3 import blake3  # noqa: E402


def _load(name, path):
    spec = importlib.util.spec_from_file_location(name, path)
    m = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(m)
    return m


IPV = _load("int_payload_verify", os.path.join(WAI, "int-payload-conformance", "int_payload_verify.py"))
OPS = _load("wai_ops_check", os.path.join(WAI, "tools", "wai_ops_check.py"))

EXTENSION = os.path.join(WAI, "extensions", "splat-exact.md")

SIGN_DOMAIN = b"wai:splat-exact\x01"            # §2
LINK_DOMAIN = b"wai:splat-exact-link\x01"       # §2
OBJECT_DOMAIN = b"moq-jwp:object\x01"           # jwp-receipts §2.1
SYN_SIGN = b"wai:synthesis\x01"                 # synthesis §2
SYN_LINK = b"wai:synthesis-link\x01"
RENDERER = "wai-splat-render:i1"
BASE = "KHR_gaussian_splatting"
EXT = "EXT_gaussian_splatting_compression_wai"
CAPABILITY = "wai.splat.int_codec"
ATTRIBUTES = [("POSITION", "VEC3"), ("KHR_gaussian_splatting:ROTATION", "VEC4"), ("KHR_gaussian_splatting:SCALE", "VEC3"),
              ("KHR_gaussian_splatting:OPACITY", "SCALAR"), ("KHR_gaussian_splatting:SH_DEGREE_0_COEF_0", "VEC3")]
FLOAT = 5126
PLANES = 14
MAX_INT = (1 << 53) - 1
MAX_COORD = 1 << 28
CLASSES = ["captured", "reconstructed", "assistive", "synthesized"]
STEP_CLASS = {"capture": "captured", "convert": "reconstructed", "assist": "assistive", "generate": "synthesized"}

# §5's constants.
NEAR, DILATION, MAX_RADIUS = 13107, 19661, 1023
CUTOFF, E_MAX = 9 << 16, 3 << 24
ALPHA_MAX, ALPHA_MIN, T_MIN = 64880, 257, 7
TILE = 16

WHITE_SPACE = set("\t\n\x0b\x0c\r \x85\xa0            "
                  "

  　")

# The checks this verifier makes, by name; `DROPPED` leaves one out (the self-falsification).
DROPPED = set()
# Changed arithmetic for the renderer's falsification; empty for the verification itself.
MUTANT = set()


def on(name):
    return name not in DROPPED


class Refuse(Exception):
    def __init__(self, reason, at=None):
        super().__init__(reason)
        self.reason, self.at = reason, at


class Bad(Exception):
    pass


def need(ok):
    if not ok:
        raise Bad()


# ---------------------------------------------------------------- §2 records

def strict_json(text):
    def pairs(ps):
        d = {}
        for k, v in ps:
            if k in d:
                raise Bad()
            d[k] = v
        return d

    def no_float(s):
        raise Bad()

    def integer(s):
        if s.startswith("-"):
            raise Bad()
        return int(s)

    try:
        v = json.loads(text, object_pairs_hook=pairs, parse_float=no_float, parse_int=integer, parse_constant=no_float)
    except (ValueError, RecursionError):
        raise Bad()

    def walk(x, depth):
        if isinstance(x, bool):
            raise Bad()
        if isinstance(x, (dict, list)):
            need(depth <= 32)
            for y in (x.values() if isinstance(x, dict) else x):
                walk(y, depth + 1)
        elif isinstance(x, str):
            need(not any(0xD800 <= ord(c) <= 0xDFFF for c in x))
        elif isinstance(x, int):
            need(0 <= x <= MAX_INT)
    walk(v, 1)
    return v


def jcs(v):
    return json.dumps(v, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")


def hexb(v, n):
    need(isinstance(v, str) and re.fullmatch(r"[0-9a-f]*", v) is not None and len(v) == 2 * n)
    return bytes.fromhex(v)


def integer(v):
    need(isinstance(v, int) and not isinstance(v, bool) and 0 <= v <= MAX_INT)
    return v


def term(v):
    need(isinstance(v, str) and 0 < len(v.encode("utf-8")) <= 2048
         and not any(c in WHITE_SPACE or unicodedata.category(c) == "Cc" for c in v))
    return v


def date_time(s):
    m = re.fullmatch(r"([0-9]{4})-([0-9]{2})-([0-9]{2})[Tt]([0-9]{2}):([0-9]{2}):([0-9]{2})(\.[0-9]+)?([Zz]|[+-]([0-9]{2}):([0-9]{2}))",
                     s if isinstance(s, str) else "")
    need(m is not None)
    y, mo, d, h, mi, sec = (int(m.group(i)) for i in range(1, 7))
    leap = y % 4 == 0 and (y % 100 != 0 or y % 400 == 0)
    days = [31, 29 if leap else 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
    need(1 <= mo <= 12 and 1 <= d <= days[mo - 1] and h <= 23 and mi <= 59 and sec <= 60)
    if m.group(9) is not None:
        need(int(m.group(9)) <= 23 and int(m.group(10)) <= 59)


def members(v, names):
    need(isinstance(v, dict) and set(v) == set(names))
    return v


def sint(v):
    """§2: `0`, or an optional `-` and a digit 1-9 followed by digits; magnitude at most 2^28."""
    need(isinstance(v, str) and re.fullmatch(r"0|-?[1-9][0-9]*", v) is not None)
    n = int(v)
    need(abs(n) <= MAX_COORD)
    return n


def read_pose(v):
    members(v, ["eye", "target", "up", "focal", "width", "height"])
    vec = []
    for k in ("eye", "target", "up"):
        need(isinstance(v[k], list) and len(v[k]) == 3)
        vec.append([sint(x) for x in v[k]])
    f, w, h = integer(v["focal"]), integer(v["width"]), integer(v["height"])
    need(1 <= f <= 4096 and 1 <= w <= 1024 and 1 <= h <= 1024 and w * h <= 1 << 20)
    return {"eye": vec[0], "target": vec[1], "up": vec[2], "focal": f, "width": w, "height": h}


def pose_key(p):
    return json.dumps(p, sort_keys=True)


def signed(v, domain):
    body = dict(v)
    del body["sig"]
    return domain + jcs(body)


def read_record(text, kind, names, fields):
    """§2: strict JSON, exactly the kind's members, then the kind's fields, then the signature."""
    try:
        v = strict_json(text)
        members(v, names)
        need(v["kind"] == kind)
        signer = hexb(v["signer_pubkey"], 32)
        sig = hexb(v["sig"], 64)
        date_time(v["issued_at"])
        r = fields(v)
    except (Bad, KeyError, TypeError, AttributeError):
        raise Refuse("record_malformed")
    if on("signature"):
        bad = ED.verify(signer, signed(v, SIGN_DOMAIN), sig)
        if bad:
            raise Refuse(bad)
    r["link"] = blake3(LINK_DOMAIN + jcs(v))
    return r


COMMITMENT = ["kind", "object", "set", "attributes", "renderer", "class", "origin", "views", "issued_at", "signer_pubkey", "sig"]
STATEMENT = ["kind", "session", "commitment", "object", "pose", "canonical", "frame", "renderer", "class", "at", "issued_at",
             "signer_pubkey", "sig"]


def commitment_fields(v):
    need(isinstance(v["views"], list) and 1 <= len(v["views"]) <= 64)
    views, seen = [], set()
    for e in v["views"]:
        members(e, ["pose", "frame"])
        p = read_pose(e["pose"])
        need(pose_key(p) not in seen)
        seen.add(pose_key(p))
        views.append({"pose": p, "frame": hexb(e["frame"], 32)})
    need(v["class"] in CLASSES)
    return {"object": hexb(v["object"], 32), "set": hexb(v["set"], 32), "attributes": hexb(v["attributes"], 32),
            "renderer": term(v["renderer"]), "class": v["class"],
            "origin": None if v["origin"] is None else hexb(v["origin"], 32), "views": views}


def statement_fields(v):
    need(v["class"] in CLASSES)
    return {"session": hexb(v["session"], 32), "commitment": hexb(v["commitment"], 32), "object": hexb(v["object"], 32),
            "pose": read_pose(v["pose"]), "canonical": None if v["canonical"] is None else integer(v["canonical"]),
            "frame": hexb(v["frame"], 32), "renderer": term(v["renderer"]), "class": v["class"], "at": integer(v["at"])}


# ---------------------------------------------------------------- §5 the renderer

def rdiv(a, b):
    """a / b to the nearest integer, a tie toward +infinity (b > 0)."""
    if "rdiv_floor" in MUTANT:
        return a // b
    return (2 * a + b) // (2 * b)


def rhu(p, s):
    """R_half_up(p, s) of integer-operators §2."""
    return p if s == 0 else (p + (1 << (s - 1))) >> s


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


def exp_neg_q16(x):
    """exp_neg_q16 of integer-operators §2, x <= 0."""
    if "exp_float" in MUTANT:
        return int(math.exp(x / 65536.0) * 65536.0)
    u = -x
    n = u // 45426
    r = x + 45426 * n
    t = r + 88670
    q = (t * t) >> 16
    y = ((23495 * q) >> 16) + 22544
    return 0 if n >= 63 else y >> n


def q16(v):
    """§4.2: v * 2^16 rounded half up, exactly, from a binary32 value."""
    return math.floor(Fraction(v) * 65536 + Fraction(1, 2))


def dot(a, b):
    return a[0] * b[0] + a[1] * b[1] + a[2] * b[2]


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


def gaussians(dec):
    """§5.2: the renderable form, or a refusal."""
    if on("layout") and dec["planes"] != PLANES:
        raise Refuse("layout_mismatch")
    lohi = []
    for i, (lo, hi) in enumerate(dec["ranges"]):
        lq, hq = q16(lo), q16(hi)
        if on("ranges") and (abs(lq) > MAX_COORD or abs(hq) > MAX_COORD or (3 <= i < 6 and lq < 0)):
            raise Refuse("range_unrenderable")
        lohi.append((lq, hq))
    cells = dec["s"] * dec["s"]
    b = dec["bytes"]
    out = []
    for k in range(dec["n"]):
        a = [lohi[p][0] + ((lohi[p][1] - lohi[p][0]) * b[p * cells + k] + 127) // 255 for p in range(PLANES)]
        out.append(a)
    return out


def camera(pose):
    """§5.3."""
    fw = [pose["target"][i] - pose["eye"][i] for i in range(3)]
    nf = math.isqrt(dot(fw, fw))
    if nf == 0:
        raise Refuse("pose_degenerate")
    f = [rdiv(x << 14, nf) for x in fw]
    rr = cross(fw, pose["up"])
    nr = math.isqrt(dot(rr, rr))
    if nr == 0:
        raise Refuse("pose_degenerate")
    r = [rdiv(x << 14, nr) for x in rr]
    u = [rhu(x, 14) for x in cross(r, f)]
    return r, u, f


def rotation(w, x, y, z):
    n = w * w + x * x + y * y + z * z
    if n == 0:
        return [[1 << 16, 0, 0], [0, 1 << 16, 0], [0, 0, 1 << 16]]

    def e(num):
        return rdiv(num << 16, n)
    return [[e(w * w + x * x - y * y - z * z), e(2 * (x * y - w * z)), e(2 * (x * z + w * y))],
            [e(2 * (x * y + w * z)), e(w * w - x * x + y * y - z * z), e(2 * (y * z - w * x))],
            [e(2 * (x * z - w * y)), e(2 * (y * z + w * x)), e(w * w - x * x - y * y + z * z)]]


def project(a, pose, cam):
    """§5.4-§5.5: the Gaussian's key and what compositing reads, or None where it is culled."""
    r, u, f = cam
    w, h, fl = pose["width"], pose["height"], pose["focal"]
    d = [a[i] - pose["eye"][i] for i in range(3)]
    xc, yc, zc = rhu(dot(r, d), 14), rhu(dot(u, d), 14), rhu(dot(f, d), 14)
    if zc < (0 if "near_zero" in MUTANT else NEAR) or zc <= 0:
        return None
    key = zc if "sort_by_z" in MUTANT else dot(d, d)
    mx = 128 * w + rdiv(256 * fl * xc, zc)
    my = 128 * h - rdiv(256 * fl * yc, zc)
    lx, ly = rdiv(13 * w * zc, 20 * fl), rdiv(13 * h * zc, 20 * fl)
    xcl, ycl = min(max(xc, -lx), lx), min(max(yc, -ly), ly)
    j00 = rdiv(fl << 32, zc)
    j02 = rdiv((-fl * xcl) << 32, zc * zc)
    j11 = -j00
    j12 = rdiv((fl * ycl) << 32, zc * zc)
    t = [[rhu(j00 * r[k] + j02 * f[k], 14) for k in range(3)], [rhu(j11 * u[k] + j12 * f[k], 14) for k in range(3)]]
    rot = rotation(a[6], a[7], a[8], a[9])
    s = a[3:6]
    m = [[rhu(rot[i][j] * s[j], 16) for j in range(3)] for i in range(3)]
    s3 = [[sum(m[i][k] * m[j][k] for k in range(3)) for j in range(3)] for i in range(3)]
    uu = [[sum(t[i][j] * s3[j][k] for j in range(3)) for k in range(3)] for i in range(2)]

    def vv(i, l):
        return rhu(sum(uu[i][k] * t[l][k] for k in range(3)), 48)
    dil = 0 if "no_dilation" in MUTANT else DILATION
    ca, cb, cc = vv(0, 0) + dil, vv(0, 1), vv(1, 1) + dil
    if ca <= 0 or cc <= 0:
        return None
    if "rect_floor" in MUTANT:
        rx, ry = math.isqrt(9 * ca) // 256, math.isqrt(9 * cc) // 256
    else:
        rx, ry = (isqrt_ceil(9 * ca) + 255) // 256, (isqrt_ceil(9 * cc) + 255) // 256
    if rx > MAX_RADIUS or ry > MAX_RADIUS:
        return None
    x0 = max(0, -((-(mx - 256 * rx - 128)) // 256))
    x1 = min(w - 1, (mx + 256 * rx - 128) // 256)
    y0 = max(0, -((-(my - 256 * ry - 128)) // 256))
    y1 = min(h - 1, (my + 256 * ry - 128) // 256)
    if x0 > x1 or y0 > y1:
        return None
    det = ca * cc - cb * cb
    if det <= 0:
        return None
    k1 = rdiv(1 << 40, math.isqrt(ca << 32))
    sq = math.isqrt((ca * det) << 16)
    k2 = rdiv(ca << 32, sq)
    k3 = rdiv((-cb) << 32, sq)
    if max(abs(k1), abs(k2), abs(k3)) > (1 << 31) - 1:
        return None

    def c16(x):
        return min(max(x, 0), 1 << 16)
    return key, (mx, my, x0, x1, y0, y1, k1, k2, k3, c16(a[13]), (c16(a[10]), c16(a[11]), c16(a[12])))


def alpha_at(g, i, j):
    """§5.6: the alpha a Gaussian gives pixel (i, j), or None where it adds nothing."""
    mx, my, _, _, _, _, k1, k2, k3, o, _ = g
    dx = 256 * i + 128 - mx
    dy = 256 * j + 128 - my
    e1 = k1 * dx
    e2 = k3 * dx + k2 * dy
    if abs(e1) > E_MAX or abs(e2) > E_MAX:
        return None
    a1, a2 = (e1 + 2048) >> 12, (e2 + 2048) >> 12
    d2 = (a1 * a1 + a2 * a2 + 128) >> 8
    if d2 > CUTOFF and "no_cutoff" not in MUTANT:
        return None
    w = exp_neg_q16(-(d2 >> 1))
    alpha = (o * w) >> 16
    if "no_alpha_max" not in MUTANT:
        alpha = min(alpha, ALPHA_MAX)
    return alpha if alpha >= ALPHA_MIN else None


def render(gs, pose):
    """§5: the frame, RGB8, rows top to bottom."""
    cam = camera(pose)
    kept = []
    for idx, a in enumerate(gs):
        p = project(a, pose, cam)
        if p is not None:
            kept.append((p[0], -idx if "ties_desc" in MUTANT else idx, p[1]))
    kept.sort(key=lambda e: (e[0], e[1]))
    w, h = pose["width"], pose["height"]
    T = [1 << 16] * (w * h)
    acc = [[0, 0, 0] for _ in range(w * h)]
    done = [False] * (w * h)
    for _, _, g in kept:
        _, _, x0, x1, y0, y1 = g[:6]
        col = g[10]
        for j in range(y0, y1 + 1):
            for i in range(x0, x1 + 1):
                p = j * w + i
                if done[p]:
                    continue
                alpha = alpha_at(g, i, j)
                if alpha is None:
                    continue
                if "contrib_floor" in MUTANT:
                    contrib = (alpha * T[p]) >> 16
                else:
                    contrib = (alpha * T[p] + (1 << 15)) >> 16
                tn = T[p] - contrib
                if tn < T_MIN and "stop_after" not in MUTANT:
                    done[p] = True
                    continue
                for ch in range(3):
                    acc[p][ch] += (col[ch] * contrib + (1 << 15)) >> 16
                T[p] = tn
                if tn < T_MIN:
                    done[p] = True
    out = bytearray()
    for p in range(w * h):
        for ch in range(3):
            out.append(min(255, (acc[p][ch] * 255 + (1 << 15)) >> 16))
    return bytes(out)


RENDERS = {}


def frame_digest(gs, gkey, pose):
    """The SHA-256 of a frame, remembered per splat set, pose and renderer rules."""
    k = (gkey, pose_key(pose), tuple(sorted(MUTANT)))
    if k not in RENDERS:
        RENDERS[k] = hashlib.sha256(render(gs, pose)).digest()
    return RENDERS[k]


# ---------------------------------------------------------------- §3 carriage and the decode

def buffer_view(g, bufs, index):
    need(isinstance(index, int) and not isinstance(index, bool) and index >= 0)
    bv = g["bufferViews"][index]
    b = g["buffers"][bv["buffer"]]
    data = bufs[b["uri"]]
    need(b["byteLength"] == len(data))
    off, n = bv.get("byteOffset", 0), bv["byteLength"]
    need(isinstance(off, int) and isinstance(n, int) and off >= 0 and n >= 0 and off + n <= len(data))
    return data[off:off + n]


def read_asset(g, bufs):
    """§3.1-§3.4, the rules that need no decode."""
    try:
        need(g["asset"]["version"] == "2.0")
        found = [p for m in g["meshes"] for p in m["primitives"] if BASE in p.get("extensions", {})]
        need(len(found) == 1)
        p = found[0]
        base = p["extensions"][BASE]
        if on("extensions_required"):
            for name in (BASE, EXT):
                need(name in g.get("extensionsUsed", []) and name in g.get("extensionsRequired", []))
        if on("primitive"):
            need(p.get("mode") == 0 and base.get("kernel") == "ellipse" and base.get("colorSpace") == "srgb_rec709_display")
            need(base.get("projection", "perspective") == "perspective" and base.get("sortingMethod", "cameraDistance") == "cameraDistance")
        attrs = p["attributes"]
        need(isinstance(attrs, dict) and len(attrs) == len(ATTRIBUTES))
        accessors = []
        for name, ty in ATTRIBUTES:
            a = g["accessors"][attrs[name]]
            need("bufferView" not in a and a.get("componentType") == FLOAT and a.get("type") == ty)
            integer(a["count"])
            accessors.append(a)
        ext = base["extensions"][EXT]
        env = buffer_view(g, bufs, ext["bufferView"])
        if on("capability") and ext["capability"] != CAPABILITY:
            if isinstance(ext["capability"], str):
                raise Refuse("capability_mismatch")
            raise Bad()
        obj = hexb(ext["object"], 32)
        set_sha = hexb(ext["set"]["sha256"], 32)
        set_inline = buffer_view(g, bufs, ext["set"]["bufferView"]) if "bufferView" in ext["set"] else None
        ctext = None
        if "commitment" in ext:
            ctext = buffer_view(g, bufs, ext["commitment"]).decode("utf-8")
    except (Bad, KeyError, TypeError, IndexError, AttributeError, UnicodeDecodeError):
        raise Refuse("asset_malformed")
    return {"envelope": env, "object": obj, "set": set_sha, "set_inline": set_inline, "commitment": ctext, "accessors": accessors}


def read_envelope(b):
    """SPEC §2: magic, manifest, payload, and nothing after it."""
    try:
        need(b[:4] == b"WAI1")
        ml = struct.unpack_from("<I", b, 4)[0]
        manifest = json.loads(b[8:8 + ml].decode("utf-8"))
        pl = struct.unpack_from("<I", b, 8 + ml)[0]
        need(12 + ml + pl == len(b) or not on("envelope_length"))
        return manifest, b[12 + ml:]
    except (Bad, ValueError, struct.error, UnicodeDecodeError):
        raise Refuse("asset_malformed")


DECODES = {}


def decode_wis1(payload, set_bytes):
    """integer-payloads §6 through the second implementations named above."""
    key = (hashlib.sha256(payload).digest(), hashlib.sha256(set_bytes).digest())
    if key in DECODES:
        return DECODES[key]
    try:
        v = IPV.parse_wis1(payload)
        need(set_bytes[:4] == b"WHM1" and struct.unpack_from("<I", set_bytes, 4)[0] == v["afrac"])
        g = OPS.load(OPS.whm1_lower(set_bytes, True))
        r = IPV.Rans(v["rans"])
        lat = []
        for ch in range(v["c"]):
            for _ in range(v["sh"] * v["sw"]):
                s = r.decode(v["cdfs"][ch], v["offsets"][ch])
                need(not r.exhausted)
                lat.append(s * (1 << v["afrac"]) + v["median_q"][ch])
        need(r.i == len(r.w) and r.x == 1 << 31)
        out = OPS.run(g, [OPS.T(v["c"], v["sh"], v["sw"], lat)], [], {})[0]
        need((out.c, out.h, out.w) == (v["attr"], v["s"], v["s"]))
        res = {"bytes": bytes(out.d), "planes": v["attr"], "s": v["s"], "n": v["n_gaussians"], "ranges": v["ranges"]}
    except Exception as e:  # noqa: BLE001 - every refusal of the decode is one code here
        if isinstance(e, (KeyboardInterrupt, SystemExit)):
            raise
        res = None
    DECODES[key] = res
    return res


def decode(c, held):
    """§3.3: the envelope, its pin against the set, the decode."""
    manifest, payload = read_envelope(c["envelope"])
    try:
        cap = manifest["model_requirement"]["capability"]
    except (KeyError, TypeError):
        raise Refuse("asset_malformed")
    if on("capability") and cap != CAPABILITY:
        raise Refuse("capability_mismatch")
    if on("object") and blake3(OBJECT_DOMAIN + c["envelope"]) != c["object"]:
        raise Refuse("object_mismatch")
    pin = (manifest.get("conditioning") or {}).get("prior") or {}
    if on("pin") and pin.get("sha256") != c["set"].hex():
        raise Refuse("set_mismatch")
    if c["set_inline"] is not None:
        set_bytes = c["set_inline"]
    else:
        found = [b for b in held if hashlib.sha256(b).digest() == c["set"]]
        if not found:
            raise Refuse("set_not_held")
        set_bytes = found[0]
    if on("set_digest") and hashlib.sha256(set_bytes).digest() != c["set"]:
        raise Refuse("set_mismatch")
    d = decode_wis1(payload, set_bytes)
    if d is None:
        raise Refuse("decode_refused")
    return d


def check_accessors(c, d):
    """§3.4: every count is the Gaussian count; POSITION's min and max are the dequantized extremes."""
    if not on("accessors"):
        return
    try:
        need(d["planes"] >= 3)
        for a in c["accessors"]:
            need(a["count"] == d["n"])
        cells = d["s"] * d["s"]
        pos = c["accessors"][0]
        need(len(pos["min"]) == 3 and len(pos["max"]) == 3)
        for axis in range(3):
            lq, hq = q16(d["ranges"][axis][0]), q16(d["ranges"][axis][1])
            qs = [lq + ((hq - lq) * d["bytes"][axis * cells + k] + 127) // 255 for k in range(d["n"])]
            need(pos["min"][axis] == min(qs) / 65536 and pos["max"][axis] == max(qs) / 65536)
    except (Bad, KeyError, TypeError, IndexError):
        raise Refuse("asset_malformed")


# ---------------------------------------------------------------- §8 the origin claim

def check_origin(text, link, obj, cls):
    """§9.2 step 6: the origin claim the commitment names (synthesis §2, §3.1, §3.2)."""
    try:
        v = strict_json(text)
        need(v.get("kind") == "wai-origin")
        need(ED.verifies(hexb(v["signer_pubkey"], 32), signed(v, SYN_SIGN), hexb(v["sig"], 64)))
        need(blake3(SYN_LINK + jcs(v)) == link)
        need(v["subject"] == {"type": "object", "content_hash": obj.hex()})
        need(v["class"] == cls)
        composed = max([CLASSES.index(STEP_CLASS[v["step"]])] + [CLASSES.index(i["class"]) for i in v["inputs"]])
        need(CLASSES[composed] == v["class"])
    except (Bad, KeyError, TypeError, ValueError, AttributeError):
        raise Refuse("origin_mismatch", "origin")


# ---------------------------------------------------------------- §9 verification

def verify_case(case, files, held_sets):
    asset = case["asset"]
    g = json.loads(files[asset])
    d0 = os.path.dirname(asset)
    bufs = {}
    for b in g.get("buffers", []):
        path = os.path.join(d0, b["uri"]) if d0 else b["uri"]
        if path in files:
            bufs[b["uri"]] = files[path]
    held = [held_sets[p] for p in case["held"]]
    try:
        c = read_asset(g, bufs)
        d = decode(c, held)
        check_accessors(c, d)
    except Refuse as r:
        raise Refuse(r.reason, "asset")
    attributes = hashlib.sha256(d["bytes"]).digest()
    report = {"object": c["object"].hex(), "attributes": attributes.hex(), "n_gaussians": d["n"], "frames": [], "statements": []}
    if c["commitment"] is None:
        return report
    try:
        cm = read_record(c["commitment"], "wai-splat-commitment", COMMITMENT, commitment_fields)
        if on("commitment_object") and cm["object"] != c["object"]:
            raise Refuse("object_mismatch")
        if on("commitment_set") and cm["set"] != c["set"]:
            raise Refuse("set_mismatch")
        if on("attributes") and cm["attributes"] != attributes:
            raise Refuse("attributes_mismatch")
        if on("renderer") and cm["renderer"] != RENDERER:
            raise Refuse("renderer_unknown")
        gs = gaussians(d)
    except Refuse as r:
        raise Refuse(r.reason, "commitment")
    gkey = attributes
    for i, v in enumerate(cm["views"]):
        try:
            got = frame_digest(gs, gkey, v["pose"])
        except Refuse as r:
            raise Refuse(r.reason, "view %d" % i)
        if on("frames") and got != v["frame"]:
            raise Refuse("frame_mismatch", "view %d" % i)
        report["frames"].append(v["frame"].hex())
    if cm["origin"] is not None and case["origin"] is not None and on("origin"):
        check_origin(case["origin"], cm["origin"], c["object"], cm["class"])
    for n, t in enumerate(case["statements"]):
        at = "statement %d" % n
        try:
            s = read_record(t, "wai-splat-view", STATEMENT, statement_fields)
            if on("statement_commitment") and s["commitment"] != cm["link"]:
                raise Refuse("commitment_mismatch")
            if on("statement_object") and s["object"] != c["object"]:
                raise Refuse("object_mismatch")
            if on("statement_renderer") and s["renderer"] != RENDERER:
                raise Refuse("renderer_unknown")
            if s["canonical"] is not None:
                i = s["canonical"]
                if on("canonical") and (i >= len(cm["views"]) or pose_key(cm["views"][i]["pose"]) != pose_key(s["pose"])):
                    raise Refuse("pose_not_canonical")
                if on("canonical_frame") and i < len(cm["views"]) and cm["views"][i]["frame"] != s["frame"]:
                    raise Refuse("frame_mismatch")
            elif on("own_frame") and frame_digest(gs, gkey, s["pose"]) != s["frame"]:
                raise Refuse("frame_mismatch")
            want = CLASSES[max(1, CLASSES.index(cm["class"]))]
            if on("class") and s["class"] != want:
                raise Refuse("class_mismatch")
        except Refuse as r:
            raise Refuse(r.reason, at)
        report["statements"].append({"canonical": s["canonical"], "class": s["class"]})
    return report


def outcome(case, files, held_sets):
    try:
        r = verify_case(case, files, held_sets)
        return dict(outcome="accepted", **r)
    except Refuse as r:
        return {"outcome": "refused", "reason": r.reason, "at": r.at}


def reasons_of_extension():
    text = open(EXTENSION, encoding="utf-8").read()
    start = text.index("## 10. Reasons")
    end = text.index("\n## 11.", start)
    return [m.group(1) for m in re.finditer(r"^\| `([a-z_]+)` \|", text[start:end], re.M)]


def run_all(files, held_sets):
    """Every case against its stated outcome: the names of the cases that differ."""
    index = json.loads(files["vectors.json"])
    bad = []
    for c in index["cases"]:
        try:
            got = outcome(c, files, held_sets)
        except (IndexError, KeyError, TypeError, ValueError, ZeroDivisionError):
            # A weakened verifier (step 4 below) can fail on an input a dropped check kept out.
            got = {"outcome": "error"}
        if got != c["expect"]:
            bad.append(c["name"])
    return bad


def png(path, w, h, rgb):
    raw = b"".join(b"\x00" + rgb[y * w * 3:(y + 1) * w * 3] for y in range(h))

    def chunk(t, d):
        return struct.pack(">I", len(d)) + t + d + struct.pack(">I", zlib.crc32(t + d) & 0xFFFFFFFF)
    with open(path, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0))
                + chunk(b"IDAT", zlib.compress(raw, 9)) + chunk(b"IEND", b""))


def main(argv):
    if len(argv) < 2:
        print(__doc__)
        return 2
    root = argv[1]
    files = {}
    for dp, _, names in os.walk(root):
        for n in names:
            p = os.path.join(dp, n)
            rel = os.path.relpath(p, root).replace(os.sep, "/")
            if rel != "README.md":
                files[rel] = open(p, "rb").read()
    index = json.loads(files["vectors.json"])
    held_sets = {}
    for c in index["cases"]:
        for p in c["held"]:
            held_sets[p] = open(os.path.join(WAI, p), "rb").read()
    fails = 0

    def check(ok, what):
        nonlocal fails
        print(("PASS " if ok else "FAIL ") + what)
        fails += 0 if ok else 1

    # 1. Every case.
    for c in index["cases"]:
        got = outcome(c, files, held_sets)
        check(got == c["expect"], "%s: %s%s" % (c["name"], got["outcome"],
                                                 "" if got["outcome"] == "accepted" else " %s at %s" % (got["reason"], got["at"])))
        if got != c["expect"]:
            print("  stated:", json.dumps(c["expect"])[:300])
            print("  got:   ", json.dumps(got)[:300])
    # 2. Every reason of §10 refuses a case, and the corpus gives no other.
    table = reasons_of_extension()
    given = {c["expect"].get("reason") for c in index["cases"]} - {None}
    check(set(table) == given, "every reason of §10 refuses a case (%d codes), and no case gives another" % len(table))
    if "--png" in argv:
        out = argv[argv.index("--png") + 1]
        os.makedirs(out, exist_ok=True)
        for c in index["cases"]:
            if c["expect"]["outcome"] != "accepted" or not c["expect"]["frames"]:
                continue
            g = json.loads(files[c["asset"]])
            bufs = {b["uri"]: files[os.path.join(os.path.dirname(c["asset"]), b["uri"])] for b in g["buffers"]}
            a = read_asset(g, bufs)
            d = decode(a, [held_sets[p] for p in c["held"]])
            cm = json.loads(a["commitment"])
            gs = gaussians(d)
            for i, v in enumerate(cm["views"]):
                pose = read_pose(v["pose"])
                png(os.path.join(out, "%s_%d.png" % (c["name"], i)), pose["width"], pose["height"], render(gs, pose))
    # 3. Mutated copies of the corpus fail.
    def mut_bytes(path, at, x=1):
        def f(fs):
            b = bytearray(fs[path])
            b[at % len(b)] ^= x
            fs[path] = bytes(b)
        return f

    def mut_index(fn):
        def f(fs):
            v = json.loads(fs["vectors.json"])
            fn(v)
            fs["vectors.json"] = json.dumps(v).encode()
        return f

    def mut_gltf(path, fn):
        def f(fs):
            g = json.loads(fs[path])
            fn(g)
            fs[path] = json.dumps(g).encode()
        return f

    def flip_frame(statements, i):
        t = statements[i]
        at = t.index('"frame":"') + len('"frame":"')
        statements[i] = t[:at] + ("1" if t[at] == "0" else "0") + t[at + 1:]

    def case_named(v, name):
        return next(c for c in v["cases"] if c["name"] == name)

    mutations = [
        ("a byte of the tiny scene's latent stream", mut_bytes("assets/tiny.bin", 900)),
        ("a byte of the cluster's envelope manifest", mut_bytes("assets/cluster.bin", 40)),
        ("a byte of the identity set inline in the cluster's buffer", mut_bytes("assets/cluster.bin", -300)),
        ("a digit of a frame in the cluster's commitment", mut_bytes("assets/cluster.commitment.json", 600)),
        ("the shell's POSITION min", mut_gltf("assets/shell.gltf", lambda g: g["accessors"][0]["min"].__setitem__(0, 0.0))),
        ("the cluster's kernel", mut_gltf("assets/cluster.gltf",
                                          lambda g: g["meshes"][0]["primitives"][0]["extensions"][BASE].__setitem__("kernel", "box"))),
        ("the shell case's held set dropped", mut_index(lambda v: case_named(v, "shell").__setitem__("held", []))),
        ("the tiny case's origin claim dropped", mut_index(lambda v: case_named(v, "statement_class").__setitem__("statements", []))),
        ("a digit of a statement's frame", mut_index(lambda v: flip_frame(case_named(v, "cluster")["statements"], 2))),
        ("a stated frame", mut_index(lambda v: case_named(v, "tiny")["expect"]["frames"].__setitem__(0, "00" * 32))),
        ("a stated reason", mut_index(lambda v: case_named(v, "cluster_reordered")["expect"].__setitem__("reason", "pose_degenerate"))),
    ]
    for what, f in mutations:
        fs = dict(files)
        f(fs)
        check(fs != files and bool(run_all(fs, held_sets)), "mutation caught: " + what)
    # 4. Each check dropped in turn: the corpus catches the weakened verifier.
    for name in ["signature", "layout", "ranges", "extensions_required", "primitive", "capability", "object", "pin",
                 "set_digest", "accessors", "commitment_object", "commitment_set", "attributes", "renderer", "frames", "origin",
                 "statement_commitment", "statement_object", "statement_renderer", "canonical", "canonical_frame", "own_frame",
                 "class", "envelope_length"]:
        DROPPED.clear()
        DROPPED.add(name)
        bad = run_all(files, held_sets)
        check(bool(bad), "check dropped and caught: %s (%s)" % (name, ", ".join(bad[:3])))
    DROPPED.clear()
    # 5. The renderer's arithmetic changed, rule by rule: the corpus's frames catch every one.
    for name in ["ties_desc", "sort_by_z", "no_dilation", "rdiv_floor", "exp_float", "contrib_floor", "stop_after",
                 "no_alpha_max", "no_cutoff", "rect_floor", "near_zero"]:
        MUTANT.clear()
        MUTANT.add(name)
        bad = run_all(files, held_sets)
        check(bool(bad), "renderer rule changed and caught: %s (%s)" % (name, ", ".join(bad[:3])))
    MUTANT.clear()
    print("%s: %d failure(s)" % ("PASS" if fails == 0 else "FAIL", fails))
    return 1 if fails else 0


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