#!/usr/bin/env python3
"""What exactness costs `wai.neural.int_hyper`: the measurement behind
research/determinism-cost.md and its record, research/determinism-cost/record.json.

The float codec is the pretrained mean-scale hyperprior `mbt2018_mean` (MSE, qualities
1-8) as its own library codes and decodes it. The integer codec is the same model exported
the way tools/wai_sota_int_synth_export.py exports quality 6 (per-channel int12 weights,
activations at 2^12, integer IGDN, the integer hyper-decoder deciding every scale bucket),
coded as a `WIH1` payload (extensions/integer-payloads.md §4) and decoded by the reference
decoder (`wai_det_cost`, which runs `int_payload::decode_int`). Every step writes files
under one work directory, so each can be rerun alone and on another machine.

    wai_determinism_cost.py select  <work> <reference-root>   pick and check the pictures
    wai_determinism_cost.py export  <work> [q ...]            pinned checkpoints -> integer sets
    wai_determinism_cost.py encode  <work> [q ...]            both codecs; the float decode on this CPU
    wai_determinism_cost.py jobs    <work> <out-tag>          the integer decode's job list
    wai_determinism_cost.py fdecode <work> <tag> [--device cpu|mps|cuda] [--threads N]
                                                              decode every float stream again
    wai_determinism_cost.py fserve  <work> <q> [--device D] [--threads N] [--ids a,b]
                                                              the timed float child of wai_det_energy.py
    wai_determinism_cost.py metrics <work>                    distortion of every decode
    wai_determinism_cost.py bd      <work>                    BD-rate, integer against float
    wai_determinism_cost.py record  <work> <record.json>      the record: raw readings and their summary

Needs the pinned stack of tools/requirements-sota-int.txt for `export`, `encode` and
`fdecode`; `select`, `jobs`, `metrics` and `bd` need numpy and Pillow only.
"""
import hashlib
import importlib.util
import io
import json
import math
import os
import struct
import sys

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
QUALITIES = list(range(1, 9))

# The published mbt2018_mean (MSE) checkpoints, by quality: the file the library's zoo
# names, and the SHA-256 of its bytes, measured when this record was made. A checkpoint
# with other bytes is refused before it is loaded. Quality 6 is the one
# wai_sota_int_synth_export.py pins, and its digest is the same.
CKPT_SHA256 = {
    1: "e522738dc4599d1538fcdb16e21aa56df9146d60e8e083542466b387c1960156",
    2: "e54a039d587afddd086b836cdd8eed4e0343d7cdd91028489ad0f3e7a72d307d",
    3: "723404a8a039371abe8a14f05160ee1614cf5e0a3fb2a4e5f7ba569b68c22601",
    4: "6dba02a30c78ff78baafd2072dd78646c27af6c4dc8416097c2c2858c1f186ce",
    5: "d504e8ebe7e6396efd0b5fd093a32136714fd7e4be03db389e6f07aee24fdb39",
    6: "a19628aba19732125e02caf5c38e233edbe2a641145a5b1779c51f4527f51dd6",
    7: "d5d441d16ba3e5a80ead947e35d2dddc12fb0d9785a8006e9b35a0cef7668d91",
    8: "8089ae3e23d9df5b7c68107c82c31da21019147aedddb03f03350de58d5385b7",
}

# The pictures: the reference manifest's PNG entries that it marks `self_test`, each one
# picture of 768x512 or 512x768 (wai/tools/wai_visual_provenance_check.py's manifest).
MANIFEST = "wai-image-references.json"
MANIFEST_SHA256 = "db33c6633f89a66d6742a37cd348208792f6d49810da739ddf26d10dc6554b66"


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


def jdump(obj, path):
    with open(path, "w") as f:
        json.dump(obj, f, indent=1, sort_keys=True)
        f.write("\n")


def jload(path):
    with open(path) as f:
        return json.load(f)


def export_module():
    spec = importlib.util.spec_from_file_location("sota_export", os.path.join(HERE, "wai_sota_int_synth_export.py"))
    m = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(m)
    return m


# ---------------------------------------------------------------- select

def cmd_select(work, root):
    raw = open(os.path.join(root, MANIFEST), "rb").read()
    if sha256(raw) != MANIFEST_SHA256:
        sys.exit(f"the manifest's SHA-256 is {sha256(raw)}, not the pinned {MANIFEST_SHA256}")
    man = json.loads(raw)
    from PIL import Image
    pics = []
    for e in man["entries"]:
        if not (e.get("self_test") and e["path"].endswith(".png") and len(e["images"]) == 1):
            continue
        im = e["images"][0]
        if (im["width"], im["height"]) not in ((768, 512), (512, 768)):
            continue
        b = open(os.path.join(root, e["path"]), "rb").read()
        if sha256(b) != e["sha256"]:
            sys.exit(f"an entry's bytes do not match the manifest ({e['sha256'][:16]})")
        px = np.array(Image.open(io.BytesIO(b)).convert("RGB"))
        h, w = px.shape[:2]
        assert h % 64 == 0 and w % 64 == 0
        pid = e["sha256"][:16]
        os.makedirs(os.path.join(work, "src"), exist_ok=True)
        px.tofile(os.path.join(work, "src", f"{pid}.rgb"))
        pics.append({"id": pid, "file_sha256": e["sha256"], "width": int(w), "height": int(h),
                     "rgb_sha256": sha256(px.tobytes())})
    jdump({"manifest_sha256": MANIFEST_SHA256,
           "rule": "PNG entries marked self_test with one picture of 768x512 or 512x768",
           "pictures": pics}, os.path.join(work, "pictures.json"))
    print(f"{len(pics)} pictures")


def pictures(work):
    return jload(os.path.join(work, "pictures.json"))["pictures"]


def source_rgb(work, p):
    return np.fromfile(os.path.join(work, "src", f"{p['id']}.rgb"), dtype=np.uint8).reshape(p["height"], p["width"], 3)


# ---------------------------------------------------------------- models

