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

FLOSS_DIR="/home/yurenh2/rrm/research/flossing"
CONDA_SH="/home/yurenh2/miniconda3/etc/profile.d/conda.sh"
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"

wait_for_pid() {
  local pid="$1"
  if [[ "${pid}" == "0" ]]; then
    return 0
  fi
  while kill -0 "${pid}" 2>/dev/null; do
    sleep 60
  done
}

activate_env() {
  source "${CONDA_SH}"
  conda activate rrm
  cd "${FLOSS_DIR}"
}

run_hrm_queue() {
  wait_for_pid "${1:-0}"
  activate_env

  CUDA_VISIBLE_DEVICES=0 python step7_interfloss.py \
    --model hrm \
    --ckpt-root "${HRM_ROOT}" \
    --ckpt-name step_26040 \
    --train-steps 10000 \
    --batch-size 8 \
    --train-lr 1e-5 \
    --floss-lr 1e-4 \
    --floss-steps 500 \
    --interfloss-at 0,500 \
    --floss-mode engelken_l2 \
    --lambda-star 0 \
    --k-lyap 8 \
    --lyap-act-steps 4 \
    --lyap-start-act 12 \
    --kl-beta 10 \
    --kl-replay-size 64 \
    --kl-batch-size 4 \
    --kl-temperature 1 \
    --seed 42 \
    --eval-every 1000 \
    --eval-n 512 \
    --eval-batch-size 32 \
    --floss-log-every 10 \
    --out step7_E_hrm_late_engelken_interfloss_kl10_start12_26040_k8_10k.json \
    > step7_E_hrm_late_engelken_interfloss_kl10_start12_26040_k8_10k.log 2>&1

  CUDA_VISIBLE_DEVICES=0 python step7_interfloss.py \
    --model hrm \
    --ckpt-root "${HRM_ROOT}" \
    --ckpt-name step_26040 \
    --train-steps 10000 \
    --batch-size 8 \
    --train-lr 1e-5 \
    --floss-lr 1e-4 \
    --floss-steps 500 \
    --interfloss-at 0,500 \
    --floss-mode volume_cf \
    --lambda-star -0.15 \
    --k-lyap 8 \
    --lyap-act-steps 4 \
    --kl-beta 10 \
    --kl-replay-size 64 \
    --kl-batch-size 4 \
    --kl-temperature 1 \
    --seed 42 \
    --eval-every 1000 \
    --eval-n 512 \
    --eval-batch-size 32 \
    --floss-log-every 10 \
    --out step7_G_hrm_volume_envelope_interfloss_kl10_lstar_neg015_26040_k8_10k.json \
    > step7_G_hrm_volume_envelope_interfloss_kl10_lstar_neg015_26040_k8_10k.log 2>&1

  CUDA_VISIBLE_DEVICES=0 python step8_basin_consistency.py \
    --model hrm \
    --ckpt-root "${HRM_ROOT}" \
    --ckpt-name step_26040 \
    --train-steps 10000 \
    --batch-size 8 \
    --lr 1e-5 \
    --consistency-beta 1 \
    --consistency-every 1 \
    --perturb-after 8 \
    --noise-std 0.02 \
    --kl-temperature 1 \
    --seed 42 \
    --eval-every 1000 \
    --eval-n 512 \
    --eval-batch-size 32 \
    --out step8_A_hrm_basin_consistency_beta1_noise002_after8_26040_10k.json \
    > step8_A_hrm_basin_consistency_beta1_noise002_after8_26040_10k.log 2>&1
}

run_trm_queue() {
  wait_for_pid "${1:-0}"
  activate_env

  CUDA_VISIBLE_DEVICES=2 python step7_interfloss.py \
    --model trm \
    --ckpt-root "${TRM_ROOT}" \
    --ckpt-name step_26041 \
    --train-steps 10000 \
    --batch-size 4 \
    --train-lr 1e-5 \
    --floss-lr 1e-4 \
    --floss-steps 500 \
    --interfloss-at 0,500 \
    --floss-mode engelken_l2 \
    --lambda-star 0 \
    --k-lyap 4 \
    --lyap-act-steps 4 \
    --lyap-start-act 12 \
    --kl-beta 10 \
    --kl-replay-size 64 \
    --kl-batch-size 4 \
    --kl-temperature 1 \
    --seed 42 \
    --eval-every 1000 \
    --eval-n 512 \
    --eval-batch-size 32 \
    --floss-log-every 10 \
    --out step7_F_trm_late_engelken_interfloss_kl10_start12_26041_k4_batch4_10k.json \
    > step7_F_trm_late_engelken_interfloss_kl10_start12_26041_k4_batch4_10k.log 2>&1

  CUDA_VISIBLE_DEVICES=2 python step7_interfloss.py \
    --model trm \
    --ckpt-root "${TRM_ROOT}" \
    --ckpt-name step_26041 \
    --train-steps 10000 \
    --batch-size 4 \
    --train-lr 1e-5 \
    --floss-lr 1e-4 \
    --floss-steps 500 \
    --interfloss-at 0,500 \
    --floss-mode volume_cf \
    --lambda-star 0.02 \
    --k-lyap 4 \
    --lyap-act-steps 4 \
    --kl-beta 10 \
    --kl-replay-size 64 \
    --kl-batch-size 4 \
    --kl-temperature 1 \
    --seed 42 \
    --eval-every 1000 \
    --eval-n 512 \
    --eval-batch-size 32 \
    --floss-log-every 10 \
    --out step7_H_trm_volume_envelope_interfloss_kl10_lstar002_26041_k4_batch4_10k.json \
    > step7_H_trm_volume_envelope_interfloss_kl10_lstar002_26041_k4_batch4_10k.log 2>&1

  CUDA_VISIBLE_DEVICES=2 python step8_basin_consistency.py \
    --model trm \
    --ckpt-root "${TRM_ROOT}" \
    --ckpt-name step_26041 \
    --train-steps 10000 \
    --batch-size 4 \
    --lr 1e-5 \
    --consistency-beta 1 \
    --consistency-every 1 \
    --perturb-after 8 \
    --noise-std 0.02 \
    --kl-temperature 1 \
    --seed 42 \
    --eval-every 1000 \
    --eval-n 512 \
    --eval-batch-size 32 \
    --out step8_B_trm_basin_consistency_beta1_noise002_after8_26041_batch4_10k.json \
    > step8_B_trm_basin_consistency_beta1_noise002_after8_26041_batch4_10k.log 2>&1
}

cmd="${1:?usage: launch_dynamics_variants_queue.sh MODE [wait_pid]}"
wait_pid="${2:-0}"

case "${cmd}" in
  hrm) run_hrm_queue "${wait_pid}" ;;
  trm) run_trm_queue "${wait_pid}" ;;
  *) echo "unknown command: ${cmd}" >&2; exit 2 ;;
esac