#!/bin/bash
#SBATCH --job-name=${DOVLA_JOB_NAME:-dovla_eval}
#SBATCH --partition=${DOVLA_PARTITION:-gpu}
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=${DOVLA_CPUS_PER_TASK:-4}
#SBATCH --gres=gpu:${DOVLA_GPUS_PER_TASK:-1}
#SBATCH --mem=${DOVLA_MEM:-32G}
#SBATCH --time=${DOVLA_TIME:-04:00:00}
#SBATCH --output=${DOVLA_LOG_DIR:-logs/slurm}/%x_%j.out
#SBATCH --error=${DOVLA_LOG_DIR:-logs/slurm}/%x_%j.err

set -euo pipefail

PROJECT_DIR="${PROJECT_DIR:-$PWD}"
VENV_PATH="${VENV_PATH:-$PROJECT_DIR/.venv}"
CHECKPOINT="${CHECKPOINT:-$PROJECT_DIR/runs/dovla_toy/best.pt}"
OUT_PATH="${OUT_PATH:-$PROJECT_DIR/runs/dovla_toy/causalstress.json}"
BACKEND="${BACKEND:-toy}"
NUM_TASKS="${NUM_TASKS:-20}"
K="${K:-16}"
SEED="${SEED:-0}"
DEVICE="${DEVICE:-auto}"

mkdir -p "${DOVLA_LOG_DIR:-logs/slurm}" "$(dirname "$OUT_PATH")"
cd "$PROJECT_DIR"

if [ -f "$VENV_PATH/bin/activate" ]; then
  # shellcheck disable=SC1091
  source "$VENV_PATH/bin/activate"
fi

python scripts/eval_causalstress.py \
  --checkpoint "$CHECKPOINT" \
  --backend "$BACKEND" \
  --out "$OUT_PATH" \
  --num-tasks "$NUM_TASKS" \
  --k "$K" \
  --seed "$SEED" \
  --device "$DEVICE"