#!/bin/bash

# Default values
DATA_DIR="data/flare_subsampled"
TASK=""
MODEL="gpt-4o"
EXP_NAME=""
SAMPLE_SIZES="100"
TEST_SETS="test"
PROMPT_BEFORE_EXEMPLARS="prompt_before_exemplars.txt"
PROMPT_AFTER_EXEMPLARS="prompt_after_exemplars_zero_shot.txt"
RESULTS_DIR="results"
SEED=42  # Default seed value
TOKEN_SUBSTITUTION_STRATEGY="many2one"
ONE2MANY_TIMESTEP="5"

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --data_dir)
            DATA_DIR="$2"
            shift 2
            ;;
        --exp_name)
            EXP_NAME="$2"
            shift 2
            ;;
        --task)
            TASK="$2"
            shift 2
            ;;
        --model)
            MODEL="$2"
            shift 2
            ;;
        --sample_sizes)
            SAMPLE_SIZES="$2"
            shift 2
            ;;
        --test_sets)
            TEST_SETS="$2"
            shift 2
            ;;
        --prompt_before_exemplars)
            PROMPT_BEFORE_EXEMPLARS="$2"
            shift 2
            ;;
        --prompt_after_exemplars)
            PROMPT_AFTER_EXEMPLARS="$2"
            shift 2
            ;;
        --results_dir)
            RESULTS_DIR="$2"
            shift 2
            ;;
        --seed)
            SEED="$2"
            shift 2
            ;;
        --token_substitution_strategy)
            TOKEN_SUBSTITUTION_STRATEGY="$2"
            shift 2
            ;;
        --one2many_timestep)
            ONE2MANY_TIMESTEP="$2"
            shift 2
            ;;
        *)
            echo "Unknown parameter: $1"
            exit 1
            ;;
    esac
done

# Check required parameters
if [[ -z "$EXP_NAME" ]]; then
    echo "Error: exp_name is required"
    exit 1
fi

if [[ -z "$TASK" || -z "$MODEL" ]]; then
    echo "Error: Required parameters missing"
    echo "Usage: $0 --task <task> --model <model> [options]"
    echo "Options:"
    echo "  --exp_name <name>          Experiment name (required)"
    echo "  --data_dir <dir>           Data directory (default: data/flare)"
    echo "  --sample_sizes <sizes>     Comma-separated list of sample sizes (default: 100)"
    echo "  --test_sets <sets>         Comma-separated list of test sets (default: test)"
    echo "  --prompt_before_exemplars <file> Prompt template file before examples (default: base_prompt_before_exemplars.txt)"
    echo "  --prompt_after_exemplars <file> Prompt template file after examples (default: base_prompt_after_exemplars_zero_shot.txt)"
    echo "  --results_dir <dir>        Results directory (default: results)"
    echo "  --seed <seed>              Random seed for reproducibility (default: 42)"
    echo "  --token_substitution_strategy <strategy> Token substitution strategy (default: many2one)"
    echo "  --one2many_timestep <timestep> Timestep for one2many token substitution (default: 5)"
    exit 1
fi

# Create job files for each combination
IFS=',' read -ra SIZES <<< "$SAMPLE_SIZES"
IFS=',' read -ra TEST_SET_ARRAY <<< "$TEST_SETS"

for size in "${SIZES[@]}"; do
    for test_set in "${TEST_SET_ARRAY[@]}"; do
        # Create unique job file for each combination
        JOB_FILE="batch_jobs/${EXP_NAME}/${TASK}_${MODEL// /_}_${size}_${test_set}.job"
        mkdir -p "batch_jobs/${EXP_NAME}"

        # Construct command for this combination
        CMD="python src/openai/batch_experiment.py \
            --exp_name $EXP_NAME \
            --data_dir $DATA_DIR \
            --task $TASK \
            --model $MODEL \
            --sample_sizes $size \
            --test_sets $test_set \
            --prompt_before_exemplars $PROMPT_BEFORE_EXEMPLARS \
            --prompt_after_exemplars $PROMPT_AFTER_EXEMPLARS \
            --seed $SEED \
            --token_substitution_strategy $TOKEN_SUBSTITUTION_STRATEGY \
            --one2many_timestep $ONE2MANY_TIMESTEP"

        # Save command to job file
        echo "$CMD" > "$JOB_FILE"
        echo "Created job file: $JOB_FILE"
    done
done
