#!/usr/bin/env bash
# Regenerate the int_hyper integer-codec vectors (the pretrained mbt2018_mean,
# exported to integers) that the Rust tests read from WAI_SOTA_INT_DIR
# (tests/int_hyper_synth.rs, tests/int_video_sota_keyframe.rs).
#
#   tools/wai_sota_vectors.sh [out-dir]        # default /tmp/sota_int
#
# Needs the pinned Python stack (tools/requirements-sota-int.txt) on PATH as
# `python` (or set PYTHON), and cargo. The pretrained mbt2018_mean checkpoint is
# downloaded into the torch.hub cache on first use and checked against its pinned
# SHA-256 (CKPT_SHA256 in wai_sota_int_synth_export.py) before it is loaded.
#
# The chain, each step consuming the one before:
#   1. the pinned public-domain source photograph      → source.png
#   2. the integer export of mbt2018_mean              → model, bitstream, hashes
#   3. the Rust integer decode of that bitstream       → frame0.rgb (the keyframe)
#   4. the P-frame encoder panning that keyframe        → video bitstream, hashes
# Every step must succeed and every file the tests read must exist at the end, or
# this exits non-zero: a partial directory would let a test skip instead of fail.
#
# The output is not portable, and part of it is not repeatable run to run on
# every machine. Step 2 runs the float model (torch, CPU) to encode the
# photograph, on the pinned stack, which each platform installs as its own build.
# - Between an arm64 macOS machine and an x86_64 Linux runner,
#   int_synth_sota_input.json (the synthesis-only latent, which the float
#   encoder and hyper-decoder compute) differed, and with it the rgb_hash in
#   int_synth_sota.json. Each side wrote the same latent in all of its runs.
# - Two runs on that one runner wrote different cond.json and
#   int_synth_sota.json, most likely in psnr_float, the float decode's PSNR
#   (the runner's files were not kept, so the field is inferred). Repeated runs
#   on the Mac wrote identical files. Only the quality test reads psnr_float,
#   against thresholds (the 1 dB bar and a 30 dB floor).
# - source.png differed between the Mac and the runner in its compressed bytes.
#   Step 1 writes it with the PNG encoder from the pinned pixels, before any
#   model runs, and its pixels (source.rgb) were identical, so this is the
#   encoder's output, not float computation.
# The other files matched in these comparisons, which does not make them
# portable. So check a decode against the vectors of the same run, and do not
# pin a hash from these files for another run, on this machine or another: it
# can differ with both decoders correct. The decode's cross-architecture check
# uses the committed tiny vector (wai-rs/tests/vectors/int_hyper/), whose hashes
# are pinned.
set -euo pipefail

here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
wai="$(dirname "$here")"
out="${1:-/tmp/sota_int}"
py="${PYTHON:-python}"

mkdir -p "$out"
out="$(cd "$out" && pwd)"

"$py" "$here/wai_sota_source_image.py" "$out/source.png"
"$py" "$here/wai_sota_int_synth_export.py" "$out/source.png" "$out"
(cd "$wai/wai-rs" && cargo run --release --no-default-features --features neural_int \
  --example int_hyper_decode_frame -- "$out" "$out/frame0.rgb")
"$py" "$here/wai_int_video_sota_keyframe_encode.py" "$out"

missing=0
for f in source.rgb \
         int_synth_sota.json int_synth_sota_model.json int_synth_sota_model.bin int_synth_sota_input.json \
         int_hs.json int_hs_model.json int_hs_model.bin int_hs_input.json \
         cond.json cond_entropy.json cond_bitstream.json \
         frame0.rgb video_sota.json video_sota_bitstream.json; do
  if [ ! -s "$out/$f" ]; then
    echo "error: $out/$f was not produced" >&2
    missing=1
  fi
done
[ "$missing" -eq 0 ] || exit 1
echo "int_hyper vectors complete in $out"