#!/bin/bash
# Source and weights are pinned; no training or upload starts during setup.
set -euo pipefail
cd "$(dirname "$0")"
STAGE="${1:-help}"
RUNTIME="${FORMA3D_V2_RUNTIME:-$HOME/Library/Application Support/Forma3D-tiny-2/venv}"
PY="$RUNTIME/bin/python"
export HF_HUB_DISABLE_TELEMETRY=1
export OMP_NUM_THREADS=4
export VECLIB_MAXIMUM_THREADS=4
REV=e80b8d895cc2e405894473abe4ce40b3472eaca7
case "$STAGE" in
  setup)
    if ! command -v uv >/dev/null; then
      echo 'Install uv first: brew install uv'; exit 1
    fi
    [ -x "$PY" ] || uv venv --python 3.11 "$RUNTIME"
    uv pip install --python "$PY" 'huggingface_hub[cli]>=0.34'
    "$RUNTIME/bin/hf" download Lucasgitguy-co/forma3d-tiny --revision "$REV" --local-dir upstream
    uv pip install --python "$PY" -r upstream/requirements-tested-macos.txt
    uv pip install --python "$PY" './upstream[data]'
    "$PY" -m image3d_student.cli doctor
    "$PY" -c 'import torch; assert torch.backends.mps.is_available(), "Metal unavailable; run in macOS Terminal"'
    mkdir -p runs
    uv pip freeze --python "$PY" > runs/environment.txt
    ;;
  prepare)
    [ -x "$PY" ] || { echo 'Run bash posttrain.sh setup first'; exit 1; }
    "$PY" prepare_v2.py --out data/forma3d-v2 --cache data/downloads \
      --scenario 128 --objaverse 128 --objaverse-shards 16 --resume \
      --grid 16 --views 4 --surface-points 4096 --seed 43 \
      --max-file-mb 16 --max-download-mb 4096 --reserve-gb 8
    ;;
  train)
    [ -x "$PY" ] || { echo 'Run bash posttrain.sh setup first'; exit 1; }
    "$PY" audit_data.py data/forma3d-v2 upstream/ATTRIBUTIONS.json
    # A fresh output directory is required. Ctrl-C saves last.pt, not optimizer resume.
    caffeinate -i "$PY" -m image3d_student.cli train \
      --manifest data/forma3d-v2/manifest.jsonl --init-from upstream/model.pt \
      --out runs/forma3d-v2 --device mps --width 64 \
      --steps 8000 --batch-size 2 --lr 0.00005 \
      --eval-every 200 --validation-fraction 0.2 --seed 43
    ;;
  evaluate)
    "$PY" stage_v2.py --run runs/forma3d-v2 --data data/forma3d-v2 \
      --base upstream/model.pt --release release/forma3d-tiny-2 --device mps
    ;;
  publish)
    # Explicit command only. Creates a private repo; makes it public after upload.
    "$PY" publish_v2.py release/forma3d-tiny-2
    ;;
  *)
    echo 'Run stages in order: setup, prepare, train, evaluate, publish (optional).'
    echo 'Example: bash posttrain.sh setup'
    ;;
esac