#!/bin/bash
CONDA_ENV_NAME='vwebarena'          # the name of your conda environment
TOTAL_PROCESSES=5                # the number of processes to run in parallel   
test_config_base_dir='config_files/vwa_not_vague/test_shopping'
TASKS_FILE=''
SKIP_COOKIES=true

# Parse optional arguments (-t for tasks file, -s to skip cookies, -c for config dir)
while getopts ":t:c:n:a:s" opt; do
    case $opt in
        t) TASKS_FILE="$OPTARG" ;;
        s) SKIP_COOKIES=true ;;
        c) test_config_base_dir="$OPTARG" ;;
        n) TOTAL_PROCESSES="$OPTARG" ;;
        a) agent_config="$OPTARG" ;;
        \?) echo "Invalid option -$OPTARG" >&2 ;;
    esac
done
# Example usage:
# ./p_run.sh -t evaluation_harness/task_subsets/shopping.txt -c config_files/vwa_not_vague/test_shopping


# Determine the start and end IDs based on files in the config base directory
start_id=$(ls "$test_config_base_dir" | grep -oP '^\d+' | sort -n | head -n 1)
end_id=$(ls "$test_config_base_dir" | grep -oP '^\d+' | sort -n | tail -n 1)

# Tmux commands for creating panes
tmux_commands=(
    'tmux split-window -h'  # 2 panes
    'tmux select-pane -t 0; tmux split-window -v' # 3 panes
    'tmux select-pane -t 0; tmux split-window -v' # 4 panes
    'tmux select-pane -t 3; tmux split-window -v' # 5 panes
    'tmux select-pane -t 3; tmux split-window -v' # 6 panes
)

# Fill with NUM_TOTAL_PROCESSES strings if want results to be saved in specific directories
results_dirs=(
    '' '' ''
)

api_keys=(
""
""
""
""
""
""
""
""
""
""
""
""
""
""
""
""
""
""
""
)
# Tells which device the captioner model is running on
agent_captioning_model_devices=(
    'server-cuda'
    'server-cuda'
    'server-cuda'
    'server-cuda'
    'server-cuda'
    'server-cuda'
    'server-cuda'
    'server-cuda'
    'server-cuda'
    'server-cuda'
)


# Copy api_keys.json to api_keys_copy.json if it exists
if [[ -f "api_keys.json" ]]; then
    cp api_keys.json api_keys_copy.json
fi

#--------------------------------------------------
# Function to run a single job in a tmux pane.
run_job() {
    pane_index=$1; start_id=$2; end_id=$3; index=$4; task_file=$5;

    if [[ -z "$task_file" ]]; then
        task_file="''"
    fi

    tmux select-pane -t $pane_index

    if [[ -n "${results_dirs[$index]}" ]]; then
        cmd="conda activate ${CONDA_ENV_NAME}; ./scripts/runs/run.sh -s $start_id -e $end_id -c ${test_config_base_dir} -k \"${api_keys[$index]}\" -d ${results_dirs[$index]} -t $task_file -m ${agent_captioning_model_devices[$index]} -a \"${agent_config}\" -r; exit"
        # cmd="conda activate ${CONDA_ENV_NAME}; ./scripts/runs/run.sh -s $start_id -e $end_id -c ${test_config_base_dir} -k \"${api_keys[$index]}\" -d ${results_dirs[$index]} -t $task_file -m ${agent_captioning_model_devices[$index]} -a \"${agent_config}\" -r"
    else
        cmd="conda activate ${CONDA_ENV_NAME}; ./scripts/runs/run.sh -s $start_id -e $end_id -c ${test_config_base_dir} -k \"${api_keys[$index]}\" -t $task_file -m ${agent_captioning_model_devices[$index]} -a \"${agent_config}\" -r; exit"
        # cmd="conda activate ${CONDA_ENV_NAME}; ./scripts/runs/run.sh -s $start_id -e $end_id -c ${test_config_base_dir} -k \"${api_keys[$index]}\" -t $task_file -m ${agent_captioning_model_devices[$index]} -a \"${agent_config}\" -r"
    fi
    tmux send-keys "$cmd" ENTER
    sleep 3
}

#--------------------------------------------------
# Function to run batches of jobs
run_batch() {
    if [[ -n "$TASKS_FILE" ]]; then
        num_jobs=$TOTAL_PROCESSES
        for ((i=0; i<num_jobs; i++)); do
            pane_index=$i
            task_file="tasks$((i+1)).txt"
            run_job $pane_index 0 1 $i $task_file
        done
    else
        args=("$@")  # all arguments as an array
        num_jobs=$(( ${#args[@]} / 2 ))
        for ((i=0; i<num_jobs; i++)); do
            pane_index=$i
            start_id=${args[$((i*2))]}
            end_id=${args[$((i*2+1))]}
            run_job $pane_index $start_id $end_id $i ""
        done
    fi
}

#--------------------------------------------------
# Create a new tmux session and the required panes
TMUX_SESSION=$(tmux new-session -d -P -F "#{session_id}")
for ((i=0; i<$TOTAL_PROCESSES-1; i++)); do
    eval ${tmux_commands[$i]}
done

# Create autologin cookies if not skipped
# Obs.: the `run.sh` calls don't create cookies; it is only created here.
if [[ "$SKIP_COOKIES" == false ]]; then
    echo "Creating autologin cookies"
    ./scripts/environments/autologin_cookies.sh 'local_vwebarena'
fi

#--------------------------------------------------
# Distribute tasks
if [[ -n "$TASKS_FILE" ]]; then
    mapfile -t all_tasks < "$TASKS_FILE"
    total_tasks=${#all_tasks[@]}
    tasks_per_process=$(( (total_tasks + TOTAL_PROCESSES - 1) / TOTAL_PROCESSES ))

    for ((i=0; i<TOTAL_PROCESSES; i++)); do
        start_index=$(( i * tasks_per_process ))
        end_index=$(( (i + 1) * tasks_per_process - 1 ))
        (( end_index >= total_tasks )) && end_index=$(( total_tasks - 1 ))

        task_file="tasks$((i+1)).txt"
        > "$task_file"
        for ((j=start_index; j<=end_index && j<total_tasks; j++)); do
            echo "${all_tasks[j]}" >> "$task_file"
        done
    done

    run_batch
    sleep 5
    # Remove temporary task files
    for ((i=0; i<TOTAL_PROCESSES; i++)); do
        rm -f "tasks$((i+1)).txt"
    done
else
    total_jobs=$((end_id - start_id + 1))
    batch_size=$((total_jobs / TOTAL_PROCESSES))
    remaining_jobs=$((total_jobs % TOTAL_PROCESSES))
    
    job_args=()
    for ((i=0; i<TOTAL_PROCESSES; i++)); do
        batch_start=$(( start_id + i * batch_size + (i < remaining_jobs ? i : remaining_jobs) ))
        batch_end=$(( batch_start + batch_size - 1 + (i < remaining_jobs ? 1 : 0) ))
        job_args+=($batch_start $batch_end)
    done

    for ((i=0; i<TOTAL_PROCESSES; i++)); do
        echo "Process $i: Start ID = ${job_args[$((i*2))]}, End ID = ${job_args[$((i*2+1))]}"
    done

    run_batch "${job_args[@]}"
fi

echo "Experiments running on tmux session $TMUX_SESSION"