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

DATA_ROOT="${DATA_ROOT:-./datasets/cifar10}"
SAVE_DIR="${SAVE_DIR:-./runs/cifar10_rtm}"
NUM_GPUS="${NUM_GPUS:-1}"

mkdir -p "$SAVE_DIR"

CKPT_DIR="$SAVE_DIR/train"
RESTORE_ARGS=()
if [ -f "$CKPT_DIR/latest-model.th" ]; then
    echo "[train_cifar10] resuming from $CKPT_DIR/latest-model.th"
    RESTORE_ARGS+=(--restore_path "$CKPT_DIR/latest-model.th")
    [ -f "$CKPT_DIR/latest-model-ema.th" ] && RESTORE_ARGS+=(--restore_ema_path "$CKPT_DIR/latest-model-ema.th")
    [ -f "$CKPT_DIR/latest-opt.th" ]       && RESTORE_ARGS+=(--restore_optimizer_path "$CKPT_DIR/latest-opt.th")
    [ -f "$CKPT_DIR/latest-sched.th" ]     && RESTORE_ARGS+=(--restore_scheduler_path "$CKPT_DIR/latest-sched.th")
    [ -f "$CKPT_DIR/latest-log.jsonl" ]    && RESTORE_ARGS+=(--restore_log_path "$CKPT_DIR/latest-log.jsonl")
fi

torchrun --nnodes=1 --nproc_per_node="$NUM_GPUS" --standalone train.py \
    --hps cifar10 \
    --save_dir "$SAVE_DIR" \
    --data_root "$DATA_ROOT" \
    --num_epochs 200 \
    --fid_freq 1000 \
    --use_se True \
    --width 256 \
    --dec_blocks '1x1,4m1,4x2,8m4,8x2,16m8,16x2,32m16,32x2' \
    --force_factor 5 \
    --imle_force_resample 5 \
    --lr 0.0008 \
    --search_type lpips \
    --n_batch 256 \
    --imle_batch 1024 \
    --iters_per_save 1000 \
    --iters_per_images 5000 \
    --iters_per_ckpt 100000 \
    --latent_dim 128 \
    --use_rtm True \
    --H_cycles 4 --L_cycles 1 --refinement_steps 4 \
    --num_tokens 4 \
    --rtm_hidden_size 128 \
    "${RESTORE_ARGS[@]}"