#!/bin/bash
#SBATCH --job-name=${DOVLA_JOB_NAME:-dovla_cil_gen}
#SBATCH --partition=${DOVLA_PARTITION:-compute}
#SBATCH --array=${DOVLA_ARRAY:-0-9}
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=${DOVLA_CPUS_PER_TASK:-8}
#SBATCH --gres=gpu:${DOVLA_GPUS_PER_TASK:-0}
#SBATCH --mem=${DOVLA_MEM:-32G}
#SBATCH --time=${DOVLA_TIME:-12:00:00}
#SBATCH --output=${DOVLA_LOG_DIR:-logs/slurm}/%x_%A_%a.out
#SBATCH --error=${DOVLA_LOG_DIR:-logs/slurm}/%x_%A_%a.err

set -euo pipefail

PROJECT_DIR="${PROJECT_DIR:-$PWD}"
VENV_PATH="${VENV_PATH:-$PROJECT_DIR/.venv}"
TASKS_PATH="${TASKS_PATH:-$PROJECT_DIR/data/tasks.jsonl}"
OUT_ROOT="${OUT_ROOT:-$PROJECT_DIR/data/cil_array}"
BACKEND="${BACKEND:-toy}"
NUM_WORKERS="${NUM_WORKERS:-4}"
STATES_PER_TASK="${STATES_PER_TASK:-1000}"
K="${K:-32}"
SHARD_SIZE="${SHARD_SIZE:-10000}"
SEED_BASE="${SEED_BASE:-0}"
RAY_ADDRESS="${RAY_ADDRESS:-}"
RESUME_FLAG="${RESUME_FLAG:-}"

mkdir -p "${DOVLA_LOG_DIR:-logs/slurm}" "$OUT_ROOT"
cd "$PROJECT_DIR"

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

export OPENCLAUDE_BASE_URL="${OPENCLAUDE_BASE_URL:-https://open-claude.com/v1}"
export OPENCLAUDE_MODEL="${OPENCLAUDE_MODEL:-<model>}"
# Set OPENCLAUDE_API_KEY in the job environment or scheduler secret store. Do not echo it.

SEED=$((SEED_BASE + SLURM_ARRAY_TASK_ID))
OUT_DIR="$OUT_ROOT/part_${SLURM_ARRAY_TASK_ID}"

CMD=(
  python scripts/generate_cil_distributed.py
  --backend "$BACKEND"
  --tasks "$TASKS_PATH"
  --out "$OUT_DIR"
  --num-workers "$NUM_WORKERS"
  --num-states-per-task "$STATES_PER_TASK"
  --k "$K"
  --seed "$SEED"
  --shard-size "$SHARD_SIZE"
)

if [ -n "$RAY_ADDRESS" ]; then
  CMD+=(--ray-address "$RAY_ADDRESS")
fi
if [ -n "$RESUME_FLAG" ]; then
  CMD+=(--resume)
fi

"${CMD[@]}"