BEGIN_TIME=$(date +%s)

CONFIG_FILE="input/config.json"
echo "CONFIG_FILE_PATH: $CONFIG_FILE"

# Read values directly from the config.json using python - no more nested key error by using a helper script
TASK_ID=$(python3 read_json.py "$CONFIG_FILE" "task_id")
DATA_DIR=$(python3 read_json.py "$CONFIG_FILE" "data_dir")
OUTPUT_DIR=$(python3 read_json.py "$CONFIG_FILE" "output_dir")
EPISODE_NUM=$(python3 read_json.py "$CONFIG_FILE" "episode_num")
GPU=$(python3 read_json.py "$CONFIG_FILE" "gpu")
T5_PATH="/weights/t5-v1_1-xxl"
NO_LANGUAGE=$(python3 read_json.py "$CONFIG_FILE" "no_language")

# For the camera keys, extract them in a way that avoids the error about 'images_info.key.*' not found
CAM_HIGH_KEY=$(python3 -c "import json; print(json.load(open('$CONFIG_FILE'))['images_info']['key'].get('cam_high', ''))")
CAM_RIGHT_WRIST_KEY=$(python3 -c "import json; print(json.load(open('$CONFIG_FILE'))['images_info']['key'].get('cam_right_wrist', ''))")

# create output path
if [ ! -d "$OUTPUT_DIR/$TASK_ID" ]; then
  mkdir -p "$OUTPUT_DIR/$TASK_ID"
  echo "Created output directory: $OUTPUT_DIR/$TASK_ID"
else
  echo "Output directory already exists: $OUTPUT_DIR/$TASK_ID"
fi

if [ "$NO_LANGUAGE" = "true" ]; then
    python3 lerobot2rdt.py \
        --data_dir $DATA_DIR \
        --output_dir $OUTPUT_DIR/$TASK_ID \
        --episode_num $EPISODE_NUM \
        --gpu $GPU \
        --t5_path $T5_PATH \
        --cam_high_key $CAM_HIGH_KEY \
        --cam_right_wrist_key $CAM_RIGHT_WRIST_KEY \
        --no_language
    status=$?
else
    python3 lerobot2rdt.py \
        --data_dir $DATA_DIR \
        --output_dir $OUTPUT_DIR/$TASK_ID \
        --episode_num $EPISODE_NUM \
        --gpu $GPU \
        --t5_path $T5_PATH \
        --cam_high_key $CAM_HIGH_KEY \
        --cam_right_wrist_key $CAM_RIGHT_WRIST_KEY
    status=$?
fi

END_TIME=$(date +%s)
echo "END_TIME: $END_TIME"
echo "TOTAL_TIME: $((END_TIME - BEGIN_TIME))"

if [ $status -eq 0 ]; then
    python3 generate_output.py $CONFIG_FILE $((END_TIME - BEGIN_TIME))
else
    echo "lerobot2rdt.py exited with status $status, skipping generate_output.py"
fi