#!/usr/bin/env bash
set -euo pipefail
cd ~/robot-benchmark
echo "[$(date +%T)] creating venv"
uv venv --python 3.12 .venv
PY=.venv/bin/python
echo "[$(date +%T)] installing torch (cu128 for Blackwell sm_120)"
uv pip install --python $PY torch --index-url https://download.pytorch.org/whl/cu128
echo "[$(date +%T)] installing core deps"
uv pip install --python $PY numpy datasets huggingface_hub tqdm pillow imageio imageio-ffmpeg opencv-python-headless matplotlib scipy
echo "[$(date +%T)] GPU smoke test"
$PY - <<PYEOF
import torch
print("torch", torch.__version__, "cuda_build", torch.version.cuda)
print("cuda_available", torch.cuda.is_available())
if torch.cuda.is_available():
    print("device", torch.cuda.get_device_name(0))
    print("capability", torch.cuda.get_device_capability(0))
    a=torch.randn(4096,4096,device="cuda"); b=torch.randn(4096,4096,device="cuda")
    c=(a@b); torch.cuda.synchronize()
    print("matmul_ok mean=", float(c.mean()))
PYEOF
echo "[$(date +%T)] SETUP_DONE"