import warnings


def check_consecutive_ranks_on_same_nodes(mpi):
    """
    Check that consecutive MPI ranks are on the same nodes, and
    return a corresponding boolean. Print a warning message if
    it is not the case, this this can affect performance.
    """
    # Skip this function if there is only one MPI rank
    if mpi.COMM_WORLD.size == 1:
        return True

    # Check that the decomposition is standard, raise a warning otherwise
    nodes = mpi.COMM_WORLD.gather( mpi.Get_processor_name(), root=0 )
    standard_decomp = None
    if mpi.COMM_WORLD.rank == 0:
        # Check that consecutive MPI ranks are on the same node
        standard_decomp = True
        unique_nodes = [ nodes[0] ]  # Initialize list of unique nodes
        for i in range(1, len(nodes)):
            if nodes[i] != nodes[i-1]: # Consecutive ranks not on the same node
                if nodes[i] not in unique_nodes:
                    unique_nodes.append( nodes[i] )
                else:
                    # This node was seen before ; this means that several
                    # MPI ranks selected it, but that they are not consecutive.
                    standard_decomp = False
                    break
        # Print a corresponding warning
        if not standard_decomp:
            warnings.warn(
            "It seems that the distribution of MPI ranks on compute nodes is such\n"
            "that consecutive MPI ranks are not located on the same node.\n"
            "(See the FBPIC output with `verbose_level = 2` for more details.)\n"
            "This type of MPI distribution can degrade the performance of FBPIC.\n"
            "Please check the options of your MPI launcher (e.g. `mpirun`, `srun`)\n"
            "in order to use a different MPI distribution.")
    standard_decomp = mpi.COMM_WORLD.bcast( standard_decomp, root=0 )
    return standard_decomp