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

ROOT="/home/yurenh2/rrm"
PY="/home/yurenh2/miniconda3/envs/rrm/bin/python"
OUT_DIR="${ROOT}/research/flossing/flossing_suite/results/toy_rnn"
mkdir -p "${OUT_DIR}"

GPU_BASE="${GPU_BASE:-0}"
GPU_PREFLOSS="${GPU_PREFLOSS:-1}"
GPU_INTER="${GPU_INTER:-3}"

write_and_launch() {
  local gpu="$1"
  local name="$2"
  local pre_epochs="$3"
  local max_inter="$4"
  local cmd="${OUT_DIR}/${name}.cmd.sh"
  local log="${OUT_DIR}/${name}.log"
  local pid="${OUT_DIR}/${name}.pid"

  cat > "${cmd}" <<EOF
#!/usr/bin/env bash
set -euo pipefail
cd "${ROOT}"
export CUDA_VISIBLE_DEVICES="${gpu}"
export PYTHONUNBUFFERED=1
exec "${PY}" research/flossing/engelken_python_flossing.py \\
  --hidden-size 80 \\
  --n-lyap 40 \\
  --train-epochs 1000 \\
  --inter-period 100 \\
  --inter-epochs 100 \\
  --batch-size 16 \\
  --input-dim 1 \\
  --train-steps 300 \\
  --lyap-steps 55 \\
  --floss-input-steps 300 \\
  --seed-ic 1 \\
  --seed-input 1 \\
  --seed-net 1 \\
  --seed-ons 1 \\
  --lr 0.001 \\
  --beta1 0.9 \\
  --beta2 0.999 \\
  --init-type 1 \\
  --recurrent-gain 1.0 \\
  --recurrent-mean-gain 0.0 \\
  --input-scale 1.0 \\
  --delay 10 \\
  --ws-std 1.0 \\
  --ws-mean 0.0 \\
  --wr-std 1.0 \\
  --wr-mean 0.0 \\
  --b-std 0.1 \\
  --b-mean 0.0 \\
  --task -1 \\
  --lyap-target 0.0 \\
  --eval-every 100 \\
  --eval-batches 4 \\
  --log-every-floss 25 \\
  --device cuda \\
  --pre-epochs "${pre_epochs}" \\
  --max-inter-episodes "${max_inter}" \\
  --out "${OUT_DIR}/${name}.json"
EOF
  chmod +x "${cmd}"
  setsid bash "${cmd}" > "${log}" 2>&1 < /dev/null &
  echo $! > "${pid}"
  echo "${name}: pid $(cat "${pid}") on GPU ${gpu}"
}

write_and_launch "${GPU_BASE}" "toy_baseline_no_floss_N80_k40_E1000" 0 0
write_and_launch "${GPU_PREFLOSS}" "toy_prefloss_N80_k40_E1000" 100 0
write_and_launch "${GPU_INTER}" "toy_pre_inter_N80_k40_E1000" 100 2

echo "queued toy RNN official-style suite in ${OUT_DIR}"