#!/usr/bin/env bash
# DAgger 纠偏（2026-10-18 首次正式预训练后新增）：bash ml/dagger/run.sh <头标签> <起点 ckpt>
#   每轮：策略闭环采起点（ml.dagger.collect）→ 老师从这些起点重新规划（gen_dataset.cjs --tasks，NPROC 并行）
#         → 打包成样本（build_dataset.cjs，与主数据同口径）→ 主数据 + 历轮 DAgger 样本一起微调 DAG_STEPS 步。
#   环境变量：DAG_ROUNDS（默认 2）、DAG_N（每轮采样题数，默认 8000，每题最多 2 个起点）、DAG_STEPS（每轮微调步数，默认 15000）、
#            DAG_LR（默认 3e-4）、NPROC、MIN_GAP_MM。
#   产物：ml/data/dag_<tag>_r<k>/（起点与老师 shard）、ml/data/samples_dag_<tag>_r<k>/、ml/runs/<tag>_dag<k>/；
#         最终模型路径写进 ml/runs/<tag>_dagger.ckpt（m5 读它）。
#   断点续跑：采集以 collect.json 为完成标记；老师 shard 先写 .part 再改名、未写完的从断点续写；打包以 .tasks.jsonl 为完成标记；
#            微调用该轮 last.pt 续训；整轮完成打 ml/runs/<tag>_dag<k>/done。
set -e -o pipefail
cd "$(dirname "$0")/../.."
PY=${PYTHON:-python}
TAG=$1; CK=$2
[ -n "$TAG" ] && [ -f "$CK" ] || { echo "usage: run.sh <tag> <ckpt>（$CK 不存在）"; exit 1; }
R=${DAG_ROUNDS:-2}; N=${DAG_N:-8000}; STEPS=${DAG_STEPS:-15000}; LR=${DAG_LR:-3e-4}; P=${NPROC:-$(nproc)}
CFGF=ml/configs/$TAG.json
SAMPLES="ml/data/samples"
for ((r=1; r<=R; r++)); do
  D=ml/data/dag_${TAG}_r$r; SD=ml/data/samples_dag_${TAG}_r$r; OUT=ml/runs/${TAG}_dag$r
  SAMPLES="$SAMPLES $SD"
  if [ -f $OUT/done ]; then echo "[dagger] 第 $r 轮已完成，跳过"; CK=$OUT/best.pt; continue; fi
  if [ -f $D/collect.json ]; then echo "[dagger] 第 $r 轮起点已采集，跳过"
  else echo "[dagger] 第 $r 轮：用 $CK 闭环采起点（$N 题）"; $PY -m ml.dagger.collect --ckpt $CK --samples ml/data/samples --n $N --seed $r --parts $P --out $D; fi
  cat $D/collect.json
  pids=()
  for f in $D/tasks_*.jsonl; do
    [ -e "$f" ] || continue
    k=$(basename $f .jsonl); k=${k#tasks_}; sh=$D/shard_dag_${TAG}_r${r}_$k.jsonl
    [ -f $sh ] && continue
    ( node ml/gen/gen_dataset.cjs --tasks $f --seed $((9000 + 100 * r + 10#$k)) --out $sh.part 2>> $D/log_$k.txt && mv $sh.part $sh ) &
    pids+=($!)
  done
  if [ ${#pids[@]} -gt 0 ]; then
    echo "[dagger] 第 $r 轮：老师重新规划（${#pids[@]} 进程）"
    t0=$(date +%s)
    while :; do
      alive=0; for p in "${pids[@]}"; do kill -0 $p 2>/dev/null && alive=$((alive+1)); done
      [ $alive = 0 ] && break; sleep 30
      echo "[dagger 节拍] $(( ($(date +%s) - t0) / 60 )) 分钟  运行中 $alive/${#pids[@]}  已写 $(cat $D/*.part $D/shard_*.jsonl 2>/dev/null | wc -l) 条"
    done
    for p in "${pids[@]}"; do wait $p || true; done   # 判定按产物（同 run_parallel.sh：Git Bash 长跑后 wait 返回码不可靠）
    for f in $D/tasks_*.jsonl; do
      k=$(basename $f .jsonl); k=${k#tasks_}; sh=$D/shard_dag_${TAG}_r${r}_$k.jsonl
      if [ ! -f $sh ]; then
        n=0; [ -f $sh.part ] && n=$(wc -l < $sh.part); need=$(wc -l < $f)
        if [ "$n" -ge "$need" ]; then mv $sh.part $sh; else echo "[dagger] $sh 未写完（$n/$need），见 $D/log_$k.txt；重跑本阶段从断点续写"; exit 6; fi
      fi
    done
  fi
  mkdir -p $SD
  todo=(); for sh in $D/shard_dag_*.jsonl; do [ -e "$sh" ] || continue; b=$(basename $sh .jsonl); [ -f $SD/$b.tasks.jsonl ] || todo+=("$sh"); done
  [ ${#todo[@]} = 0 ] || printf '%s\n' "${todo[@]}" | xargs -P "$P" -I{} node ml/prep/build_dataset.cjs {} $SD
  $PY -m ml.data --index $SD | tail -n 3
  res=""; [ -f $OUT/last.pt ] && res="--resume $OUT/last.pt"
  echo "[dagger] 第 $r 轮：微调 $STEPS 步（样本：$SAMPLES）"
  $PY -m ml.train --samples $SAMPLES --out $OUT --config $CFGF --init $CK --steps $STEPS --lr $LR --warmup 500 $res
  touch $OUT/done; CK=$OUT/best.pt
done
echo "$CK" > ml/runs/${TAG}_dagger.ckpt
echo "[dagger] 完成：后续阶段使用 $CK"