open-standards/wai/int-payload-conformance/int_payload_verify.py

#!/usr/bin/env python3
"""Independent verifier for the integer payload corpus (extensions/integer-payloads.md).

Python standard library only; it shares no code with wai-rs. Written from the
extension's text, it:

  - parses every committed payload (WIV1, WIH1, WIA1, WIS1) with the rules and in the
    check order of the extension's §2-§6, and requires each to parse;
  - converts each payload's JSON source independently (§10) and requires the result to
    equal the committed payload byte for byte, and the payload's SHA-256 to be the one
    the index states;
  - decodes the WIV1 payloads with keyframe 0 in full - the rANS decoder with its
    bypass escape, its stop at the first exhausted symbol and its canonical end,
    whole-pel and half-pel motion compensation - and requires every frame's SHA-256
    and the clip's to be the ones the vector's sha256.json states;
  - parses every refusal fixture and requires a parse-layer fixture to be refused with
    its code, and a decode-layer fixture to parse; a decode-layer fixture of WIV1
    keyframe 0 is decoded and must be refused with its code (stream_exhausted,
    stream_noncanonical).

Not decoded independently: WIH1, WIV1 with keyframe 1 (its keyframe is a WIH1), WIA1
and WIS1, whose decodes run integer networks this script does not implement. Their
decoded outputs are checked by wai-rs against the same sidecars, on every leg of the
cross-architecture matrix for WIV1 and WIH1. The decode-layer fixtures of those
formats are checked here to parse, and by wai-rs for their codes.

Frame streams (extensions/frame-stream.md, the corpus frame-stream/), from that text and
integer-payloads.md §16:

  - splits every clip a stream names again, with its own rANS encoder (the encoder of
    integer-payloads.md §2, state from 2^31), and requires each committed frame payload
    byte for byte; splits the neural clip and joins its frames, and requires both;
  - parses every WAI1 envelope and makes frame-stream §5's checks over every stream, in
    order, with the BLAKE3 content hash of jwp-receipts §2.1 (../claim-conformance/
    _blake3.py), decoding every WIF1 frame of keyframe 0 (its own stream, its canonical
    end, the prediction from the previous decoded frame) and requiring each frame's
    SHA-256 and, for a clean split, the whole clip's frames byte for byte; a WIF1 frame of
    keyframe 1 and a WNF1 frame are parsed and held to their group's geometry, not decoded;
  - refuses every payload fixture with its code at its layer, every split and join
    refusal with its code, and checks every catalog case (moq-streaming-format §4.3).

NMSF tracks (extensions/frame-stream.md §9, the corpus frame-stream/nmsf/): it reads every
NMSF object (the 26-byte header of draft-herz-moq-nmsf-01 §3.2, its components §3.6) and
catalog entry with its own parser, maps every rewrap stream and requires each committed frame
object's WNF1 payload and manifest (its frame member, pin and capability), admits the frame
objects with its own frame-stream §5 checks (number and capture_ms included), unwraps them
and requires the NMSF objects back byte for byte, and checks every binding, catalog, unwrap
and links case, and that every code of §9.6 is given by one.

Conditional inter-frame video (integer-payloads.md §17, the corpus int-cond/), from that
text and integer-operators.md §13.2:

  - parses the WCB1 bundle (its WHM1 networks and its tables, read as JSON that holds only
    integers), plans each network on a frame's shapes, and decodes the committed WIC1 clip
    in full: the context network over the previous frame's feature (zero for frame 0), the
    entropy context network, the rANS decode of each latent element against its σ's bucket,
    the conditional decoder and the synthesis with its RGB8 step, every network's integer
    arithmetic and the check ranges of whm1/1 written here; requires every frame's SHA-256
    to be the sidecar's;
  - splits the clip and joins its frames, makes frame-stream §5's checks over every stream
    of wai.video.int_cond.frame and decodes its frames, refuses every fixture with its code
    at its layer, every set fixture and every join refusal.

    python3 int-payload-conformance/int_payload_verify.py int-payload-conformance
    python3 int-payload-conformance/int_payload_verify.py --frame-stream int-payload-conformance
    python3 int-payload-conformance/int_payload_verify.py --int-cond int-payload-conformance

Exit status 1 on any failure.
"""
import base64
import hashlib
import json
import math
import os
import struct
import sys

MAGICS = (b"WIV1", b"WIH1", b"WIA1", b"WIS1", b"WIF1", b"WCF1", b"WIC1")
MAX_TENSOR = 1 << 26
MAX_CLIP = 1 << 30
MAX_FRAMES = 1 << 20
MAX_MEDIAN_Q = 1 << 47
CAP_MAGIC = {
    "wai.video.int_motion": b"WIV1",
    "wai.video.int_hyper": b"WIV1",
    "wai.neural.int_hyper": b"WIH1",
    "wai.audio.int_codec": b"WIA1",
    "wai.splat.int_codec": b"WIS1",
}


class Refused(Exception):
    def __init__(self, code, why):
        super().__init__(f"{code}: {why}")
        self.code = code


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

    def left(self):
        return len(self.b) - self.p

    def take(self, n):
        if n > self.left():
            raise Refused("stream_length", "the payload ends inside a section")
        s = self.b[self.p:self.p + n]
        self.p += n
        return s

    def u8(self):
        return self.take(1)[0]

    def u16(self):
        return struct.unpack("<H", self.take(2))[0]

    def u32(self):
        return struct.unpack("<I", self.take(4))[0]

    def i32(self):
        return struct.unpack("<i", self.take(4))[0]

    def cdf(self, n):
        t = list(struct.unpack(f"<{n}I", self.take(4 * n)))
        check_cdf(t)
        return t

    def stream(self):
        n = self.u32()
        if n % 4 or n < 8:
            raise Refused("stream_length", "a rANS stream's length is not a multiple of 4, or is under 8")
        return self.take(n)

    def end(self):
        if self.left():
            raise Refused("trailing", "bytes follow the last section")


def check_cdf(t):
    if not 3 <= len(t) <= 4097:
        raise Refused("cdf_malformed", "length outside 3..=4097")
    if t[0] != 0 or t[-1] != 65536:
        raise Refused("cdf_malformed", "not from 0 to 65536")
    if any(a >= b for a, b in zip(t, t[1:])):
        raise Refused("cdf_malformed", "not strictly increasing")


def check_magic(p, want):
    m = p[:4]
    if len(m) == 4 and m == want:
        return
    if len(m) == 4 and m in MAGICS:
        raise Refused("format_capability_mismatch", "another integer format")
    raise Refused("bad_magic", "not the capability's magic")


def opening(p, magic, header):
    check_magic(p, magic)
    if len(p) < 5:
        raise Refused("stream_length", "no version byte")
    if p[4] != 1:
        raise Refused("version", "a version this revision does not define")
    if len(p) < header:
        raise Refused("stream_length", "the payload ends inside its fixed header")
    return Reader(p, 5)


def too_large(dims, cap):
    n = 1
    for d in dims:
        n *= d
    if n > cap:
        raise Refused("too_large", "over the cap")


