#!/bin/bash

tasks=(    
#"adjust_bottle"
#"beat_block_hammer"
#"blocks_ranking_rgb"
#"blocks_ranking_size"
#"click_alarmclock"
#"click_bell"
#"dump_bin_bigbin"
#"grab_roller"
#"handover_block"
#"handover_mic"

#"hanging_mug"
#"lift_pot"
#"move_can_pot"
#"move_pillbottle_pad"
#"move_playingcard_away"
#"move_stapler_pad"
#"open_laptop"
#"open_microwave"
#"pick_dual_bottles"
#"pick_diverse_bottles"


#"place_a2b_left"
#"place_a2b_right"
#"place_bread_basket"
#"place_bread_skillet"
#"place_can_basket"
#"place_cans_plasticbox"
#"place_container_plate"
#"place_dual_shoes"
#"place_empty_cup"
#"place_fan"

#"place_burger_fries"
#"place_mouse_pad"
#"place_object_basket"
#"place_object_scale"
#"place_object_stand"
#"place_phone_stand"
#"place_shoe"
#"press_stapler"
#"put_bottles_dustbin"
#"put_object_cabinet"

#"rotate_qrcode"
#"shake_bottle"
#"shake_bottle_horizontally"
#"stack_blocks_three"
#"stack_blocks_two"
#"stack_bowls_three"
#"stack_bowls_two"
#"stamp_seal"
#"turn_switch"
)

# directory where processed data will be saved
SAVE_ROOT_DIR="/yourpath/RoboTwin/policy/LLaVA-VLA/training_data"
PYTHON_SCRIPT="/yourpath/RoboTwin/policy/LLaVA-VLA/llava/process_data/process_data.py"

# check if the script is run with sufficient arguments
if [ $# -lt 1 ]; then
    echo "Usage: $0 <task_name> <task_config> [<future_step_num>] | $0 all <task_config> [<future_step_num>]"
    exit 1
fi

# process all tasks or a task
if [ "$1" == "all" ]; then
    if [ $# -lt 2 ]; then
        echo "Usage for all tasks: $0 all <task_config> [<future_step_num>]"
        exit 1
    fi
    task_config="$2"
    future_step_num=${3:-5}  # default value is 5

    echo "📦 Found ${#tasks[@]} tasks: ${tasks[*]}"
    echo "Task configuration: $task_config"
    for task in "${tasks[@]}"; do
        DATA_ROOT_DIR="/yourpath/RoboTwin/data/$task/$task_config/data/"
        INSTRUCTION_DIR="/yourpath/RoboTwin/data/$task/$task_config/instructions"
        echo " Starting task: $task"
        python "$PYTHON_SCRIPT" \
            --task_name "$task" \
            --task_config "$task_config" \
            --data_root_dir "$DATA_ROOT_DIR" \
            --instruction_dir "$INSTRUCTION_DIR" \
            --image_root_path "/yourpath/RoboTwin/policy/LLaVA-VLA/pictures" \
            --save_root_dir "$SAVE_ROOT_DIR" \
            --future_step_num "$future_step_num"
        if [ $? -eq 0 ]; then
            echo "Task $task completed successfully"
        else
            echo "Task $task failed"
        fi
        echo "----------------------------------------"
    done

    echo " All tasks completed!"
else
    # task-specific processing
    task_name="$1"
    task_config="$2"
    future_step_num=${3:-5}  # default value is 5
    DATA_ROOT_DIR="/yourpath/RoboTwin/data/$task_name/$task_config/data/"
    INSTRUCTION_DIR="/yourpath/RoboTwin/data/$task_name/$task_config/instructions"

    echo "Start processing task: $task_name"
    python "$PYTHON_SCRIPT" \
        --task_name "$task_name" \
        --task_config "$task_config" \
        --data_root_dir "$DATA_ROOT_DIR" \
        --instruction_dir "$INSTRUCTION_DIR" \
        --image_root_path "/yourpath/RoboTwin/pictures" \
        --save_root_dir "$SAVE_ROOT_DIR" \
        --future_step_num "$future_step_num"
    if [ $? -eq 0 ]; then
        echo " Task $task_name completed successfully"
    else
        echo " Task $task_name failed"
    fi
fi
