#!/usr/bin/env bash
set -euo pipefail

ROOT=${ROOT:-/workspace/mp_20_pxrdnet}
ENV_ACTIVATE=${ENV_ACTIVATE:-/workspace/miniforge/bin/activate}
ENV_NAME=${ENV_NAME:-pxrdnet}

MODEL_PATH=${MODEL_PATH:-${ROOT}/hydra/singlerun/2026-07-20/pxrdgen_raw512_run02}
OUTPUT_ROOT=${OUTPUT_ROOT:-${ROOT}/paper_results_pxrdgen_match_only/cpu_multinode/full_test}
TOTAL_MATERIALS=${TOTAL_MATERIALS:-9046}
CHUNK_SIZE=${CHUNK_SIZE:?Set CHUNK_SIZE to the best single-node material count.}
NODE_INDEX=${NODE_INDEX:-${SLURM_ARRAY_TASK_ID:-0}}

FIRST_IDX=$((NODE_INDEX * CHUNK_SIZE))
if (( FIRST_IDX >= TOTAL_MATERIALS )); then
  echo "NODE_INDEX=${NODE_INDEX} starts at ${FIRST_IDX}, beyond TOTAL_MATERIALS=${TOTAL_MATERIALS}; nothing to do."
  exit 0
fi
REMAINING=$((TOTAL_MATERIALS - FIRST_IDX))
if (( REMAINING < CHUNK_SIZE )); then
  NUM_MATERIALS=${REMAINING}
else
  NUM_MATERIALS=${CHUNK_SIZE}
fi

OUTPUT_DIR=${OUTPUT_DIR:-${OUTPUT_ROOT}/node_$(printf "%05d" "${NODE_INDEX}")}
CPU_WORKERS=${CPU_WORKERS:-0}
TORCH_THREADS_PER_WORKER=${TORCH_THREADS_PER_WORKER:-0}
NUM_STARTING_POINTS=${NUM_STARTING_POINTS:-1}
NUM_GRADIENT_STEPS=${NUM_GRADIENT_STEPS:-5000}
N_STEP_EACH=${N_STEP_EACH:-100}
PROGRESS_LOG_INTERVAL=${PROGRESS_LOG_INTERVAL:-500}
PROGRESS_MININTERVAL=${PROGRESS_MININTERVAL:-10}

source "${ENV_ACTIVATE}" "${ENV_NAME}"
cd "${ROOT}"

export PROJECT_ROOT="${ROOT}"
export HYDRA_JOBS="${ROOT}/hydra"
export WABDB_DIR="${ROOT}/wabdb"
export PYTHONPATH="${ROOT}:${ROOT}/scripts:${PYTHONPATH:-}"
export WANDB_MODE=disabled
export CUDA_VISIBLE_DEVICES=""
export PIP_CACHE_DIR="${PIP_CACHE_DIR:-/workspace/.cache/pip}"
export TMPDIR="${TMPDIR:-/workspace/tmp}"

mkdir -p "${OUTPUT_DIR}"

echo "Running node chunk:"
echo "  NODE_INDEX=${NODE_INDEX}"
echo "  FIRST_IDX=${FIRST_IDX}"
echo "  NUM_MATERIALS=${NUM_MATERIALS}"
echo "  CHUNK_SIZE=${CHUNK_SIZE}"
echo "  OUTPUT_DIR=${OUTPUT_DIR}"

python -W ignore "${ROOT}/cpu_generation_benchmark/run_cpu_generation_benchmark.py" \
  --model-path "${MODEL_PATH}" \
  --output-dir "${OUTPUT_DIR}" \
  --first-idx "${FIRST_IDX}" \
  --num-materials "${NUM_MATERIALS}" \
  --num-starting-points "${NUM_STARTING_POINTS}" \
  --workers "${CPU_WORKERS}" \
  --torch-threads-per-worker "${TORCH_THREADS_PER_WORKER}" \
  --num-gradient-steps "${NUM_GRADIENT_STEPS}" \
  --n-step-each "${N_STEP_EACH}" \
  --progress-log-interval "${PROGRESS_LOG_INTERVAL}" \
  --progress-mininterval "${PROGRESS_MININTERVAL}" \
  2>&1 | tee "${OUTPUT_DIR}/benchmark.log"