#!/bin/bash
# Instantiate and submit SLURM benchmark jobs.
#
# Usage: bash slurm/submit.sh [--papi] [--nspins N] [--params LIST] [--variants LIST] [--node NODE]
#
# Examples:
#   bash slurm/submit.sh
#   bash slurm/submit.sh --papi --nspins 500
#   bash slurm/submit.sh --variants "ref avx2" --params "512 1024"
#   bash slurm/submit.sh --node node2334   # pin all jobs to a specific node

set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
TEMPLATE="${REPO_ROOT}/slurm/templates/bench_mlkem.sh.tmpl"

# ── Defaults ─────────────────────────────────────────────────────────────────
NSPINS=1000
WITH_PAPI=OFF
PARAMS="512 768 1024"
VARIANTS="ref avx2 refnv refo0"
BENCH_NODE=""

while [[ $# -gt 0 ]]; do
  case "$1" in
    --papi)     WITH_PAPI=ON ;;
    --nspins)   shift; NSPINS="$1" ;;
    --params)   shift; PARAMS="$1" ;;
    --variants) shift; VARIANTS="$1" ;;
    --node)     shift; BENCH_NODE="$1" ;;
    *) echo "unknown flag: $1" >&2; exit 1 ;;
  esac
  shift
done

# Build directory created by build.sh.
BUILD_DIR="${REPO_ROOT}/harness/build-hpc"

if [[ ! -d "$BUILD_DIR" ]]; then
  echo "ERROR: $BUILD_DIR not found — run slurm/build.sh first" >&2
  exit 1
fi

echo "=== pqc-bench submit ==="
echo "NSPINS   : $NSPINS"
echo "WITH_PAPI: $WITH_PAPI"
echo "PARAMS   : $PARAMS"
echo "VARIANTS : $VARIANTS"
echo "NODE     : ${BENCH_NODE:-any}"
echo ""

JOBS_SUBMITTED=0

for PARAM in $PARAMS; do
  for VARIANT in $VARIANTS; do
    BINARY="${BUILD_DIR}/bench_mlkem${PARAM}_${VARIANT}"
    if [[ ! -x "$BINARY" ]]; then
      echo "SKIP  bench_mlkem${PARAM}_${VARIANT} — binary not found"
      continue
    fi

    # Output goes into data/raw/kyber/mlkem{PARAM}/{VARIANT}/ so the aggregation
    # tool infers algorithm and variant from the directory structure.
    OUTPUT_DIR="${REPO_ROOT}/data/raw/kyber/mlkem${PARAM}/${VARIANT}"
    mkdir -p "$OUTPUT_DIR"

    # Instantiate template.
    JOB_SCRIPT="$(mktemp /tmp/bench_mlkem${PARAM}_${VARIANT}.XXXXXX.sh)"
    export PARAM VARIANT NSPINS BUILD_DIR OUTPUT_DIR WITH_PAPI BENCH_NODE
    envsubst '${PARAM} ${VARIANT} ${NSPINS} ${BUILD_DIR} ${OUTPUT_DIR} ${WITH_PAPI} ${BENCH_NODE}' \
      < "$TEMPLATE" > "$JOB_SCRIPT"
    chmod +x "$JOB_SCRIPT"

    SBATCH_ARGS="--parsable"
    if [[ -n "$BENCH_NODE" ]]; then
      SBATCH_ARGS="$SBATCH_ARGS --nodelist=$BENCH_NODE"
    fi
    JOB_ID=$(sbatch $SBATCH_ARGS "$JOB_SCRIPT")
    echo "SUBMIT bench_mlkem${PARAM}_${VARIANT}  job=${JOB_ID}  out=${OUTPUT_DIR}/${JOB_ID}.out"
    JOBS_SUBMITTED=$((JOBS_SUBMITTED + 1))
  done
done

echo ""
echo "Submitted $JOBS_SUBMITTED jobs."