"""Reproducible end-to-end validation harness for the elic + mlicpp decoders.

elic (`wai.neural.elic`) and mlicpp (`wai.neural.mlicpp`) have no pretrained
model in the CompressAI zoo, so we can't cross-check them against a reference
codec the way bmshj2018 / mbt2018_mean are checked. Instead we build *minimal
CompressAI-compatible* models (real EntropyBottleneck + GaussianConditional,
random-but-non-trivial weights), run them through the SAME export + encode logic
the real tools use, and emit the `.wai` envelope + a reference reconstruction.
The Rust tests `elic_matches_reference` / `mlicpp_matches_reference` then decode
the envelope and must match the reference within ORT-version g_s rounding.

This makes those two validations reproducible from the repo (the artifacts the
Rust tests consume otherwise live only in /tmp). Run before the gated tests:

  python tools/wai_validate_synthetic.py            # writes /tmp artifacts
  cargo test --features neural --test neural_native_decode _matches_reference

Needs torch + compressai + onnxruntime + numpy + Pillow.
"""
from __future__ import annotations

import argparse
import json
import math
import struct
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import onnxruntime as ort
from PIL import Image
from compressai.models import CompressionModel
from compressai.entropy_models import GaussianConditional, EntropyBottleneck
from compressai.ans import BufferedRansEncoder


def scale_table():
    return torch.exp(torch.linspace(math.log(0.11), math.log(256), 64))


def _convs(M, N):
    dn = lambda ci, co: nn.Conv2d(ci, co, 5, 2, 2)
    up = lambda ci, co: nn.ConvTranspose2d(ci, co, 5, 2, 2, output_padding=1)
    g_a = nn.Sequential(dn(3, M), nn.ReLU(), dn(M, M), nn.ReLU(), dn(M, M), nn.ReLU(), dn(M, M))
    g_s = nn.Sequential(up(M, M), nn.ReLU(), up(M, M), nn.ReLU(), up(M, M), nn.ReLU(), up(M, 3))
    h_a = nn.Sequential(dn(M, N), nn.ReLU(), dn(N, N))
    h_s = nn.Sequential(up(N, M), nn.ReLU(), up(M, 2 * M))
    return g_a, g_s, h_a, h_s


class MiniElic(CompressionModel):
    def __init__(self, M=48, N=48):
        super().__init__()
        self.entropy_bottleneck = EntropyBottleneck(N)
        self.gaussian_conditional = GaussianConditional(None)
        self.g_a, self.g_s, self.h_a, self.h_s = _convs(M, N)
        self.context = nn.Conv2d(M, 2 * M, 5, 1, 2)


class MiniMlic(CompressionModel):
    def __init__(self, M=48, N=48):
        super().__init__()
        self.entropy_bottleneck = EntropyBottleneck(N)
        self.gaussian_conditional = GaussianConditional(None)
        self.g_a, self.g_s, self.h_a, self.h_s = _convs(M, N)
        self.entropy = nn.Conv2d(3 * M, 2 * M, 5, 1, 2)


def build(model_cls, seed):
    torch.manual_seed(seed)
    net = model_cls().eval()
    for p in net.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p, gain=3.0)   # gain 3 → real latent entropy
    if hasattr(net, "gaussian_conditional"):
        net.gaussian_conditional.update_scale_table(scale_table())
    net.update(force=True)
    return net


def fblob(eb):
    cl = eb._cdf_length.int().tolist()
    return {"channels": len(cl), "cdf_length": cl, "offset": eb._offset.int().tolist(),
            "medians": eb._get_medians().detach().reshape(-1).tolist(),
            "cdf": [r[:cl[i]] for i, r in enumerate(eb._quantized_cdf.int().tolist())]}


def gblob(g):
    cl = g._cdf_length.int().tolist()
    return {"scale_table": g.scale_table.reshape(-1).tolist(), "cdf_length": cl,
            "offset": g._offset.int().tolist(),
            "cdf": [r[:cl[i]] for i, r in enumerate(g._quantized_cdf.int().tolist())]}


