#!/usr/bin/env bash
set -euo pipefail
# Canonical MNIST IDX bytes, served by the CVDF mirror with the authors' permission.
# https://yann.lecun.org/exdb/mnist/  https://github.com/cvdfoundation/mnist
# Checksums are fixed; downloads are kept outside source/distribution packages.
output=${1:-build/benchmarks/mnist/data}
mkdir -p "$output"
checksum() { if command -v sha256sum >/dev/null; then sha256sum "$1" | cut -d' ' -f1; else shasum -a 256 "$1" | cut -d' ' -f1; fi; }
while read -r name compressed raw; do
  archive="$output/$name.gz"
  if [[ ! -f "$archive" ]]; then
    temporary=$(mktemp "$output/.mnist-download.XXXXXX")
    trap 'rm -f "$temporary"' EXIT
    curl --fail --location --retry 2 --connect-timeout 15 --max-time 180 "https://storage.googleapis.com/cvdf-datasets/mnist/$name.gz" --output "$temporary"
    [[ $(checksum "$temporary") == "$compressed" ]] || { echo "Checksum mismatch: $name.gz" >&2; exit 1; }
    mv "$temporary" "$archive"
    trap - EXIT
  fi
  [[ $(checksum "$archive") == "$compressed" ]] || { echo "Checksum mismatch: $name.gz" >&2; exit 1; }
  if [[ ! -f "$output/$name" ]]; then
    temporary=$(mktemp "$output/.mnist-idx.XXXXXX")
    trap 'rm -f "$temporary"' EXIT
    gzip -dc "$archive" > "$temporary"
    [[ $(checksum "$temporary") == "$raw" ]] || { echo "IDX checksum mismatch: $name" >&2; exit 1; }
    mv "$temporary" "$output/$name"
    trap - EXIT
  fi
  [[ $(checksum "$output/$name") == "$raw" ]] || { echo "IDX checksum mismatch: $name" >&2; exit 1; }
  echo "$name verified"
done <<'CHECKSUMS'
train-images-idx3-ubyte 440fcabf73cc546fa21475e81ea370265605f56be210a4024d2ca8f203523609 ba891046e6505d7aadcbbe25680a0738ad16aec93bde7f9b65e87a2fc25776db
train-labels-idx1-ubyte 3552534a0a558bbed6aed32b30c495cca23d567ec52cac8be1a0730e8010255c 65a50cbbf4e906d70832878ad85ccda5333a97f0f4c3dd2ef09a8a9eef7101c5
t10k-images-idx3-ubyte 8d422c7b0a1c1c79245a5bcf07fe86e33eeafee792b84584aec276f5a2dbc4e6 0fa7898d509279e482958e8ce81c8e77db3f2f8254e26661ceb7762c4d494ce7
t10k-labels-idx1-ubyte f7ae60f92e00ec6debd23a6088c31dbd2371eca3ffa0defaefb259924204aec6 ff7bcfd416de33731a308c3f266cc351222c34898ecbeaf847f06e48f7ec33f2
CHECKSUMS