#!/bin/bash
# Wait for the W-runner (X→Y→Z→W pipeline) to finish, then launch SRM-AOL training.
set -e
cd /home/yurenh2/rrm/srm
source /home/yurenh2/miniconda3/etc/profile.d/conda.sh
conda activate rrm
export CUDA_VISIBLE_DEVICES=2

W_RUNNER_PID=$1
LOG_DIR=/home/yurenh2/rrm/research/flossing
echo "[$(date '+%H:%M:%S')] SRM-runner waiting for W-runner PID $W_RUNNER_PID..." >> $LOG_DIR/step4_runner.log
while kill -0 "$W_RUNNER_PID" 2>/dev/null; do
  sleep 30
done
echo "[$(date '+%H:%M:%S')] W done. Starting SRM-AOL from-scratch training (Sudoku 1k, hidden=512, 3000 steps)" >> $LOG_DIR/step4_runner.log

# Make sure log/ckpt dirs exist
mkdir -p runs ckpts

# SRM-AOL training: same setup as XYZ baselines but with SRM model + Lyapunov diag every 25 steps
python scripts/train_srm.py \
  --n-steps 3000 --batch-size 8 \
  --hidden-size 512 --n-iters 12 --n-aol-layers 2 \
  --kappa 0.9 --eta 1.0 --alpha 1.0 \
  --warmup-steps 200 \
  --k-lyap 2 --lyap-iters 8 --lyap-every 50 \
  --eval-every 300 --eval-n 512 --eval-batch-size 32 \
  --out runs/srm_aol_v1_sudoku1k_3k.json \
  --save-ckpt ckpts/srm_aol_v1_3k.pt \
  > $LOG_DIR/srm_aol_run.log 2>&1

echo "[$(date '+%H:%M:%S')] SRM-AOL training complete" >> $LOG_DIR/step4_runner.log