def export_onnx(mod, dummy, path):
    torch.onnx.export(mod, (dummy,), str(path), input_names=["i"], output_names=["o"],
                      dynamic_axes={"i": {2: "h", 3: "w"}, "o": {2: "h", 3: "w"}},
                      opset_version=17, dynamo=False)


def test_image(path, size):
    if path.exists():
        return
    x, y = np.meshgrid(np.linspace(0, 1, size), np.linspace(0, 1, size))
    r = (x * 255).astype(np.uint8); g = (y * 255).astype(np.uint8)
    b = (((np.sin(x * 12) * np.cos(y * 12)) * 0.5 + 0.5) * 255).astype(np.uint8)
    img = np.stack([r, g, b], -1)
    yy, xx = np.ogrid[:size, :size]
    img[(yy - size // 2) ** 2 + (xx - size // 2) ** 2 < (size // 5) ** 2] = [200, 40, 120]
    Image.fromarray(img).save(path)


def write_wai(path, cap, kind, header, z, y, H, W):
    manifest = {"wai": "1.0", "media": "image", "intent": "replicate",
                "model_requirement": {"capability": cap}, "conditioning": {"kind": kind},
                "target": {"w": W, "h": H}}
    mb = json.dumps(manifest, separators=(",", ":")).encode()
    payload = header + z + y
    path.write_bytes(b"WAI1" + struct.pack("<I", len(mb)) + mb + struct.pack("<I", len(payload)) + payload)


def _session(path):
    """Open an ONNX model exactly as the Rust decoder does.

    `wai-rs/src/ort_session.rs::load_deterministic` evaluates every model with one
    intra-op thread, one inter-op thread and full graph optimisation, because float
    ONNX inference is only reproducible that way. The fixtures must be ENCODED under
    the same settings. With onnxruntime's defaults (every core) a convolution's
    reduction is split across threads, the summation order changes and the float
    output shifts. Where a scale then straddles a threshold of the Gaussian scale
    table, encoder and decoder select adjacent CDFs, and the arithmetic-coded stream
    desyncs from that symbol onward.

    Measured, not hypothetical: under default sessions the MLIC++ fixture diverged at
    symbol 416 of 12,288 -- encoded with scale 0.78787208 against CDF 16, decoded
    with 0.78788924 against CDF 17, either side of scale_table[16] = 0.78788900 --
    and decoded to max|delta|=255. Under these settings it decodes exactly, and
    DCVC's residual max|delta|=1 goes to 0.
    """
    so = ort.SessionOptions()
    so.intra_op_num_threads = 1
    so.inter_op_num_threads = 1
    so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    return ort.InferenceSession(str(path), sess_options=so, providers=["CPUExecutionProvider"])


def onx(sess, t):
    return torch.from_numpy(sess.run(None, {sess.get_inputs()[0].name: t.numpy().astype(np.float32)})[0])


def roundtrip_elic(net, x, mdir, size):
    mdir.mkdir(parents=True, exist_ok=True)
    gc = net.gaussian_conditional
    export_onnx(net.h_s, torch.zeros(1, 48, 4, 4), mdir / "hyper_synthesis.onnx")
    export_onnx(net.context, torch.zeros(1, 48, 16, 16), mdir / "context.onnx")
    export_onnx(net.g_s, torch.zeros(1, 48, 16, 16), mdir / "synthesis.onnx")
    (mdir / "hyper_cdfs_q3.json").write_text(json.dumps(fblob(net.entropy_bottleneck), separators=(",", ":")))
    (mdir / "gauss_cdfs_q3.json").write_text(json.dumps(gblob(gc), separators=(",", ":")))
    hs = _session(mdir / "hyper_synthesis.onnx")
    cx = _session(mdir / "context.onnx")
    gs = _session(mdir / "synthesis.onnx")
    with torch.no_grad():
        y = net.g_a(x); z = net.h_a(y)
        zs = net.entropy_bottleneck.compress(z); z_hat = net.entropy_bottleneck.decompress(zs, z.size()[-2:])
        base = onx(hs, z_hat); sb, mb = base.chunk(2, 1)
        _, Cy, Sy, _ = y.shape
        yy, xx = torch.meshgrid(torch.arange(Sy), torch.arange(Sy), indexing="ij")
        anc = ((yy + xx) % 2 == 0).numpy()
        syb = torch.round(y - mb).to(torch.int32)
        ya = ((syb.float() + mb) * torch.from_numpy(anc)).to(y.dtype)
        ctx = onx(cx, ya); sc, mc = ctx.chunk(2, 1); syc = torch.round(y - mc).to(torch.int32)
        ib = gc.build_indexes(sb)[0]; ic = gc.build_indexes(sc)[0]; syb, syc = syb[0], syc[0]
        syms, idxs = [], []
        for isa in (True, False):
            st, it = (syb, ib) if isa else (syc, ic)
            for ch in range(Cy):
                for r in range(Sy):
                    for c in range(Sy):
                        if bool(anc[r, c]) == isa:
                            syms.append(int(st[ch, r, c])); idxs.append(int(it[ch, r, c]))
        enc = BufferedRansEncoder()
        enc.encode_with_indexes(syms, idxs, gc._quantized_cdf.tolist(), gc._cdf_length.tolist(), gc._offset.tolist())
        yb = enc.flush()
        yhat = (syb.float() + mb) * torch.from_numpy(anc) + (syc.float() + mc) * torch.from_numpy(~anc)
        img = onx(gs, yhat if yhat.dim() == 4 else yhat.unsqueeze(0))
    _, Cz, Sz, _ = z.shape
    hdr = struct.pack("<IIIIHHBBH", size, size, len(zs[0]), len(yb), Cz, Cy, 2, 4, 3)
    write_wai(Path("/tmp/test.elic.wai"), "wai.neural.elic", "elic", hdr, zs[0], yb, size, size)
    (np.clip(img[0].numpy(), 0, 1).transpose(1, 2, 0) * 255).round().astype(np.uint8).tofile("/tmp/ref_elic.rgb")
    return len(yb), len(syms)


def roundtrip_mlicpp(net, x, mdir, size, K=4):
    mdir.mkdir(parents=True, exist_ok=True)
    gc = net.gaussian_conditional
    export_onnx(net.h_s, torch.zeros(1, 48, 4, 4), mdir / "hyper_synthesis.onnx")
    export_onnx(net.entropy, torch.zeros(1, 144, 16, 16), mdir / "entropy.onnx")
    export_onnx(net.g_s, torch.zeros(1, 48, 16, 16), mdir / "synthesis.onnx")
    (mdir / "hyper_cdfs_q3.json").write_text(json.dumps(fblob(net.entropy_bottleneck), separators=(",", ":")))
    (mdir / "gauss_cdfs_q3.json").write_text(json.dumps(gblob(gc), separators=(",", ":")))
    hs = _session(mdir / "hyper_synthesis.onnx")
    en = _session(mdir / "entropy.onnx")
    gs = _session(mdir / "synthesis.onnx")
    with torch.no_grad():
        y = net.g_a(x); z = net.h_a(y)
        zs = net.entropy_bottleneck.compress(z); z_hat = net.entropy_bottleneck.decompress(zs, z.size()[-2:])
        hyper = onx(hs, z_hat); _, Cy, Sy, _ = y.shape; sw = Cy // K
        yy, xx = torch.meshgrid(torch.arange(Sy), torch.arange(Sy), indexing="ij")
        anc = ((yy + xx) % 2 == 0).numpy()
        y_dec = torch.zeros(1, Cy, Sy, Sy); y0 = y[0]; syms, idxs = [], []
        for k in range(K):
            c0, c1 = k * sw, (k + 1) * sw
            for isa in (True, False):
                params = onx(en, torch.cat([y_dec, hyper], 1)); sc, mc = params.chunk(2, 1)
                idx = gc.build_indexes(sc)[0]; m0 = mc[0]
                for ch in range(c0, c1):
                    for r in range(Sy):
                        for c in range(Sy):
                            if bool(anc[r, c]) == isa:
                                sym = int(round(float(y0[ch, r, c] - m0[ch, r, c])))
                                syms.append(sym); idxs.append(int(idx[ch, r, c]))
                                y_dec[0, ch, r, c] = sym + m0[ch, r, c]
        enc = BufferedRansEncoder()
        enc.encode_with_indexes(syms, idxs, gc._quantized_cdf.tolist(), gc._cdf_length.tolist(), gc._offset.tolist())
        yb = enc.flush(); img = onx(gs, y_dec)
    _, Cz, Sz, _ = z.shape
    hdr = struct.pack("<IIIIHHBBHH", size, size, len(zs[0]), len(yb), Cz, Cy, 2, 4, 3, K)
    write_wai(Path("/tmp/test.mlic.wai"), "wai.neural.mlicpp", "mlicpp", hdr, zs[0], yb, size, size)
    (np.clip(img[0].numpy(), 0, 1).transpose(1, 2, 0) * 255).round().astype(np.uint8).tofile("/tmp/ref_mlic.rgb")
    return len(yb), len(syms)


class SynthSnac(nn.Module):
    """SNAC's multi-rate-sum decode interface (3 levels, ×4/×2/×1 → finest),
    with clean ONNX ops only — the real SNAC decoder's noise/attention blocks
    don't export (a sink-side concern). Validates snac.rs's multi-level wire
    parse + dynamic codes_0..codes_{N-1} dispatch."""
    HOP = 512
    def forward(self, c0, c1, c2):
        f0 = c0.float().repeat_interleave(4, dim=1)
        f1 = c1.float().repeat_interleave(2, dim=1)
        z = (f0 + f1 + c2.float()) / 12000.0
        return z.unsqueeze(1).repeat_interleave(self.HOP, dim=2)


def roundtrip_snac(mdir: Path):
    mdir.mkdir(parents=True, exist_ok=True)
    sr, lengths = 24_000, [12, 24, 48]
    nsamp = lengths[-1] * SynthSnac.HOP
    rng = np.random.RandomState(7)
    codes = [torch.from_numpy(rng.randint(0, 4096, size=(1, L)).astype(np.int64)) for L in lengths]
    names = [f"codes_{i}" for i in range(3)]
    torch.onnx.export(SynthSnac().eval(), tuple(codes), str(mdir / "decoder.onnx"),
                      input_names=names, output_names=["audio_values"],
                      dynamic_axes={**{f"codes_{i}": {1: f"t{i}"} for i in range(3)},
                                    "audio_values": {2: "s"}}, opset_version=17, dynamo=False)
    sess = _session(mdir / "decoder.onnx")
    ref = sess.run(None, {f"codes_{i}": codes[i].numpy() for i in range(3)})[0]
    hdr = struct.pack("<III", sr, nsamp, 3) + b"".join(struct.pack("<I", L) for L in lengths)
    body = b"".join(codes[i][0].numpy().astype(np.uint16).tobytes() for i in range(3))
    man = (b'{"wai":"1.0","media":"audio","intent":"replicate","model_requirement":'
           b'{"capability":"wai.neural.snac"},"conditioning":{"kind":"snac"}}')
    payload = hdr + body
    Path("/tmp/test.snac.wai").write_bytes(
        b"WAI1" + struct.pack("<I", len(man)) + man + struct.pack("<I", len(payload)) + payload)
    ref.astype(np.float32).tofile("/tmp/ref_snac.f32")
    return nsamp


class MiniDCVC(CompressionModel):
    """Hyperprior entropy + a temporal decoder (latents, ref_feature) →
    (image, feature). Validates the inter-frame frame-loop orchestration."""
    def __init__(self, M=48, N=48, Cf=48):
        super().__init__()
        self.entropy_bottleneck = EntropyBottleneck(N)
        self.gaussian_conditional = GaussianConditional(None)
        self.g_a, self.g_s, self.h_a, self.h_s = _convs(M, N)
        self.t_feat = nn.Conv2d(M + Cf, Cf, 3, 1, 1)
        up = lambda ci, co: nn.ConvTranspose2d(ci, co, 5, 2, 2, output_padding=1)
        self.t_img = nn.Sequential(up(M + Cf, M), nn.ReLU(), up(M, M), nn.ReLU(),
                                   up(M, M), nn.ReLU(), up(M, 3))


class _TDec(nn.Module):
    def __init__(s, m): super().__init__(); s.m = m
    def forward(s, latents, ref_feature):
        x = torch.cat([latents, ref_feature], 1)
        return s.m.t_img(x), torch.relu(s.m.t_feat(x))


def roundtrip_dcvc(net, x, mdir: Path, size, Cf=48):
    mdir.mkdir(parents=True, exist_ok=True)
    gc = net.gaussian_conditional
    export_onnx(net.h_s, torch.zeros(1, 48, 4, 4), mdir / "hyper_synthesis.onnx")
    torch.onnx.export(_TDec(net).eval(), (torch.zeros(1, 48, 16, 16), torch.zeros(1, Cf, 16, 16)),
                      str(mdir / "temporal_decoder.onnx"), input_names=["latents", "ref_feature"],
                      output_names=["image", "feature"],
                      dynamic_axes={"latents": {2: "h", 3: "w"}, "ref_feature": {2: "h", 3: "w"},
                                    "image": {2: "h", 3: "w"}, "feature": {2: "h", 3: "w"}},
                      opset_version=17, dynamo=False)
    (mdir / "hyper_cdfs_q3.json").write_text(json.dumps(fblob(net.entropy_bottleneck), separators=(",", ":")))
    (mdir / "gauss_cdfs_q3.json").write_text(json.dumps(gblob(gc), separators=(",", ":")))
    hs = _session(mdir / "hyper_synthesis.onnx")
    gs = _session(mdir / "temporal_decoder.onnx")
    frames = [np.roll(np.asarray(x[0].permute(1, 2, 0)), shift=8 * k, axis=1) for k in range(3)]
    F = torch.zeros(1, Cf, 16, 16); fcodes = []; ref = []
    for fr in frames:
        xi = torch.from_numpy(np.ascontiguousarray(fr)).permute(2, 0, 1).unsqueeze(0)
        with torch.no_grad():
            y = net.g_a(xi); z = net.h_a(y); zs = net.entropy_bottleneck.compress(z)
            z_hat = net.entropy_bottleneck.decompress(zs, z.size()[-2:])
            params = onx(hs, z_hat); sc, mn = params.chunk(2, 1)
            ys = gc.compress(y, gc.build_indexes(sc), means=mn); y_hat = gc.decompress(ys, gc.build_indexes(sc), means=mn)
            img, F = onx(gs, [y_hat, F]) if False else [torch.from_numpy(t) for t in gs.run(None, {"latents": y_hat.numpy(), "ref_feature": F.numpy()})]
        fcodes.append((zs[0], ys[0])); ref.append((np.clip(img[0].numpy(), 0, 1).transpose(1, 2, 0) * 255).round().astype(np.uint8))
    _, Cz, Sz, _ = z.shape; _, Cy, Sy, _ = y.shape
    hdr = struct.pack("<IIIIHHBBHH", size, size, 3, 30000, Cy, Cz, 4, 2, 3, Cf)
    table = b"".join(struct.pack("<II", len(zb), len(yb)) for zb, yb in fcodes)
    body = b"".join(zb + yb for zb, yb in fcodes)
    write_wai(Path("/tmp/test.dcvc.wai"), "wai.neural.dcvc_rt", "dcvc_rt", hdr, b"", table + body, size, size)
    np.concatenate([r.reshape(-1) for r in ref]).astype(np.uint8).tofile("/tmp/ref_dcvc.rgb")
    return 3


from compressai.entropy_models.entropy_models import pmf_to_quantized_cdf  # noqa: E402
_RL = 1 << 31


def _rans_cat_encode(symbols, cdfs):
    """Bypass-free categorical rANS matching wai-rs decode_categorical."""
    state, out = _RL, []
    for sym, cdf in zip(reversed(symbols), reversed(cdfs)):
        start, freq = cdf[sym], cdf[sym + 1] - cdf[sym]
        x_max = ((_RL >> 16) << 32) * freq
        while state >= x_max:
            out.append(state & 0xFFFFFFFF); state >>= 32
        state = ((state // freq) << 16) + (state % freq) + start
    words = [state & 0xFFFFFFFF, (state >> 32) & 0xFFFFFFFF] + list(reversed(out))
    return b"".join(struct.pack("<I", w) for w in words)


class MiniGLC(CompressionModel):
    """Categorical/VQ entropy: hyper + checkerboard context predict per-position
    logits over a K-codebook; synthesis maps the index grid → image. Logits are
    ROUNDED in-graph so encoder/decoder ORT builds agree (categorical CDFs are
    razor-sensitive — the determinism contract, sharper than Gaussian)."""
    def __init__(self, K=16, Cz=32):
        super().__init__()
        self.entropy_bottleneck = EntropyBottleneck(Cz)
        up = lambda ci, co: nn.ConvTranspose2d(ci, co, 5, 2, 2, output_padding=1)
        self.hyper_entropy = nn.Sequential(up(Cz, 64), nn.ReLU(), up(64, K))
        self.context_entropy = nn.Conv2d(1, K, 5, 1, 2)
        self.synthesis = nn.Sequential(up(1, 32), nn.ReLU(), up(32, 32), nn.ReLU(),
                                       up(32, 32), nn.ReLU(), up(32, 3))


class _Round(nn.Module):
    def __init__(s, m): super().__init__(); s.m = m
    def forward(s, x): return torch.round(s.m(x))


def roundtrip_glc(net, mdir: Path, size, K=16, Cz=32, Sz=4, Sy=16):
    mdir.mkdir(parents=True, exist_ok=True)
    torch.onnx.export(_Round(net.hyper_entropy), (torch.zeros(1, Cz, Sz, Sz),), str(mdir / "hyper_entropy.onnx"),
                      input_names=["i"], output_names=["o"], dynamic_axes={"i": {2: "h", 3: "w"}, "o": {2: "h", 3: "w"}}, opset_version=17, dynamo=False)
    torch.onnx.export(_Round(net.context_entropy), (torch.zeros(1, 1, Sy, Sy),), str(mdir / "context_entropy.onnx"),
                      input_names=["i"], output_names=["o"], dynamic_axes={"i": {2: "h", 3: "w"}, "o": {2: "h", 3: "w"}}, opset_version=17, dynamo=False)
    export_onnx(net.synthesis, torch.zeros(1, 1, Sy, Sy), mdir / "synthesis.onnx")
    (mdir / "hyper_cdfs_q3.json").write_text(json.dumps(fblob(net.entropy_bottleneck), separators=(",", ":")))
    he = _session(mdir / "hyper_entropy.onnx"); ce = _session(mdir / "context_entropy.onnx")
    gs = _session(mdir / "synthesis.onnx")
    rng = np.random.RandomState(7)
    sm = lambda l: (lambda e: e / e.sum())(np.exp(l - l.max()))
    cdf_of = lambda l: pmf_to_quantized_cdf(torch.tensor(sm(l), dtype=torch.float32), 16).tolist()
    z = torch.randn(1, Cz, Sz, Sz); zs = net.entropy_bottleneck.compress(z); z_hat = net.entropy_bottleneck.decompress(zs, (Sz, Sz))
    base = he.run(None, {"i": z_hat.numpy()})[0][0]
    idx = rng.randint(0, K, size=(Sy, Sy)).astype(np.int64)
    anc = np.array([[(y + x) % 2 == 0 for x in range(Sy)] for y in range(Sy)])
    ctx = ce.run(None, {"i": np.where(anc, idx, 0).astype(np.float32)[None, None]})[0][0]
    syms, cdfs = [], []
    for isa in (True, False):
        for y in range(Sy):
            for x in range(Sy):
                if anc[y, x] == isa:
                    lg = base[:, y, x] if isa else base[:, y, x] + ctx[:, y, x]
                    syms.append(int(idx[y, x])); cdfs.append(cdf_of(lg))
    yb = _rans_cat_encode(syms, cdfs)
    ref = gs.run(None, {"i": idx.astype(np.float32)[None, None]})[0]
    (np.clip(ref[0], 0, 1).transpose(1, 2, 0) * 255).round().astype(np.uint8).tofile("/tmp/ref_glc.rgb")
    hdr = struct.pack("<IIIIHHBBH", size, size, len(zs[0]), len(yb), Cz, K, 2, 4, 3)
    write_wai(Path("/tmp/test.glc.wai"), "wai.neural.glc", "glc", hdr, zs[0], yb, size, size)
    return len(syms)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--size", type=int, default=256)
    args = ap.parse_args()
    img_path = Path("/tmp/wai_test.png")
    test_image(img_path, args.size)
    img = Image.open(img_path).convert("RGB").resize((args.size, args.size), Image.LANCZOS)
    x = torch.from_numpy(np.asarray(img, np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0)

    # Each harness is independent; a failure in one must not block the others
    # (e.g. a synthetic model that only traces at certain sizes). Report a
    # per-decoder PASS/FAIL summary, and exit non-zero unless EVERY harness succeeded.
    harnesses = [
        ("elic", lambda: roundtrip_elic(build(MiniElic, 0), x, Path("/tmp/wai_elic"), args.size)),
        ("mlicpp", lambda: roundtrip_mlicpp(build(MiniMlic, 1), x, Path("/tmp/wai_mlic"), args.size)),
        ("snac", lambda: roundtrip_snac(Path("/tmp/wai_snac"))),
        ("dcvc", lambda: roundtrip_dcvc(build(MiniDCVC, 2), x, Path("/tmp/wai_dcvc"), args.size)),
        ("glc", lambda: roundtrip_glc(build(MiniGLC, 3), Path("/tmp/wai_glc"), args.size)),
    ]
    ok = 0
    for name, fn in harnesses:
        try:
            fn()
            print(f"  ✓ {name}")
            ok += 1
        except Exception as e:  # noqa: BLE001 — report and continue
            print(f"  ✗ {name}: {type(e).__name__}: {e}")
    print(f"generated {ok}/{len(harnesses)} synthetic neural fixtures into /tmp")
    print("now: cargo test --features neural --test neural_native_decode -- _matches_reference")
    # Every harness must succeed, not merely one of them.
    #
    # `0 if ok else 1` reported success whenever a single harness worked, and the
    # loop above deliberately catches per-harness exceptions to keep going. So a
    # 1-of-5 run exited 0, the CI step went green, and the four decoders whose
    # fixtures were never written then SKIPPED in the Rust suite — where a skip is
    # reported as a pass. The job came out fully green having checked one decoder
    # out of five, and nothing in the output said so.
    #
    # Partial generation is still useful interactively: the per-harness ✓/✗ lines
    # are printed either way, and a caller who wants one fixture can ignore the
    # exit code. What is not acceptable is CI reading a partial run as a pass.
    return 0 if ok == len(harnesses) else 1


if __name__ == "__main__":
    raise SystemExit(main())