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

: "${ROOT:?Set ROOT to the output parent folder for generated VariousSpeed datasets.}"
: "${BASE:?Set BASE to the folder containing *_no_noops_1.0.0_lerobot source datasets.}"
RUN_TAG="${RUN_TAG:-four_suite_v2}"
TIMESTAMP="${TIMESTAMP:-$(date +%Y%m%d)}"
SPEEDS="${SPEEDS:-0.5 0.75 1.0 1.25 2.0}"
WRITE_VIDEOS="${WRITE_VIDEOS:-1}"
OVERWRITE="${OVERWRITE:-0}"
MAX_EPISODES="${MAX_EPISODES:-}"

export UV_CACHE_DIR="${UV_CACHE_DIR:-/tmp/uv-cache}"
export HF_HOME="${HF_HOME:-/tmp/hf-cache}"
export HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-/tmp/hf-cache/datasets}"

suites=(
  libero_spatial
  libero_object
  libero_goal
  libero_10
)

log() {
  printf '[%(%Y-%m-%d %H:%M:%S)T] %s\n' -1 "$*"
}

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

  read -r -a speed_args <<< "$SPEEDS"
  common_args=(
    --dst "$ROOT"
    --auto-name
    --run-tag "$RUN_TAG"
    --timestamp "$TIMESTAMP"
    --speeds "${speed_args[@]}"
    --clean-transl-eps 1e-4
    --clean-rot-eps 1e-4
    --min-segment-len 1
  )

  if [[ "$WRITE_VIDEOS" == "1" ]]; then
    common_args+=(--write-videos)
  fi
  if [[ "$OVERWRITE" == "1" ]]; then
    common_args+=(--overwrite)
  fi
  if [[ -n "$MAX_EPISODES" ]]; then
    common_args+=(--max-episodes "$MAX_EPISODES")
  fi

  log "building four LIBERO speed datasets"
  log "root=$ROOT base=$BASE run_tag=$RUN_TAG timestamp=$TIMESTAMP speeds=$SPEEDS"

  for suite in "${suites[@]}"; do
    log "processing $suite"
    uv run python scripts/build_libero_speed_dataset.py \
      --src "$BASE/${suite}_no_noops_1.0.0_lerobot" \
      --task-suite-name "$suite" \
      "${common_args[@]}"
  done

  log "done"
}

main "$@"