# File containing common commands for NetKet Test infrastructure

import pytest
import netket as nk

skipif_mpi = pytest.mark.skipif(nk.utils.mpi.n_nodes > 1, reason="Only run without MPI")
"""Use as a decorator to mark a test to be skipped when running under MPI.
For example:

Example:
>>> @skipif_mpi
>>> def test_my_serial_function():
>>>     your_serial_test()

"""

onlyif_mpi = pytest.mark.skipif(nk.utils.mpi.n_nodes < 2, reason="Only run with MPI")
"""Use as a decorator to mark a test to be executed only when running with at least 2 MPI
nodes.

It can be used in combination with the fixtures defined in test/conftest.py, namely
_mpi_rank, mpi_size and _mpi_comm, that retrieve the number of mpi rank, size and comm
used by netket.

If you need to trick netket into executing some code without MPI when running under MPI,
you can use the class netket_disable_mpi below

Example:

>>> @onlyif_mpi
>>> def test_my_parallel_function():
>>>     x = compute_some_stuff_with_mpi
>>>     with netket_disable_mpi():
>>>         x_serial = compute_some_stuff_withjout_mpi

"""


class netket_disable_mpi:
    """
    Temporarily disables MPI functions inside of NetKet, tricking NetKet
    into executing all code on every rank independently.

    Example:

    >>> with netket_disable_mpi():
    >>>     run_code

    """

    def __enter__(self):
        self._orig_nodes = nk.utils.mpi.n_nodes
        nk.utils.mpi.n_nodes = 1
        nk.utils.mpi.mpi.n_nodes = 1
        nk.utils.mpi.primitives.n_nodes = 1

    def __exit__(self, exc_type, exc_value, traceback):
        nk.utils.mpi.n_nodes = self._orig_nodes
        nk.utils.mpi.mpi.n_nodes = self._orig_nodes
        nk.utils.mpi.primitives.n_nodes = self._orig_nodes
