#!/usr/bin/env bash
set -euo pipefail

PROJECT_DIR="${PROJECT_DIR:-/workspace/protomorph_dinov3_runpod}"
cd "$PROJECT_DIR"

python3.11 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip wheel setuptools

# PyTorch 2.4.0 does not have official CUDA 13 wheels. Use cu124 on RunPod/RTX 3090.
pip install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu124
pip install -r requirements-core.txt

# Map RunPod variable names to standard Hugging Face names for download/upload tools.
if [[ -n "${hf_key:-}" && -z "${HF_TOKEN:-}" ]]; then
  export HF_TOKEN="$hf_key"
fi
if [[ -n "${hf_repo:-}" && -z "${HF_REPO_ID:-}" ]]; then
  export HF_REPO_ID="$hf_repo"
fi

mkdir -p "${HF_HOME:-/workspace/hf_cache}"
python scripts/create_random_head.py --num-classes 10 --out-dir checkpoints
python scripts/smoke_test_head_only.py

echo
echo "Setup complete."
echo "To start Jupyter: bash runpod/start_jupyter.sh"
echo "To upload to Hugging Face: bash runpod/upload_to_hf.sh"