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

# Run one physical node's chunk for a given wave.
#
# Mapping:
#   chunk_index = WAVE_INDEX * ACTUAL_NODES + NODE_RANK
#
# This lets a smaller number of real nodes process a larger plan in multiple
# waves without changing the underlying generation code.

ACTUAL_NODES=${ACTUAL_NODES:?Set ACTUAL_NODES to the number of real compute nodes available in this wave.}
NODE_RANK=${NODE_RANK:-${SLURM_ARRAY_TASK_ID:-0}}
WAVE_INDEX=${WAVE_INDEX:?Set WAVE_INDEX, starting from 0.}
TOTAL_MATERIALS=${TOTAL_MATERIALS:-9046}
CHUNK_SIZE=${CHUNK_SIZE:?Set CHUNK_SIZE to the best single-node material count.}

if (( ACTUAL_NODES <= 0 )); then
  echo "ERROR: ACTUAL_NODES must be positive." >&2
  exit 2
fi
if (( NODE_RANK < 0 || NODE_RANK >= ACTUAL_NODES )); then
  echo "ERROR: NODE_RANK=${NODE_RANK} must be in [0, ACTUAL_NODES)." >&2
  exit 2
fi

NODE_INDEX=$((WAVE_INDEX * ACTUAL_NODES + NODE_RANK))
FIRST_IDX=$((NODE_INDEX * CHUNK_SIZE))

if (( FIRST_IDX >= TOTAL_MATERIALS )); then
  echo "Wave node has no work:"
  echo "  WAVE_INDEX=${WAVE_INDEX}"
  echo "  ACTUAL_NODES=${ACTUAL_NODES}"
  echo "  NODE_RANK=${NODE_RANK}"
  echo "  NODE_INDEX=${NODE_INDEX}"
  echo "  FIRST_IDX=${FIRST_IDX}"
  echo "  TOTAL_MATERIALS=${TOTAL_MATERIALS}"
  exit 0
fi

export NODE_INDEX TOTAL_MATERIALS CHUNK_SIZE

ROOT=${ROOT:-/workspace/mp_20_pxrdnet}
bash "${ROOT}/cpu_multinode_generation/run_node_chunk.sh"