#!/bin/bash
script_name=`basename "$0"`
id=${script_name%.*}

dataset=${dataset:-'none'}
gpus=${gpus:-'0,1,2,3,4,5,6,7'}
start_pid=${start_pid:-0}
N=${N:-40}  # total number of prompts
task=${task:-'neg-opp-lib-v2'}
thresh=${thresh:-0.85}

while [ $# -gt 0 ]; do
    if [[ $1 == *"--"* ]]; then
        param="${1/--/}"
        declare $param="$2"
    fi
    shift
done

#### common
path=$dataset/usps-loss-above=$thresh-100-improve-gpt-4.json
version='v1'
num_seeds=2

#### round 2
## change group!!!
# cmd="bash improve-h2080.sh --group exp-r4-$dataset-$task --path $path --wandb 1 --task $task --num_seeds 3 --num_seeds_train $num_seeds --transfer 1 --override y --version $version"
## combine
# cmd="bash improve-h2080.sh --group exp-r4-$dataset-$task-combine --path $path --wandb 1 --task $task --num_seeds 3 --num_seeds_train $num_seeds --transfer 1 --override y --version $version --loss spherical"
## opp v2
# path=$dataset/usps-loss-above=0.85-100-improve-gpt-4-v2.json
# cmd="bash improve-h2080.sh --group exp-r4-$dataset-$task-combine-opp-v2 --path $path --wandb 1 --task $task --num_seeds 3 --num_seeds_train $num_seeds --transfer 1 --override y --version $version --loss spherical"

#### FINAL
cmd="bash improve-h2080.sh --group exp-filter-$dataset-$task --path $path --wandb 1 --task $task --num_seeds 2 --num_seeds_train $num_seeds --override y --version $version --loss spherical"




IFS=',' read -ra gpus <<< "$gpus"
M=${#gpus[@]}   # Set your desired number of groups
session_name=exp-$dataset

#### create session
tmux kill-session -t $session_name
tmux new-session -d -s $session_name


#### distribute jobs
numbers_per_group=$(( (N + 1) / M ))
# Loop through each group
for ((group=0; group<M; group++)); do
    ## obtain divided prompt_ids
    start=$((start_pid + group * numbers_per_group))
    end=$((start + numbers_per_group - 1))
    group_numbers=""
    for ((num=start; num<=end; num++)); do
        if [ "$num" -eq "$end" ]; then
            group_numbers+="$num"
        else
            group_numbers+="$num,"
        fi
    done

    echo ${group_numbers[@]}

    ## commands
    # echo "${group_numbers[@]}"
    tmux select-window -t $i

    # Send the command to the current window
    gpu=${gpus[$group]}
    tmux send-keys "$cmd --prompt_id ${group_numbers[@]} --gpu $gpu" C-m

    # Create a new window with a specific name
    tmux new-window #-n "Window $i"

done

tmux attach-session -t $session_name