#!/usr/bin/env bash
set -uo pipefail
cd "$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"

# ===== 环境（与原消融一致，保证可复现）=====
export TRANSFORMERS_OFFLINE=1 HF_HUB_OFFLINE=1 PYTHONUNBUFFERED=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True TOKENIZERS_PARALLELISM=false
export PYTHONHASHSEED=${PYTHONHASHSEED:-42}

# ===== 参数（与 s1_both 一致）=====
SEED=${SEED:-42}
GPU=${GPU:-0}
QBATCH=${QBATCH:-8}
N_OUTER=${N_OUTER:-5}; N_INNER=${N_INNER:-3}
N_TRIALS=${N_TRIALS:-20}; EPOCHS=${EPOCHS:-20}; PATIENCE=${PATIENCE:-5}
INPUT=${INPUT:-data/interim/internal.csv}
CHEM_CACHE=data/processed/chemeleon_embeddings.npz
QWEN=models/qwen2.5-7b-instruct

SRC=models/abl_full/s1_both/seed${SEED}
DST=models/abl_full/s1_both_chemeleon/seed${SEED}

# ===== 方案A：只拷 best_params + epoch_mean（切勿拷 test_metrics.json！）=====
for i in $(seq 0 $((N_OUTER-1))); do
  mkdir -p "$DST/outer_fold_$i"
  cp "$SRC/outer_fold_$i/best_params.json" "$DST/outer_fold_$i/"
  cp "$SRC/outer_fold_$i/epoch_mean.json"  "$DST/outer_fold_$i/"
done

# ===== 自动续跑：崩溃/被抢占后重试；靠 --resume-dir 跳过已完成的 fold =====
MAX_RETRY=${MAX_RETRY:-100}     # 最多重试次数
POLL=${POLL:-60}                # 重试/等待间隔(秒)
NEED=${NEED:-14000}             # 启动所需空闲显存(MiB)，Qwen QLoRA 约 14GB
mkdir -p logs

wait_free(){                    # 抢占期显存不够就等，避免一起来就 OOM
  while :; do
    free=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits -i "$GPU" 2>/dev/null)
    [ "${free:-0}" -ge "$NEED" ] && return 0
    echo "[$(date '+%F %T')] GPU${GPU} 空闲 ${free}MiB (<${NEED})，等待 ${POLL}s..."
    sleep "$POLL"
  done
}

n=1
while :; do
  wait_free
  echo "[$(date '+%F %T')] >>> s1_both_chemeleon try $n on GPU${GPU}"
  if CUDA_VISIBLE_DEVICES=$GPU python -m lnp_ml.modeling.nested_cv_optuna \
      --input-path "$INPUT" \
      --output-dir models/abl_full/s1_both_chemeleon --resume-dir "$DST" \
      --seed $SEED --device cuda \
      --n-outer-folds $N_OUTER --n-inner-folds $N_INNER \
      --n-trials $N_TRIALS --epochs-per-trial $EPOCHS --inner-patience $PATIENCE \
      --batch-size $QBATCH \
      --use-chemeleon --chemeleon-cache $CHEM_CACHE \
      --use-moe \
      --use-llm --use-rag --rag-top-k 4 --use-soft-prompt --no-llm-freeze \
      --llm-use-qlora --llm-model-path $QWEN; then
    echo "[$(date '+%F %T')] <<< s1_both_chemeleon DONE"; break
  fi
  (( n >= MAX_RETRY )) && { echo "[$(date '+%F %T')] !!! FAILED x${MAX_RETRY}"; exit 1; }
  echo "[$(date '+%F %T')] 崩溃/被抢占，${POLL}s 后断点续跑 (attempt $n)"; sleep "$POLL"; ((n++))
done