#!/usr/bin/env bash
# One-shot bootstrap for the head-to-head Gradio benchmark on a RunPod
# (or any Ubuntu 22.04 + CUDA 12.1 L40S box).
#
# Usage on the pod (after you've uploaded your three .nii.gz files into
# /workspace/test_volumes/):
#
#   bash prepare_runpod.sh <hf_space_or_repo_url>
#
# Example:
#   bash prepare_runpod.sh https://huggingface.co/spaces/HarshithReddy01/voxpixel
#
# What it does:
#   1. Clones the VoxPixel repo into /workspace/voxpixel
#   2. Installs system deps + Python deps + builds mamba_ssm and
#      selective_scan_cuda_oflex (this is the slow step, ~10-15 min)
#   3. Runs the 9-call x 3-variant benchmark
#   4. Tars the CSV + log into /workspace/results.tar.gz so you can
#      download it from the Jupyter file browser

set -euo pipefail

REPO_URL="${1:-https://huggingface.co/spaces/HarshithReddy01/voxpixel}"
WORKDIR="/workspace/voxpixel"
VOLDIR="/workspace/test_volumes"
RESULTS="/workspace/gradio_baseline_results.csv"
LOG="/workspace/bench.log"

echo "===> [1/5] Sanity checks"
nvidia-smi | head -n 20
test -d "$VOLDIR" || { echo "ERROR: $VOLDIR missing. Upload your .nii.gz volumes there first."; exit 1; }
ls -lh "$VOLDIR"

echo "===> [2/5] Cloning repo into $WORKDIR"
if [ ! -d "$WORKDIR" ]; then
    git clone "$REPO_URL" "$WORKDIR"
fi
cd "$WORKDIR"
git pull --rebase || true

echo "===> [3/5] Installing system + Python deps + CUDA extensions"
apt-get update -qq
apt-get install -y -qq build-essential ninja-build git libgl1 libglib2.0-0 >/dev/null

python -m pip install --upgrade pip wheel setuptools packaging ninja >/dev/null

if ! python -c "import torch" 2>/dev/null; then
    pip install --index-url https://download.pytorch.org/whl/cu121 \
        torch torchvision torchaudio
fi

pip install --no-cache-dir -r requirements.txt

if ! python -c "import mamba_ssm" 2>/dev/null; then
    echo "  -> Building mamba-ssm (~5-10 min)..."
    pip install "mamba-ssm>=2.2.2" --no-build-isolation
fi

if ! python -c "import selective_scan_cuda_oflex" 2>/dev/null; then
    echo "  -> Building selective_scan_cuda_oflex (~3-5 min)..."
    cd SRMA-Mamba/selective_scan
    pip install --no-build-isolation -e . -v
    cd "$WORKDIR"
fi

echo "===> [4/5] Verifying the stack"
python - <<'PY'
import torch, mamba_ssm, selective_scan_cuda_oflex
print("torch:", torch.__version__, "cuda:", torch.cuda.is_available(),
      "device:", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu")
print("mamba_ssm OK:", mamba_ssm.__file__)
print("selective_scan_cuda_oflex OK:", selective_scan_cuda_oflex.__file__)
PY

echo "===> [5/5] Running 9-call benchmark across V1, V2, V3"
VOLS=( "$VOLDIR"/56.nii.gz "$VOLDIR"/58.nii.gz "$VOLDIR"/60.nii.gz )
for v in "${VOLS[@]}"; do
    test -f "$v" || { echo "ERROR: missing volume $v"; exit 1; }
done

python bench_gradio_baseline.py \
    --volumes "${VOLS[@]}" \
    --trials 3 \
    --modality T1 \
    --out "$RESULTS" 2>&1 | tee "$LOG"

echo "===> Done. Bundling results."
tar -czf /workspace/results.tar.gz \
    -C /workspace \
    "$(basename "$RESULTS")" \
    "$(basename "$LOG")"
ls -lh /workspace/results.tar.gz "$RESULTS" "$LOG"
echo
echo "Download /workspace/results.tar.gz from the Jupyter file browser,"
echo "then 'Stop' or 'Terminate' this pod from the RunPod dashboard."