#!/bin/bash
#SBATCH --job-name=prop_bench
#SBATCH --time=01:00:00
#SBATCH --nodes=1
#SBATCH --ntasks=48
#SBATCH --cpus-per-task=1
#SBATCH --mem-per-cpu=1G
#SBATCH --output=bench_%j.out
#SBATCH --error=bench_%j.out
set -u
echo "host=$(hostname)  start=$(date '+%F %T')"

# ---- environment -----------------------------------------------------------
source /cluster/apps/local/env2lmod.sh 2>/dev/null || true
module purge 2>/dev/null || true
module load stack/2024-06 gcc/12.2.0 openmpi/4.1.6 openfoam/2306
if ! command -v pimpleFoam >/dev/null 2>&1; then
  [ -n "${FOAM_ETC:-}" ] && { set +u; source "$FOAM_ETC/bashrc"; set -u; }
fi
command -v pimpleFoam || { echo "FATAL: pimpleFoam not on PATH"; exit 1; }
echo "OpenFOAM: ${WM_PROJECT_VERSION:-unknown}"

CASE=$SLURM_SUBMIT_DIR/case
cd "$CASE" || exit 1
T0=0.0133348
DT=2.9761904761904763e-06

for N in 12 24 48; do
  echo "================ ranks = $N ================"
  rm -rf processor*
  sed -i "s/^numberOfSubdomains.*/numberOfSubdomains $N;/" system/decomposeParDict
  # only the restart time; keep every other time out of the way
  /usr/bin/time -f "decomposePar %e s" decomposePar -force -time $T0 > log.decomp.$N 2>&1
  tail -1 log.decomp.$N
  grep -q "FOAM FATAL" log.decomp.$N && { echo "decomposePar FAILED for $N"; tail -20 log.decomp.$N; continue; }

  # 20 steps only
  END=$(python3 -c "print(repr($T0 + 20*$DT))")
  sed -i "s/^endTime.*/endTime         $END;/" system/controlDict

  S=$(date +%s)
  srun -n $N pimpleFoam -parallel > log.run.$N 2>&1
  RC=$?; E=$(date +%s)
  if grep -q "FOAM FATAL" log.run.$N; then
    echo "SOLVER FAILED at $N ranks:"; grep -A12 "FOAM FATAL" log.run.$N | head -25; continue
  fi
  NS=$(grep -c "^Time = " log.run.$N)
  echo "ranks=$N  rc=$RC  wall=$((E-S)) s  steps=$NS"
  python3 - "$N" <<'PY'
import re,sys
n=sys.argv[1]
t=[float(x) for x in re.findall(r"ExecutionTime = ([0-9.]+) s", open(f"log.run.{n}").read())]
if len(t)>6:
    d=[t[i+1]-t[i] for i in range(len(t)-1)][5:]      # skip startup
    d.sort()
    print(f"  steady s/step (median of {len(d)}): {d[len(d)//2]:.2f}")
PY
done
echo "end=$(date '+%F %T')"