import shutil


def mpi_executable(preferred_executable=None):
    """
    Return an mpi executable found on the current system.

    Depending on your MPI implementation, the executable name
    to run an MPI application may differ. This function will
    check which one is available and return the first valid one
    it finds. Valid in this case means that it can be found with
    methods like `which`.
    To override which executable to check first you can pass
    your preferred executable as an argument.

    Parameters
    ----------
    preferred_executable : str, optional
        The first executable to check for on the system. If it
        isn't found, will continue with the regular search for
        valid MPI executables.

    Returns
    -------
    str
        The name of a valid MPI executable.

    Raises
    ------
    RuntimeError
        If no valid MPI executable could be found, a RuntimeError
        is raised.

    """
    if preferred_executable:
        if shutil.which(preferred_executable):
            return preferred_executable
        else:
            raise RuntimeError(
                f"The given preferred mpi executable `{preferred_executable}` "
                "was not found on this system"
            )

    executables = ["mpirun", "mpiexec", "srun"]
    for executable in executables:
        if shutil.which(executable):
            return executable

    raise RuntimeError(
        "Could not find an mpi installation. Make sure your PATH is set correctly."
    )