#!/bin/bash
# pymufem — launch a μfem case under MPI.
#
# Resolves libmufem (which ships mpiexec and the runtime libs) at run time so
# the script works under any Python minor version and any venv layout.

set -e

# pymufem lives in <env>/bin/. The sibling python interpreter is what we use
# to locate libmufem inside this env's site-packages.
BIN_DIR=$(dirname "$(realpath "$0")")
PYTHON=$BIN_DIR/python
[ -x "$PYTHON" ] || PYTHON=$BIN_DIR/python3
[ -x "$PYTHON" ] || { echo "pymufem: no python interpreter next to $0" >&2; exit 1; }

# Locate the libmufem package WITHOUT importing it: __init__.py preloads the
# engine + every third-party .so via ctypes, which is unnecessary just to find
# mpiexec, and would mask the real cause if any of those loads failed.
LIBMUFEM_DIR=$("$PYTHON" -c \
    'import importlib.util, os, sys
s = importlib.util.find_spec("libmufem")
if s is None or not s.origin:
    sys.exit("pymufem: libmufem package not found; install it via '"'"'pip install libmufem'"'"'")
print(os.path.dirname(s.origin))') || exit 1

export LD_LIBRARY_PATH="$LIBMUFEM_DIR/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
export OPAL_PREFIX="$LIBMUFEM_DIR"

MPIEXEC=$LIBMUFEM_DIR/bin/mpiexec
[ -x "$MPIEXEC" ] || { echo "pymufem: $MPIEXEC missing or not executable" >&2; exit 1; }

# Pin each rank to a physical core, round-robin across sockets so ranks land on
# separate NUMA domains first. Fall back to --oversubscribe when the requested
# rank count over-commits the machine (pinning and oversubscription conflict).
# Mirrors .scripts/run_case.sh — duplicated by design, keep them in sync.
PHYSICAL_CORES=$(lscpu 2>/dev/null | awk '/^Core\(s\) per socket:/ {c=$4} /^Socket\(s\):/ {s=$2} END{if(c && s) print c*s}')
PHYSICAL_CORES=${PHYSICAL_CORES:-8}

if [[ $1 =~ ^[0-9]+$ ]]; then
    NP=$1; shift
    if (( NP <= PHYSICAL_CORES )); then
        "$MPIEXEC" --allow-run-as-root --bind-to core --map-by socket -n "$NP" "$PYTHON" "$@"
    else
        printf '\033[5;31mWARNING: %d ranks > %d physical cores — oversubscribing; ranks share cores and performance will degrade.\033[0m\n' \
            "$NP" "$PHYSICAL_CORES" >&2
        "$MPIEXEC" --allow-run-as-root --oversubscribe -n "$NP" "$PYTHON" "$@"
    fi
else
    "$MPIEXEC" --allow-run-as-root --bind-to core --map-by socket "$PYTHON" "$@"
fi
