#!/usr/bin/env bash
# Purpose: create the course's Python environment on one machine: a uv virtual environment
#          holding the PyTorch build for its track (and MLX on a Mac), torchvision, JupyterLab
#          and matplotlib, then print the versions and the accelerator PyTorch can see
# Platform: spark (the aarch64 CUDA wheel; the NGC container is the other Spark path),
#           strix, mac, nvidia (inside WSL2 on Windows)
# Minimum memory: 8 GB
# Assumes: uv is installed and on PATH; internet access to download wheels; run from the
#          course directory (~/llm-course), which will hold .venv
#
# Usage: TRACK=<spark|strix|mac|nvidia> bash setup-env.sh
#   Optional environment variables:
#     PYTHON=3.12            Python version for the environment (uv downloads it if missing)
#     VENV=.venv             where to create it
#     FRESH=1                replace an existing .venv instead of reusing it
#     TORCH_BACKEND=<value>  override the track's PyTorch index: cu130, cu126, rocm7.2,
#                            cpu or auto (values uv pip install --torch-backend accepts;
#                            cu130 needs an NVIDIA driver of 580 or later, cu126 one of
#                            560.28.03 or later; cu128 resolves to torch 2.11.0, not 2.14.0,
#                            so this course does not use it)
#     ROCM_NIGHTLY=1         Track X only: use AMD's documented nightly index with --pre
#                            instead of the stable rocm7.2 index
#
# Track defaults: spark -> the cu130 index (has aarch64 wheels); nvidia -> auto (uv queries
# the installed driver and picks the most compatible CUDA index); strix -> rocm7.2;
# mac -> the PyPI wheel, which is the MPS build on Apple silicon, plus mlx.
set -euo pipefail

TRACK="${TRACK:-}"
PYTHON="${PYTHON:-3.12}"
VENV="${VENV:-.venv}"
FRESH="${FRESH:-0}"
TORCH_BACKEND="${TORCH_BACKEND:-}"
ROCM_NIGHTLY="${ROCM_NIGHTLY:-0}"

die() { echo "setup-env: $*" >&2; exit 1; }
command -v uv >/dev/null || die "uv is not installed; see https://docs.astral.sh/uv/getting-started/installation/ and rerun"
case "$TRACK" in
  spark|strix|mac|nvidia) ;;
  "") die "set TRACK=spark, TRACK=strix, TRACK=mac or TRACK=nvidia" ;;
  *) die "unknown TRACK '$TRACK' (expected spark, strix, mac or nvidia)" ;;
esac
if [[ "$TRACK" == "mac" && "$(uname -m)" != "arm64" ]]; then
  die "TRACK=mac needs an Apple silicon Mac; uname -m printed $(uname -m)"
fi

echo "==> uv $(uv --version | awk '{print $2}')"
if [[ -f "$VENV/pyvenv.cfg" && "$FRESH" != "1" ]]; then
  echo "==> Reusing the existing environment at $VENV (set FRESH=1 to replace it)"
elif [[ -f "$VENV/pyvenv.cfg" ]]; then
  echo "==> Replacing the existing environment at $VENV"
  uv venv --clear --python "$PYTHON" "$VENV"
else
  echo "==> Creating $VENV with Python $PYTHON"
  uv venv --python "$PYTHON" "$VENV"
fi
# shellcheck disable=SC1091
source "$VENV/bin/activate"

case "$TRACK" in
  spark)
    BACKEND="${TORCH_BACKEND:-cu130}"
    echo "==> Installing PyTorch (CUDA build, aarch64 wheel) and torchvision from the $BACKEND index"
    uv pip install torch torchvision --torch-backend="$BACKEND"
    ;;
  nvidia)
    BACKEND="${TORCH_BACKEND:-auto}"
    echo "==> Installing PyTorch (CUDA build) and torchvision with --torch-backend=$BACKEND"
    uv pip install torch torchvision --torch-backend="$BACKEND"
    ;;
  strix)
    if [[ "$ROCM_NIGHTLY" == "1" ]]; then
      echo "==> Installing PyTorch (ROCm build) and torchvision from AMD's documented nightly index"
      uv pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm7.2
    else
      BACKEND="${TORCH_BACKEND:-rocm7.2}"
      echo "==> Installing PyTorch (ROCm build) and torchvision with --torch-backend=$BACKEND"
      uv pip install torch torchvision --torch-backend="$BACKEND"
    fi
    ;;
  mac)
    echo "==> Installing PyTorch (the macOS arm64 wheel; MPS is built in), torchvision and MLX"
    uv pip install torch torchvision mlx
    ;;
esac

echo "==> Installing JupyterLab and matplotlib"
uv pip install jupyterlab matplotlib

echo "==> Verifying"
python - <<'PY'
import platform
import torch
print(f"python {platform.python_version()} ({platform.machine()})")
print(f"torch {torch.__version__}   cuda build: {torch.version.cuda}   hip build: {torch.version.hip}")
mps = getattr(torch.backends, "mps", None)
if torch.cuda.is_available():
    print(f"accelerator: cuda / {torch.cuda.get_device_name(0)}")
elif mps is not None and mps.is_available():
    print("accelerator: mps (Apple silicon GPU)")
else:
    built = "built" if (mps is not None and mps.is_built()) else "not built"
    print(f"accelerator: none visible; the lab will run on the CPU (mps support {built})")
PY
if [[ "$TRACK" == "mac" ]]; then
  python -c 'import mlx.core as mx; print("mlx", mx.__version__, "default device:", mx.default_device())'
fi
echo "jupyterlab $(jupyter lab --version)"
python -c 'import matplotlib; print("matplotlib", matplotlib.__version__)'
echo "==> Done. Copy the lines above into the Environment section of labbook.md."
echo "    Activate in every new terminal with: source $VENV/bin/activate"