#!/usr/bin/env bash
#SBATCH --job-name=run_nemo_est05
#SBATCH --qos=np
#SBATCH --nodes=3
#SBATCH --ntasks=257                 # 1 xios + 256 nemo
#SBATCH --hint=nomultithread
#SBATCH --time=00:30:00
#SBATCH --output=%x_%j.out
#SBATCH --error=%x_%j.err

set -euo pipefail

# ---------- Modules ----------
if command -v module >/dev/null 2>&1; then
  module purge
  module load prgenv/gnu
  module load openmpi hdf5 netcdf4
fi

ld_path_parts=()
[[ -n "${NETCDF4_DIR:-}" ]] && ld_path_parts+=("${NETCDF4_DIR}/lib")
[[ -n "${HDF5_DIR:-}" ]] && ld_path_parts+=("${HDF5_DIR}/lib")
[[ -n "${N421_EXTRA_LD_LIBRARY_PATH:-}" ]] && ld_path_parts+=("${N421_EXTRA_LD_LIBRARY_PATH}")
[[ -n "${LD_LIBRARY_PATH:-}" ]] && ld_path_parts+=("${LD_LIBRARY_PATH}")

if (( ${#ld_path_parts[@]} > 0 )); then
  export LD_LIBRARY_PATH="$(IFS=:; echo "${ld_path_parts[*]}")"
fi

# ---------- Helpers ----------
ts(){ TZ=EET date '+%Y-%m-%d %H:%M:%S (utc%z)'; }
log(){ printf "[%s] %s\n" "$(ts)" "$*"; }

# ---------- Inputs ----------
: "${model_run_dir:?model_run_dir env not set}"
cd "${model_run_dir}"

XIOS=./xios_server.exe
NEMO=./nemo.exe
[[ -x "${XIOS}" ]] || { echo "Missing or non-executable: ${XIOS}"; exit 2; }
[[ -x "${NEMO}" ]] || { echo "Missing or non-executable: ${NEMO}"; exit 2; }

log "model_run_dir: ${model_run_dir}"
log "Hostlist: ${SLURM_NODELIST:-unknown}"
log "Tasks: ${SLURM_NTASKS:-?}  Nodes: ${SLURM_NNODES:-?}"

# ---------- Rank layout ----------
XIOS_TASKS=1
XIOS_MEM_PER_CPU=15gb
NEMO_TASKS=256
NEMO_MEM_PER_CPU=1gb

# ---------- Run ----------
echo "$(date +%s)" > model_last_start
start_sec=$(date +%s)

log "Launching: ${XIOS_TASKS} x ${XIOS}  +  ${NEMO_TASKS} x ${NEMO}"
set -x
srun \
  -n "${XIOS_TASKS}" -N 1 --mem-per-cpu="${XIOS_MEM_PER_CPU}" "${XIOS}" \
  : \
  -n "${NEMO_TASKS}"        --mem-per-cpu="${NEMO_MEM_PER_CPU}" "${NEMO}"
set +x

# ---------- Timing ----------
echo "$(date +%s)" > model_last_end
end_sec=$(date +%s)
run_s=$(( end_sec - start_sec ))
log "NEMO run time: $(date -u -d @${run_s} +%H:%M:%S)"