def load_model(q, allow_new=False):
    """mbt2018_mean (MSE) at quality q from the published checkpoint, its bytes checked
    against CKPT_SHA256 before they are loaded, as tensors only."""
    import torch
    from compressai.zoo import image as zoo_image
    from compressai.zoo.pretrained import load_pretrained
    url = zoo_image.model_urls["mbt2018-mean"]["mse"][q]
    path = os.path.join(torch.hub.get_dir(), "checkpoints", os.path.basename(url))
    if not os.path.exists(path):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        torch.hub.download_url_to_file(url, path, progress=False)
    blob = open(path, "rb").read()
    got = sha256(blob)
    # The zoo's file names carry the first eight hex digits of the digest.
    if not os.path.basename(url).split("-")[-1].startswith(got[:8]):
        sys.exit(f"q{q}: {path} has SHA-256 {got}, which does not match the prefix in its name")
    want = CKPT_SHA256[q]
    if want is None and not allow_new:
        sys.exit(f"q{q}: no pinned digest; measured {got}")
    if want is not None and got != want:
        sys.exit(f"q{q}: {path} has SHA-256 {got}, pinned {want}")
    state = torch.load(io.BytesIO(blob), map_location="cpu", weights_only=True)
    m = zoo_image.model_architectures["mbt2018-mean"].from_state_dict(load_pretrained(state)).eval()
    return m, {"quality": q, "url_name": os.path.basename(url), "sha256": got}


# ---------------------------------------------------------------- integer export

def int_net(E, seq, kind):
    """One integer network from a float Sequential, quantised as E (the quality-6 export
    tool) quantises g_s and h_s."""
    net = {"afrac": E.AFRAC, "gf": E.GF, "shift": E.SHIFT, "layers": []}
    import torch
    for layer in seq.children():
        t = type(layer).__name__
        if t == "ConvTranspose2d":
            wi, sc = E.q_weight_pc(layer.weight.detach().numpy())
            net["layers"].append({"type": "convT", "in_ch": int(wi.shape[0]), "out_ch": int(wi.shape[1]),
                                  "k": int(wi.shape[2]), "stride": int(layer.stride[0]), "pad": int(layer.padding[0]),
                                  "opad": int(layer.output_padding[0]), "weight": wi.reshape(-1).tolist(),
                                  "mult": np.round(sc * (2 ** E.SHIFT)).astype(np.int64).tolist(),
                                  "bias": np.round(layer.bias.detach().numpy() * (2 ** E.AFRAC)).astype(np.int64).tolist()})
        elif t == "Conv2d":
            wi, sc = E.q_weight_conv(layer.weight.detach().numpy())
            net["layers"].append({"type": "conv", "in_ch": int(wi.shape[1]), "out_ch": int(wi.shape[0]),
                                  "k": int(wi.shape[2]), "stride": int(layer.stride[0]), "pad": int(layer.padding[0]),
                                  "weight": wi.reshape(-1).tolist(),
                                  "mult": np.round(sc * (2 ** E.SHIFT)).astype(np.int64).tolist(),
                                  "bias": np.round(layer.bias.detach().numpy() * (2 ** E.AFRAC)).astype(np.int64).tolist()})
        elif t == "GDN":
            with torch.no_grad():
                beta = layer.beta_reparam(layer.beta).numpy()
                C = beta.shape[0]
                gamma = layer.gamma_reparam(layer.gamma).reshape(C, C).numpy()
            net["layers"].append({"type": "igdn", "c": int(C),
                                  "gamma": np.round(gamma * (2 ** E.GF)).astype(np.int64).reshape(-1).tolist(),
                                  "beta": np.round(beta * (2 ** (2 * E.AFRAC))).astype(np.int64).tolist()})
        elif t == "ReLU":
            net["layers"].append({"type": "relu"})
        elif t == "LeakyReLU":
            net["layers"].append({"type": "leaky", "slope_q": 655, "slope_frac": 16})
        else:
            sys.exit(f"{kind}: no integer form for a {t} layer")
    return net


def requant(raw, mult, bias, shift):
    """((raw * mult) >> shift) + bias per output channel, exact: in int64 where the
    product provably fits, else in Python integers."""
    if np.abs(raw).max(initial=0) * np.abs(mult).max(initial=0) < 2 ** 62:
        return ((raw * mult[:, None, None]) >> shift) + bias[:, None, None]
    return ((raw.astype(object) * mult[:, None, None]) >> shift).astype(np.int64) + bias[:, None, None]


def conv_t_fast(x, L, shift):
    """The integer transposed convolution of `int_hyper_synth` (weight [in,out,k,k]),
    tap by tap: out[:, iy*s+ky, ix*s+kx] += w[:, :, ky, kx]^T x[:, iy, ix]. Integer sums
    are exact in any order, so this equals the export tool's per-pixel scatter."""
    cin, co, k, s, p, op = L["in_ch"], L["out_ch"], L["k"], L["stride"], L["pad"], L["opad"]
    w = np.array(L["weight"], dtype=np.int64).reshape(cin, co, k, k)
    _, H, W = x.shape
    Ho, Wo = (H - 1) * s + k + op - 2 * p, (W - 1) * s + k + op - 2 * p
    raw = np.zeros((co, Ho + 2 * p, Wo + 2 * p), dtype=np.int64)
    xf = x.reshape(cin, -1)
    assert np.abs(x).max(initial=0) * np.abs(w).max(initial=0) * cin * k * k < 2 ** 63
    for ky in range(k):
        for kx in range(k):
            c = (w[:, :, ky, kx].T @ xf).reshape(co, H, W)
            raw[:, ky:ky + s * H:s, kx:kx + s * W:s] += c
    raw = raw[:, p:p + Ho, p:p + Wo]
    return requant(raw, np.array(L["mult"], dtype=np.int64), np.array(L["bias"], dtype=np.int64), shift)


