#!/usr/bin/env bash
# reproduce.sh: the cross-simulator corpus, regenerated and judged from nothing but this
# repository, a Rust toolchain and the network.
#
#   bash wai/quantum-cross-simulator/reproduce.sh [--tier ci|full|gpu] [--out FILE]
#        [--jobs N] [--inject FAULT] [--self-test] [--perf] [--no-provider] [--memory-gb N]
#
# It needs macOS on arm64 or Linux on x86_64, curl, tar, a C compiler and rustup (or cargo at
# the version rust-toolchain.toml names). No account and no key: every download is pinned by
# version and SHA-256.
#
#   1. uv, the pinned release, unless that version is already on PATH;
#   2. CPython 3.12.16 and the engines of env/requirements.lock, every wheel checked against its
#      hash, in a virtual environment of their own;
#   3. the engine under test, built from this checkout with its lockfile (--locked);
#   4. tools/compare.py: every circuit on every engine, each answer held to its reference within
#      its stated tolerance; it prints each failure and PASS or FAIL;
#   5. provider/test.sh: the provider template, added, held to the protocol's conformance vectors,
#      run and its receipt verified.
#
# --self-test runs each fault compare.py can inject on one circuit (on two for the gpu tier's
# small circuits refused), and passes only if every one of those runs fails, on the engine and the
# check that fault breaks. --perf runs tools/perf.py.
# --tier gpu builds the engine with the platform's GPU methods (Metal on macOS) and the fabric's
# driver (wai-quantum's example fabric_state, with every GPU member this platform has), and holds
# each member's state to the CPU's bit for bit; with --self-test, its faults must be caught (a
# member's bit flipped; every member gone; every circuit under 3 qubits refused and every other
# answered: the last two only what QCORPUS_GPU_EXPECT declares catches, since a member it declares
# must pass every circuit it is asked on).
# --memory-gb stops a job whose processes hold more than that, as a time-out stops it.
# Downloads and builds go to $QCORPUS_CACHE (default ~/.cache/wai-quantum-cross-simulator).
# Exit status: 0 every check passed, 1 a check failed, 2 the run could not be made.
set -euo pipefail

HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO="$(cd "$HERE/../.." && pwd)"
TIER=ci
OUT=""
JOBS=2
INJECT=""
SELFTEST=0
PERF=0
PROVIDER=1
MEMORY_GB=""
while [ $# -gt 0 ]; do
  case "$1" in
    --tier) TIER="$2"; shift 2 ;;
    --out) OUT="$2"; shift 2 ;;
    --jobs) JOBS="$2"; shift 2 ;;
    --inject) INJECT="$2"; shift 2 ;;
    --self-test) SELFTEST=1; shift ;;
    --perf) PERF=1; shift ;;
    --no-provider) PROVIDER=0; shift ;;
    --memory-gb) MEMORY_GB="$2"; shift 2 ;;
    -h|--help) sed -n '2,32p' "$0"; exit 0 ;;
    *) echo "reproduce.sh: unknown argument $1" >&2; exit 2 ;;
  esac
done

die() { echo "reproduce.sh: $*" >&2; exit 2; }
# The gpu tier's faults member_absent and small_refused are caught by what the host declares it
# has: without a fabric member declared there is nothing to catch them, and the self-test cannot
# be judged.
if [ $SELFTEST -eq 1 ] && [ "$TIER" = gpu ] && [[ "${QCORPUS_GPU_EXPECT:-}" != *:* ]]; then
  die "the gpu tier's self-test needs QCORPUS_GPU_EXPECT to name a fabric member this host must have (family:key, as the gpu tier's results name it)"
fi
CACHE="${QCORPUS_CACHE:-${XDG_CACHE_HOME:-$HOME/.cache}/wai-quantum-cross-simulator}"
mkdir -p "$CACHE"
for t in curl tar; do command -v "$t" >/dev/null || die "needs $t"; done

# 1. uv, pinned.
UV_VERSION=0.13.6
# The two platforms every pinned engine publishes wheels for (env/requirements.in says which do not
# elsewhere). On another, the run stops here rather than build a third party's code from source.
case "$(uname -s)-$(uname -m)" in
  Darwin-arm64) UV_TARGET=aarch64-apple-darwin; UV_SHA=60fbd5c8a1ed681845eea28d9d622fcde03c51e91a43465b64a1e17db1ed3472 ;;
  Linux-x86_64) UV_TARGET=x86_64-unknown-linux-gnu; UV_SHA=f29c3e94ffabd7e6f3ba521dba169a9b7ec12100bf3257087f0f24f6123f3645 ;;
  *) die "the pinned engines publish no wheels for $(uname -s) $(uname -m): run it on macOS on arm64 or Linux on x86_64" ;;
