# -*- coding: utf-8 -*-

# PyNX - Python tools for Nano-structures Crystallography
#   (c) 2020-present : ESRF-European Synchrotron Radiation Facility
#       authors:
#         Vincent Favre-Nicolin, favre@esrf.fr

import os

# Import mpi4py - use environment variables to test if it is really
# necessary, but if not avoid the import.
try:
    rank_gt_1 = None
    if 'OMPI_COMM_WORLD_SIZE' in os.environ:
        if int(os.environ['OMPI_COMM_WORLD_SIZE']) > 1:
            rank_gt_1 = True
        else:
            rank_gt_1 = False
    if 'PMI_SIZE' in os.environ and rank_gt_1 is not True:
        if int(os.environ['PMI_SIZE']) > 1:
            rank_gt_1 = True
    # A bit conservative here - no import if rank is 1, but test anyway if no rank was found
    # Not sure what environment variable we are *guaranteed* for the different MPI flavours
    if ('PMI_RANK' in os.environ or 'PMIX_RANK' in os.environ) and (rank_gt_1 is None or rank_gt_1 is True):
        from mpi4py import MPI

        mpic = MPI.COMM_WORLD
    else:
        MPI = None
except ImportError:
    MPI = None