def conv_fast(x, L, shift):
    cin, co, k, s, p = L["in_ch"], L["out_ch"], L["k"], L["stride"], L["pad"]
    w = np.array(L["weight"], dtype=np.int64).reshape(co, cin, k, k)
    _, H, W = x.shape
    Ho, Wo = (H + 2 * p - k) // s + 1, (W + 2 * p - k) // s + 1
    xp = np.zeros((cin, H + 2 * p, W + 2 * p), dtype=np.int64)
    xp[:, p:p + H, p:p + W] = x
    assert np.abs(x).max(initial=0) * np.abs(w).max(initial=0) * cin * k * k < 2 ** 63
    raw = np.zeros((co, Ho, Wo), dtype=np.int64)
    for ky in range(k):
        for kx in range(k):
            sub = xp[:, ky:ky + s * Ho:s, kx:kx + s * Wo:s].reshape(cin, -1)
            raw += (w[:, :, ky, kx] @ sub).reshape(co, Ho, Wo)
    return requant(raw, np.array(L["mult"], dtype=np.int64), np.array(L["bias"], dtype=np.int64), shift)


def run_hs(net, x):
    """The integer hyper-decoder (convolutions and ReLUs only), as `int_hyper_synth`
    runs it."""
    for L in net["layers"]:
        t = L["type"]
        if t == "convT":
            x = conv_t_fast(x, L, net["shift"])
        elif t == "conv":
            x = conv_fast(x, L, net["shift"])
        elif t == "relu":
            x = np.maximum(x, 0)
        elif t == "leaky":
            x = np.where(x >= 0, x, (x * L["slope_q"]) >> L["slope_frac"])
        else:
            sys.exit(f"h_s has a {t} layer")
    return x


def cmd_export(work, qs, allow_new=False):
    """Per quality, the integer set (integer-payloads §9: entropy.json, g_s.whm, h_s.whm)
    and the model's identity."""
    import torch
    from compressai.entropy_models import GaussianConditional
    torch.use_deterministic_algorithms(True)
    torch.set_num_threads(1)
    E = export_module()
    out = {}
    for q in qs:
        m, ident = load_model(q, allow_new)
        g_s = int_net(E, m.g_s, "g_s")
        h_s = int_net(E, m.h_s, "h_s")
        eb = m.entropy_bottleneck
        eb.update()
        median = eb._get_medians().reshape(-1).detach().numpy().astype(np.float64)
        scale_table = list(torch.exp(torch.linspace(np.log(0.11), np.log(256.0), 64)).numpy().astype(float))
        gc = GaussianConditional(scale_table)
        gc.update()
        ent = {"afrac": E.AFRAC, "frac": E.AFRAC,
               "scale_edges": np.floor(np.array(scale_table) * (2 ** E.AFRAC)).astype(np.int64).tolist(),
               "cdfs": gc._quantized_cdf.tolist(), "cdf_lengths": gc._cdf_length.reshape(-1).int().tolist(),
               "offsets": gc._offset.reshape(-1).int().tolist(),
               "z_cdfs": eb._quantized_cdf.tolist(), "z_cdf_lengths": eb._cdf_length.reshape(-1).int().tolist(),
               "z_offsets": eb._offset.reshape(-1).int().tolist(),
               "z_median_q": np.round(median * (2 ** E.AFRAC)).astype(np.int64).tolist()}
        d = os.path.join(work, "sets", f"q{q}")
        os.makedirs(d, exist_ok=True)
        open(os.path.join(d, "g_s.whm"), "wb").write(E.model_to_bin(g_s))
        open(os.path.join(d, "h_s.whm"), "wb").write(E.model_to_bin(h_s))
        with open(os.path.join(d, "entropy.json"), "w") as f:
            json.dump(ent, f)
        jdump(h_s, os.path.join(d, "h_s.json"))  # the encoder's copy of h_s
        ident["set"] = {n: sha256(open(os.path.join(d, n), "rb").read()) for n in ("entropy.json", "g_s.whm", "h_s.whm")}
        ident["N"], ident["M"] = int(m.N), int(m.M)
        out[q] = ident
        print(f"q{q}: N={m.N} M={m.M} checkpoint {ident['sha256'][:16]}… set g_s {ident['set']['g_s.whm'][:16]}…")
    path = os.path.join(work, "models.json")
    prev = jload(path) if os.path.exists(path) else {}
    prev.update({str(k): v for k, v in out.items()})
    jdump(prev, path)


# ---------------------------------------------------------------- encode

def to_rgb8(x_hat):
    """A float reconstruction in [0, 1] as 8-bit samples: x·255 rounded half to even, as
    the float library's own tensor arithmetic gives it on the device it ran on."""
    import torch
    return (x_hat.clamp(0, 1) * 255).round().to(torch.uint8)[0].permute(1, 2, 0).contiguous().cpu().numpy()


def float_container(strings, zshape):
    y, z = strings[0][0], strings[1][0]
    return b"DETF" + struct.pack("<IIII", zshape[0], zshape[1], len(y), len(z)) + y + z


def read_float_container(b):
    assert b[:4] == b"DETF"
    zh, zw, ly, lz = struct.unpack("<IIII", b[4:20])
    y = b[20:20 + ly]
    z = b[20 + ly:20 + ly + lz]
    assert 20 + ly + lz == len(b)
    return [[y], [z]], (zh, zw)


def float_decode(m, strings, zshape):
    """The float library's own decode of a mean-scale stream (its `decompress`), with the
    decoded indexes and symbols kept, so a mismatch can be placed: in the entropy decode
    (other indexes or symbols) or after it."""
    import torch
    with torch.no_grad():
        z_hat = m.entropy_bottleneck.decompress(strings[1], zshape)
        scales_hat, means_hat = m.h_s(z_hat).chunk(2, 1)
        indexes = m.gaussian_conditional.build_indexes(scales_hat)
        symbols = m.gaussian_conditional.decompress(strings[0], indexes)
        y_hat = symbols + means_hat
        x_hat = m.g_s(y_hat).clamp_(0, 1)
    idx = indexes.cpu().to(torch.int32).numpy().tobytes()
    sym = symbols.cpu().round().to(torch.int32).numpy().tobytes()
    return to_rgb8(x_hat), sha256(idx), sha256(sym)