def parse_wiv1(p, want_keyframe=None):
    r = opening(p, b"WIV1", 36)
    mvp, kf, res0 = r.u8(), r.u8(), r.u8()
    n_frames, h, w, block, q = r.u32(), r.u32(), r.u32(), r.u32(), r.u32()
    offset, cdf_len, res1 = r.i32(), r.u16(), r.u16()
    if res0 or res1:
        raise Refused("reserved_nonzero", "reserved")
    if mvp not in (1, 2):
        raise Refused("field_range", "mv_precision")
    if kf > 1:
        raise Refused("field_range", "keyframe")
    if 0 in (n_frames, h, w):
        raise Refused("field_range", "n_frames, h or w")
    if block == 0:
        raise Refused("field_range", "block")
    if not 1 <= q <= 65535:
        raise Refused("field_range", "q")
    if not 3 <= cdf_len <= 4097:
        raise Refused("field_range", "cdf_len")
    if n_frames == 1 and (mvp != 1 or block != 1):
        raise Refused("field_range", "a one-frame clip has mv_precision 1 and block 1")
    if want_keyframe is not None and want_keyframe != kf:
        raise Refused("keyframe_capability_mismatch", "keyframe")
    too_large([h, w, 3], MAX_TENSOR)
    too_large([n_frames, h, w, 3], MAX_CLIP)
    too_large([n_frames], MAX_FRAMES)
    nb = -(-h // block) * -(-w // block)
    too_large([2, n_frames - 1, nb], MAX_TENSOR)
    cdf = r.cdf(cdf_len)
    raw = r.take((n_frames - 1) * nb * 4)
    mvs = [struct.unpack_from("<hh", raw, 4 * i) for i in range((n_frames - 1) * nb)]
    rans = r.stream()
    wih1 = None
    if kf == 1:
        k = r.u32()
        if k != r.left():
            raise Refused("nested_length", "K is not the rest")
        wih1 = r.take(k)
        parse_wih1(wih1)
    else:
        r.end()
    return dict(mv_precision=mvp, keyframe=kf, n_frames=n_frames, h=h, w=w, block=block, q=q,
                offset=offset, cdf=cdf, mvs=mvs, rans=rans, wih1=wih1)


def parse_wih1(p):
    r = opening(p, b"WIH1", 32)
    res = r.take(3)
    m, sh, sw, n_z, zh, zw = (r.u32() for _ in range(6))
    if any(res):
        raise Refused("reserved_nonzero", "reserved")
    if 0 in (m, sh, sw, n_z, zh, zw):
        raise Refused("field_range", "a dimension")
    too_large([n_z, zh, zw], MAX_TENSOR)
    too_large([2, m, sh, sw], MAX_TENSOR)
    z, y = r.stream(), r.stream()
    r.end()
    return dict(m=m, sh=sh, sw=sw, n_z=n_z, zh=zh, zw=zw, z_rans=z, rans=y)


def parse_wia1(p):
    r = opening(p, b"WIA1", 32)
    afrac, res0 = r.u8(), r.u16()
    sr, n_ch, lat_len, t_out = r.u32(), r.u32(), r.u32(), r.u32()
    offset, cdf_len, res1 = r.i32(), r.u16(), r.u16()
    if res0 or res1:
        raise Refused("reserved_nonzero", "reserved")
    if not 1 <= afrac <= 30:
        raise Refused("field_range", "afrac")
    if 0 in (sr, n_ch, lat_len, t_out):
        raise Refused("field_range", "a dimension")
    if not 3 <= cdf_len <= 4097:
        raise Refused("field_range", "cdf_len")
    too_large([n_ch, lat_len], MAX_TENSOR)
    too_large([t_out], MAX_TENSOR)
    med = list(struct.unpack(f"<{n_ch}q", r.take(8 * n_ch)))
    if any(abs(x) > MAX_MEDIAN_Q for x in med):
        raise Refused("field_range", "median")
    cdfs = [r.cdf(cdf_len) for _ in range(n_ch)]
    rans = r.stream()
    r.end()
    return dict(afrac=afrac, sr=sr, n_ch=n_ch, lat_len=lat_len, t_out=t_out, offset=offset,
                median_q=med, cdfs=cdfs, rans=rans)


def parse_wis1(p):
    r = opening(p, b"WIS1", 32)
    afrac, res = r.u8(), r.u16()
    c, sh, sw, s, attr, n_g = (r.u32() for _ in range(6))
    if res:
        raise Refused("reserved_nonzero", "reserved")
    if not 1 <= afrac <= 30:
        raise Refused("field_range", "afrac")
    if 0 in (c, sh, sw, s, attr, n_g):
        raise Refused("field_range", "a dimension")
    if n_g > s * s:
        raise Refused("field_range", "n_gaussians")
    too_large([c, sh, sw], MAX_TENSOR)
    too_large([attr, s, s], MAX_TENSOR)
    raw = r.take(8 * attr)
    ranges = [struct.unpack_from("<ff", raw, 8 * i) for i in range(attr)]
    for lo, hi in ranges:
        if not (math.isfinite(lo) and math.isfinite(hi)) or lo > hi:
            raise Refused("field_range", "a range")
    lens = list(struct.unpack(f"<{c}H", r.take(2 * c)))
    if any(not 3 <= n <= 4097 for n in lens):
        raise Refused("field_range", "a CDF length")
    offs = list(struct.unpack(f"<{c}i", r.take(4 * c)))
    med = list(struct.unpack(f"<{c}q", r.take(8 * c)))
    if any(abs(x) > MAX_MEDIAN_Q for x in med):
        raise Refused("field_range", "median")
    cdfs = [r.cdf(n) for n in lens]
    rans = r.stream()
    r.end()
    return dict(afrac=afrac, c=c, sh=sh, sw=sw, s=s, attr=attr, n_gaussians=n_g, ranges=ranges,
                cdfs=cdfs, offsets=offs, median_q=med, rans=rans)


def parse(cap, p):
    m = CAP_MAGIC[cap]
    if m == b"WIV1":
        return parse_wiv1(p, 1 if cap == "wai.video.int_hyper" else 0)
    return {b"WIH1": parse_wih1, b"WIA1": parse_wia1, b"WIS1": parse_wis1}[m](p)


# ---- writers: the JSON sources converted independently (extensions/integer-payloads.md §10) ----

def put_stream(b):
    return struct.pack("<I", len(b)) + b


class NotConvertible(Exception):
    """A JSON source the converters of §10 refuse. Raised explicitly, never through
    `assert`, so `python3 -O` checks the same rules."""


def require(ok, why):
    if not ok:
        raise NotConvertible(why)


def wiv1_from_json(j, keyframe=None):
    nf, h, w, block = j["n_frames"], j["h"], j["w"], j["block"]
    mvp = j.get("mv_precision", 1)
    require(mvp in (1, 2), "mv_precision is neither 1 nor 2")
    if nf == 1:
        # A one-frame clip's header fixes mv_precision 1 and block 1 (§3, §10).
        mvp, block = 1, 1
    nb = -(-h // block) * -(-w // block)
    mvs = j["mvs"]
    require(len(mvs) == nf and not mvs[0], "mvs is not one list per frame, the first empty")
    out = b"".join([b"WIV1", bytes([1, mvp, 1 if keyframe else 0, 0]),
                    struct.pack("<IIIIIiHH", nf, h, w, block, j["q"], j["offset"], len(j["cdf"]), 0),
                    struct.pack(f"<{len(j['cdf'])}I", *j["cdf"])])
    for fr in mvs[1:]:
        require(isinstance(fr, list), "a frame's motion vectors are not a list")
        require(len(fr) <= nb, "a frame carries more motion vectors than blocks")
        for dy, dx in list(fr) + [[0, 0]] * (nb - len(fr)):
            out += struct.pack("<hh", dy, dx)
    out += put_stream(base64.b64decode(j["rans_b64"]))
    if keyframe:
        out += put_stream(keyframe)
    return out


def wih1_from_json(j):
    return (b"WIH1" + bytes([1, 0, 0, 0]) + struct.pack("<6I", j["m"], j["sh"], j["sw"], j["n_z"], j["zh"], j["zw"])
            + put_stream(bytes(j["z_rans"])) + put_stream(bytes(j["rans"])))


def wia1_from_json(j, sr):
    n = len(j["cdfs"][0])
    require(all(len(t) == n for t in j["cdfs"]), "the channels' tables are not of one length")
    return (b"WIA1" + bytes([1, j["afrac"], 0, 0])
            + struct.pack("<IIIIiHH", sr, j["n_ch"], j["lat_len"], j["t_out"], j["lo"], n, 0)
            + struct.pack(f"<{j['n_ch']}q", *j["median_q"])
            + b"".join(struct.pack(f"<{n}I", *t) for t in j["cdfs"])
            + put_stream(base64.b64decode(j["rans_b64"])))


def f32(x):
    """`x` as IEEE 754 binary32, refused unless it is exactly a binary32 value (§10)."""
    require(isinstance(x, (int, float)) and math.isfinite(x), f"{x!r} is not a finite number")
    try:
        b = struct.pack("<f", x)
    except OverflowError:
        raise NotConvertible(f"{x!r} is not a binary32 value")
    require(struct.unpack("<f", b)[0] == x, f"{x!r} is not a binary32 value")
    return b


def wis1_from_json(j):
    c = j["c"]
    return (b"WIS1" + bytes([1, j["afrac"], 0, 0])
            + struct.pack("<6I", c, j["sh"], j["sw"], j["s"], j["attr"], j["n_gaussians"])
            + b"".join(f32(lo) + f32(hi) for lo, hi in zip(j["lo"], j["hi"]))
            + struct.pack(f"<{c}H", *(len(t) for t in j["cdfs"]))
            + struct.pack(f"<{c}i", *j["offsets"])
            + struct.pack(f"<{c}q", *j["median_q"])
            + b"".join(struct.pack(f"<{len(t)}I", *t) for t in j["cdfs"])
            + put_stream(base64.b64decode(j["rans_b64"])))


# ---- the rANS decoder (precision 16, bypass 4) and the WIV1 keyframe-0 decode ----

I32_MIN, I32_MAX = -(1 << 31), (1 << 31) - 1


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


class Rans:
    def __init__(self, b):
        self.w = [struct.unpack_from("<I", b, i)[0] for i in range(0, len(b), 4)]
        self.i, self.x, self.exhausted = 2, (self.w[1] << 32) | self.w[0], False

    def renorm(self):
        if self.x < (1 << 31):
            if self.i >= len(self.w):
                self.exhausted = True
                return
            self.x = (self.x << 32) | self.w[self.i]
            self.i += 1

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

    def decode(self, cdf, offset):
        n = len(cdf)
        max_value = n - 2
        cum = self.x & 0xFFFF
        lo, hi = 1, n - 1
        while lo < hi:
            mid = (lo + hi) // 2
            if cdf[mid] > cum:
                hi = mid
            else:
                lo = mid + 1
        s = lo - 1
        self.x = (cdf[s + 1] - cdf[s]) * (self.x >> 16) + (self.x & 0xFFFF) - cdf[s]
        self.renorm()
        value = s
        if value == max_value:
            val = self.bits(4)
            count, groups = val, 0
            while val == 15 and groups < 8 and not self.exhausted:
                val = self.bits(4)
                count += val
                groups += 1
            raw = 0
            for j in range(min(count, 8)):
                raw |= self.bits(4) << (4 * j)
            value = raw >> 1
            value = sat(-value - 1) if raw & 1 else sat(value + max_value)
        return sat(value + offset)


def decode_wiv1_k0(v):
    h, w, block, q, prec = v["h"], v["w"], v["block"], v["q"], v["mv_precision"]
    nbx = -(-w // block)
    nb = nbx * -(-h // block)
    d = Rans(v["rans"])
    frames, prev = [], None
    for f in range(v["n_frames"]):
        mvs = v["mvs"][(f - 1) * nb:f * nb] if f else None
        # The frame grows a sample at a time as its symbols decode: nothing is
        # allocated ahead of the stream that fills it.
        cur = bytearray()
        for y in range(h):
            for x in range(w):
                dy, dx = mvs[(y // block) * nbx + x // block] if f else (0, 0)
                for c in range(3):
                    if f == 0:
                        pred = 128
                    else:
                        def p(yy, xx):
                            return prev[(min(max(yy, 0), h - 1) * w + min(max(xx, 0), w - 1)) * 3 + c]
                        if prec == 1:
                            pred = p(y + dy, x + dx)
                        else:
                            iy, fy, ix, fx = y + (dy >> 1), dy & 1, x + (dx >> 1), dx & 1
                            if not fy and not fx:
                                pred = p(iy, ix)
                            elif not fy:
                                pred = (p(iy, ix) + p(iy, ix + 1) + 1) >> 1
                            elif not fx:
                                pred = (p(iy, ix) + p(iy + 1, ix) + 1) >> 1
                            else:
                                pred = (p(iy, ix) + p(iy, ix + 1) + p(iy + 1, ix) + p(iy + 1, ix + 1) + 2) >> 2
                    r = d.decode(v["cdf"], v["offset"])
                    if d.exhausted:
                        # §2: the first symbol during whose decode the stream runs out.
                        raise Refused("stream_exhausted", "the residual stream ran out")
                    cur.append(min(max(pred + r * q, 0), 255))
        frames.append(bytes(cur))
        prev = cur
    if d.i != len(d.w) or d.x != 1 << 31:
        # §2: after the last symbol every word has been read and the state is 2^31.
        raise Refused("stream_noncanonical", "the residual stream does not end canonically")
    return frames


def main(cdir, frame_stream_only=False, int_cond_only=False):
    wai = os.path.normpath(os.path.join(cdir, ".."))
    if frame_stream_only or int_cond_only:
        fs_fails = frame_stream_main(wai) if frame_stream_only else int_cond_main(wai)
        for x in fs_fails:
            print("FAIL", x)
        return 1 if fs_fails else 0

    def rd(rel):
        with open(os.path.join(wai, rel), "rb") as f:
            return f.read()

    def jl(rel):
        return json.loads(rd(rel))

    idx = jl("int-payload-conformance/vectors.json")
    fails, decoded, parsed_only = [], 0, []
    for v in idx["vectors"]:
        name, cap = v["name"], v["capability"]
        p = rd(v["payload"])
        try:
            fields = parse(cap, p)
        except Refused as e:
            fails.append(f"{name}: does not parse: {e}")
            continue
        if hashlib.sha256(p).hexdigest() != v["payload_sha256"]:
            fails.append(f"{name}: payload SHA-256 differs from the index")
        src = v["converted_from"]
        j = jl(src["json"])
        m = CAP_MAGIC[cap]
        try:
            if m == b"WIV1":
                conv = wiv1_from_json(j, rd(src["keyframe"]) if "keyframe" in src else None)
            elif m == b"WIH1":
                conv = wih1_from_json(j)
            elif m == b"WIA1":
                conv = wia1_from_json(j, jl(src["sample_rate"])["sample_rate"])
            else:
                conv = wis1_from_json(j)
        except NotConvertible as e:
            fails.append(f"{name}: {src['json']} is not convertible: {e}")
            continue
        if conv != p:
            fails.append(f"{name}: the independent conversion of {src['json']} differs from the payload")
        if m == b"WIV1" and fields["keyframe"] == 0:
            frames = decode_wiv1_k0(fields)
            side = jl(v["sidecar"])[name]
            digests = [hashlib.sha256(f).digest() for f in frames]
            if [d.hex() for d in digests] != side["frames_sha256"] or hashlib.sha256(b"".join(digests)).hexdigest() != side["clip_sha256"]:
                fails.append(f"{name}: the independent decode differs from {v['sidecar']}")
            else:
                decoded += 1
                print(f"ok   {name}: parsed, converted, decoded to its sidecar's SHA-256 ({len(frames)} frames)")
        else:
            parsed_only.append(name)
            print(f"ok   {name}: parsed and converted (not decoded independently)")
    checked, decode_layer = 0, []
    for f in idx["refusals"]:
        name, cap, code, layer = f["name"], f["capability"], f["code"], f["layer"]
        p = rd(f["payload"])
        try:
            fields = parse(cap, p)
            got = None
        except Refused as e:
            got = e.code
        if layer == "parse":
            if got != code:
                fails.append(f"{name}: parsed with {got}, want {code}")
                continue
        else:
            if got is not None:
                fails.append(f"{name}: refused at parse with {got}, want a decode-layer {code}")
                continue
            if CAP_MAGIC[cap] == b"WIV1" and fields["keyframe"] == 0:
                try:
                    decode_wiv1_k0(fields)
                    fails.append(f"{name}: decoded, want {code}")
                    continue
                except Refused as e:
                    if e.code != code:
                        fails.append(f"{name}: decode refused with {e.code}, want {code}")
                        continue
            else:
                decode_layer.append(name)
        checked += 1
        print(f"ok   {name}: {code}" + ("" if layer == "parse" or name not in decode_layer else " (parses; the code is the decode's)"))
    # wiv1_end_state codes three symbols from the state 2^31 + 1, as the words [1, 4]; its
    # canonical twin, the words [0, 4], decodes to one grey pixel.
    twin = bytearray(rd("int-payload-conformance/refusals/wiv1_end_state.bin"))
    twin[-8:-4] = struct.pack("<I", 0)
    try:
        if decode_wiv1_k0(parse("wai.video.int_motion", bytes(twin))) != [bytes([128, 128, 128])]:
            fails.append("wiv1_end_state's canonical twin does not decode to one grey pixel")
        else:
            print("ok   wiv1_end_state's canonical twin decodes to one grey pixel")
    except Refused as e:
        fails.append(f"wiv1_end_state's canonical twin is refused: {e}")
    print(f"vectors: {decoded} decoded independently; parsed and converted only: {', '.join(parsed_only)}")
    print(f"refusal fixtures checked: {checked}; decode-layer codes left to wai-rs: {', '.join(decode_layer)}")
    if os.path.isdir(os.path.join(cdir, "frame-stream")):
        fails += frame_stream_main(wai)
    if os.path.isdir(os.path.join(cdir, "int-cond")):
        fails += int_cond_main(wai)
    for x in fails:
        print("FAIL", x)
    return 1 if fails else 0


# ---- frame streams (extensions/frame-stream.md; integer-payloads.md §16) ----

HERE = os.path.dirname(os.path.abspath(__file__))
FS = "int-payload-conformance/frame-stream"
MAX_INDEX = (1 << 53) - 1
INT_FRAME_CAPS = {"wai.video.int_motion.frame": 0, "wai.video.int_hyper.frame": 1}
NEURAL_FRAME_CAPS = ("wai.neural.dcvc_rt.frame", "wai.neural.dcvc_fm.frame")
FRAME_CAPS = tuple(INT_FRAME_CAPS) + NEURAL_FRAME_CAPS


def _content_hash(env):
    """jwp-receipts §2.1: BLAKE3 over the object domain and the envelope."""
    sys.path.insert(0, os.path.join(HERE, "..", "claim-conformance"))
    from _blake3 import blake3  # written from its specification; standard library only
    return blake3(b"moq-jwp:object\x01" + env)


def parse_wif1(p, want_keyframe=None):
    """integer-payloads.md §16, its checks in §2's order."""
    r = opening(p, b"WIF1", 32)
    kind, kf, mvp = r.u8(), r.u8(), r.u8()
    h, w, block, q = r.u32(), r.u32(), r.u32(), r.u32()
    offset, cdf_len, res = r.i32(), r.u16(), r.u16()
    if res:
        raise Refused("reserved_nonzero", "reserved")
    if kind > 1:
        raise Refused("field_range", "kind")
    intra = kind == 0
    if kf > 1:
        raise Refused("field_range", "keyframe")
    if (intra and mvp != 0) or (not intra and mvp not in (1, 2)):
        raise Refused("field_range", "mv_precision")
    if 0 in (h, w):
        raise Refused("field_range", "h or w")
    if (intra and block != 0) or (not intra and block == 0):
        raise Refused("field_range", "block")
    if (intra and not 1 <= q <= 65535) or (not intra and (q != 0 or offset != 0)):
        raise Refused("field_range", "q or offset")
    if (intra and not 3 <= cdf_len <= 4097) or (not intra and cdf_len != 0):
        raise Refused("field_range", "cdf_len")
    if want_keyframe is not None and want_keyframe != kf:
        raise Refused("keyframe_capability_mismatch", "keyframe")
    too_large([h, w, 3], MAX_TENSOR)
    f = dict(kind=kind, keyframe=kf, h=h, w=w)
    if intra:
        f.update(q=q, offset=offset, cdf=r.cdf(cdf_len), rans=None, wih1=None)
        if kf == 1:
            k = r.u32()
            if k != r.left():
                raise Refused("nested_length", "K is not the rest")
            f["wih1"] = r.take(k)
            parse_wih1(f["wih1"])
        else:
            f["rans"] = r.stream()
            r.end()
    else:
        nb = -(-h // block) * -(-w // block)
        too_large([2, nb], MAX_TENSOR)
        raw = r.take(nb * 4)
        f.update(mv_precision=mvp, block=block, mvs=[struct.unpack_from("<hh", raw, 4 * i) for i in range(nb)])
        f["rans"] = r.stream()
        r.end()
    return f


def write_wif1(f):
    if f["kind"] == 0:
        out = (b"WIF1" + bytes([1, 0, f["keyframe"], 0])
               + struct.pack("<IIIIiHH", f["h"], f["w"], 0, f["q"], f["offset"], len(f["cdf"]), 0)
               + struct.pack(f"<{len(f['cdf'])}I", *f["cdf"]))
        return out + (put_stream(f["rans"]) if f["keyframe"] == 0 else put_stream(f["wih1"]))
    return (b"WIF1" + bytes([1, 1, f["keyframe"], f["mv_precision"]])
            + struct.pack("<IIIIiHH", f["h"], f["w"], f["block"], 0, 0, 0, 0)
            + b"".join(struct.pack("<hh", dy, dx) for dy, dx in f["mvs"]) + put_stream(f["rans"]))


def parse_wnf1(p):
    """frame-stream.md §4, its checks in order."""
    if p[:4] != b"WNF1":
        raise Refused("bad_magic", "not WNF1")
    if len(p) < 5:
        raise Refused("stream_length", "no version byte")
    if p[4] != 1:
        raise Refused("version", "version")
    if len(p) < 40:
        raise Refused("stream_length", "the fixed header")
    r = Reader(p, 5)
    kind, res0 = r.u8(), r.u16()
    h, w, fps, cy, cz = r.u32(), r.u32(), r.u32(), r.u16(), r.u16()
    sy, sz, q, cf, res1, lz, ly = r.u8(), r.u8(), r.u16(), r.u16(), r.u16(), r.u32(), r.u32()
    if res0 or res1:
        raise Refused("reserved_nonzero", "reserved")
    if kind > 1:
        raise Refused("field_range", "kind")
    header = (h, w, fps, cy, cz, sy, sz, q, cf)
    check_wnf1_header(header)
    z, y = r.take(lz), r.take(ly)
    r.end()
    return dict(kind=kind, header=header, z=z, y=y)


def check_wnf1_header(hd):
    h, w, fps, cy, cz, sy, sz, q, cf = hd
    if 0 in (h, w, cy, cz, cf):
        raise Refused("field_range", "a dimension")
    if sy > 15 or sz > 15:
        raise Refused("field_range", "a side's logarithm")
    too_large([h, w, 3], MAX_TENSOR)
    too_large([cy, 1 << sy, 1 << sy], MAX_TENSOR)
    too_large([cz, 1 << sz, 1 << sz], MAX_TENSOR)
    too_large([cf, 1 << sy, 1 << sy], MAX_TENSOR)


def write_wnf1(kind, hd, z, y):
    h, w, fps, cy, cz, sy, sz, q, cf = hd
    return (b"WNF1" + bytes([1, kind, 0, 0]) + struct.pack("<IIIHHBBHHHII", h, w, fps, cy, cz, sy, sz, q, cf, 0, len(z), len(y))
            + z + y)


def parse_frame(cap, p):
    if cap in INT_FRAME_CAPS:
        return parse_wif1(p, INT_FRAME_CAPS[cap])
    if cap in NEURAL_FRAME_CAPS:
        return parse_wnf1(p)
    raise Refused("unsupported_capability", cap)


# The rANS encoder of integer-payloads.md §2 (precision 16, bypass 4), state from 2^31.

def _value_syms(value, cdf, offset):
    n = len(cdf)
    max_value = n - 2
    v = value - offset
    if v < 0:
        sym, raw = max_value, -2 * v - 1
    elif v >= max_value:
        sym, raw = max_value, 2 * (v - max_value)
    else:
        sym, raw = v, None
    out = [("t", cdf[sym], cdf[sym + 1] - cdf[sym])]
    if raw is not None:
        nb = 0
        while raw >> (nb * 4):
            nb += 1
        count = nb
        while count >= 15:
            out.append(("b", 15))
            count -= 15
        out.append(("b", count))
        out.extend(("b", (raw >> (4 * j)) & 15) for j in range(nb))
    return out


def rans_encode(values, cdf, offset):
    syms = [s for v in values for s in _value_syms(v, cdf, offset)]
    state, words = 1 << 31, []
    for s in reversed(syms):
        if s[0] == "t":
            start, freq = s[1], s[2]
            x_max = ((1 << 31 >> 16) << 32) * freq
            while state >= x_max:
                words.append(state & 0xFFFFFFFF)
                state >>= 32
            state = ((state // freq) << 16) + (state % freq) + start
        else:
            x_max = ((1 << 31 >> 16) << 32) * (1 << 12)
            if state >= x_max:
                words.append(state & 0xFFFFFFFF)
                state >>= 32
            state = (state << 4) | s[1]
    return struct.pack("<II", state & 0xFFFFFFFF, state >> 32) + b"".join(struct.pack("<I", x) for x in reversed(words))


def split_wiv1(cap, wiv1):
    """frame-stream.md §6: one group of WIF1 frames, each frame's symbols the clip's."""
    kf = INT_FRAME_CAPS[cap]
    v = parse_wiv1(wiv1, kf)
    per, nb = v["h"] * v["w"] * 3, -(-v["h"] // v["block"]) * -(-v["w"] // v["block"])
    d = Rans(v["rans"])
    syms = []
    for _ in range(kf, v["n_frames"]):
        fr = []
        for _ in range(per):
            fr.append(d.decode(v["cdf"], v["offset"]))
            if d.exhausted:
                raise Refused("stream_exhausted", "the residual stream ran out")
        syms.append(fr)
    if d.i != len(d.w) or d.x != 1 << 31:
        raise Refused("stream_noncanonical", "the residual stream does not end canonically")
    base = dict(keyframe=kf, h=v["h"], w=v["w"])
    out = [write_wif1(dict(base, kind=0, q=v["q"], offset=v["offset"], cdf=v["cdf"],
                           rans=rans_encode(syms[0], v["cdf"], v["offset"]) if kf == 0 else None, wih1=v["wih1"]))]
    for f in range(1, v["n_frames"]):
        out.append(write_wif1(dict(base, kind=1, mv_precision=v["mv_precision"], block=v["block"],
                                   mvs=v["mvs"][(f - 1) * nb:f * nb], rans=rans_encode(syms[f - kf], v["cdf"], v["offset"]))))
    return out


def decode_frame(rans, h, w, q, offset, cdf, prev=None, prec=1, block=1, mvs=()):
    """One WIF1 frame from its own stream (integer-payloads.md §16)."""
    d = Rans(rans)
    nbx = -(-w // block)
    cur = bytearray()
    for y in range(h):
        for x in range(w):
            dy, dx = mvs[(y // block) * nbx + x // block] if prev is not None else (0, 0)
            for c in range(3):
                if prev is None:
                    pred = 128
                else:
                    def p(yy, xx):
                        return prev[(min(max(yy, 0), h - 1) * w + min(max(xx, 0), w - 1)) * 3 + c]
                    if prec == 1:
                        pred = p(y + dy, x + dx)
                    else:
                        iy, fy, ix, fx = y + (dy >> 1), dy & 1, x + (dx >> 1), dx & 1
                        if not fy and not fx:
                            pred = p(iy, ix)
                        elif not fy:
                            pred = (p(iy, ix) + p(iy, ix + 1) + 1) >> 1
                        elif not fx:
                            pred = (p(iy, ix) + p(iy + 1, ix) + 1) >> 1
                        else:
                            pred = (p(iy, ix) + p(iy, ix + 1) + p(iy + 1, ix) + p(iy + 1, ix + 1) + 2) >> 2
                r = d.decode(cdf, offset)
                if d.exhausted:
                    raise Refused("stream_exhausted", "the frame's stream ran out")
                cur.append(min(max(pred + r * q, 0), 255))
    if d.i != len(d.w) or d.x != 1 << 31:
        raise Refused("stream_noncanonical", "the frame's stream does not end canonically")
    return bytes(cur)


def read_envelope(b):
    """SPEC §2: the manifest and the payload, or envelope_malformed."""
    try:
        if b[:4] != b"WAI1":
            raise ValueError("magic")
        ml = struct.unpack_from("<I", b, 4)[0]
        m = json.loads(b[8:8 + ml].decode("utf-8"))
        pl = struct.unpack_from("<I", b, 8 + ml)[0]
        payload = b[12 + ml:12 + ml + pl]
        if len(payload) != pl or not isinstance(m, dict):
            raise ValueError("length")
        return m, payload
    except (ValueError, struct.error, UnicodeDecodeError) as e:
        raise Refused("envelope_malformed", str(e))


def _is_hash(s):
    return isinstance(s, str) and len(s) == 64 and all(c in "0123456789abcdef" for c in s)


def read_link(m):
    f = m.get("frame")
    if not isinstance(f, dict):
        raise Refused("frame_link_missing", "no frame object")
    i, prev = f.get("index"), f.get("prev", 0)
    if not (type(i) is int and 0 <= i <= MAX_INDEX):
        raise Refused("frame_link_malformed", "index")
    if not (prev is None or _is_hash(prev)):
        raise Refused("frame_link_malformed", "prev")
    # number and capture_ms, OPTIONAL: present, each an integer in its range.
    for k, low in (("number", 0), ("capture_ms", 1)):
        if k in f and not (type(f[k]) is int and low <= f[k] <= MAX_INDEX):
            raise Refused("frame_link_malformed", k)
    return i, prev, f.get("number")


class Stream:
    """frame-stream.md §5. `decode`: decode WIF1 keyframe-0 frames; otherwise check each
    frame's geometry against its group's intra frame."""

    def __init__(self, cap, decode):
        self.cap, self.decode, self.last, self.state = cap, decode, None, None

    def push(self, env):
        try:
            return self._admit(env)
        except Refused:
            self.last = self.state = None
            raise

    def _admit(self, env):
        m, payload = read_envelope(env)
        if (m.get("model_requirement") or {}).get("capability") != self.cap:
            raise Refused("frame_capability_mismatch", "another capability")
        index, prev, number = read_link(m)
        f = parse_frame(self.cap, payload)
        intra = f["kind"] == 0
        if intra and not (index == 0 and prev is None):
            raise Refused("frame_link_mismatch", "an intra frame with a link")
        if not intra and not (index >= 1 and prev is not None):
            raise Refused("frame_link_mismatch", "an inter frame without one")
        h = _content_hash(env).hex()
        prior = (m.get("conditioning") or {}).get("prior")
        out = b""
        if intra:
            self.last = self.state = None
            if self.cap in INT_FRAME_CAPS:
                geom = (f["h"], f["w"])
                if self.decode and f["keyframe"] == 0:
                    out = decode_frame(f["rans"], f["h"], f["w"], f["q"], f["offset"], f["cdf"])
                self.state = dict(geom=geom, q=f.get("q"), offset=f.get("offset"), cdf=f.get("cdf"), ref=out)
            else:
                self.state = dict(geom=f["header"])
        else:
            if self.last is None:
                raise Refused("frame_no_reference", "no reference")
            if prev != self.last["hash"]:
                raise Refused("frame_gap" if index > self.last["index"] + 1 else "frame_prev_mismatch", "prev")
            if index != self.last["index"] + 1:
                raise Refused("frame_index_mismatch", "index")
            if number != (None if self.last["number"] is None else self.last["number"] + 1):
                raise Refused("frame_number_mismatch", "number")
            if prior != self.last["prior"]:
                raise Refused("frame_prior_changed", "pin")
            geom = (f["h"], f["w"]) if self.cap in INT_FRAME_CAPS else f["header"]
            if geom != self.state["geom"]:
                raise Refused("geometry_mismatch", "geometry")
            if self.decode and self.cap in INT_FRAME_CAPS and f["keyframe"] == 0:
                st = self.state
                out = decode_frame(f["rans"], f["h"], f["w"], st["q"], st["offset"], st["cdf"], st["ref"],
                                   f["mv_precision"], f["block"], f["mvs"])
                st["ref"] = out
        self.last = dict(hash=h, index=index, prior=prior, number=number)
        return out


def split_dcvc(clip):
    if len(clip) < 26:
        raise Refused("stream_length", "header")
    h, w, n, fps, cy, cz, sy, sz, q, cf = struct.unpack_from("<IIIIHHBBHH", clip, 0)
    if n == 0:
        raise Refused("field_range", "no frames")
    hd = (h, w, fps, cy, cz, sy, sz, q, cf)
    check_wnf1_header(hd)
    at = 26 + 8 * n
    if at > len(clip):
        raise Refused("stream_length", "the frame table")
    out = []
    for i in range(n):
        lz, ly = struct.unpack_from("<II", clip, 26 + 8 * i)
        if at + lz + ly > len(clip):
            raise Refused("stream_length", "a bitstream")
        out.append(write_wnf1(int(i > 0), hd, clip[at:at + lz], clip[at + lz:at + lz + ly]))
        at += lz + ly
    if at != len(clip):
        raise Refused("trailing", "after the last bitstream")
    return out


def join_dcvc(frames):
    fs = [parse_wnf1(f) for f in frames]
    if not fs or any(f["kind"] != int(i > 0) for i, f in enumerate(fs)):
        raise Refused("field_range", "a group is its intra frame, then inter frames")
    if any(f["header"] != fs[0]["header"] for f in fs):
        raise Refused("geometry_mismatch", "a header")
    h, w, fps, cy, cz, sy, sz, q, cf = fs[0]["header"]
    return (struct.pack("<IIIIHHBBHH", h, w, len(fs), fps, cy, cz, sy, sz, q, cf)
            + b"".join(struct.pack("<II", len(f["z"]), len(f["y"])) for f in fs)
            + b"".join(f["z"] + f["y"] for f in fs))


def catalog_read(entry):
    caps = entry.get("wai_capabilities")
    caps = caps if isinstance(caps, list) else []
    if any(not isinstance(c, str) for c in caps):
        raise Refused("catalog_frame_stream_malformed", "wai_capabilities")
    member = entry.get("wai_frame_stream")
    if member is None:
        if any(c in FRAME_CAPS for c in caps):
            raise Refused("catalog_frame_stream_malformed", "a frame capability with no frame stream")
        return None
    if not isinstance(member, dict):
        raise Refused("catalog_frame_stream_malformed", "not an object")
    cap = member.get("capability")
    if not isinstance(cap, str) or cap not in FRAME_CAPS:
        raise Refused("catalog_frame_stream_malformed", "capability")
    if caps != [cap]:
        raise Refused("catalog_frame_stream_malformed", "wai_capabilities")
    return cap


def catalog_check(cap, env, object_id):
    m, _ = read_envelope(env)
    if (m.get("model_requirement") or {}).get("capability") != cap:
        raise Refused("catalog_frame_mismatch", "capability")
    f = m.get("frame")
    if not isinstance(f, dict):
        raise Refused("catalog_frame_mismatch", "no frame")
    if object_id is not None and not (type(f.get("index")) is int and f["index"] == object_id):
        raise Refused("catalog_frame_mismatch", "index")


def frame_stream_main(wai):
    def rd(rel):
        with open(os.path.join(wai, rel), "rb") as f:
            return f.read()

    fails = []
    idx = json.loads(rd(f"{FS}/index.json"))
    for s in idx["streams"]:
        name, cap, mode = s["name"], s["capability"], s["mode"]
        objs = s["objects"]
        envs = [rd(o["envelope"]) for o in objs]
        decode = mode == "full" and cap in INT_FRAME_CAPS and INT_FRAME_CAPS[cap] == 0
        st, outs, bad = Stream(cap, decode), [], []
        for k, (o, e) in enumerate(zip(objs, envs)):
            try:
                out = st.push(e)
                got = "ok"
            except Refused as r:
                got, out = r.code, None
            if got != o["expect"]:
                bad.append(f"object {k}: {got}, want {o['expect']}")
                continue
            if got == "ok":
                if _content_hash(e).hex() != o["content_hash"]:
                    bad.append(f"object {k}: content hash")
                if decode:
                    outs.append(out)
                    if "frame_sha256" in o and hashlib.sha256(out).hexdigest() != o["frame_sha256"]:
                        bad.append(f"object {k}: frame SHA-256")
        if s["split_from"]:
            frames = []
            for src in s["split_from"]:
                frames += split_wiv1(cap, rd(src["payload"]))
            payloads = [read_envelope(e)[1] for e in envs]
            if frames != payloads:
                bad.append("the independent split differs from the committed frame payloads")
            if decode:
                whole, side = [], []
                for src in s["split_from"]:
                    whole += decode_wiv1_k0(parse_wiv1(rd(src["payload"]), 0))
                    side += json.loads(rd(src["sidecar"]))[src["name"]]["frames_sha256"]
                if outs != whole:
                    bad.append("the frames are not the whole clips' frames byte for byte")
                if [hashlib.sha256(f).hexdigest() for f in outs] != side:
                    bad.append("the frames are not the sidecars' SHA-256")
        for b in bad:
            fails.append(f"frame stream {name}: {b}")
        if not bad:
            how = "decoded" if decode else "checked to its links and geometry"
            print(f"ok   frame stream {name} ({cap}): {len(objs)} objects as expected, {how}"
                  + (", split again byte for byte" if s["split_from"] else ""))
    for p in idx["payloads"]:
        name, cap, code, layer = p["name"], p["capability"], p["code"], p["layer"]
        b = rd(p["payload"])
        try:
            f = parse_frame(cap, b)
            got = None
        except Refused as e:
            got = e.code
        if layer == "decode" and got is None:
            try:
                decode_frame(f["rans"], f["h"], f["w"], f["q"], f["offset"], f["cdf"])
            except Refused as e:
                got = e.code
        if got != code:
            fails.append(f"frame payload {name}: {got}, want {code} at {layer}")
        else:
            print(f"ok   frame payload {name}: {code} ({layer})")
    d = idx["dcvc"]
    clip = rd(d["clip"])
    frames = [rd(f) for f in d["frames"]]
    if split_dcvc(clip) != frames or join_dcvc(frames) != clip:
        fails.append("dcvc: the split and the join are not each other's inverse over the committed clip")
    else:
        print(f"ok   dcvc: the clip splits into its {len(frames)} frames and they join into it")
    for r in d["split_refusals"]:
        try:
            split_dcvc(rd(r["clip"]))
            got = None
        except Refused as e:
            got = e.code
        (print(f"ok   dcvc {r['name']}: {got}") if got == r["code"] else fails.append(f"dcvc {r['name']}: {got}, want {r['code']}"))
    for r in d["join_refusals"]:
        try:
            join_dcvc([rd(f) for f in r["frames"]])
            got = None
        except Refused as e:
            got = e.code
        (print(f"ok   dcvc {r['name']}: {got}") if got == r["code"] else fails.append(f"dcvc {r['name']}: {got}, want {r['code']}"))
    cases = json.loads(rd(idx["catalog"]))
    for c in cases["cases"]:
        try:
            cap = catalog_read(c["entry"])
            got = "ok" if cap is not None else "none"
        except Refused as e:
            cap, got = None, e.code
        bad = [] if got == c["read"] else [f"read {got}, want {c['read']}"]
        if cap is not None:
            for k, o in enumerate(c["objects"]):
                try:
                    catalog_check(cap, rd(o["envelope"]), o.get("object_id"))
                    g = "ok"
                except Refused as e:
                    g = e.code
                if g != o["expect"]:
                    bad.append(f"object {k}: {g}, want {o['expect']}")
        for b in bad:
            fails.append(f"catalog {c['name']}: {b}")
        if not bad:
            print(f"ok   catalog {c['name']}: {got}")
    if os.path.isfile(os.path.join(wai, NMSF, "index.json")):
        fails += nmsf_main(wai)
    return fails


# ---- NMSF tracks (extensions/frame-stream.md §9; draft-herz-moq-nmsf-01 §3) ----

NMSF = FS + "/nmsf"
NMSF_CODECS = {"dcvc-rt": ("wai.neural.dcvc_rt.frame", "dcvc_rt"), "dcvc-fm": ("wai.neural.dcvc_fm.frame", "dcvc_fm")}
NMSF_CODES = ("nmsf_catalog_malformed", "nmsf_track_mismatch", "nmsf_codec_unmapped", "nmsf_model_unbound",
              "nmsf_binding_malformed", "nmsf_object_malformed", "nmsf_reserved", "nmsf_group_order",
              "nmsf_value_unmappable", "nmsf_shape_unmappable", "nmsf_group_header_changed", "nmsf_no_form")
_ABSENT = object()


def _num(v):
    return type(v) in (int, float)


def nmsf_object(b):
    """§9.1: the 26-byte header (the draft's field sizes, big-endian), then the payload."""
    if len(b) < 26:
        raise Refused("nmsf_object_malformed", "the header")
    ft, qp, fn, pts, w, h, plen = struct.unpack_from(">BBIQIII", b, 0)
    if plen != len(b) - 26:
        raise Refused("nmsf_object_malformed", "payload_len")
    if ft > 1:
        raise Refused("nmsf_reserved", "frame_type")
    if qp > 63:
        raise Refused("nmsf_reserved", "qp")
    return (ft, qp, fn, pts, w, h), b[26:]


def nmsf_components(p, n):
    at, out = 0, []
    for _ in range(n):
        if len(p) - at < 16:
            raise Refused("nmsf_object_malformed", "a component's header")
        c, h, w, dl = struct.unpack_from(">IIII", p, at)
        if dl > len(p) - at - 16:
            raise Refused("nmsf_object_malformed", "data_len")
        out.append((c, h, w, p[at + 16:at + 16 + dl]))
        at += 16 + dl
    if at != len(p):
        raise Refused("nmsf_object_malformed", "after the last component")
    return out


def nmsf_write(hd, payload):
    return struct.pack(">BBIQIII", *hd, len(payload)) + payload


def nmsf_component(c, h, w, data):
    return struct.pack(">IIII", c, h, w, len(data)) + data


def nmsf_shape(comp):
    c, h, w, _ = comp
    if h != w or h < 1 or h & (h - 1) or h.bit_length() - 1 > 15:
        raise Refused("nmsf_shape_unmappable", "a side")
    if not 1 <= c <= 65535:
        raise Refused("nmsf_shape_unmappable", "channels")
    return c, h.bit_length() - 1


def nmsf_binding(v):
    """§9.3."""
    def bad(why):
        return Refused("nmsf_binding_malformed", why)
    if not isinstance(v, dict):
        raise bad("not an object")
    if not isinstance(v.get("codec"), str):
        raise bad("codec")
    mv = v.get("modelVersion")
    if mv is not None and not isinstance(mv, str):
        raise bad("modelVersion")
    cf = v.get("cf")
    if not (type(cf) is int and 1 <= cf <= 65535):
        raise bad("cf")
    p = v.get("prior")
    if p is not None:
        if not isinstance(p, dict) or not all(isinstance(p.get(k), str) for k in ("id", "sha256", "determinism")):
            raise bad("prior")
        if not _is_hash(p["sha256"]) or any(p.get(k) is not None and not isinstance(p[k], str) for k in ("form", "tier")):
            raise bad("prior")
    return dict(codec=v["codec"], mv=mv, cf=cf, prior=p)


def fps_milli(f):
    m = round(f * 1000.0)
    if 1 <= m <= 0xFFFFFFFF and m / 1000.0 == f:
        return int(m)
    raise Refused("nmsf_value_unmappable", "framerate")


def nmsf_fields(e):
    def m(why):
        return Refused("nmsf_catalog_malformed", why)
    if not isinstance(e, dict):
        raise m("an entry")
    if not (isinstance(e.get("name"), str) and e["name"]):
        raise m("name")
    if e.get("packaging") != "nvc":
        raise m("packaging")
    if not isinstance(e.get("isLive"), bool):
        raise m("isLive")
    for k in ("codec", "colorspace"):
        if not isinstance(e.get(k), str):
            raise m(k)
    for k in ("width", "height", "gopSize"):
        if not (type(e.get(k)) is int and e[k] >= 1):
            raise m(k)
    fr = e.get("framerate")
    if not (_num(fr) and math.isfinite(fr) and fr > 0):
        raise m("framerate")
    if "priority" in e and not _num(e["priority"]):
        raise m("priority")
    for k in ("nvcRole", "depends"):
        if k in e and not isinstance(e[k], str):
            raise m(k)
    mv = None
    if "nvc" in e:
        if not isinstance(e["nvc"], dict):
            raise m("nvc")
        if "modelVersion" in e["nvc"]:
            mv = e["nvc"]["modelVersion"]
            if not isinstance(mv, str):
                raise m("nvc.modelVersion")
    return mv


def nmsf_read(entries):
    """§9.2: one entry or two, then the codec and the framerate."""
    mvs = [nmsf_fields(e) for e in entries]
    if len(entries) == 1:
        if "nvcRole" in entries[0]:
            raise Refused("nmsf_catalog_malformed", "a single track with a role")
        mode = 1
    elif len(entries) == 2:
        roles = [e.get("nvcRole") for e in entries]
        if sorted(r for r in roles if r is not None) != ["hyperprior", "latent"] or None in roles:
            raise Refused("nmsf_catalog_malformed", "roles")
        hi = roles.index("hyperprior")
        hy, la = entries[hi], entries[1 - hi]
        if hy["name"] == la["name"] or la.get("depends") != hy["name"]:
            raise Refused("nmsf_catalog_malformed", "depends")
        for k in ("codec", "isLive", "width", "height", "framerate", "colorspace", "gopSize"):
            a, b = hy.get(k, _ABSENT), la.get(k, _ABSENT)
            same = float(a) == float(b) if (_num(a) and _num(b) and not isinstance(a, bool) and not isinstance(b, bool)) else a == b
            if not same:
                raise Refused("nmsf_track_mismatch", k)
        if mvs[0] != mvs[1]:
            raise Refused("nmsf_track_mismatch", "modelVersion")
        mode = 2
    else:
        raise Refused("nmsf_catalog_malformed", "one track or two")
    codec = entries[0]["codec"]
    if codec not in NMSF_CODECS:
        raise Refused("nmsf_codec_unmapped", codec)
    cap, kind = NMSF_CODECS[codec]
    return dict(mode=mode, entries=entries, codec=codec, mv=mvs[0], cap=cap, kind=kind, fps=fps_milli(entries[0]["framerate"]))


def nmsf_bind(t, binds):
    for b in binds:
        if b["codec"] == t["codec"] and b["mv"] == t["mv"]:
            return b
    raise Refused("nmsf_model_unbound", t["codec"])


def nmsf_wai_entry(t, b, name):
    first, cap = t["entries"][0], t["cap"]
    e = {"name": name, "packaging": "wai", "isLive": first["isLive"], "role": "video", "width": first["width"],
         "height": first["height"], "framerate": first["framerate"], "wai_capabilities": [cap],
         "wai_frame_stream": {"capability": cap}}
    if b["prior"] is not None:
        pe = {"capability": cap, "sha256": b["prior"]["sha256"]}
        for k in ("form", "tier"):
            if b["prior"].get(k) is not None:
                pe[k] = b["prior"][k]
        e["wai_priors"] = [pe]
    e["wai_nmsf"] = t["entries"]
    return e


def nmsf_read_wai(entry):
    """§9.2, a WAI entry: its frame stream, wai_nmsf, and the pin it lists."""
    cap = catalog_read(entry)
    if cap is None:
        raise Refused("nmsf_no_form", "no frame stream")
    if entry.get("wai_nmsf") is None:
        raise Refused("nmsf_no_form", "no wai_nmsf")
    if not isinstance(entry["wai_nmsf"], list):
        raise Refused("nmsf_catalog_malformed", "wai_nmsf")
    t = nmsf_read(entry["wai_nmsf"])
    if t["cap"] != cap:
        raise Refused("nmsf_catalog_malformed", "the codec's capability")
    pins = []
    wp = entry.get("wai_priors")
    if wp is not None:
        if not isinstance(wp, list):
            raise Refused("nmsf_catalog_malformed", "wai_priors")
        for p in wp:
            if not isinstance(p, dict):
                raise Refused("nmsf_catalog_malformed", "a wai_priors entry")
            if p.get("capability") != cap:
                continue
            if not _is_hash(p.get("sha256")) or (p.get("form") is not None and not isinstance(p["form"], str)):
                raise Refused("nmsf_catalog_malformed", "a pin")
            pins.append({"sha256": p["sha256"], "form": p.get("form")})
    if len(pins) > 1:
        raise Refused("nmsf_no_form", "several pins")
    return t, (pins[0] if pins else None)


class NmsfRewrap:
    """§9.4. `expect` gives the frame object's payload and `frame` member, which the
    caller holds the committed object to; `commit` makes it the group's last."""

    def __init__(self, t, b):
        self.t, self.b, self.g = t, b, None

    def expect(self, oid, objs):
        try:
            return self._map(oid, objs)
        except Refused:
            self.g = None
            raise

    def _map(self, oid, objs):
        if len(objs) != self.t["mode"]:
            raise Refused("nmsf_track_mismatch", "objects")
        parsed = [nmsf_object(o) for o in objs]
        hd = parsed[0][0]
        if any(p[0] != hd for p in parsed):
            raise Refused("nmsf_track_mismatch", "headers")
        if self.t["mode"] == 1:
            z, y = nmsf_components(parsed[0][1], 2)
        else:
            z, y = nmsf_components(parsed[0][1], 1)[0], nmsf_components(parsed[1][1], 1)[0]
        ft, qp, fn, pts, w, h = hd
        if ft == 0:
            self.g = None
            if oid is not None and oid != 0:
                raise Refused("nmsf_group_order", "an Intra frame past Object ID 0")
            index, prev = 0, None
        else:
            g = self.g
            if g is None or fn != g["number"] + 1 or (oid is not None and oid != g["index"] + 1):
                raise Refused("nmsf_group_order", "the group")
            index, prev = g["index"] + 1, g["hash"]
        if pts > MAX_INDEX:
            raise Refused("nmsf_value_unmappable", "pts_ms")
        cz, sz = nmsf_shape(z)
        cy, sy = nmsf_shape(y)
        header = (h, w, self.t["fps"], cy, cz, sy, sz, qp, self.b["cf"])
        try:
            check_wnf1_header(header)
        except Refused as e:
            raise Refused("nmsf_shape_unmappable", e.code)
        if ft == 1 and header != self.g["header"]:
            raise Refused("nmsf_group_header_changed", "the header")
        member = {"index": index, "prev": prev, "number": fn}
        if pts:
            member["capture_ms"] = pts
        self.pending = dict(index=index, number=fn, header=header)
        return write_wnf1(ft, header, z[3], y[3]), member

    def commit(self, env):
        self.g = dict(self.pending, hash=_content_hash(env).hex())


class NmsfUnwrap:
    """§9.5: a sink's checks (frame-stream §5), then the NMSF objects."""

    def __init__(self, t, pin):
        self.t, self.pin, self.st = t, pin, Stream(t["cap"], False)

    def push(self, env):
        self.st.push(env)
        try:
            return self._map(env)
        except Refused:
            self.st.last = self.st.state = None
            raise

    def _map(self, env):
        m, payload = read_envelope(env)
        f = parse_wnf1(payload)
        if f["kind"] == 0:
            p = (m.get("conditioning") or {}).get("prior")
            got = None if p is None else {"sha256": p.get("sha256"), "form": p.get("form")}
            if got != self.pin:
                raise Refused("nmsf_model_unbound", "the pin")
        fr = m["frame"]
        if "number" not in fr or fr["number"] > 0xFFFFFFFF:
            raise Refused("nmsf_value_unmappable", "number")
        h, w, fps, cy, cz, sy, sz, q, cf = f["header"]
        if q > 63 or fps != self.t["fps"]:
            raise Refused("nmsf_value_unmappable", "q or fps")
        hd = (f["kind"], q, fr["number"], fr.get("capture_ms", 0), w, h)
        zc = nmsf_component(cz, 1 << sz, 1 << sz, f["z"])
        yc = nmsf_component(cy, 1 << sy, 1 << sy, f["y"])
        return [nmsf_write(hd, zc + yc)] if self.t["mode"] == 1 else [nmsf_write(hd, zc), nmsf_write(hd, yc)]


def _code(f, *a):
    try:
        return "ok", f(*a)
    except Refused as e:
        return e.code, None


def nmsf_main(wai):
    def rd(rel):
        with open(os.path.join(wai, rel), "rb") as f:
            return f.read()

    fails = []
    idx = json.loads(rd(f"{NMSF}/index.json"))
    binds = [nmsf_binding(b) for b in idx["bindings"]]
    for c in idx["binding_cases"]:
        got, _ = _code(nmsf_binding, c["binding"])
        (print(f"ok   nmsf binding {c['name']}: {got}") if got == c["expect"] and got != "ok"
         else fails.append(f"nmsf binding {c['name']}: {got}, want {c['expect']}"))
    nmsf_cases, wai_entries = {}, {}
    for c in idx["nmsf_catalog"]:
        nmsf_cases[c["name"]] = c

        def m(c=c):
            t = nmsf_read(c["nmsf"])
            return nmsf_wai_entry(t, nmsf_bind(t, binds), c["wai_name"])
        got, w = _code(m)
        if got != c["expect"] or (got == "ok" and w != c["wai"]):
            fails.append(f"nmsf catalog {c['name']}: {got}, want {c['expect']}")
        else:
            print(f"ok   nmsf catalog {c['name']}: {got}")
        if w is not None:
            wai_entries[c["name"]] = w
    tracks = {}
    for c in idx["wai_catalog"]:
        got, tp = _code(nmsf_read_wai, c["wai"])
        good = got == c["expect"] and (got != "ok" or (tp[0]["entries"] == c["nmsf"] and tp[1] == c["pin"]))
        (print(f"ok   nmsf wai catalog {c['name']}: {got}") if good else fails.append(f"nmsf wai catalog {c['name']}: {got}, want {c['expect']}"))
        if tp is not None:
            tracks[c["name"]] = tp
    for s in idx["rewrap"]:
        case = nmsf_cases[s["catalog"]]
        t = nmsf_read(case["nmsf"])
        rw, bad, pairs = NmsfRewrap(t, nmsf_bind(t, binds)), [], []
        for k, o in enumerate(s["objects"]):
            objs = [rd(r) for r in o["nmsf"]]
            got, exp = _code(rw.expect, o.get("object_id"), objs)
            if got != o["expect"]:
                bad.append(f"object {k}: {got}, want {o['expect']}")
                continue
            if got != "ok":
                continue
            env = rd(o["wai"])
            mm, payload = read_envelope(env)
            cond = mm.get("conditioning") or {}
            if (payload != exp[0] or mm.get("frame") != exp[1] or mm.get("media") != "video" or mm.get("intent") != "perceptual"
                    or (mm.get("model_requirement") or {}).get("capability") != t["cap"] or cond.get("kind") != t["kind"]
                    or cond.get("prior") != rw.b["prior"] or _content_hash(env).hex() != o["content_hash"]):
                bad.append(f"object {k}: the frame object is not the mapping of its NMSF objects")
            rw.commit(env)
            pairs.append((env, objs))
        # The frame objects are a stream a sink admits, and unwrap to the NMSF objects.
        st = Stream(t["cap"], False)
        uw = NmsfUnwrap(*nmsf_read_wai(wai_entries[s["catalog"]]))
        for k, (env, objs) in enumerate(pairs):
            got, _ = _code(st.push, env)
            if got != "ok":
                bad.append(f"frame object {k}: a sink gives {got}")
            got, back = _code(uw.push, env)
            if back != objs:
                bad.append(f"frame object {k}: unwraps to {got}, not its NMSF objects")
        for b in bad:
            fails.append(f"nmsf rewrap {s['name']}: {b}")
        if not bad:
            print(f"ok   nmsf rewrap {s['name']}: {len(s['objects'])} objects as expected, "
                  f"{len(pairs)} frame objects mapped and unwrapped back byte for byte")
    for s in idx["unwrap"]:
        uw, bad = NmsfUnwrap(*tracks[s["catalog"]]), []
        for k, o in enumerate(s["objects"]):
            got, back = _code(uw.push, rd(o["wai"]))
            if got != o["expect"] or (got == "ok" and back != [rd(r) for r in o["nmsf"]]):
                bad.append(f"object {k}: {got}, want {o['expect']}")
        for b in bad:
            fails.append(f"nmsf unwrap {s['name']}: {b}")
        if not bad:
            print(f"ok   nmsf unwrap {s['name']}: {len(s['objects'])} objects as expected")
    for s in idx["links"]:
        st, bad = Stream(s["capability"], False), []
        for k, o in enumerate(s["objects"]):
            got, _ = _code(st.push, rd(o["wai"]))
            if got != o["expect"]:
                bad.append(f"object {k}: {got}, want {o['expect']}")
        for b in bad:
            fails.append(f"nmsf links {s['name']}: {b}")
        if not bad:
            print(f"ok   nmsf links {s['name']}")
    text = json.dumps(idx)
    missing = [c for c in NMSF_CODES if f'"{c}"' not in text]
    if missing:
        fails.append(f"nmsf: no case gives {', '.join(missing)}")
    else:
        print(f"ok   nmsf: every one of the {len(NMSF_CODES)} codes is given by a case")
    return fails


# ---- conditional inter-frame video (integer-payloads.md §17) ----

IC = "int-payload-conformance/int-cond"
COND_CAP, COND_FRAME_CAP = "wai.video.int_cond", "wai.video.int_cond.frame"
I64_MAX, I64_MIN = (1 << 63) - 1, -(1 << 63)
I128_MAX = (1 << 127) - 1


def parse_wcf1(p):
    """§17.1, its checks in §2's order."""
    r = opening(p, b"WCF1", 32)
    kind, res = r.u8(), r.u16()
    dims = [r.u32() for _ in range(6)]
    if res:
        raise Refused("reserved_nonzero", "reserved")
    if kind > 1:
        raise Refused("field_range", "kind")
    if 0 in dims:
        raise Refused("field_range", "a dimension is zero")
    h, w, m, sh, sw, cf = dims
    too_large([h, w, 3], MAX_TENSOR)
    too_large([2, m, sh, sw], MAX_TENSOR)
    too_large([cf, sh, sw], MAX_TENSOR)
    rans = r.stream()
    r.end()
    return dict(kind=kind, header=tuple(dims), rans=rans)


def write_wcf1(kind, header, rans):
    return b"WCF1" + struct.pack("<BBH6I", 1, kind, 0, *header) + struct.pack("<I", len(rans)) + rans


def parse_wic1(p):
    """§17.2: the clip header, then each frame, in order."""
    r = opening(p, b"WIC1", 12)
    res = r.take(3)
    n = r.u32()
    if res != b"\0\0\0":
        raise Refused("reserved_nonzero", "reserved")
    if n == 0:
        raise Refused("field_range", "no frames")
    if n > MAX_FRAMES:
        raise Refused("too_large", "frames")
    frames, raw = [], []
    for k in range(n):
        b = r.take(r.u32())
        f = parse_wcf1(b)
        if (f["kind"] == 0) != (k == 0):
            raise Refused("field_range", "frame 0 intra, the rest inter")
        if k == 0:
            h, w = f["header"][:2]
            too_large([n, h, w, 3], MAX_CLIP)
        elif f["header"] != frames[0]["header"]:
            raise Refused("geometry_mismatch", "a frame's header is not frame 0's")
        frames.append(f)
        raw.append(b)
    r.end()
    return frames, raw


def write_wic1(frames):
    return b"WIC1" + struct.pack("<B3xI", 1, len(frames)) + b"".join(struct.pack("<I", len(f)) + f for f in frames)


class ModelMismatch(Exception):
    pass


def parse_whm1(b):
    """integer-payloads.md §9: a WHM1 network. Tags 0, 1, 3 and 4; an inverse GDN (tag
    2) is outside what this verifier decodes and refuses the set here."""
    def need(cond):
        if not cond:
            raise ModelMismatch("WHM1")
    need(b[:4] == b"WHM1" and len(b) >= 20)
    afrac, gf, shift, n = struct.unpack_from("<4I", b, 4)
    p, layers = 20, []
    for _ in range(n):
        need(p < len(b))
        tag = b[p]
        p += 1
        if tag in (0, 1):
            nf = 6 if tag == 0 else 5
            need(p + 4 * nf <= len(b))
            f = struct.unpack_from(f"<{nf}I", b, p)
            p += 4 * nf
            cin, cout, k = f[0], f[1], f[2]
            nw = cin * cout * k * k
            need(p + 2 * nw + 16 * cout <= len(b))
            wts = list(struct.unpack_from(f"<{nw}h", b, p))
            p += 2 * nw
            mult = list(struct.unpack_from(f"<{cout}q", b, p))
            bias = list(struct.unpack_from(f"<{cout}q", b, p + 8 * cout))
            p += 16 * cout
            L = dict(tag=tag, cin=cin, cout=cout, k=k, stride=f[3], pad=f[4], opad=f[5] if tag == 0 else 0, w=wts, mult=mult, bias=bias)
            layers.append(L)
        elif tag == 3:
            layers.append(dict(tag=3))
        elif tag == 4:
            need(p + 12 <= len(b))
            q, fr = struct.unpack_from("<qI", b, p)
            p += 12
            layers.append(dict(tag=4, q=q, f=fr))
        else:
            raise ModelMismatch("a layer this verifier does not decode")
    need(p == len(b))
    return dict(afrac=afrac, gf=gf, shift=shift, layers=layers)


def _int_json(text):
    """A JSON text read strictly: no member named twice, no fraction or exponent."""
    def pairs(kv):
        d = {}
        for k, v in kv:
            if k in d:
                raise ModelMismatch("a member named twice")
            d[k] = v
        return d

    def no_float(s):
        raise ModelMismatch("not an integer")

    def no_const(s):
        raise ModelMismatch("not JSON")
    try:
        return json.loads(text.decode("utf-8"), object_pairs_hook=pairs, parse_float=no_float, parse_constant=no_const)
    except (ValueError, UnicodeDecodeError) as e:
        raise ModelMismatch(str(e))


def parse_wcb1(b):
    """§17.3: four WHM1 networks and the tables; any rule broken is model_mismatch."""
    try:
        if b[:4] != b"WCB1":
            raise ModelMismatch("magic")
        p, parts = 4, []
        for _ in range(5):
            if p + 4 > len(b):
                raise ModelMismatch("a section's length")
            n = struct.unpack_from("<I", b, p)[0]
            if p + 4 + n > len(b):
                raise ModelMismatch("a section past the bundle")
            parts.append(b[p + 4:p + 4 + n])
            p += 4 + n
        if p != len(b):
            raise ModelMismatch("trailing")
        nets = [parse_whm1(x) for x in parts[:4]]
        afrac = nets[0]["afrac"]
        if any(n["afrac"] != afrac for n in nets) or not 1 <= afrac <= 30:
            raise ModelMismatch("afrac")
        t = _int_json(parts[4])
        isint = lambda v: type(v) is int
        edges, rows, lens, offs = t.get("scale_edges"), t.get("cdfs"), t.get("cdf_lengths"), t.get("offsets")
        if t.get("afrac") != afrac or not all(type(x) is list for x in (edges, rows, lens, offs)):
            raise ModelMismatch("tables")
        n = len(edges)
        if not 1 <= n <= 65536 or len(rows) != n or len(lens) != n or len(offs) != n:
            raise ModelMismatch("table counts")
        if not all(isint(e) and I64_MIN <= e <= I64_MAX for e in edges) or not all(isint(o) and I32_MIN <= o <= I32_MAX for o in offs):
            raise ModelMismatch("an edge or offset")
        cdfs = []
        for row, ln in zip(rows, lens):
            if type(row) is not list or not isint(ln) or not 0 <= ln <= len(row) or not all(isint(v) and I32_MIN <= v <= I32_MAX for v in row):
                raise ModelMismatch("a table")
            try:
                check_cdf(row[:ln])
            except Refused:
                raise ModelMismatch("a table is not a CDF")
            cdfs.append(row[:ln])
        return dict(nets=dict(zip(("ctx", "prior", "dec", "recon"), nets)), afrac=afrac, edges=edges, cdfs=cdfs, offsets=offs)
    except (ModelMismatch, struct.error) as e:
        raise Refused("model_mismatch", str(e))


def plan_net(net, c, h, w):
    """integer-payloads.md §7, *Planning the networks*, for the layers this verifier
    decodes: (out_c, out_h, out_w), or the refusal."""
    cap = MAX_TENSOR
    if 0 in (c, h, w):
        raise Refused("model_mismatch", "an empty input")
    if c * h * w > cap:
        raise Refused("too_large", "a network's input")
    for L in net["layers"]:
        t = L["tag"]
        if t in (0, 1):
            co, k, s, p = L["cout"], L["k"], L["stride"], L["pad"]
            if co == 0 or k == 0 or s == 0 or L["cin"] != c or net["shift"] > 127:
                raise Refused("model_mismatch", "a layer does not fit")
            if t == 0:
                fh, fw = (h - 1) * s + k + L["opad"], (w - 1) * s + k + L["opad"]
                if fh > cap or fw > cap or co * fh * fw > cap:
                    raise Refused("too_large", "a transposed convolution")
                if fh <= 2 * p or fw <= 2 * p:
                    raise Refused("model_mismatch", "no output")
                c, h, w = co, fh - 2 * p, fw - 2 * p
            else:
                hp, wp = h + 2 * p, w + 2 * p
                if hp > cap or wp > cap:
                    raise Refused("too_large", "a padded side")
                if hp < k or wp < k:
                    raise Refused("model_mismatch", "a kernel past its input")
                ho, wo = (hp - k) // s + 1, (wp - k) // s + 1
                if co * ho * wo > cap:
                    raise Refused("too_large", "a convolution's output")
                c, h, w = co, ho, wo
        elif t == 4 and L["f"] > 63:
            raise Refused("model_mismatch", "a leaky shift")
    return c, h, w


def _check(net, L, c, x):
    """whm1/1's check range on a layer's input (integer-operators.md §13.2): P(max|x|)."""
    X = max((max(abs(v) for v in pl) for pl in x), default=0)
    if L["tag"] in (0, 1):
        n = c * L["k"] * L["k"] * max(abs(v) for v in L["w"]) * X
        ok = n <= I64_MAX and ((n * max(abs(v) for v in L["mult"])) >> net["shift"]) + 1 + max(abs(v) for v in L["bias"]) <= I64_MAX
    else:
        ok = X * abs(L["q"]) <= I64_MAX
    if not ok:
        raise Refused("field_range", "a value outside a check range of the network")


def run_net(net, x, c, h, w, rgb8=False):
    """A WHM1 network's integers (integer-operators.md §13.2) on x, c planes of h*w."""
    shift = net["shift"]
    for L in net["layers"]:
        t = L["tag"]
        if t in (0, 1, 4):
            _check(net, L, c, x)
        if t == 1:
            co, k, s, p, W = L["cout"], L["k"], L["stride"], L["pad"], L["w"]
            ho, wo = (h + 2 * p - k) // s + 1, (w + 2 * p - k) // s + 1
            out = []
            for oc in range(co):
                acc = [0] * (ho * wo)
                for ic in range(c):
                    xp = x[ic]
                    for ky in range(k):
                        for kx in range(k):
                            wv = W[((oc * c + ic) * k + ky) * k + kx]
                            if wv == 0:
                                continue
                            for oy in range(ho):
                                iy = oy * s + ky - p
                                if iy < 0 or iy >= h:
                                    continue
                                lo = max(0, -(-(p - kx) // s))
                                hi = min(wo, (w - 1 + p - kx) // s + 1)
                                if lo >= hi:
                                    continue
                                a0 = iy * w + lo * s + kx - p
                                seg = xp[a0:a0 + (hi - lo - 1) * s + 1:s]
                                b0 = oy * wo
                                acc[b0 + lo:b0 + hi] = [u + wv * v for u, v in zip(acc[b0 + lo:b0 + hi], seg)]
                m, b = L["mult"][oc], L["bias"][oc]
                out.append([((a * m) >> shift) + b for a in acc])
            x, c, h, w = out, co, ho, wo
        elif t == 0:
            co, k, s, p, W = L["cout"], L["k"], L["stride"], L["pad"], L["w"]
            fh, fw = (h - 1) * s + k + L["opad"], (w - 1) * s + k + L["opad"]
            ho, wo = fh - 2 * p, fw - 2 * p
            out = []
            for oc in range(co):
                full = [0] * (fh * fw)
                for ic in range(c):
                    xp = x[ic]
                    for ky in range(k):
                        for kx in range(k):
                            wv = W[((ic * co + oc) * k + ky) * k + kx]
                            if wv == 0:
                                continue
                            for iy in range(h):
                                d0 = (iy * s + ky) * fw + kx
                                row = xp[iy * w:(iy + 1) * w]
                                full[d0:d0 + (w - 1) * s + 1:s] = [u + wv * v for u, v in zip(full[d0:d0 + (w - 1) * s + 1:s], row)]
                m, b = L["mult"][oc], L["bias"][oc]
                out.append([((full[(oy + p) * fw + ox + p] * m) >> shift) + b for oy in range(ho) for ox in range(wo)])
            x, c, h, w = out, co, ho, wo
        elif t == 3:
            x = [[v if v > 0 else 0 for v in pl] for pl in x]
        elif t == 4:
            q, f = L["q"], L["f"]
            x = [[v if v >= 0 else (v * q) >> f for v in pl] for pl in x]
    if rgb8:
        one, half = 1 << net["afrac"], 1 << (net["afrac"] - 1)
        x = [[(min(max(v, 0), one) * 255 + half) >> net["afrac"] for v in pl] for pl in x]
    return x, c, h, w


def cond_plan(model, hd):
    """§17.4 step 3: every network planned, then the shapes the frame declares."""
    h, w, m, sh, sw, cf = hd
    n = model["nets"]
    ctx = plan_net(n["ctx"], cf, sh, sw)
    prior = plan_net(n["prior"], ctx[0], sh, sw)
    dec = plan_net(n["dec"], m + ctx[0], sh, sw)
    recon = plan_net(n["recon"], cf, sh, sw)
    if ctx[1:] != (sh, sw) or prior != (2 * m, sh, sw) or dec != (cf, sh, sw):
        raise Refused("geometry_mismatch", "a network does not give the shape the frame declares")
    if recon[0] != 3:
        raise Refused("model_mismatch", "recon does not give three channels")
    if recon[1:] != (h, w):
        raise Refused("geometry_mismatch", "recon does not give h × w")


def cond_frame(model, hd, rans, ref):
    """§17.4 steps 5-9: one frame from its stream and the reference feature (None: zero)."""
    h, w, m, sh, sw, cf = hd
    n, af = model["nets"], model["afrac"]
    feat = ref if ref is not None else [[0] * (sh * sw) for _ in range(cf)]
    c, cc, _, _ = run_net(n["ctx"], feat, cf, sh, sw)
    pr, _, _, _ = run_net(n["prior"], c, cc, sh, sw)
    sigma, mu = [v for pl in pr[:m] for v in pl], [v for pl in pr[m:] for v in pl]
    d = Rans(rans)
    edges, cdfs, offs = model["edges"], model["cdfs"], model["offsets"]
    e0, last = edges[0], edges[:-1]
    y = []
    for s, u in zip(sigma, mu):
        sq = max(s, e0)
        b = len(edges) - 1 - sum(1 for e in last if sq <= e)
        r = d.decode(cdfs[b], offs[b])
        if d.exhausted:
            raise Refused("stream_exhausted", "the latent stream ran out")
        v = (r << af) + u
        if not I64_MIN <= v <= I64_MAX:
            raise Refused("field_range", "a latent value leaves i64")
        y.append(v)
    if d.i != len(d.w) or d.x != 1 << 31:
        raise Refused("stream_noncanonical", "the latent stream does not end canonically")
    plane = sh * sw
    x = [y[i * plane:(i + 1) * plane] for i in range(m)] + c
    f, fc, _, _ = run_net(n["dec"], x, m + cc, sh, sw)
    out, oc, oh, ow = run_net(n["recon"], f, fc, sh, sw, rgb8=True)
    return bytes(out[ch][i] for i in range(oh * ow) for ch in range(3)), f


def cond_clip(clip, set_bytes):
    frames, _ = parse_wic1(clip)
    if set_bytes is None:
        raise Refused("set_incomplete", "model.wcb")
    model = parse_wcb1(set_bytes)
    hd = frames[0]["header"]
    cond_plan(model, hd)
    out, ref = [], None
    for f in frames:
        rgb, ref = cond_frame(model, hd, f["rans"], ref)
        out.append(rgb)
    return out


class CondStream:
    """frame-stream.md §5 over wai.video.int_cond.frame, decoding every frame."""

    def __init__(self, model, set_sha):
        self.model, self.set_sha, self.last, self.state = model, set_sha, None, None

    def push(self, env):
        try:
            return self._admit(env)
        except Refused:
            self.last = self.state = None
            raise

    def _admit(self, env):
        m, payload = read_envelope(env)
        if (m.get("model_requirement") or {}).get("capability") != COND_FRAME_CAP:
            raise Refused("frame_capability_mismatch", "another capability")
        index, prev, number = read_link(m)
        f = parse_wcf1(payload)
        intra = f["kind"] == 0
        if intra and not (index == 0 and prev is None):
            raise Refused("frame_link_mismatch", "an intra frame with a link")
        if not intra and not (index >= 1 and prev is not None):
            raise Refused("frame_link_mismatch", "an inter frame without one")
        h = _content_hash(env).hex()
        prior = (m.get("conditioning") or {}).get("prior")
        if intra:
            self.last = self.state = None
            ref = None
        else:
            if self.last is None:
                raise Refused("frame_no_reference", "no reference")
            if prev != self.last["hash"]:
                raise Refused("frame_gap" if index > self.last["index"] + 1 else "frame_prev_mismatch", "prev")
            if index != self.last["index"] + 1:
                raise Refused("frame_index_mismatch", "index")
            if number != (None if self.last["number"] is None else self.last["number"] + 1):
                raise Refused("frame_number_mismatch", "number")
            if prior != self.last["prior"]:
                raise Refused("frame_prior_changed", "pin")
            if f["header"] != self.state["header"]:
                raise Refused("geometry_mismatch", "an inter frame's header is not its group's")
            ref = self.state["feature"]
        if not prior or prior.get("sha256") != self.set_sha:
            raise Refused("prior_unverified", "the corpus pins every frame to its set")
        cond_plan(self.model, f["header"])
        rgb, feat = cond_frame(self.model, f["header"], f["rans"], ref)
        self.state = dict(header=f["header"], feature=feat)
        self.last = dict(hash=h, index=index, number=number, prior=prior)
        return rgb


def int_cond_main(wai):
    def rd(rel):
        with open(os.path.join(wai, rel), "rb") as f:
            return f.read()

    fails = []
    idx = json.loads(rd(f"{IC}/index.json"))
    set_bytes = rd(idx["set"]["model.wcb"])
    set_sha = hashlib.sha256(set_bytes).hexdigest()
    if set_sha != idx["set_sha256"] or idx["prior"]["sha256"] != set_sha:
        fails.append("int-cond: the set's SHA-256 is not the index's or the pin's")
    model = parse_wcb1(set_bytes)
    clip = rd(idx["clip"]["payload"])
    side = json.loads(rd(idx["clip"]["sidecar"]))["clip"]
    whole = cond_clip(clip, set_bytes)
    digests = [hashlib.sha256(f).digest() for f in whole]
    if [d.hex() for d in digests] != side["frames_sha256"] or hashlib.sha256(b"".join(digests)).hexdigest() != side["clip_sha256"]:
        fails.append("int-cond clip: the independent decode differs from its sidecar")
    else:
        print(f"ok   int-cond clip: decoded independently to its sidecar's SHA-256 ({len(whole)} frames)")
    frames = [rd(f) for f in idx["clip"]["frames"]]
    if parse_wic1(clip)[1] != frames or write_wic1(frames) != clip:
        fails.append("int-cond clip: the split and the join are not the committed frames and clip")
    else:
        print(f"ok   int-cond clip: splits into its {len(frames)} frames and they join into it")
    for s in idx["streams"]:
        st, outs, bad = CondStream(model, set_sha), [], []
        for k, o in enumerate(s["objects"]):
            e = rd(o["envelope"])
            try:
                out, got = st.push(e), "ok"
            except Refused as r:
                out, got = None, r.code
            if got != o["expect"]:
                bad.append(f"object {k}: {got}, want {o['expect']}")
                continue
            if got == "ok":
                if _content_hash(e).hex() != o["content_hash"]:
                    bad.append(f"object {k}: content hash")
                if hashlib.sha256(out).hexdigest() != o.get("frame_sha256"):
                    bad.append(f"object {k}: frame SHA-256")
                outs.append(out)
        if s["name"] == "clip" and outs != whole:
            bad.append("its frames are not the clip's, byte for byte")
        fails += [f"int-cond stream {s['name']}: {b}" for b in bad]
        if not bad:
            print(f"ok   int-cond stream {s['name']}: {len(s['objects'])} objects as expected, decoded")
    for f in idx["refusals"]:
        name, cap, code, layer = f["name"], f["capability"], f["code"], f["layer"]
        b = rd(f["payload"])
        parse = parse_wic1 if cap == COND_CAP else parse_wcf1
        try:
            parse(b)
            got = None
        except Refused as e:
            got = e.code
        if layer == "decode":
            if got is not None:
                fails.append(f"int-cond {name}: refused at parse with {got}, want a decode-layer {code}")
                continue
            try:
                cond_clip(b, set_bytes)
            except Refused as e:
                got = e.code
        if got != code:
            fails.append(f"int-cond {name}: {got}, want {code} at {layer}")
        else:
            print(f"ok   int-cond {name}: {code} ({layer})")
    for sc in idx["sets"]:
        path = sc["set"].get("model.wcb")
        try:
            cond_clip(clip, rd(path) if path else None)
            got = None
        except Refused as e:
            got = e.code
        (print(f"ok   int-cond set {sc['name']}: {got}") if got == sc["code"] else fails.append(f"int-cond set {sc['name']}: {got}, want {sc['code']}"))
    for j in idx["join_refusals"]:
        try:
            parse_wic1(write_wic1([rd(x) for x in j["frames"]]))
            got = None
        except Refused as e:
            got = e.code
        (print(f"ok   int-cond join {j['name']}: {got}") if got == j["code"] else fails.append(f"int-cond join {j['name']}: {got}, want {j['code']}"))
    return fails


if __name__ == "__main__":
    args = sys.argv[1:]
    only = "--frame-stream" in args
    cond = "--int-cond" in args
    args = [a for a in args if a not in ("--frame-stream", "--int-cond")]
    sys.exit(main(args[0] if args else os.path.dirname(os.path.abspath(__file__)), frame_stream_only=only, int_cond_only=cond))