#!/usr/bin/env bash
# 全量数据 final 模型：3-fold Optuna 调参 + 全量固定 epoch 重训
# MPNN + MoE + Qwen2.5-7B QLoRA + soft-RAG
# 被 OOM / 抢占 / SSH 断连杀掉后自动等显存并续跑（Optuna 从 sqlite 恢复）
set -uo pipefail

GPU=${GPU:-0}
SEED=${SEED:-42}
OUT=${OUT:-models/final}
N_TRIALS=${N_TRIALS:-20}
EPOCHS=${EPOCHS:-20}          # 与 nested CV 的 EPOCHS 保持一致
PATIENCE=${PATIENCE:-5}       # 与 nested CV 的 PATIENCE 保持一致
N_FOLDS=${N_FOLDS:-3}
BATCH=${BATCH:-8}
REG_BYPASS=${REG_BYPASS:-off}
FREEZE=${FREEZE:-3}
PRETRAIN=${PRETRAIN:-models/pretrain/mpnn/pretrain_delivery.pt}
MIN_FREE_MB=${MIN_FREE_MB:-14000}
MAX_RETRY=${MAX_RETRY:-50}
RETRY_WAIT=${RETRY_WAIT:-300}
MIN_OK_SEC=${MIN_OK_SEC:-180}   # 存活不足这么久就挂 → 判为配置/代码错误，立即停止重试

LOG="${OUT}/train.log"
STUDY="${OUT}/optuna_study.sqlite3"

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

mkdir -p "${OUT}"

# 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
}

# Optuna 的 study.optimize(n_trials=N) 在续跑时语义是"再跑 N 个"，不是"补到 N 个"。
# 所以每次重启前先数已完成的 trial 只补差额，否则被抢占几次就会多跑几十个 trial。
remaining_trials() {
  if [ ! -f "${STUDY}" ]; then echo "${N_TRIALS}"; return; fi
  python - "${STUDY}" "${N_TRIALS}" <<'PY' 2>/dev/null || echo "${N_TRIALS}"
import sys, optuna
optuna.logging.set_verbosity(optuna.logging.WARNING)
try:
    st = optuna.load_study(study_name="final_optuna_cv", storage=f"sqlite:///{sys.argv[1]}")
    done = sum(1 for t in st.trials if t.state == optuna.trial.TrialState.COMPLETE)
except Exception:
    done = 0
print(max(0, int(sys.argv[2]) - done))
PY
}

EXTRA=()
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
  NT=$(remaining_trials)
  echo "[$(date '+%F %T')] ATTEMPT ${attempt}/${MAX_RETRY} gpu=${GPU} 本次补 ${NT} 个 trial（目标 ${N_TRIALS}）reg_bypass=${REG_BYPASS} batch=${BATCH}" >>"${LOG}"
  t0=${SECONDS}

  python -u -m lnp_ml.modeling.final_train_optuna_cv \
    --input-path data/interim/internal.csv \
    --output-dir "${OUT}" \
    --seed ${SEED} \
    --n-folds ${N_FOLDS} \
    --n-trials ${NT} \
    --epochs-per-trial ${EPOCHS} \
    --patience ${PATIENCE} \
    --batch-size ${BATCH} \
    --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 \
    --freeze-backbone-epochs ${FREEZE} \
    --device cuda \
    ${EXTRA[@]+"${EXTRA[@]}"} \
    >>"${LOG}" 2>&1
  rc=$?
  dt=$((SECONDS - t0))

  if [ ${rc} -eq 0 ]; then
    echo "[$(date '+%F %T')] DONE（用时 ${dt}s）→ ${OUT}/model.pt" >>"${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 后续跑" >>"${LOG}"
  sleep "${RETRY_WAIT}"
done

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