#!/usr/bin/env bash
# 完整 nested CV：MPNN + MoE + Qwen2.5-7B QLoRA + soft-RAG
# 单卡一个分片。被 OOM / 抢占 / 断连杀掉后自动等显存并续跑。
set -uo pipefail

GPU=${GPU:?必须指定，例如 GPU=0}
SEED=${SEED:-42}
FOLDS=${FOLDS:-}                    # 空=跑全部 5 折；"0,1,2"=只跑这三折
REG_BYPASS=${REG_BYPASS:-on}        # on=delivery/size 绕开 MoE+LLM；off=接入
BATCH=${BATCH:-8}                   # 显存紧就设 4
N_OUTER=${N_OUTER:-5}
N_INNER=${N_INNER:-3}
N_TRIALS=${N_TRIALS:-20}
EPOCHS=${EPOCHS:-20}
PATIENCE=${PATIENCE:-5}
MIN_FREE_MB=${MIN_FREE_MB:-14000}
MAX_RETRY=${MAX_RETRY:-50}
RETRY_WAIT=${RETRY_WAIT:-300}
MIN_OK_SEC=${MIN_OK_SEC:-180}    # 存活不足这么久就挂 → 判为配置/代码错误，立即停止重试
PRETRAIN=${PRETRAIN:-}           # 非空且文件存在才注入 --init-from-pretrain

RUN_DIR="models/final_cv/${REG_BYPASS}_seed${SEED}"
TAG="gpu${GPU}$([ -n "${FOLDS}" ] && echo "_folds$(echo "${FOLDS}" | tr -d ',')")"
LOG="${RUN_DIR}/${TAG}.log"

export TRANSFORMERS_OFFLINE=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export TOKENIZERS_PARALLELISM=false
export CUDA_VISIBLE_DEVICES=${GPU}

mkdir -p "${RUN_DIR}"

# nvidia-smi 不受 CUDA_VISIBLE_DEVICES 影响，-i 用物理卡号
wait_free() {
  while :; do
    local free
    free=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits -i "${GPU}")
    [ "${free}" -ge "${MIN_FREE_MB}" ] && return 0
    echo "[$(date '+%F %T')] GPU${GPU} 仅空 ${free}MiB < ${MIN_FREE_MB}MiB，60s 后重试" >>"${LOG}"
    sleep 60
  done
}

EXTRA=()
if [ -n "${FOLDS}" ]; then
  # 注意：需先给 nested_cv_optuna.main 增加 --only-folds 选项，否则 Typer 直接报错
  EXTRA+=(--only-folds "${FOLDS}")
fi
if [ -n "${PRETRAIN}" ] && [ -f "${PRETRAIN}" ]; then
  EXTRA+=(--init-from-pretrain "${PRETRAIN}")
elif [ -n "${PRETRAIN}" ]; then
  echo "[$(date '+%F %T')] 警告：${PRETRAIN} 不存在，跳过预训练初始化" >>"${LOG}"
fi

for attempt in $(seq 1 "${MAX_RETRY}"); do
  wait_free
  echo "[$(date '+%F %T')] ATTEMPT ${attempt}/${MAX_RETRY} gpu=${GPU} seed=${SEED} folds=${FOLDS:-all} reg_bypass=${REG_BYPASS} batch=${BATCH} trials=${N_TRIALS}" >>"${LOG}"
  t0=${SECONDS}

  python -u -m lnp_ml.modeling.nested_cv_optuna \
    --input-path data/interim/internal.csv \
    --output-dir models/final_cv \
    --resume-dir "${RUN_DIR}" \
    --seed ${SEED} \
    --use-mpnn \
    --use-moe \
    --reg-bypass ${REG_BYPASS} \
    --use-llm --use-rag --rag-top-k 4 \
    --use-soft-prompt --llm-use-qlora \
    --llm-model-path models/qwen2.5-7b-instruct \
    --no-llm-freeze \
    --llm-max-length 1536 \
    --n-outer-folds ${N_OUTER} --n-inner-folds ${N_INNER} \
    --n-trials ${N_TRIALS} --epochs-per-trial ${EPOCHS} \
    --inner-patience ${PATIENCE} \
    --batch-size ${BATCH} --device cuda \
    ${EXTRA[@]+"${EXTRA[@]}"} \
    >>"${LOG}" 2>&1
  rc=$?
  dt=$((SECONDS - t0))

  if [ ${rc} -eq 0 ]; then
    echo "[$(date '+%F %T')] DONE gpu=${GPU} folds=${FOLDS:-all}（用时 ${dt}s）" >>"${LOG}"
    exit 0
  fi
  if [ ${dt} -lt ${MIN_OK_SEC} ]; then
    echo "[$(date '+%F %T')] 仅存活 ${dt}s 即以 ${rc} 退出，判定为配置/代码错误而非抢占，停止重试" >>"${LOG}"
    exit ${rc}
  fi
  echo "[$(date '+%F %T')] 运行 ${dt}s 后以 ${rc} 退出，${RETRY_WAIT}s 后带 --resume-dir 续跑" >>"${LOG}"
  sleep "${RETRY_WAIT}"
done

echo "[$(date '+%F %T')] 超过 MAX_RETRY 次仍失败，放弃" >>"${LOG}"
exit 1