def tables_sha256(m):
    """The entropy tables a float decode reads (both models' quantised CDFs, lengths and
    offsets): a decode on another machine is comparable only under the same tables."""
    import torch
    h = hashlib.sha256()
    for em in (m.entropy_bottleneck, m.gaussian_conditional):
        for t in (em._quantized_cdf, em._cdf_length, em._offset):
            h.update(t.cpu().to(torch.int32).numpy().tobytes())
    return h.hexdigest()


def wih1(m_ch, sh, sw, n_z, zh, zw, z_rans, rans):
    o = b"WIH1" + bytes([1, 0, 0, 0]) + struct.pack("<6I", m_ch, sh, sw, n_z, zh, zw)
    return o + struct.pack("<I", len(z_rans)) + z_rans + struct.pack("<I", len(rans)) + rans


def cmd_encode(work, qs):
    import torch
    from compressai.ans import RansEncoder
    torch.use_deterministic_algorithms(True)
    torch.set_num_threads(int(os.environ.get("DET_THREADS", "1")))
    E = export_module()
    pics = pictures(work)
    rows = {}
    for q in qs:
        m, _ = load_model(q)
        m.update()
        tables = tables_sha256(m)
        d = os.path.join(work, "sets", f"q{q}")
        hs = jload(os.path.join(d, "h_s.json"))
        ent = jload(os.path.join(d, "entropy.json"))
        edges = np.array(ent["scale_edges"], dtype=np.int64)
        mq = np.array(ent["z_median_q"], dtype=np.int64)
        for sub in ("float", "int", "dec"):
            os.makedirs(os.path.join(work, sub, f"q{q}"), exist_ok=True)
        for p in pics:
            px = source_rgb(work, p)
            x = torch.from_numpy(px.astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0)
            with torch.no_grad():
                # the float codec, as its library codes and decodes it
                c = m.compress(x)
                fb = float_container(c["strings"], tuple(c["shape"]))
                open(os.path.join(work, "float", f"q{q}", f"{p['id']}.detf"), "wb").write(fb)
                rgb, ih, sh_ = float_decode(m, c["strings"], tuple(c["shape"]))
                rgb.tofile(os.path.join(work, "dec", f"q{q}", f"{p['id']}.float.rgb"))
                # the integer codec: the float analysis, then everything the decoder reads in integers
                y = m.g_a(x)
                z = m.h_a(y)
            y_np = y[0].numpy()
            eb_median = m.entropy_bottleneck._get_medians().reshape(-1).detach().numpy().astype(np.float64)
            z_sym = np.round(z[0].numpy() - eb_median[:, None, None]).astype(np.int64)
            Cz, zh, zw = z_sym.shape
            z_bits = RansEncoder().encode_with_indexes(
                z_sym.reshape(-1).tolist(), np.repeat(np.arange(Cz), zh * zw).tolist(),
                ent["z_cdfs"], ent["z_cdf_lengths"], ent["z_offsets"])
            zi = z_sym * (2 ** E.AFRAC) + mq[:, None, None]
            sm = run_hs(hs, zi)
            M = sm.shape[0] // 2
            sigma, mu = sm[:M], sm[M:]
            resid = np.round(y_np - mu.astype(np.float64) / (2 ** E.AFRAC)).astype(np.int64)
            idx = E.det_build_index_vec(sigma, edges)
            y_bits = RansEncoder().encode_with_indexes(resid.reshape(-1).tolist(), idx.reshape(-1).tolist(),
                                                       ent["cdfs"], ent["cdf_lengths"], ent["offsets"])
            pay = wih1(M, sigma.shape[1], sigma.shape[2], Cz, zh, zw, z_bits, y_bits)
            open(os.path.join(work, "int", f"q{q}", f"{p['id']}.wih1"), "wb").write(pay)
            rows[f"{p['id']}/q{q}"] = {
                "float_stream_bytes": len(c["strings"][0][0]) + len(c["strings"][1][0]),
                "float_stream_sha256": sha256(fb),
                "float_tables_sha256": tables,
                "float_cpu": {"rgb_sha256": sha256(rgb.tobytes()), "indexes_sha256": ih, "symbols_sha256": sh_},
                "int_stream_bytes": len(z_bits) + len(y_bits),
                "int_payload_bytes": len(pay),
                "int_payload_sha256": sha256(pay),
            }
            row = rows[f"{p['id']}/q{q}"]
            print(f"q{q} {p['id']}: float {row['float_stream_bytes']} B, integer {row['int_stream_bytes']} B", flush=True)
    path = os.path.join(work, "encoded.json")
    prev = jload(path) if os.path.exists(path) else {}
    prev.update(rows)
    jdump(prev, path)


def cmd_jobs(work, tag):
    """One line per integer payload: its set, the payload, where the decode goes."""
    lines = []
    for q in QUALITIES:
        d = os.path.join(work, "int", f"q{q}")
        if not os.path.isdir(d):
            continue
        os.makedirs(os.path.join(work, "dec", f"q{q}"), exist_ok=True)
        for f in sorted(os.listdir(d)):
            pid = f.split(".")[0]
            out = os.path.join(work, "dec", f"q{q}", f"{pid}.int.{tag}.rgb") if tag != "-" else "-"
            lines.append(f"{os.path.join(work, 'sets', f'q{q}')}\t{os.path.join(d, f)}\t{out}")
    sys.stdout.write("\n".join(lines) + "\n")


