#!/usr/bin/env bash
# Phase Annotate: sequentially annotate all four LIBERO suites.
# Usage:
#   bash examples/LIBERO/train_files/run_phase_annotate.sh
#   OUTPUT_ROOT=/custom/path bash examples/LIBERO/train_files/run_phase_annotate.sh

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"

# ── Configurable parameters ────────────────────────────────────────────────────
LIBERO_DATA_ROOT="${REPO_ROOT}/LIBERO-DATA"
OUTPUT_ROOT="${OUTPUT_ROOT:-${REPO_ROOT}/LIBERO-DATA-phase}"
SIGNAL_SOURCE="${SIGNAL_SOURCE:-action}"
MEDIAN_FILTER_SIZE="${MEDIAN_FILTER_SIZE:-5}"
NUM_WORKERS="${NUM_WORKERS:-1}"

SUITES=(
    "libero_10_no_noops_1.0.0_lerobot"
    "libero_goal_no_noops_1.0.0_lerobot"
    "libero_object_no_noops_1.0.0_lerobot"
    "libero_spatial_no_noops_1.0.0_lerobot"
)
# ──────────────────────────────────────────────────────────────────────────────

echo "========================================================"
echo "  LIBERO Phase Annotation"
echo "========================================================"
echo "  LIBERO_DATA_ROOT : ${LIBERO_DATA_ROOT}"
echo "  OUTPUT_ROOT      : ${OUTPUT_ROOT}"
echo "  signal_source    : ${SIGNAL_SOURCE}"
echo "  num_workers      : ${NUM_WORKERS}"
echo "========================================================"

cd "${REPO_ROOT}"

FAILED=()
for SUITE in "${SUITES[@]}"; do
    SRC_DIR="${LIBERO_DATA_ROOT}/${SUITE}"
    OUT_DIR="${OUTPUT_ROOT}/${SUITE}"

    echo ""
    echo "-------- [$(date '+%H:%M:%S')] Processing: ${SUITE} --------"

    if [[ ! -d "${SRC_DIR}/data" ]]; then
        echo "  SKIP: ${SRC_DIR}/data not found (parquet not downloaded yet)"
        continue
    fi

    mkdir -p "${OUT_DIR}"

    if python -m phaseaware_data_construction.phase_annotate \
        --src_dir "${SRC_DIR}" \
        --output_dir "${OUT_DIR}" \
        --signal_source "${SIGNAL_SOURCE}" \
        --median_filter_size "${MEDIAN_FILTER_SIZE}" \
        --num_workers "${NUM_WORKERS}"; then
        echo "  OK: annotation written to ${OUT_DIR}"
    else
        echo "  FAILED: ${SUITE}"
        FAILED+=("${SUITE}")
    fi
done

# ── Validate all annotated suites ─────────────────────────────────────────────
echo ""
echo "========================================================"
echo "  Validation"
echo "========================================================"
for SUITE in "${SUITES[@]}"; do
    SRC_DIR="${LIBERO_DATA_ROOT}/${SUITE}"
    OUT_DIR="${OUTPUT_ROOT}/${SUITE}"

    if [[ ! -d "${OUT_DIR}/data" ]]; then
        echo "  SKIP validate: ${SUITE} (no output yet)"
        continue
    fi

    echo ""
    echo "-------- [$(date '+%H:%M:%S')] Validating: ${SUITE} --------"
    python -m phaseaware_data_construction.phase_annotate \
        --src_dir "${SRC_DIR}" \
        --output_dir "${OUT_DIR}" \
        --validate || true   # validation failures are non-fatal
done

# ── Final report ──────────────────────────────────────────────────────────────
echo ""
echo "========================================================"
if [[ ${#FAILED[@]} -eq 0 ]]; then
    echo "  All suites completed successfully."
else
    echo "  The following suites FAILED:"
    for s in "${FAILED[@]}"; do echo "    - ${s}"; done
    exit 1
fi
echo "========================================================"