# Path: gym3/util.py
# def call_func(fn_path, **kwargs):
#     """
#     Look up a function for a function path, e.g. fn_path='gym3:call_func' will
#     return this function, then call it with the specified kwargs.
#     """
#     assert isinstance(fn_path, str)
#     module_path, fn_name = fn_path.split(":")
#     module = importlib.import_module(module_path)
#     return getattr(module, fn_name)(**kwargs)


import base64
import os
import pickle
import subprocess
import sys
from gym3.util import call_func



def run_test_with_mpi(fn_path, kwargs=None, nproc=2, timeout=30):
    if kwargs is None:
        kwargs = {}
    serialized_fn = base64.b64encode(pickle.dumps((fn_path, kwargs)))
    subprocess.check_call(
        [
            "mpiexec",
            "-n",
            str(nproc),
            sys.executable,
            "-m",
            "gym3.internal.test_with_mpi",
            serialized_fn,
        ],
        env=os.environ,
        timeout=timeout,
    )


if __name__ == "__main__":
    fn_path, kwargs = pickle.loads(base64.b64decode(sys.argv[1]))

# Path: gym3/util.py
# def call_func(fn_path, **kwargs):
#     """
#     Look up a function for a function path, e.g. fn_path='gym3:call_func' will
#     return this function, then call it with the specified kwargs.
#     """
#     assert isinstance(fn_path, str)
#     module_path, fn_name = fn_path.split(":")
#     module = importlib.import_module(module_path)
#     return getattr(module, fn_name)(**kwargs)


# Path: gym3/internal/test_with_mpi.py
import base64
import os
import pickle
import subprocess
import sys
from gym3.util import call_func




def run_test_with_mpi(fn_path, kwargs=None, nproc=2, timeout=30):
    if kwargs is None:
        kwargs = {}
    serialized_fn = base64.b64encode(pickle.dumps((fn_path, kwargs)))
    subprocess.check_call(
        [
            "mpiexec",
            "-n",
            str(nproc),
            sys.executable,
            "-m",
            "gym3.internal.test_with_mpi",
            serialized_fn,
        ],
        env=os.environ,
        timeout=timeout,
    )


if __name__ == "__main__":
    fn_path, kwargs = pickle.loads(base64.b64decode(sys.argv[1]))

    call_func(fn_path, **kwargs)