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

ROOT_DIR="$(cd "$(dirname "$0")" && pwd)"
cd "$ROOT_DIR"

export UV_CACHE_DIR="${UV_CACHE_DIR:-$ROOT_DIR/.uv-cache}"
export UV_PYTHON_INSTALL_DIR="${UV_PYTHON_INSTALL_DIR:-$ROOT_DIR/.uv-python}"

uv venv --python 3.9 .venv
source .venv/bin/activate

uv pip install --upgrade pip wheel setuptools==80.9.0
uv pip install --python .venv/bin/python \
  --index-url https://download.pytorch.org/whl/cu128 \
  torch==2.8.0+cu128 torchvision==0.23.0+cu128 torchaudio==2.8.0+cu128
uv pip install --python .venv/bin/python -r requirements-pusht-5090.txt
uv pip install --python .venv/bin/python -e .

echo "uv environment ready at $ROOT_DIR/.venv"