esac
UV=""
if command -v uv >/dev/null && [ "$(uv --version | awk '{print $2}')" = "$UV_VERSION" ]; then
  UV="$(command -v uv)"
elif [ -x "$CACHE/uv-$UV_VERSION/uv" ]; then
  UV="$CACHE/uv-$UV_VERSION/uv"
else
  echo "== uv $UV_VERSION ($UV_TARGET)" >&2
  f="$CACHE/uv-$UV_TARGET.tar.gz"
  curl --proto '=https' --tlsv1.2 -sSfL -o "$f" "https://github.com/astral-sh/uv/releases/download/$UV_VERSION/uv-$UV_TARGET.tar.gz"
  got="$( (command -v sha256sum >/dev/null && sha256sum "$f" || shasum -a 256 "$f") | awk '{print $1}')"
  [ "$got" = "$UV_SHA" ] || die "uv's archive has SHA-256 $got, not $UV_SHA"
  mkdir -p "$CACHE/uv-$UV_VERSION"
  tar -xzf "$f" -C "$CACHE/uv-$UV_VERSION" --strip-components 1
  UV="$CACHE/uv-$UV_VERSION/uv"
fi
export UV_CACHE_DIR="$CACHE/uv-cache" UV_PYTHON_INSTALL_DIR="$CACHE/python"

# 2. The interpreter and the engines, every wheel against its hash.
PY_VERSION=3.12.16
LOCK="$HERE/env/requirements.lock"
LOCK_ID="$( (command -v sha256sum >/dev/null && sha256sum "$LOCK" || shasum -a 256 "$LOCK") | cut -c1-16)"
VENV="$CACHE/venv-$PY_VERSION-$LOCK_ID"
if [ ! -x "$VENV/bin/python" ] || [ ! -f "$VENV/.complete" ]; then
  echo "== CPython $PY_VERSION and the pinned engines" >&2
  rm -rf "$VENV"
  "$UV" venv --quiet --python "$PY_VERSION" --python-preference only-managed "$VENV"
  "$UV" pip install --quiet --python "$VENV/bin/python" --require-hashes --only-binary :all: -r "$LOCK"
  touch "$VENV/.complete"
fi
PY="$VENV/bin/python"

# 3. The engine under test, from this checkout.
command -v cargo >/dev/null || die "needs cargo: install rustup (https://rustup.rs); rust-toolchain.toml picks the version"
# The gpu tier's builds: the platform's single-precision GPU methods in the engine (Metal, on
# macOS), and the fabric's driver with every GPU member the platform can have.
CLI_FEATURES=""
FABRIC_FEATURES="quantum_fabric_webgpu"
case "$(uname -s)" in
  Darwin) CLI_FEATURES="metal"; FABRIC_FEATURES="$FABRIC_FEATURES,quantum_fabric_metal" ;;
  Linux) FABRIC_FEATURES="$FABRIC_FEATURES,quantum_fabric_cuda" ;;
esac
[ "$TIER" = gpu ] || CLI_FEATURES=""
echo "== the engine under test (cargo build --locked --release${CLI_FEATURES:+ --features $CLI_FEATURES})" >&2
( cd "$REPO/wai/wai-quantum-cli" && cargo build --locked --release --quiet --bin wai-quantum --example cross_simulator ${CLI_FEATURES:+--features "$CLI_FEATURES"} ) || die "the build failed"
TARGET="${CARGO_TARGET_DIR:-$REPO/wai/wai-quantum-cli/target}"
export WAI_CROSS_SIMULATOR="$TARGET/release/examples/cross_simulator" WAI_QUANTUM="$TARGET/release/wai-quantum"
[ -x "$WAI_CROSS_SIMULATOR" ] || die "no driver at $WAI_CROSS_SIMULATOR"
if [ "$TIER" = gpu ]; then
  echo "== the fabric's driver (cargo build --locked --release --features $FABRIC_FEATURES)" >&2
  ( cd "$REPO/wai/wai-quantum" && cargo build --locked --release --quiet --example fabric_state --features "$FABRIC_FEATURES" ) || die "the fabric driver's build failed"
  export WAI_FABRIC_DRIVER="${CARGO_TARGET_DIR:-$REPO/wai/wai-quantum/target}/release/examples/fabric_state"
  [ -x "$WAI_FABRIC_DRIVER" ] || die "no fabric driver at $WAI_FABRIC_DRIVER"
