#!/usr/bin/env bash
# Stage-A orchestration: wait for continue-pretrain to exit, baseline eval, then SFT.
set -u
cd "$(dirname "$0")"
export PYTHONPATH="$PWD"
BASE_CKPT=ckpt/tiny18m2/model_best.pt
EVAL_LOG=logs/baseline_eval_tiny18m2.log
SFT_LOG=logs/stage_a_sft.log
SFT_DIR=ckpt/tiny18m2_sft_a

echo "[stage_a] $(date) waiting for continue-pretrain (train_lm.py on tiny18m2) to exit..."
while pgrep -f "train/train_lm.py --resume ckpt/tiny18m2" >/dev/null 2>&1; do
  sleep 60
done
echo "[stage_a] $(date) training exited; last train log tail:"
tail -3 logs/train_phase2b_r3.log 2>/dev/null || true
# refresh best.pt pointer
BASE=$(ls -t ckpt/tiny18m2/model_best.pt 2>/dev/null || echo "$BASE_CKPT")
echo "[stage_a] $(date) baseline eval on $BASE -> $EVAL_LOG"
.venv/bin/python research/eval.py --ckpt "$BASE" --probes data/eval_probes.jsonl > "$EVAL_LOG" 2>&1
echo "=== probes_researcher ===" >> "$EVAL_LOG"
.venv/bin/python research/eval.py --ckpt "$BASE" --probes data/probes_researcher.jsonl >> "$EVAL_LOG" 2>&1
echo "[stage_a] $(date) baseline eval done (floor to beat)."

if [ -s "$SFT_LOG" ] && [ -d "$SFT_DIR" ]; then
  echo "[stage_a] $(date) Stage-A already started (sentinel present); not relaunching."
  echo "  resuming/continuing manually; see $SFT_LOG"
  exit 0
fi
echo "[stage_a] $(date) launching Stage-A SFT (conservative, KL-anchored) -> $SFT_LOG"
nohup .venv/bin/python train/train_sft_v4.py \
  --base "$BASE" \
  --data data/stage_a_gold.jsonl \
  --tok data/tokenizer.json \
  --ckpt "$SFT_DIR" \
  --val-bin data/valid_mix.bin \
  --epochs 3 --batch 4 --seq 512 --lr 8e-6 \
  --kl 0.05 --train-scope last2 --threads 8 \
  > "$SFT_LOG" 2>&1 &
echo "[stage_a] $(date) SFT pid $! logged to $SFT_LOG"