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

cd "$(dirname "$0")/.."

DATA_DIR="${DATA_DIR:-/path/to/cifar-10}"
SAVE_DIR="${SAVE_DIR:-./runs/cifar10_rtm}"
NUM_GPUS="${NUM_GPUS:-1}"
CFG_FILE="${CFG_FILE:-src/configs/CIFAR10/StyleGAN2-ADA-RTM.yaml}"

mkdir -p "$SAVE_DIR"

RESUME_ARGS=()
if [ -d "$SAVE_DIR/checkpoints" ]; then
    prev_run=$(find "$SAVE_DIR/checkpoints" -mindepth 1 -maxdepth 1 -type d 2>/dev/null \
                 | sort | tail -n1)
    if [ -n "$prev_run" ] && ls "$prev_run"/model=*.pth >/dev/null 2>&1; then
        echo "[train_cifar10] resuming from $prev_run"
        RESUME_ARGS=(-ckpt "$prev_run")
    fi
fi

CMD=(python src/main.py -t
     --cfg_file "$CFG_FILE"
     --data_dir "$DATA_DIR"
     --save_dir "$SAVE_DIR"
     -metrics fid is prdc
     --num_workers 4
     --print_freq 100
     --save_freq 2000
     --seed 0
     "${RESUME_ARGS[@]}")

if [ "$NUM_GPUS" -gt 1 ]; then
    "${CMD[@]}" -DDP
else
    "${CMD[@]}"
fi