#!/bin/bash

task_name=${1}
task_config=${2}
expert_data_num=${3}
seed=${4}
action_dim=${5}
gpu_id=${6}
shift 6

head_camera_type=D435

DEBUG=False
save_ckpt=True

alg_name=robot_dp_$action_dim
config_name=${alg_name}
addition_info=train
exp_name=${task_name}-robot_dp-${addition_info}
run_dir="data/outputs/${exp_name}_seed${seed}"
if [ "${action_dim}" = "8" ]; then
    zarr_path="data/${task_name}-${task_config}-8d-${expert_data_num}.zarr"
else
    zarr_path="data/${task_name}-${task_config}-${expert_data_num}.zarr"
fi
extra_args=()
for override in "$@"; do
    case "${override}" in
        task.dataset.zarr_path=*)
            zarr_path="${override#task.dataset.zarr_path=}"
            ;;
        *)
            extra_args+=("${override}")
            ;;
    esac
done
if [[ "${zarr_path}" = /* ]]; then
    zarr_dir="${zarr_path}"
else
    zarr_dir="./${zarr_path}"
fi

echo -e "\033[33mgpu id (to use): ${gpu_id}\033[0m"


if [ $DEBUG = True ]; then
    wandb_mode=offline
    # wandb_mode=online
    echo -e "\033[33mDebug mode!\033[0m"
    echo -e "\033[33mDebug mode!\033[0m"
    echo -e "\033[33mDebug mode!\033[0m"
else
    wandb_mode=online
    echo -e "\033[33mTrain mode\033[0m"
fi

export HYDRA_FULL_ERROR=1 
export CUDA_VISIBLE_DEVICES=${gpu_id}

if [ ! -d "${zarr_dir}" ]; then
    process_args=("${task_name}" "${task_config}" "${expert_data_num}" --save-dir "${zarr_path}")
    if [ "${action_dim}" = "8" ]; then
        process_args+=(--right-arm-only)
    fi
    bash process_data.sh "${process_args[@]}"
fi

python train.py --config-name=${config_name}.yaml \
                            task.name=${task_name} \
                            task.dataset.zarr_path="${zarr_path}" \
                            training.debug=$DEBUG \
                            training.seed=${seed} \
                            training.device="cuda:0" \
                            exp_name=${exp_name} \
                            logging.mode=${wandb_mode} \
                            setting=${task_config} \
                            expert_data_num=${expert_data_num} \
                            head_camera_type=$head_camera_type \
                            "${extra_args[@]}"
                            # checkpoint.save_ckpt=${save_ckpt}
                            # hydra.run.dir=${run_dir} \
