#!/usr/bin/env bash
# 样本级切分（StratifiedKFold）下 baseline vs MoE+LLM 对比实验。
# 串行跑两个分支（24GB 单卡装不下两个 7B）。SSH 断连、GPU 抢占、OOM
# 都靠 --resume-dir 的折级断点续跑恢复，已完成的折不会重算。
set -u

GPU=${GPU:-0}
SEED=${SEED:-42}
ROOT=${ROOT:-models/cv5_sample}
FIX_HP=${FIX_HP:-models/final/best_params.json}
FIX_EM=${FIX_EM:-11}
BATCH=${BATCH:-8}
N_OUTER=${N_OUTER:-5}
N_INNER=${N_INNER:-3}

MIN_FREE_MB=${MIN_FREE_MB:-12000}   # 显存不够就等，不硬上
WAIT_STEP=${WAIT_STEP:-60}
MAX_RETRY=${MAX_RETRY:-30}
RETRY_WAIT=${RETRY_WAIT:-120}
MIN_OK_SEC=${MIN_OK_SEC:-120}       # 存活不足这么久且无进度 -> 判为配置错误，停止重试

mkdir -p "${ROOT}"
LOG="${ROOT}/run_seed${SEED}.log"
log() { echo "[$(date '+%F %T')] $*" >>"${LOG}"; }

done_folds() {   # $1 = 分支目录
  ls "$1"/outer_fold_*/test_metrics.json 2>/dev/null | wc -l
}

wait_free() {
  while :; do
    local free
    free=$(nvidia-smi --id="${GPU}" --query-gpu=memory.free \
             --format=csv,noheader,nounits 2>/dev/null || echo 0)
    [ "${free}" -ge "${MIN_FREE_MB}" ] && break
    log "GPU${GPU} 空闲 ${free}MiB < ${MIN_FREE_MB}MiB，${WAIT_STEP}s 后重试"
    sleep "${WAIT_STEP}"
  done
}

# $1 = 分支名，其余参数为该分支特有的 CLI 开关
run_variant() {
  local name=$1; shift
  local dir="${ROOT}/${name}/seed${SEED}"
  mkdir -p "${dir}"

  local before after rc dt
  before=$(done_folds "${dir}")
  if [ "${before}" -ge "${N_OUTER}" ]; then
    log "${name}: ${before}/${N_OUTER} 折已完成，整个分支跳过"
    return 0
  fi

  local attempt
  for attempt in $(seq 1 "${MAX_RETRY}"); do
    wait_free
    before=$(done_folds "${dir}")
    log "${name} ATTEMPT ${attempt}/${MAX_RETRY}（已完成 ${before}/${N_OUTER} 折）"
    local t0=${SECONDS}

    CUDA_VISIBLE_DEVICES=${GPU} TRANSFORMERS_OFFLINE=1 TOKENIZERS_PARALLELISM=false \
    PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
    python -u -m lnp_ml.modeling.nested_cv_optuna \
      --input-path data/interim/internal.csv \
      --output-dir "${ROOT}" \
      --resume-dir "${dir}" \
      --seed "${SEED}" \
      --use-mpnn \
      --fix-hparams-json "${FIX_HP}" \
      --fix-epoch-mean "${FIX_EM}" \
      --n-outer-folds "${N_OUTER}" --n-inner-folds "${N_INNER}" \
      --batch-size "${BATCH}" --device cuda \
      ${@+"$@"} >>"${LOG}" 2>&1
    rc=$?
    dt=$((SECONDS - t0))

    if [ ${rc} -eq 0 ]; then
      log "${name} DONE（本次用时 ${dt}s，共 $(done_folds "${dir}")/${N_OUTER} 折）"
      return 0
    fi

    after=$(done_folds "${dir}")
    if [ ${dt} -lt ${MIN_OK_SEC} ] && [ "${after}" -le "${before}" ]; then
      log "${name} 仅存活 ${dt}s 且无进度（rc=${rc}），判为配置/代码错误而非抢占，停止重试"
      log "${name} 请查看本日志末尾的 traceback"
      return ${rc}
    fi
    log "${name} 运行 ${dt}s 后以 ${rc} 退出（进度 ${before}->${after}），${RETRY_WAIT}s 后续跑"
    sleep "${RETRY_WAIT}"
  done

  log "${name} 重试 ${MAX_RETRY} 次仍未完成，放弃"
  return 1
}

log "===== 开始：样本级切分对比实验 seed=${SEED} gpu=${GPU} ====="

run_variant baseline || { log "baseline 失败，不再继续 moe_llm"; exit 1; }

run_variant moe_llm \
  --use-moe --moe-n-experts 4 --moe-top-k 2 \
  --reg-bypass off \
  --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 \
  || { log "moe_llm 失败"; exit 1; }

log "===== 全部完成，两个分支各 ${N_OUTER} 折 ====="