#!/bin/bash
# setup-torch.sh — install PyTorch (CPU or CUDA) + the NVIDIA driver on one node.
#
# *** THIS IS THE ONE PLACE TO CHANGE GPU / CUDA / DRIVER / TORCH VERSIONS. ***
# setup-common.sh calls this after creating the venv; edit here to retarget CUDA.
#
# A physical GPU (lspci) is the trigger; the GPU hint is informational (a hint=true
# node with no GPU is mis-scheduled -> CPU). CUDA_TAG picks the CUDA build (cu121 is
# proven with nvidia-driver-535; newer drivers pair with cu124/cu126). CUDA_TAG=cpu
# forces the CPU wheel even on a GPU node.
#
# Config — env vars (what the notebook sets) or positional args (manual runs):
#   FL_USE_GPU        auto|true|false        ($1)  GPU hint
#   FL_CUDA_TAG       cu121|cu124|cu126|cpu  ($2)  PyTorch CUDA wheel; 'cpu' forces CPU
#   FL_NVIDIA_DRIVER  535|550|...                  apt NVIDIA driver version
#   FL_TORCH_SPEC     torch|torch==2.5.1|...        pip spec for the torch install
#   FL_VENV           path                          venv to install into (default ~/fl-venv)
set -uo pipefail
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
source "$SCRIPT_DIR/setup-lib.sh"

USE_GPU_HINT="${1:-${FL_USE_GPU:-auto}}"
CUDA_TAG="${2:-${FL_CUDA_TAG:-cu121}}"
NVIDIA_DRIVER="${FL_NVIDIA_DRIVER:-535}"
TORCH_SPEC="${FL_TORCH_SPEC:-torch}"
VENV="${FL_VENV:-$HOME/fl-venv}"
PY="$VENV/bin/python"
PIP="$VENV/bin/pip"

[ -x "$PIP" ] || { echo "ERROR: [setup-torch] venv pip not found at $PIP"; exit 1; }

GPU_PRESENT=no
if lspci 2>/dev/null | grep -qi nvidia; then
    GPU_PRESENT=yes
fi
if [ "$CUDA_TAG" = "cpu" ]; then
    GPU_PRESENT=no          # explicit CPU override, regardless of hardware
fi
echo "### PROGRESS: [setup-torch] hint=$USE_GPU_HINT detected=$GPU_PRESENT "\
"cuda_tag=$CUDA_TAG driver=$NVIDIA_DRIVER torch='$TORCH_SPEC'"

if [ "$GPU_PRESENT" = "yes" ]; then
    echo "### PROGRESS: [setup-torch] installing NVIDIA driver $NVIDIA_DRIVER (may take minutes)"
    apt_retry install -y -qq "nvidia-driver-$NVIDIA_DRIVER" "nvidia-utils-$NVIDIA_DRIVER" >/dev/null 2>&1 || \
        echo "### PROGRESS: [setup-torch] driver apt install returned nonzero (continuing)"
    sudo modprobe nvidia 2>/dev/null || true
    echo "### PROGRESS: [setup-torch] installing CUDA PyTorch wheel ($CUDA_TAG)"
    if "$PIP" install -q $TORCH_SPEC --index-url "https://download.pytorch.org/whl/${CUDA_TAG}"; then
        echo "### PROGRESS: [setup-torch] CUDA torch wheel installed"
    else
        echo "### PROGRESS: [setup-torch] CUDA wheel ($CUDA_TAG) failed; falling back to CPU wheel"
        "$PIP" install -q $TORCH_SPEC --index-url "https://download.pytorch.org/whl/cpu" || {
            echo "ERROR: [setup-torch] torch install failed"; exit 1; }
    fi
else
    if [ "$USE_GPU_HINT" = "true" ] && [ "$CUDA_TAG" != "cpu" ]; then
        echo "### PROGRESS: [setup-torch] WARNING — GPU requested but none detected (lspci); "\
"installing CPU wheel (node may be mis-scheduled)"
    fi
    echo "### PROGRESS: [setup-torch] installing CPU PyTorch wheel"
    "$PIP" install -q $TORCH_SPEC --index-url "https://download.pytorch.org/whl/cpu" || {
        echo "ERROR: [setup-torch] torch install failed"; exit 1; }
fi

# If the hardware is present but CUDA isn't live, the driver usually just needs a
# reboot to load its kernel module — say so clearly (training still runs on CPU).
if [ "$GPU_PRESENT" = "yes" ]; then
    CUDA_OK="$("$PY" -c 'import torch;print(torch.cuda.is_available())' 2>/dev/null)"
    if [ "$CUDA_OK" != "True" ]; then
        echo "### PROGRESS: [setup-torch] NOTE — GPU present but torch.cuda.is_available()=False; "\
"reboot the node to load the NVIDIA driver:  sudo reboot   (training runs on CPU until then)"
    else
        echo "### PROGRESS: [setup-torch] CUDA is live (torch.cuda.is_available()=True)"
    fi
fi

echo "### PROGRESS: [setup-torch] done"