def cmd_fdecode(work, tag, device, threads):
    """Decode every float stream in a fresh process on `device`: the RGB8, the indexes and
    the symbols, by SHA-256."""
    import platform
    import torch
    torch.use_deterministic_algorithms(True)
    if threads:
        torch.set_num_threads(threads)
    out = {"tag": tag, "device": device, "threads": torch.get_num_threads(), "torch": torch.__version__,
           "machine": platform.machine(), "system": platform.system(), "results": {}}
    srcs = {p["id"]: p for p in pictures(work)}
    for q in QUALITIES:
        d = os.path.join(work, "float", f"q{q}")
        if not os.path.isdir(d):
            continue
        m, _ = load_model(q)
        m.update()
        out.setdefault("tables_sha256", {})[f"q{q}"] = tables_sha256(m)
        m = m.to(device)
        for f in sorted(os.listdir(d)):
            strings, zshape = read_float_container(open(os.path.join(d, f), "rb").read())
            rgb, ih, sh_ = float_decode(m, strings, zshape)
            pid = f.split(".")[0]
            src = source_rgb(work, srcs[pid])
            diff = rgb.astype(np.int64) - src.astype(np.int64)
            out["results"][f"{pid}/q{q}"] = {"rgb_sha256": sha256(rgb.tobytes()), "indexes_sha256": ih,
                                             "symbols_sha256": sh_, "sse": int((diff * diff).sum())}
        print(f"q{q}: {len(os.listdir(d))} decoded on {device}", flush=True)
    os.makedirs(os.path.join(work, "fdecode"), exist_ok=True)
    jdump(out, os.path.join(work, "fdecode", f"{tag}.json"))


def cmd_fserve(work, q, device, threads, ids):
    """The timed child of tools/wai_det_energy.py for the float decode: the library's own
    `decompress` of every float stream at quality q (or of the pictures `ids`) on
    `device`, then the 8-bit samples on the host. Each decode's RGB8 is hashed and
    compared with the first decode of its stream; a differing decode is counted, not
    hidden (`done <decodes> <ns> mismatches=<n>`)."""
    import time
    import torch
    torch.use_deterministic_algorithms(True, warn_only=device != "cpu")
    if threads:
        torch.set_num_threads(threads)
    m, _ = load_model(q)
    m.update()
    m = m.to(device)
    d = os.path.join(work, "float", f"q{q}")
    files = [f for f in sorted(os.listdir(d)) if not ids or f.split(".")[0] in ids]
    streams = [read_float_container(open(os.path.join(d, f), "rb").read()) for f in files]

    def one(st):
        with torch.no_grad():
            x_hat = m.decompress(st[0], st[1])["x_hat"]
        return sha256(to_rgb8(x_hat).tobytes())

    want = [one(st) for st in streams]
    print(f"ready {len(streams)}", flush=True)
    for line in sys.stdin:
        w = line.split()
        if not w or w[0] == "exit":
            return
        k = int(w[1])
        t0 = time.monotonic_ns()
        n = bad = 0
        for _ in range(k):
            for st, h in zip(streams, want):
                bad += one(st) != h
                n += 1
        print(f"done {n} {time.monotonic_ns() - t0} mismatches={bad}", flush=True)


# ---------------------------------------------------------------- metrics

GAUSS = None
MSSSIM_W = (0.0448, 0.2856, 0.3001, 0.2363, 0.1333)


def _gauss():
    global GAUSS
    if GAUSS is None:
        g = np.exp(-((np.arange(11) - 5.0) ** 2) / (2 * 1.5 ** 2))
        GAUSS = g / g.sum()
    return GAUSS


def _filt(x):
    """Separable 11-tap Gaussian (sigma 1.5), valid region only."""
    g = _gauss()
    h, w = x.shape
    t = sum(g[i] * x[i:h - 10 + i, :] for i in range(11))
    return sum(g[i] * t[:, i:w - 10 + i] for i in range(11))


