#!/usr/bin/env bash
# Set up and run the Laya-CoreML trolley-problem benchmark on Apple Silicon.
#
# Requires macOS 15+ on Apple Silicon and a Python 3.11-3.13 interpreter
# (Laya-CoreML does not yet support 3.14). Downloads three Core ML checkpoints
# (~1.5 GB total) from Hugging Face on first run, then reproduces every number
# in the post: the two full-prompt checkpoints, the compute-unit sweep, and
# the compressed-prompt Neural Engine run. All output (venv, models, results)
# is written under this directory in `.venv/`, `models/`, and `out/`, which
# stay out of the published bundle.
set -euo pipefail

cd "$(dirname "$0")"

PYTHON=""
for candidate in python3.13 python3.12 python3.11; do
  if command -v "$candidate" >/dev/null 2>&1; then
    PYTHON="$candidate"
    break
  fi
done
if [ -z "$PYTHON" ]; then
  echo "No Python 3.11-3.13 interpreter found on PATH (Laya-CoreML does not support 3.14+)." >&2
  exit 1
fi

mkdir -p out
"$PYTHON" -m venv .venv
.venv/bin/python -m pip install --quiet --upgrade pip
.venv/bin/python -m pip install --quiet 'laya-coreml==0.1.0' 'huggingface_hub==1.32.0'

.venv/bin/hf download aac6fef/laya-typed-decisions-coreml \
  --revision 28d24fa8d67a3264556b23391ec6c3fd98573056 \
  --local-dir models/typed-decisions

.venv/bin/hf download aac6fef/laya-multilingual-coreml \
  --revision 8139e9089273319512c730218903784074133187 \
  --local-dir models/multilingual

.venv/bin/hf download aac6fef/laya-multilingual-coreml-ane \
  --revision 39d6a9b3d0f67f06da74fbade6121ea134cbdb21 \
  --local-dir models/ane

# The two full-prompt checkpoints, each against the identical, hash-verified
# eleven-question state.
.venv/bin/python bench_trolley.py \
  --model-dir models/typed-decisions \
  --checkpoint-label aac6fef/laya-typed-decisions-coreml \
  --revision 28d24fa8d67a3264556b23391ec6c3fd98573056 \
  --state state.json --questions questions.json \
  --jev-response jev-response.json \
  --out out/response-typed-decisions.json \
  --timings out/timings-typed-decisions.json

.venv/bin/python bench_trolley.py \
  --model-dir models/multilingual \
  --checkpoint-label aac6fef/laya-multilingual-coreml \
  --revision 8139e9089273319512c730218903784074133187 \
  --state state.json --questions questions.json \
  --jev-response jev-response.json \
  --out out/response-multilingual.json \
  --timings out/timings-multilingual.json

# The compute-unit sweep on the `action` question alone (cpu_gpu vs cpu_ne vs
# cpu vs all), on the typed-decisions checkpoint.
.venv/bin/python bench_trolley.py --sweep \
  --model-dir models/typed-decisions \
  --checkpoint-label aac6fef/laya-typed-decisions-coreml \
  --revision 28d24fa8d67a3264556b23391ec6c3fd98573056 \
  --state state.json --questions questions.json \
  --timings out/sweep-typed-decisions.json

# The Neural Engine sidebar: a different, compressed prompt (state-96.json /
# questions-96.json) that fits the ANE bundle's 96-token budget. --skip-hash-check
# is required here because this is deliberately not the eleven-question prompt.
.venv/bin/python bench_trolley.py \
  --model-dir models/ane \
  --checkpoint-label aac6fef/laya-multilingual-coreml-ane \
  --revision 39d6a9b3d0f67f06da74fbade6121ea134cbdb21 \
  --compute-units cpu_ne \
  --state state-96.json --questions questions-96.json \
  --skip-hash-check \
  --out out/response-ane.json \
  --timings out/timings-ane.json

echo "Done. See the out/ directory for every result file."