#!/bin/bash

source config.sh
mkdir -p logs

log_message() {
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" | tee -a logs/master_HPC.log
}

count_user_jobs() {
  squeue -u $USER -h | wc -l
}

# Count total combinations
num_locations=${#locations[@]}
num_sigma=${#sigma_betas_vals[@]}
num_IDs=${#ID_vals[@]}
num_delta_weather=${#True_e_vals[@]}
num_rho_Mean=${#rho_Means[@]}
num_rho_K=${#rho_Ks[@]}
num_imports=${#import_rates[@]}
num_School_Terms=${#School_Terms[@]}

total_jobs=$((num_locations * num_sigma * num_IDs * num_delta_weather * num_rho_Mean * num_rho_K * num_imports * num_School_Terms))
batch_size=300

log_message "Total jobs to submit: $total_jobs"

start=1
batch_number=1

while [ $start -le $total_jobs ]; do
  # --- STOP flag check ---
  if [ -f STOP ]; then
    log_message "STOP file found. Exiting gracefully without submitting further jobs."
    exit 0
  fi
  # ------------------------

  end=$((start + batch_size - 1))
  [ $end -gt $total_jobs ] && end=$total_jobs

  while [ $(count_user_jobs) -gt 0 ]; do
    log_message "Waiting for current jobs to finish..."
    sleep 60

    # Also check STOP during waiting
    if [ -f STOP ]; then
      log_message "STOP file found during wait. Exiting gracefully."
      exit 0
    fi
  done

  log_message "Submitting batch $batch_number: jobs $start to $end"
  job_id=$(sbatch --parsable --array=${start}-${end} main_HPC.sh)

  if [ $? -ne 0 ]; then
    log_message "ERROR submitting batch $batch_number"
    exit 1
  fi

  log_message "Submitted batch $batch_number with job ID: $job_id"
  start=$((end + 1))
  batch_number=$((batch_number + 1))
done

log_message "All batches submitted."