def ms_ssim_channel(a, b):
    """MS-SSIM of one channel, 8-bit samples (dynamic range 255): five scales, an 11-tap
    Gaussian window of sigma 1.5, K1 0.01, K2 0.03, 2x2 mean pooling between scales
    (an odd row or column dropped), contrast-structure at the first four scales and
    luminance-contrast-structure at the fifth, weights 0.0448, 0.2856, 0.3001, 0.2363,
    0.1333, each term clamped at zero before its power."""
    C1, C2 = (0.01 * 255) ** 2, (0.03 * 255) ** 2
    a, b = a.astype(np.float64), b.astype(np.float64)
    val = 1.0
    for s in range(5):
        mu_a, mu_b = _filt(a), _filt(b)
        saa = _filt(a * a) - mu_a ** 2
        sbb = _filt(b * b) - mu_b ** 2
        sab = _filt(a * b) - mu_a * mu_b
        cs = (2 * sab + C2) / (saa + sbb + C2)
        if s < 4:
            val *= max(cs.mean(), 0.0) ** MSSSIM_W[s]
            h, w = a.shape[0] // 2 * 2, a.shape[1] // 2 * 2
            a = a[:h, :w].reshape(h // 2, 2, w // 2, 2).mean(axis=(1, 3))
            b = b[:h, :w].reshape(h // 2, 2, w // 2, 2).mean(axis=(1, 3))
        else:
            l = (2 * mu_a * mu_b + C1) / (mu_a ** 2 + mu_b ** 2 + C1)
            val *= max((l * cs).mean(), 0.0) ** MSSSIM_W[s]
    return val


def ms_ssim(a, b):
    return float(np.mean([ms_ssim_channel(a[:, :, c], b[:, :, c]) for c in range(3)]))


def cmd_metrics(work, tags):
    pics = pictures(work)
    enc = jload(os.path.join(work, "encoded.json"))
    rows = {}
    for p in pics:
        src = source_rgb(work, p)
        for q in QUALITIES:
            k = f"{p['id']}/q{q}"
            if k not in enc:
                continue
            r = {}
            for path, fname in (("float", f"{p['id']}.float.rgb"), ("int", f"{p['id']}.int.{tags[0]}.rgb")):
                rgb = np.fromfile(os.path.join(work, "dec", f"q{q}", fname), dtype=np.uint8)
                rgb = rgb.reshape(src.shape)
                d = rgb.astype(np.int64) - src.astype(np.int64)
                r[path] = {"sse": int((d * d).sum()), "samples": int(d.size), "ms_ssim": ms_ssim(rgb, src),
                           "rgb_sha256": sha256(rgb.tobytes())}
            rows[k] = r
        print(p["id"], flush=True)
    jdump(rows, os.path.join(work, "metrics.json"))


# ---------------------------------------------------------------- BD-rate

def psnr(sse, n):
    return 10 * math.log10(255.0 ** 2 * n / sse)


def ms_ssim_db(v):
    return -10 * math.log10(1 - v)


def pchip_slopes(x, y):
    """Interior slopes of the monotone piecewise-cubic Hermite interpolant (weighted
    harmonic mean of the neighbouring secants, zero at a change of sign), end slopes by
    the three-point shape-preserving formula."""
    n = len(x)
    h = [x[i + 1] - x[i] for i in range(n - 1)]
    d = [(y[i + 1] - y[i]) / h[i] for i in range(n - 1)]
    m = [0.0] * n
    for i in range(1, n - 1):
        if d[i - 1] * d[i] > 0:
            w1, w2 = 2 * h[i] + h[i - 1], h[i] + 2 * h[i - 1]
            m[i] = (w1 + w2) / (w1 / d[i - 1] + w2 / d[i])
    def end(h0, h1, d0, d1):
        s = ((2 * h0 + h1) * d0 - h0 * d1) / (h0 + h1)
        if s * d0 <= 0:
            return 0.0
        if d0 * d1 <= 0 and abs(s) > abs(3 * d0):
            return 3 * d0
        return s
    m[0] = end(h[0], h[1], d[0], d[1])
    m[-1] = end(h[-1], h[-2], d[-1], d[-2])
    return m


def pchip_integral(x, y, a, b):
    """The integral from a to b of the interpolant through (x, y): each segment's cubic
    Hermite polynomial integrated exactly."""
    m = pchip_slopes(x, y)
    total = 0.0
    for i in range(len(x) - 1):
        lo, hi = max(a, x[i]), min(b, x[i + 1])
        if hi <= lo:
            continue
        h = x[i + 1] - x[i]
        def F(t):  # antiderivative in t = (u - x_i)/h, times h
            h00 = t - t ** 3 + t ** 4 / 2
            h10 = t ** 2 / 2 - 2 * t ** 3 / 3 + t ** 4 / 4
            h01 = t ** 3 - t ** 4 / 2
            h11 = -t ** 3 / 3 + t ** 4 / 4
            return h * (h00 * y[i] + h10 * h * m[i] + h01 * y[i + 1] + h11 * h * m[i + 1])
        total += F((hi - x[i]) / h) - F((lo - x[i]) / h)
    return total


def bd_rate(r_ref, q_ref, r_test, q_test):
    """Bjontegaard delta rate (ITU-T VCEG-M33), log-rate interpolated as a monotone
    piecewise cubic of quality and averaged over the overlap of the two quality ranges:
    the percentage by which the test needs more rate than the reference at equal quality.
    None when either curve is not strictly increasing in quality."""
    def prep(r, qv):
        pts = sorted(zip(qv, (math.log10(v) for v in r)))
        xs = [p[0] for p in pts]
        if any(b <= a for a, b in zip(xs, xs[1:])):
            return None
        return xs, [p[1] for p in pts]
    A, B = prep(r_ref, q_ref), prep(r_test, q_test)
    if A is None or B is None:
        return None
    lo, hi = max(A[0][0], B[0][0]), min(A[0][-1], B[0][-1])
    if hi <= lo:
        return None
    ia = pchip_integral(A[0], A[1], lo, hi)
    ib = pchip_integral(B[0], B[1], lo, hi)
    return (10 ** ((ib - ia) / (hi - lo)) - 1) * 100


def cmd_bd(work):
    pics = pictures(work)
    enc = jload(os.path.join(work, "encoded.json"))
    met = jload(os.path.join(work, "metrics.json"))
    per = {}
    for p in pics:
        n = p["width"] * p["height"]
        rf, ri, pf, pi, sf, si = [], [], [], [], [], []
        for q in QUALITIES:
            k = f"{p['id']}/q{q}"
            rf.append(enc[k]["float_stream_bytes"] * 8 / n)
            ri.append(enc[k]["int_stream_bytes"] * 8 / n)
            pf.append(psnr(met[k]["float"]["sse"], met[k]["float"]["samples"]))
            pi.append(psnr(met[k]["int"]["sse"], met[k]["int"]["samples"]))
            sf.append(ms_ssim_db(met[k]["float"]["ms_ssim"]))
            si.append(ms_ssim_db(met[k]["int"]["ms_ssim"]))
        per[p["id"]] = {"bd_rate_psnr": bd_rate(rf, pf, ri, pi), "bd_rate_ms_ssim": bd_rate(rf, sf, ri, si)}
        print(p["id"], per[p["id"]])
    jdump(per, os.path.join(work, "bd.json"))


# ---------------------------------------------------------------- the record

def stats(v):
    v = sorted(v)
    n = len(v)
    if n == 0:
        return None
    med = v[n // 2] if n % 2 else (v[n // 2 - 1] + v[n // 2]) / 2
    return {"n": n, "median": med, "mean": sum(v) / n, "min": v[0], "max": v[-1]}


def summarise(rec):
    """Everything the record states beyond its raw readings, computed from them alone:
    tools/wai_determinism_cost_verify.py recomputes each value independently."""
    pics = rec["pictures"]["items"]
    pts = rec["points"]
    qs = sorted(int(q) for q in rec["models"])
    out = {}
    bd = {"psnr": {}, "ms_ssim": {}}
    for p in pics:
        n = p["width"] * p["height"]
        ks = [f"{p['id']}/q{q}" for q in qs]
        rf = [pts[k]["float_stream_bytes"] * 8 / n for k in ks]
        ri = [pts[k]["int_stream_bytes"] * 8 / n for k in ks]
        for metric, f in (("psnr", lambda r: psnr(r["sse"], r["samples"])), ("ms_ssim", lambda r: ms_ssim_db(r["ms_ssim"]))):
            bd[metric][p["id"]] = bd_rate(rf, [f(pts[k]["float"]) for k in ks], ri, [f(pts[k]["int"]) for k in ks])
    out["bd_rate_percent"] = {m: {"per_picture": v, **stats([x for x in v.values() if x is not None]),
                                  "undefined": sorted(k for k, x in v.items() if x is None)} for m, v in bd.items()}
    perq = {}
    for q in qs:
        ks = [f"{p['id']}/q{q}" for p in pics]
        perq[f"q{q}"] = {
            "float_bpp": stats([pts[k]["float_stream_bytes"] * 8 / (pts[k]["float"]["samples"] / 3) for k in ks]),
            "rate_ratio": stats([pts[k]["int_stream_bytes"] / pts[k]["float_stream_bytes"] for k in ks]),
            "float_psnr_db": stats([psnr(pts[k]["float"]["sse"], pts[k]["float"]["samples"]) for k in ks]),
            "delta_psnr_db": stats([psnr(pts[k]["int"]["sse"], pts[k]["int"]["samples"]) - psnr(pts[k]["float"]["sse"], pts[k]["float"]["samples"]) for k in ks]),
            "delta_ms_ssim_db": stats([ms_ssim_db(pts[k]["int"]["ms_ssim"]) - ms_ssim_db(pts[k]["float"]["ms_ssim"]) for k in ks]),
        }
    out["per_quality"] = perq
    ref_int = {k: v["int"]["rgb_sha256"] for k, v in pts.items()}
    xi = {}
    for tag, d in rec["integer_decodes"].items():
        r = d["results"]
        xi[tag] = {"decodes": len(r),
                   "identical": sum(1 for k, h in r.items() if h == ref_int.get(k)),
                   "refused": sum(1 for h in r.values() if h.startswith("refused:")),
                   "missing": sorted(k for k in ref_int if k not in r)}
        xi[tag]["differ"] = xi[tag]["decodes"] - xi[tag]["identical"] - xi[tag]["refused"]
    out["integer_cross_platform"] = xi
    out["integer_cross_platform_total"] = {"platforms": len(xi), **{f: sum(v[f] for v in xi.values()) for f in ("decodes", "identical", "differ", "refused")}}
    xf = {}
    for tag, d in rec["float_decodes"].items():
        r = d["results"]
        c = {"decodes": len(r), "identical": 0, "rgb_differ": 0, "indexes_differ": 0, "symbols_differ": 0,
             "tables_differ": sorted(q for q, h in d.get("tables_sha256", {}).items() if h != rec["float_tables_sha256"][q])}
        dps, dsym = [], []
        for k, v in r.items():
            ref = pts[k]
            same = v["rgb_sha256"] == ref["float"]["rgb_sha256"]
            c["identical"] += same
            c["rgb_differ"] += not same
            c["indexes_differ"] += v["indexes_sha256"] != ref["float_cpu_indexes_sha256"]
            c["symbols_differ"] += v["symbols_sha256"] != ref["float_cpu_symbols_sha256"]
            dp = psnr(v["sse"], ref["float"]["samples"]) - psnr(ref["float"]["sse"], ref["float"]["samples"])
            if not same:
                dps.append(dp)
            if v["symbols_sha256"] != ref["float_cpu_symbols_sha256"]:
                dsym.append(dp)
        c["delta_psnr_db_where_differ"] = stats(dps)
        c["delta_psnr_db_where_symbols_differ"] = stats(dsym)
        c["missing"] = sorted(k for k in pts if k not in r)
        xf[tag] = c
    out["float_cross_platform"] = xf
    def live(b):
        """energy-measurement §3.1: the rail advanced in two consecutive intervals, in at
        least five and in at least four of every five, over at least 300 ms."""
        return bool(b) and b["longest_run"] >= 2 and b["advanced"] >= 5 and 5 * b["advanced"] >= 4 * b["intervals"] and b["span_ns"] >= 300_000_000

    tm = {}
    for run in rec["timing"]:
        rs = run["rounds"]
        t = {"rounds": len(rs), "decodes_per_round": sorted({r["decodes"] for r in rs}),
             "ms_per_decode": stats([r["child_wall_ns"] / r["decodes"] / 1e6 for r in rs]),
             "load_avg_1min": stats([r["load_avg_before"][0] for r in rs] + [r["load_avg_after"][0] for r in rs])}
        if all("task" in r for r in rs):
            t["task_j_per_decode"] = stats([r["task"]["energy_nj"] / r["decodes"] / 1e9 for r in rs])
            t["task_cpu_ms_per_decode"] = stats([r["task"]["cpu_ns"] / r["decodes"] / 1e6 for r in rs])
            t["billed_j_per_decode"] = stats([r["task"]["billed_nj"] / r["decodes"] / 1e9 for r in rs])
        for cn in ("rapl", "ioreport-gpu"):
            if all(cn in r for r in rs):
                ok = [r for r in rs if r[cn].get("verdict") in ("clean", "relative-only")
                      and all(live(r[cn][b]) for b in ("busy", "idle_before", "idle_after"))]
                if cn == "ioreport-gpu":
                    ok = [r for r in ok if (r.get("gpu_time", {}).get("share") or 0) >= 0.99]
                t[cn] = {"verdicts": [r[cn].get("verdict") for r in rs],
                         "live_rounds": sum(all(live(r[cn][b]) for b in ("busy", "idle_before", "idle_after")) for r in rs),
                         "ratio": stats([r[cn]["ratio"] for r in rs if r[cn].get("ratio") is not None]),
                         "usable_rounds": len(ok),
                         "marginal_j_per_decode": stats([r[cn]["marginal_nj"] / r["decodes"] / 1e9 for r in ok]),
                         "publishable": len(ok) >= 7 and all(r[cn]["verdict"] == "clean" for r in ok)}
                if cn == "ioreport-gpu":
                    t[cn]["gpu_share"] = stats([r["gpu_time"]["share"] for r in rs if r.get("gpu_time", {}).get("share") is not None])
        tm[run["label"]] = t
    out["timing"] = tm
    return out


def cmd_record(work, dest):
    pics = jload(os.path.join(work, "pictures.json"))
    enc = jload(os.path.join(work, "encoded.json"))
    met = jload(os.path.join(work, "metrics.json"))
    models = jload(os.path.join(work, "models.json"))
    points, tables = {}, {}
    for k, e in sorted(enc.items()):
        q = k.split("/")[1]
        if q in tables and tables[q] != e["float_tables_sha256"]:
            sys.exit(f"{k}: other float tables than its quality's other pictures")
        tables[q] = e["float_tables_sha256"]
        points[k] = {"float_stream_bytes": e["float_stream_bytes"], "int_stream_bytes": e["int_stream_bytes"],
                     "int_payload_bytes": e["int_payload_bytes"], "float_stream_sha256": e["float_stream_sha256"],
                     "int_payload_sha256": e["int_payload_sha256"],
                     "float_cpu_indexes_sha256": e["float_cpu"]["indexes_sha256"],
                     "float_cpu_symbols_sha256": e["float_cpu"]["symbols_sha256"],
                     "float": met[k]["float"], "int": met[k]["int"]}
        if met[k]["float"]["rgb_sha256"] != e["float_cpu"]["rgb_sha256"]:
            sys.exit(f"{k}: the measured float decode is not the one the encoder wrote")
    rec = {"format": "wai-determinism-cost/1", "procedure": "wai/research/determinism-cost.md",
           "runs": "wai/research/determinism-cost/runs/", "measured": os.environ.get("DET_MEASURED", "2026-10-22"),
           "pictures": {"manifest_sha256": pics["manifest_sha256"], "rule": pics["rule"],
                        "items": [{kk: p[kk] for kk in ("id", "file_sha256", "width", "height", "rgb_sha256")} for p in pics["pictures"]]},
           "models": {q: {"checkpoint_sha256": v["sha256"], "checkpoint_name": v["url_name"], "N": v["N"], "M": v["M"], "set_sha256": v["set"]}
                      for q, v in sorted(models.items(), key=lambda kv: int(kv[0]))},
           "float_tables_sha256": tables, "points": points,
           "integer_decodes": {}, "float_decodes": {}, "timing": []}
    xd = os.path.join(work, "xplat")
    for f in sorted(os.listdir(xd)) if os.path.isdir(xd) else []:
        if f.endswith(".tsv"):
            tag = f[:-4]
            meta = jload(os.path.join(xd, tag + ".json"))
            res = {}
            for line in open(os.path.join(xd, f)):
                c = line.rstrip("\n").split("\t")
                k = f"{os.path.basename(c[0]).split('.')[0]}/{os.path.basename(os.path.dirname(c[0]))}"
                res[k] = f"refused:{c[2]}" if c[1] == "REFUSED" else c[1]
            rec["integer_decodes"][tag] = {"platform": meta, "results": dict(sorted(res.items()))}
    fd = os.path.join(work, "fdecode")
    for f in sorted(os.listdir(fd)) if os.path.isdir(fd) else []:
        if not f.endswith(".json"):
            continue
        d = jload(os.path.join(fd, f))
        meta = jload(os.path.join(fd, f[:-5] + ".meta")) if os.path.exists(os.path.join(fd, f[:-5] + ".meta")) else {}
        rec["float_decodes"][d["tag"]] = {"platform": {**meta, "device": d["device"], "threads": d["threads"], "torch": d["torch"],
                                                       "machine": d["machine"], "system": d["system"]},
                                          "tables_sha256": d.get("tables_sha256", {}), "results": d["results"]}
    ed = os.path.join(work, "energy")
    for f in sorted(os.listdir(ed)) if os.path.isdir(ed) else []:
        if f.endswith(".json"):
            run = jload(os.path.join(ed, f))
            # The child's command line, relative to each host's work area (no home or volume paths).
            run["child"] = [a.split("/wai-build/", 1)[1] if "/wai-build/" in a else a for a in run["child"]]
            rec["timing"].append(run)
    rec["summary"] = summarise(rec)
    jdump(rec, dest)
    print(json.dumps(rec["summary"]["bd_rate_percent"]["psnr"] | {"per_picture": None}, indent=1))


def main(argv):
    if len(argv) < 3:
        sys.exit(__doc__)
    cmd, work = argv[1], argv[2]
    rest = argv[3:]
    if cmd == "select":
        cmd_select(work, rest[0])
    elif cmd == "export":
        allow = "--allow-new" in rest
        cmd_export(work, [int(a) for a in rest if a != "--allow-new"] or QUALITIES, allow)
    elif cmd == "encode":
        cmd_encode(work, [int(a) for a in rest] or QUALITIES)
    elif cmd == "jobs":
        cmd_jobs(work, rest[0] if rest else "-")
    elif cmd == "fdecode":
        tag = rest[0]
        dev = rest[rest.index("--device") + 1] if "--device" in rest else "cpu"
        th = int(rest[rest.index("--threads") + 1]) if "--threads" in rest else 0
        cmd_fdecode(work, tag, dev, th)
    elif cmd == "fserve":
        q = int(rest[0])
        dev = rest[rest.index("--device") + 1] if "--device" in rest else "cpu"
        th = int(rest[rest.index("--threads") + 1]) if "--threads" in rest else 0
        ids = rest[rest.index("--ids") + 1].split(",") if "--ids" in rest else []
        cmd_fserve(work, q, dev, th, ids)
    elif cmd == "metrics":
        cmd_metrics(work, rest or ["native"])
    elif cmd == "bd":
        cmd_bd(work)
    elif cmd == "record":
        cmd_record(work, rest[0])
    else:
        sys.exit(__doc__)


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