model="gpt-4o-mini"
tasks="even-pairs,majority,unmarked-reversal,marked-copy,binary-addition,binary-multiplication,bucket-sort,compute-sqrt,cycle-navigation,dyck-2-3,first,marked-reversal,missing-duplicate-string,modular-arithmetic-simple,odds-first,parity,repeat-01,stack-manipulation"
prompt_type="io_prompt"
sample_size="100"
test_set="test"
substitution_strategy='one2one'
seed=42

# Check for mode argument
if [ $# -eq 0 ]; then
    echo "Usage: $0 <mode>"
    echo "Available modes: submit, resubmit, process"
    exit 1
fi

run_mode="$1"

exp_name="${model}_${prompt_type}_${substitution_strategy}_seed${seed}_sample${sample_size}_${test_set}"
data_dir="data/flare_subsampled"
prompt_before_exemplars="prompts/prompt_before_exemplars.txt"
prompt_after_exemplars="prompts/${prompt_type}.txt"

if [[ "$run_mode" == "submit" ]]; then
    # First block: create
    poetry run  bash src/openai/run_batch_workflow.sh \
        --mode create \
        --tasks $tasks \
        --data_dir $data_dir \
        --test_sets $test_set \
        --sample_sizes $sample_size \
        --exp_name $exp_name \
        --prompt_before_exemplars $prompt_before_exemplars \
        --prompt_after_exemplars $prompt_after_exemplars \
        --model $model \
        --token_substitution_strategy $substitution_strategy

    # Second block: submit
    poetry run  bash src/openai/run_batch_workflow.sh \
        --mode submit \
        --tasks $tasks \
        --data_dir $data_dir \
        --test_sets $test_set \
        --sample_sizes $sample_size \
        --exp_name $exp_name \
        --prompt_before_exemplars $prompt_before_exemplars \
        --prompt_after_exemplars $prompt_after_exemplars \
        --model $model \
        --token_substitution_strategy $substitution_strategy
elif [[ "$run_mode" == "resubmit" ]]; then
    # Third block: resubmit
    poetry run  bash src/openai/run_batch_workflow.sh \
        --mode resubmit \
        --tasks $tasks \
        --data_dir $data_dir \
        --test_sets $test_set \
        --sample_sizes $sample_size \
        --exp_name $exp_name \
        --prompt_before_exemplars $prompt_before_exemplars \
        --prompt_after_exemplars $prompt_after_exemplars \
        --model $model \
        --token_substitution_strategy $substitution_strategy
elif [[ "$run_mode" == "process" ]]; then
    # Fourth block: process
    poetry run  bash src/openai/run_batch_workflow.sh \
        --mode process \
        --tasks $tasks \
        --data_dir $data_dir \
        --test_sets $test_set \
        --sample_sizes $sample_size \
        --exp_name $exp_name \
        --prompt_before_exemplars $prompt_before_exemplars \
        --prompt_after_exemplars $prompt_after_exemplars \
        --model $model \
        --token_substitution_strategy $substitution_strategy
else
    echo "Unknown mode: $run_mode"
    echo "Available modes: submit, resubmit, process"
    exit 1
fi
