#!/usr/bin/env bash
set -eo pipefail

MODEL="${1:-hrm}"
GPU="${2:-0}"

source /home/yurenh2/miniconda3/etc/profile.d/conda.sh
conda activate rrm
cd /home/yurenh2/rrm/research/flossing

HRM_ROOT="/home/yurenh2/rrm/hrm/checkpoints/Sudoku-extreme-1k-aug-1000 ACT-torch/HierarchicalReasoningModel_ACTV1 righteous-python"
TRM_ROOT="/home/yurenh2/rrm/trm/checkpoints/Sudoku-extreme-1k-aug-1000-ACT-torch/pretrain_mlp_t_sudoku_singleGPU"

run_one() {
  local tag="$1"
  shift
  echo "[$(date -Is)] START ${tag}"
  CUDA_VISIBLE_DEVICES="${GPU}" python step9_trajectory_perturb_train.py "$@" \
    > "${tag}.log" 2>&1
  echo "[$(date -Is)] DONE ${tag}"
}

if [[ "${MODEL}" == "hrm" ]]; then
  run_one step9_A_hrm_single_perturb_sigma1e-3_26040_10k \
    --model hrm \
    --ckpt-root "${HRM_ROOT}" \
    --ckpt-name step_26040 \
    --mode single_perturbed_ce \
    --train-steps 10000 \
    --batch-size 8 \
    --lr 1e-5 \
    --noise-std 0.001 \
    --perturb both \
    --seed 42 \
    --eval-every 1000 \
    --eval-n 512 \
    --eval-batch-size 32 \
    --out step9_A_hrm_single_perturb_sigma1e-3_26040_10k.json

  run_one step9_B_hrm_multi4_perturb_sigma1e-3_26040_10k \
    --model hrm \
    --ckpt-root "${HRM_ROOT}" \
    --ckpt-name step_26040 \
    --mode multi_perturbed_ce \
    --n-trajectories 4 \
    --train-steps 10000 \
    --batch-size 8 \
    --lr 1e-5 \
    --noise-std 0.001 \
    --perturb both \
    --seed 42 \
    --eval-every 1000 \
    --eval-n 512 \
    --eval-batch-size 32 \
    --out step9_B_hrm_multi4_perturb_sigma1e-3_26040_10k.json
elif [[ "${MODEL}" == "trm" ]]; then
  run_one step9_C_trm_single_perturb_sigma1e-3_26041_batch4_10k \
    --model trm \
    --ckpt-root "${TRM_ROOT}" \
    --ckpt-name step_26041 \
    --mode single_perturbed_ce \
    --train-steps 10000 \
    --batch-size 4 \
    --lr 1e-5 \
    --noise-std 0.001 \
    --perturb both \
    --seed 42 \
    --eval-every 1000 \
    --eval-n 512 \
    --eval-batch-size 32 \
    --out step9_C_trm_single_perturb_sigma1e-3_26041_batch4_10k.json

  run_one step9_D_trm_multi4_perturb_sigma1e-3_26041_batch4_10k \
    --model trm \
    --ckpt-root "${TRM_ROOT}" \
    --ckpt-name step_26041 \
    --mode multi_perturbed_ce \
    --n-trajectories 4 \
    --train-steps 10000 \
    --batch-size 4 \
    --lr 1e-5 \
    --noise-std 0.001 \
    --perturb both \
    --seed 42 \
    --eval-every 1000 \
    --eval-n 512 \
    --eval-batch-size 32 \
    --out step9_D_trm_multi4_perturb_sigma1e-3_26041_batch4_10k.json
else
  echo "unknown model: ${MODEL}" >&2
  exit 2
fi