#!/bin/bash

MODEL_LIST="models.txt"
SUPERVISOR_LIST="models_supervisor.txt"
CODER_LIST="models_coder.txt"
NUM_TESTS=10
OUTDIR="/global/cfs/projectdirs/atlas/llm4hep/oct_11_tests/"

usage() {
    echo "Usage: $0 [--mode identical|pairwise]"
    echo "  --mode identical : Use the same model for both supervisor and coder (from models.txt) [default]"
    echo "  --mode pairwise  : Use all pairs (from models_supervisor.txt and models_coder.txt)"
    exit 1
}

# Default mode
MODE="identical"

# Parse arguments
while [[ $# -gt 0 ]]; do
    case "$1" in
        --mode)
            MODE="$2"
            shift 2
            ;;
        *)
            usage
            ;;
    esac
done

if [[ "$MODE" == "identical" ]]; then
    # One model for both supervisor and coder
    while IFS= read -r model; do
        model=$(echo "$model" | xargs)
        [ -z "$model" ] && continue
        echo "Supervisor & Coder: $model"
        sbatch --job-name="${model}_${model}" jobs/run_tests.sh "$model" "$model" "$NUM_TESTS" "$OUTDIR"
    done < "$MODEL_LIST"
elif [[ "$MODE" == "pairwise" ]]; then
    # Different models for supervisor and coder
    while IFS= read -r supervisor; do
        supervisor=$(echo "$supervisor" | xargs)
        [ -z "$supervisor" ] && continue
        while IFS= read -r coder; do
            coder=$(echo "$coder" | xargs)
            [ -z "$coder" ] && continue
            echo "Supervisor: $supervisor, Coder: $coder"
            sbatch --job-name="${supervisor}_${coder}" jobs/run_tests.sh "$supervisor" "$coder" "$NUM_TESTS" "$OUTDIR"
        done < "$CODER_LIST"
    done < "$SUPERVISOR_LIST"
else
    usage
fi