fi
# What was built and run, as git tree ids of the files as they are on disk, committed or not:
# `git rev-parse <commit>:wai/wai-quantum` gives the same id for any commit that holds the same
# engine. A tree that differs from HEAD's is marked "(uncommitted)", and outside a git checkout
# the id is "unknown".
INDEX="$(mktemp)"
tree() {
  local id head
  id=$( (cd "$REPO" && cp "$(git rev-parse --git-path index)" "$INDEX" && GIT_INDEX_FILE="$INDEX" git add -A -- "$1" >/dev/null 2>&1 && GIT_INDEX_FILE="$INDEX" git write-tree --prefix="$1/") 2>/dev/null ) || { echo unknown; return; }
  head=$( (cd "$REPO" && git rev-parse "HEAD:$1") 2>/dev/null || true)
  [ "$id" = "$head" ] && echo "$id" || echo "$id (uncommitted)"
}
# The toolchain cargo built with: rustc as rust-toolchain.toml picks it in the checkout.
export WAI_BUILD_INFO="wai/wai-quantum $(tree wai/wai-quantum), wai/wai-quantum-qpu $(tree wai/wai-quantum-qpu), wai/wai-quantum-cli $(tree wai/wai-quantum-cli), wai/quantum-cross-simulator/tools $(tree wai/quantum-cross-simulator/tools), wai/quantum-cross-simulator/corpus $(tree wai/quantum-cross-simulator/corpus); $(cd "$REPO/wai/wai-quantum-cli" && rustc --version 2>/dev/null || true)"
rm -f "$INDEX"

if [ $SELFTEST -eq 1 ]; then
  # Each fault must fail the run on the check it breaks, on the engine it was put into, and
  # nothing but that engine may fail. The gpu tier's fault runs with the gpu tier alone.
  status=0
  ran=0
  while read -r fault eng tier only timeout what; do
    if [ "$TIER" = gpu ]; then [ "$tier" = gpu ] || continue; else [ "$tier" != gpu ] || continue; fi
    ran=$((ran + 1))
    log="$CACHE/selftest-$fault.log"
    if "$PY" -I "$HERE/tools/compare.py" --tier "$tier" --jobs "$JOBS" --inject "$fault" --only "$only" --timeout "$timeout" > "$log" 2>&1; then
      echo "self-test: the fault '$fault' passed unnoticed" >&2; status=1
    elif ! grep -q -E "^FAIL +[^ ]+ +$eng +.*$what" "$log"; then
      echo "self-test: the fault '$fault' failed the run, but not $eng's check '$what':" >&2; grep "^FAIL " "$log" | head -3 >&2; tail -2 "$log" >&2; status=1
    elif grep "^FAIL " "$log" | grep -v -E "^FAIL +[^ ]+ +$eng " | grep -q .; then
      echo "self-test: the fault '$fault' failed another engine as well:" >&2; grep "^FAIL " "$log" | grep -v -E "^FAIL +[^ ]+ +$eng " | head -3 >&2; status=1
    else
      echo "self-test: the fault '$fault' was caught: $(grep -m1 -E "^FAIL +[^ ]+ +$eng +.*$what" "$log" | cut -c1-200)"
    fi
  done <<'FAULTS'
