#!/usr/bin/env bash
set -euo pipefail

# One-shot validation that prefers Apple Silicon MPS.

bash oss_llm/testNanoGPT/10_clone_nanogpt.sh
bash oss_llm/testNanoGPT/00_setup_env.sh
bash oss_llm/testNanoGPT/20_prepare_shakespeare.sh

# Train with MPS if available (script falls back to CPU)
bash oss_llm/testNanoGPT/31_train_mps_quick.sh

# Sample with MPS if available, else CPU
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
# shellcheck source=./_common.sh
source "${SCRIPT_DIR}/_common.sh"
require_venv
cd_nanogpt

DEVICE="mps"
if ! "${VENV_PY}" -c "import torch; import sys; sys.exit(0 if torch.backends.mps.is_available() else 1)"; then
  DEVICE="cpu"
fi

if [[ ! -f "out-shakespeare-char/ckpt.pt" ]]; then
  echo "ERROR: ckpt.pt not found after training" >&2
  exit 1
fi

"${VENV_PY}" sample.py --out_dir=out-shakespeare-char --device="${DEVICE}" --dtype=float32 --max_new_tokens=200

printf '\nOK: nanoGPT MPS smoke test completed (device=%s).\n' "${DEVICE}"