#!/bin/bash
set -euo pipefail
cd "$(dirname "$0")"
export HF_HOME="$PWD/.cache/huggingface"
export UV_CACHE_DIR="$PWD/.cache/uv"
export UV_LINK_MODE=copy
export PYTHONUNBUFFERED=1
export CACTUS_TELEMETRY=0
export PYTORCH_ENABLE_MPS_FALLBACK=0
mkdir -p logs
case "${1:-help}" in
  setup)
    uv run --no-project --python 3.12 python setup.py "${2:-all}"
    ;;
  fleet)
    test -x .venvs/fleet/bin/python || { echo "Run: bash run.sh setup fleet"; exit 1; }
    caffeinate -i .venvs/fleet/bin/python fleet/run.py all 2>&1 | tee -a logs/fleet.log
    ;;
  forma)
    test -x .venvs/forma/bin/python || { echo "Run: bash run.sh setup forma"; exit 1; }
    caffeinate -i .venvs/forma/bin/python forma3d/run.py all 2>&1 | tee -a logs/forma.log
    ;;
  all)
    bash run.sh fleet
    bash run.sh forma
    ;;
  report)
    uv run --no-project --python 3.12 python report.py
    ;;
  test)
    .venvs/forma/bin/python -m unittest discover -s tests -v
    if test -x .venvs/fleet/bin/python; then
      JAX_PLATFORMS=cpu .venvs/fleet/bin/python -m unittest discover -s tests -p test_fleet.py -v
    fi
    ;;
  *)
    echo "Usage: bash run.sh {setup [all|fleet|forma]|fleet|forma|all|report|test}"
    ;;
esac