amplitude aer.statevector ci qft_n5_x19 900 distance from the exact state
bit_order wai.statevector ci bit_order_product 900 distance from the exact state
counts stim.tableau ci reset_after_superposition 900 largest deviation of an outcome's count
counts_spread aer.statevector full qft_counts_n8 900 L1 distance
dem wai.qec ci surface_z_d3 900 mechanisms
qec_rate wai.qec ci surface_z_d3 900 detection and flip rates
stabilizer stim.tableau ci clifford_random_n65 900 reduced stabilizer group
fixed_point wai.exact ci ghz_n3 900 state hash
drop_gate wai.exact ci small_angle_k32 900 the op-log is the circuit
drop_phase wai.statevector full below_grid_k48 900 holds the phase
refuse wai.mps ci mps_ghz_n30 900 gave no answer
crash wai.auto ci ghz_n3 900 gave no answer
hang wai.qec ci surface_z_d3 60 gave no answer
accept wai.exact ci rotations_n6 900 accepts a circuit it must refuse
qec_words wai.qec ci qec_inverted_and_tagged 900 gave no answer
qec_crash wai.qec ci qec_refused_random_detector 900 gave no answer
exact_words wai.exact ci rotations_n6 900 gave no answer
qec_short wai.qec ci surface_z_d3 900 detectors, observables and shots returned
exact_turn wai.exact full qft_wide_n20 900 64 amplitudes against the closed form
phase_flip wai.statevector full repeated_phase_k19_quarter 900 distance from the exact state, with its phase
mps_phase aer.mps full mps_qft_n40 900 amplitudes against the closed form
counts_collapse stim.tableau ci clifford_counts_n40 900 the most frequent outcome's count
pauli_bound wai.pauli ci pauli_rotations_n6 900 within its truncation bound and rounding
stabrank_amp wai.stabrank ci stabrank_clifford_t_n5 900 amplitudes of the exact sum
cutrank_amp wai.cutrank ci cutrank_n8 900 amplitudes against the integer state, 2-norm
trajectory_mean wai.trajectories ci noisy_ghz_n3 900 the mean of
density_value wai.density_matrix ci noisy_ghz_n3 900 under noise against the integer density matrix
frame_sign wai.stabilizer ci noisy_ghz_n3 900 the mean of
zne_raw wai.mitigate ci mitigate_zne_ghz_n3 900 extrapolated, against the Richardson
readout_raw wai.mitigate ci mitigate_readout_n3 900 every quasi-probability
vd_purity wai.mitigate full mitigate_vd_n2 900 the purity
pulse_p1 wai.pulse ci pulse_rabi 900 against the closed form
decode_flip wai.matching ci decode_repetition_d3 900 not a least-weight correction
uf_flip wai.union_find ci decode_repetition_d3 900 single edges of the graph
member_bit wai\.fabric\[(webgpu|metal|cuda):[-a-z0-9]+\]? gpu ghz_n5 900 against the CPU's bit-exact method
member_absent wai\.fabric\[(webgpu|metal|cuda):[-a-z0-9]+\]? gpu ghz_n5 900 checks passed by what this host declares it has
small_refused (wai\.statevector_single|wai\.fabric\[(webgpu|metal|cuda):[-a-z0-9]+\]?) gpu ghz_n2,ghz_n5 900 checks passed by what this host declares it has
FAULTS
  # A self-test that ran no fault proves nothing.
  [ $ran -gt 0 ] || { echo "self-test: no fault ran" >&2; status=1; }
  [ $status -eq 0 ] && echo "self-test: every one of the $ran injected faults was caught" || echo "self-test: FAIL" >&2
  exit $status
fi

if [ $PERF -eq 1 ]; then
  # QCORPUS_PERF_ARGS passes options to tools/perf.py: --repeat N, --quick, --only PREFIX.
  # shellcheck disable=SC2086
  exec "$PY" -I "$HERE/tools/perf.py" ${OUT:+--out "$OUT"} ${QCORPUS_PERF_ARGS:-}
fi

# 4. The corpus.
[ -n "$OUT" ] || OUT="$CACHE/run-$TIER-$(uname -s)-$(uname -m).json"
echo "== the corpus, tier $TIER" >&2
set +e
"$PY" -I "$HERE/tools/compare.py" --tier "$TIER" --jobs "$JOBS" --out "$OUT" ${INJECT:+--inject "$INJECT"} ${MEMORY_GB:+--memory-gb "$MEMORY_GB"} 2> "$CACHE/compare-progress.log"
corpus=$?
set -e
[ $corpus -le 1 ] || die "the comparison could not be made (see $CACHE/compare-progress.log)"
echo "results: $OUT"

# 5. The provider template.
provider=0
if [ $PROVIDER -eq 1 ] && [ "$TIER" != gpu ]; then
  echo "== the provider template" >&2
  bash "$HERE/provider/test.sh" "$WAI_QUANTUM" "$PY" || provider=1
fi
[ $corpus -eq 0 ] && [ $provider -eq 0 ] && { echo "reproduce: PASS"; exit 0; }
echo "reproduce: FAIL"
exit 1