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

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd -- "${SCRIPT_DIR}/.." && pwd)"
PYTHON_BIN="${PYTHON_BIN:-python}"
CONFIG="${CONFIG:-${SCRIPT_DIR}/configs/stead_empty8_local.yaml}"
TRAIN_H5="${PROJECT_ROOT}/Train/data/processed/from_raw/stead/train/stead_100hz_60s_train.h5"
NUM_GPUS="${NUM_GPUS:-8}"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --num-gpus)
      [[ $# -ge 2 ]] || { echo "--num-gpus requires a value" >&2; exit 2; }
      NUM_GPUS="$2"
      shift 2
      ;;
    --num-gpus=*)
      NUM_GPUS="${1#*=}"
      shift
      ;;
    *)
      echo "Unknown argument: $1" >&2
      echo "Usage: bash Train/run_train_local.sh [--num-gpus N]" >&2
      exit 2
      ;;
  esac
done
[[ "${NUM_GPUS}" =~ ^[1-9][0-9]*$ ]] || { echo "NUM_GPUS must be a positive integer" >&2; exit 2; }

cd "${PROJECT_ROOT}"
[[ -f "${TRAIN_H5}" ]] || {
  echo "Missing processed STEAD data: ${TRAIN_H5}" >&2
  echo "Run: bash Train/run_download_prepare.sh" >&2
  exit 1
}

export PYTHONUNBUFFERED=1
echo "Training with ${NUM_GPUS} GPU(s); global batch size is read from ${CONFIG}"
if [[ "${NUM_GPUS}" -eq 1 ]]; then
  exec "${PYTHON_BIN}" Train/training/train.py --config "${CONFIG}"
fi
exec "${PYTHON_BIN}" -m torch.distributed.run \
  --standalone \
  --nproc_per_node="${NUM_GPUS}" \
  Train/training/train.py --config "${CONFIG}"