#!/bin/bash
set -euo pipefail

# Submit short-budget sweep jobs for successive-halving style exploration.
# Usage:
#   slurm/submit_gguf_sweep.sh
#
# Env overrides:
#   RUN_TAG=sweepA
#   SWEEP_TIME=02:00:00
#   MAX_CASES_LIST=10,20,40
#   K_LIST=1,2
#   CONTEXT_LIST=2048,4096
#   JOBS_LIST=2,4
#   REQUIRE_SMOKE=1

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd -P)"
REPO_DIR="$(cd "$SCRIPT_DIR/.." && pwd -P)"

RUN_TAG="${RUN_TAG:-sweep-$(date +%Y%m%dT%H%M%S)}"
SWEEP_TIME="${SWEEP_TIME:-02:00:00}"
MAX_CASES_LIST="${MAX_CASES_LIST:-10,20,40}"
K_LIST="${K_LIST:-1,2}"
CONTEXT_LIST="${CONTEXT_LIST:-2048,4096}"
JOBS_LIST="${JOBS_LIST:-2,4}"
REQUIRE_SMOKE="${REQUIRE_SMOKE:-1}"

if [[ "${REQUIRE_SMOKE,,}" =~ ^(true|yes|on)$ ]]; then
    REQUIRE_SMOKE=1
elif [[ "${REQUIRE_SMOKE,,}" =~ ^(false|no|off)$ ]]; then
    REQUIRE_SMOKE=0
fi

csv_to_array() {
    local csv="$1"
    local -n out_ref="$2"
    IFS=',' read -r -a out_ref <<< "$csv"
}

cd "$REPO_DIR"

if (( REQUIRE_SMOKE == 1 )); then
    echo "Running mandatory smoke test before sweep..."
    WAIT_FOR_COMPLETION=1 slurm/run_gguf_smoke.sh
fi

csv_to_array "$MAX_CASES_LIST" MAX_CASES_ARR
csv_to_array "$K_LIST" K_ARR
csv_to_array "$CONTEXT_LIST" CONTEXT_ARR
csv_to_array "$JOBS_LIST" JOBS_ARR

echo "Submitting sweep jobs (run tag: ${RUN_TAG})"
echo "time=${SWEEP_TIME}"
printf '%-14s %-45s %-30s\n' "JOB_ID" "RUN_ID" "PARAMS"
printf '%-14s %-45s %-30s\n' "------" "------" "------"

for max_cases in "${MAX_CASES_ARR[@]}"; do
    for k in "${K_ARR[@]}"; do
        for ctx in "${CONTEXT_ARR[@]}"; do
            for jobs in "${JOBS_ARR[@]}"; do
                run_id="${RUN_TAG}-m${max_cases}-k${k}-c${ctx}-j${jobs}"
                job_id="$(sbatch --parsable \
                    --time="$SWEEP_TIME" \
                    --export=ALL,SMOKE_TEST_MODE=0,EVAL_RUN_ID="$run_id",EVAL_MAX_CASES="$max_cases",EVAL_K="$k",CONTEXT_SIZE="$ctx",EVAL_JOBS="$jobs" \
                    slurm/eval_gguf_eval.sh)"

                printf '%-14s %-45s %-30s\n' \
                    "$job_id" "$run_id" "max=${max_cases},k=${k},ctx=${ctx},jobs=${jobs}"
            done
        done
    done
done
