#!/bin/bash
#SBATCH --job-name=dovla_lerobot_export
#SBATCH --account=def-yalda_gpu
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=4
#SBATCH --mem=24G
#SBATCH --time=01:00:00
#SBATCH --output=outputs/hpc/logs/%x_%j.out
#SBATCH --error=outputs/hpc/logs/%x_%j.err

set -euo pipefail

PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
DATASET="${DATASET:?Set DATASET to a CIL dataset or collection directory}"
OUT="${OUT:?Set OUT to the LeRobot-style export directory}"
PYTHON="${PYTHON:-$PROJECT_DIR/.venv/bin/python}"
SELECTION="${SELECTION:-best}"
GROUP_SAMPLING="${GROUP_SAMPLING:-sequential}"
SEED="${SEED:-0}"
SPLIT="${SPLIT:-train}"
MAX_GROUPS="${MAX_GROUPS:-}"
NO_IMAGES="${NO_IMAGES:-0}"

cd "$PROJECT_DIR"
mkdir -p "$OUT" outputs/hpc/logs

ARGS=(
  scripts/export_lerobot_dataset.py
  --dataset "$DATASET"
  --out "$OUT"
  --split "$SPLIT"
  --selection "$SELECTION"
  --group-sampling "$GROUP_SAMPLING"
  --seed "$SEED"
)

if [[ -n "$MAX_GROUPS" ]]; then
  ARGS+=(--max-groups "$MAX_GROUPS")
fi
if [[ "$NO_IMAGES" == "1" ]]; then
  ARGS+=(--no-images)
fi

"$PYTHON" "${ARGS[@]}"