#!/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}
mkdir -p logs models/abl_full reports

# ===== 参数 =====
SEED=${SEED:-42}
BATCH=${BATCH:-16}; 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
UNIMOL_CACHE=data/processed/unimol_embeddings.npz
MOLT5=models/molt5-base; BIOT5=models/biot5-plus-base; QWEN=models/qwen2.5-7b-instruct
LLM_COMMON="--use-llm --use-rag --rag-top-k 4 --use-soft-prompt --no-llm-freeze"

GPU_POOL=${GPU_POOL:-0,1}   # 我们可用的卡
MAX_RETRY=${MAX_RETRY:-5}
POLL=${POLL:-60}            # 轮询/重试间隔(秒)
IFS=',' read -r -a POOL <<< "$GPU_POOL"

ts(){ awk '{ print strftime("%F %T"), $0; fflush() }'; }

record_env(){ local dir=$1; shift
  { echo "date: $(date -Is)"
    echo "git_commit: $(git rev-parse HEAD 2>/dev/null)"
    echo "git_dirty_files: $(git status --porcelain 2>/dev/null | wc -l)"
    echo "seed: $SEED  pythonhashseed: $PYTHONHASHSEED"
    echo "cmd: $*"; } > "${dir}/repro.txt"
  git diff  > "${dir}/local_uncommitted.patch" 2>/dev/null
  pip freeze > "${dir}/pip_freeze.txt" 2>/dev/null
}

wait_free(){ local gpu=$1 need=$2 free
  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..." >&2
    sleep "$POLL"
  done
}

run_one(){ local gpu=$1 name=$2 need=$3; shift 3
  local rundir="models/abl_full/${name}/seed${SEED}"
  if ls "${rundir}"/summary.json "${rundir}"/*/summary.json >/dev/null 2>&1; then
    echo "[$(date '+%F %T')] SKIP ${name}（已完成）"; return 0; fi
  mkdir -p "${rundir}"; local n=1
  while :; do
    wait_free "$gpu" "$need"
    echo "[$(date '+%F %T')] >>> ${name} try $n on GPU${gpu}" | tee -a "${rundir}/run.log"
    record_env "${rundir}" "GPU=$gpu $*"
    if CUDA_VISIBLE_DEVICES="$gpu" python -m lnp_ml.modeling.nested_cv_optuna \
         --input-path "${INPUT}" \
         --output-dir "models/abl_full/${name}" --resume-dir "${rundir}" \
         --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} \
         "$@" 2>&1 | ts | tee -a "${rundir}/run.log"; then
      echo "[$(date '+%F %T')] <<< ${name} DONE" | tee -a "${rundir}/run.log"; return 0; fi
    (( n > MAX_RETRY )) && { echo "[$(date '+%F %T')] !!! ${name} FAILED x${MAX_RETRY}" | tee -a "${rundir}/run.log"; return 1; }
    echo "[$(date '+%F %T')] ${name} 失败, ${POLL}s 后断点续跑重试..." | tee -a "${rundir}/run.log"; sleep "$POLL"; ((n++))
  done
}

# ===== 变体队列：名称|需显存MiB|flags =====
QUEUE=$(mktemp); LOCK="${QUEUE}.lock"
add(){ echo "$*" >> "$QUEUE"; }
# Study 2：encoder 消融（MoE/LLM 全关）
add "s2_chemeleon|8000|--batch-size ${BATCH} --use-chemeleon --chemeleon-cache ${CHEM_CACHE}"
add "s2_unimol|8000|--batch-size ${BATCH} --use-unimol --unimol-cache ${UNIMOL_CACHE}"
# Study 3：LLM backbone（同集成模式，T5→LoRA）
add "s3_molt5|10000|--batch-size ${BATCH} --use-mpnn ${LLM_COMMON} --llm-use-lora --llm-model-path ${MOLT5}"
add "s3_biot5|10000|--batch-size ${BATCH} --use-mpnn ${LLM_COMMON} --llm-use-lora --llm-model-path ${BIOT5}"
# Study 1：核心模块（LLM=Qwen-RAG，Qwen→QLoRA）
add "s1_baseline|8000|--batch-size ${BATCH} --use-mpnn"
add "s1_moe|8000|--batch-size ${BATCH} --use-mpnn --use-moe"
add "s1_llm|14000|--batch-size ${QBATCH} --use-mpnn ${LLM_COMMON} --llm-use-qlora --llm-model-path ${QWEN}"
add "s1_both|14000|--batch-size ${QBATCH} --use-mpnn --use-moe ${LLM_COMMON} --llm-use-qlora --llm-model-path ${QWEN}"

pop_job(){ exec 9>"$LOCK"; flock 9
  local line; line=$(head -n1 "$QUEUE" 2>/dev/null)
  [ -n "$line" ] && sed -i '1d' "$QUEUE"
  flock -u 9; echo "$line"; }

worker(){ local gpu=$1
  while :; do
    local job; job=$(pop_job); [ -z "$job" ] && break
    local name=${job%%|*}; local r=${job#*|}; local need=${r%%|*}; local flags=${r#*|}
    run_one "$gpu" "$name" "$need" $flags
  done
}

echo "[$(date '+%F %T')] START seed=${SEED} pool=${GPU_POOL}"
for g in "${POOL[@]}"; do worker "$g" & done
wait

# ===== 汇总 =====
latest(){ ls -dt "models/abl_full/$1/seed${SEED}"/summary.json "models/abl_full/$1/seed${SEED}"/*/summary.json 2>/dev/null | head -1 | xargs -r dirname; }
sum_args=()
for name in s1_baseline s1_moe s1_llm s1_both s2_chemeleon s2_unimol s3_molt5 s3_biot5; do
  d=$(latest "$name"); [ -n "$d" ] && sum_args+=(--run "${name}=$d")
done
d=$(latest s1_baseline); [ -n "$d" ] && sum_args+=(--run "s2_mpnn=$d")
python scripts/summarize_ablation.py "${sum_args[@]}" --out reports/ablation_full_summary.csv 2>&1 | ts | tee logs/summary_full.log
echo "[$(date '+%F %T')] ALL DONE"