#!/usr/bin/env bash
set -uo pipefail
cd ~/lnp_ml
export PYTHONUNBUFFERED=1 TOKENIZERS_PARALLELISM=false
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export PYTHONHASHSEED=${PYTHONHASHSEED:-42}

GPU=${GPU:-0}
SEED=${SEED:-42}
NEED=${NEED:-8000}
POLL=${POLL:-60}
MAX_RETRY=${MAX_RETRY:-100}
NAME=s2_mole
RUNDIR=models/abl_full/${NAME}/seed${SEED}
mkdir -p "$RUNDIR" logs

n=1
while :; do
  while :; do
    free=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits -i "$GPU" 2>/dev/null)
    [ "${free:-0}" -ge "$NEED" ] && break
    echo "[$(date '+%F %T')] GPU${GPU} 剩 ${free}MiB(<${NEED}) 等 ${POLL}s..."
    sleep "$POLL"
  done

  echo "[$(date '+%F %T')] >>> try $n on GPU${GPU}"
  if CUDA_VISIBLE_DEVICES="$GPU" python -m lnp_ml.modeling.nested_cv_optuna \
       --input-path data/interim/internal.csv \
       --output-dir models/abl_full/${NAME} --resume-dir "${RUNDIR}" \
       --seed ${SEED} --device cuda \
       --n-outer-folds 5 --n-inner-folds 3 \
       --n-trials 20 --epochs-per-trial 20 --inner-patience 5 \
       --batch-size 16 \
       --use-mole --mole-cache data/processed/mole_embeddings.npz \
       2>&1 | tee -a "${RUNDIR}/run.log"; then
    echo "[$(date '+%F %T')] <<< DONE"; break
  fi
  (( n > MAX_RETRY )) && { echo "FAILED x${MAX_RETRY}"; exit 1; }
  echo "[$(date '+%F %T')] 崩溃/被抢占, ${POLL}s 后断点续跑 (try $((n+1)))..."
  sleep "$POLL"; ((n++))
done