#!/usr/bin/env bash
# Install GPU PyTorch for the surfer experiments (local RTX 3050 Ti, driver 580 → CUDA 12.x).
#
# RECOMMENDED (no sudo — installs to ~/.local, matching your existing numpy/scipy/databento):
#     bash scripts/install_torch_gpu.sh
#
# System-wide (only if you really want it):
#     sudo bash scripts/install_torch_gpu.sh
#
# Pick a different CUDA wheel tag if cu124 ever 404s (cu126 / cu128 also work on driver 580):
#     bash scripts/install_torch_gpu.sh cu126
set -euo pipefail

CUDA_TAG="${1:-cu124}"
TORCH_INDEX="https://download.pytorch.org/whl/${CUDA_TAG}"

if [ "$(id -u)" -eq 0 ]; then
  echo "[install] running as ROOT → system-wide site-packages (--break-system-packages)"
  FLAGS="--break-system-packages"
else
  echo "[install] running as USER → ~/.local (matches existing packages; no sudo needed)"
  FLAGS="--user --break-system-packages"
fi

echo "[install] torch from ${TORCH_INDEX}"
python3 -m pip install ${FLAGS} torch --index-url "${TORCH_INDEX}"

echo "[verify] importing torch + checking CUDA on the local GPU"
python3 - <<'PY'
import torch
print("  torch:", torch.__version__)
print("  cuda available:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("  device:", torch.cuda.get_device_name(0),
          "| capability:", torch.cuda.get_device_capability(0))
    x = torch.randn(1_000_000, device="cuda")
    print("  GPU tensor op OK, sum =", float(x.sum()))
else:
    raise SystemExit("CUDA NOT available to torch — check driver / try a different CUDA_TAG (cu126/cu128)")
print("  ✅ GPU PyTorch ready")
PY