#!/bin/bash
# E: α=3 λ*=0 (calculated TRM-equivalent of HRM α=10)
# F: α=2 λ*=0 (slightly weaker)
set -e
cd /home/yurenh2/rrm/research/flossing
source /home/yurenh2/miniconda3/etc/profile.d/conda.sh
conda activate rrm
export CUDA_VISIBLE_DEVICES=0

PREV_RUNNER_PID=$1
echo "[$(date '+%H:%M:%S')] EF-runner waiting for prev runner PID $PREV_RUNNER_PID..." >> step5_runner.log
while kill -0 "$PREV_RUNNER_PID" 2>/dev/null; do
  sleep 30
done
echo "[$(date '+%H:%M:%S')] Previous runner done. Starting E/F (α=3/2 calculated equivalents)" >> step5_runner.log

CKPT_ROOT="/home/yurenh2/rrm/trm/checkpoints/Sudoku-extreme-1k-aug-1000-ACT-torch/pretrain_mlp_t_sudoku_singleGPU"
SRC_CKPT="step_26041"

run_cf() {
  local TAG=$1; local ALPHA=$2; local LSTAR=$3
  local OUT=step5_${TAG}_trm_${SRC_CKPT}_alpha${ALPHA}_lstar${LSTAR}.json
  local LOG=step5_${TAG}.log
  if [ -f "$OUT" ]; then echo "skip $TAG"; return; fi
  echo "[$(date '+%H:%M:%S')] Starting $TAG (α=$ALPHA λ*=$LSTAR) from $SRC_CKPT" >> step5_runner.log
  python step5_train_trm_cf.py \
    --ckpt-root "$CKPT_ROOT" --ckpt-name "$SRC_CKPT" \
    --n-steps 500 --batch-size 8 \
    --alpha-rf $ALPHA --rf-mode fixed --lambda-star $LSTAR \
    --k-lyap 2 --lyap-act-steps 4 \
    --eval-every 100 --eval-n 512 --eval-batch-size 32 \
    --out $OUT > $LOG 2>&1
  echo "[$(date '+%H:%M:%S')] $TAG done" >> step5_runner.log
}

run_cf E 3 0
run_cf F 2 0

echo "[$(date '+%H:%M:%S')] EF complete" >> step5_runner.log