#!/bin/bash
#SBATCH --account=nvr_lpr_rvp
#SBATCH --qos=normal
#SBATCH --partition=cpu
#SBATCH --qos=cpu-long
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=40
#SBATCH --time=04:00:00
#SBATCH --job-name=groot17_prep_data
#SBATCH --output=/lustre/fsw/portfolios/nvr/users/jtremblay/yu/logs/groot17_prep_data_%j.out
#SBATCH --error=/lustre/fsw/portfolios/nvr/users/jtremblay/yu/logs/groot17_prep_data_%j.err
#SBATCH --comment=fact_off

set -euo pipefail

WORKSPACE=/lustre/fsw/portfolios/nvr/users/jtremblay
PYTHON=$WORKSPACE/conda_envs/GR00T17/bin/python
GROOT_DIR=$WORKSPACE/yu/Isaac-GR00T

SRC_ROOT=$WORKSPACE/yu/conflict_maniskill/demo_conflict
OUT_ROOT=$WORKSPACE/yu/groot17_data

export HF_HOME=$WORKSPACE/hugging_face
export PYTHONPATH=$GROOT_DIR:${PYTHONPATH:-}

mkdir -p $WORKSPACE/yu/logs "$OUT_ROOT"

echo "========================================"
echo " GR00T Data Preparation"
echo " src:  $SRC_ROOT"
echo " out:  $OUT_ROOT"
echo " time: $(date)"
echo "========================================"

$PYTHON $GROOT_DIR/scripts/prepare_conflict_data.py \
    --src_root "$SRC_ROOT" \
    --out_root "$OUT_ROOT" \
    --num_demos 300 \
    --fps 30

echo "========================================"
echo " Data preparation complete: $(date)"
echo "========================================"