export CUDA_VISIBLE_DEVICES=4,5
SEED_LIST=${SEEDS:-"42"}
LOOP_DIST=uniform
POISSON_LAMBDA=
NPROC=2
RUN_PREFIX=loop_varc_var
WANDB_PROJECT=VisionARC
WANDB_TAG=variance
OUTPUT_DIR=saves/variance_runs

DATE_TAG=$(date +%Y-%m-%d)
BASE_OUTPUT_DIR="${OUTPUT_DIR}/${DATE_TAG}"
mkdir -p "${BASE_OUTPUT_DIR}"

for SEED in ${SEED_LIST}; do
  RUN_NAME="${RUN_PREFIX}_s${SEED}_${WANDB_TAG}"
  SAVE_PREFIX="${BASE_OUTPUT_DIR}/${RUN_NAME}"

  EXTRA_POISSON_FLAG=""
  if [ -n "${POISSON_LAMBDA}" ] && [ "${LOOP_DIST}" = "poisson" ]; then
    EXTRA_POISSON_FLAG=" --loop-step-poisson-lambda ${POISSON_LAMBDA}"
  fi

  torchrun --nproc_per_node="${NPROC}" offline_train_loop_variable_ARC.py \
    --epochs 100 \
    --loop-core-depth 2 \
    --max-loop-steps 6 \
    --min-loop-steps 2 \
    --batch-size 32 \
    --image-size 64 \
    --patch-size 2 \
    --learning-rate 3e-4 \
    --weight-decay 0 \
    --embed-dim 512 \
    --mlp-dim 1024 \
    --num-heads 8 \
    --include-rearc \
    --num-colors 12 \
    --data-root "raw_data/ARC-AGI" \
    --train-split "training" \
    --wandb-project "${WANDB_PROJECT}" \
    --wandb-run-name "${RUN_NAME}" \
    --save-path "${SAVE_PREFIX}_final.pt" \
    --best-save-path "${SAVE_PREFIX}_best.pt" \
    --lr-scheduler "cosine" \
    --vis-every 50 \
    --distributed \
    --use-wandb \
    --gate-entropy-weight 5e-4 \
    --loop-penalty-weight 1e-3 \
    --train-dynamic-exit \
    --eval-dynamic-exit \
    --randomize-loop-steps \
    --loop-step-dist "${LOOP_DIST}" \
    --seed "${SEED}"${EXTRA_POISSON_FLAG}
done
