#!/usr/bin/env bash
# Clone-and-run Maze training, defaulted for 2x A6000 (48G). Repo-relative; no hardcoded paths.
# Usage:
#   bash run_maze_a6000.sh              # 2x A6000, gbs 384 (default)
#   bash run_maze_a6000.sh 2 384        # explicit
#   bash run_maze_a6000.sh 1 192        # single A6000
#   SMOKE=1 bash run_maze_a6000.sh      # 200-step smoke test first (minutes)
set -eo pipefail

REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
NGPU="${1:-2}"
GBS="${2:-384}"
DATA="$REPO/data/maze-30x30-hard-1k"
TRM="$REPO/trm"
RUN_NAME="pretrain_att_maze30x30_${NGPU}gpu_gbs${GBS}"

# conda env (default 'rrm'); override with CONDA_ENV / CONDA_SH
CONDA_SH="${CONDA_SH:-}"
[[ -z "$CONDA_SH" ]] && for p in "$HOME/miniconda3/etc/profile.d/conda.sh" \
  "$HOME/anaconda3/etc/profile.d/conda.sh" "/opt/conda/etc/profile.d/conda.sh"; do
  [[ -f "$p" ]] && CONDA_SH="$p" && break; done
[[ -f "$CONDA_SH" ]] && source "$CONDA_SH" && conda activate "${CONDA_ENV:-rrm}"

[[ -d "$DATA" ]] || { echo "FATAL: dataset missing at $DATA"; exit 1; }
[[ -f "$TRM/pretrain.py" ]] || { echo "FATAL: trm/pretrain.py missing"; exit 1; }
cd "$TRM"
export WANDB_MODE=offline

EPOCHS=50000; EVAL=5000
if [[ "${SMOKE:-0}" == "1" ]]; then EPOCHS=200; EVAL=200; RUN_NAME="smoke_${RUN_NAME}"; fi

ARGS=(
  arch=trm
  "data_paths=[$DATA]"
  "evaluators=[]"
  epochs=$EPOCHS eval_interval=$EVAL
  lr=1e-4 puzzle_emb_lr=1e-4 weight_decay=1.0 puzzle_emb_weight_decay=1.0
  global_batch_size="$GBS"
  arch.L_layers=2 arch.H_cycles=3 arch.L_cycles=4
  +run_name="$RUN_NAME" ema=True
  +checkpoint_every_eval=true
)
LOG="$REPO/maze_${RUN_NAME}.log"

echo "launching $RUN_NAME on $NGPU GPU(s), gbs=$GBS  (epochs=$EPOCHS)"
if [[ "$NGPU" -gt 1 ]]; then
  nohup torchrun --nproc-per-node "$NGPU" --rdzv_backend=c10d --rdzv_endpoint=localhost:0 \
    --nnodes=1 pretrain.py "${ARGS[@]}" > "$LOG" 2>&1 &
else
  nohup python pretrain.py "${ARGS[@]}" > "$LOG" 2>&1 &
fi
echo "pid $!   log: $LOG"
echo "watch:  tail -f $LOG | grep -E 'exact|accuracy|Traceback|Error'"
echo "ckpts:  $TRM/checkpoints/maze-30x30-hard-1k.../$RUN_NAME/  (1 per ${EVAL} epochs)"