import os

from colossalai.context.parallel_mode import ParallelMode
from colossalai.core import global_context as gpc


def get_tesseract_dim_dep_from_env():
    try:
        tesseract_dim = int(os.environ['TESSERACT_DIM'])
        tesseract_dep = int(os.environ['TESSERACT_DEP'])
        assert tesseract_dim > 0, 'TESSERACT_DIM must be larger than zero'
        assert tesseract_dep > 0, 'TESSERACT_DEP must be larger than zero'
        return tesseract_dim, tesseract_dep

    except KeyError as e:
        raise EnvironmentError('TESSERACT_DIM or TESSERACT_DEP is not found in the current environment, '
                               'please make sure that you have used the correct process group initializer')


def assert_tesseract_initialization():
    assert gpc.is_initialized(ParallelMode.PARALLEL_2P5D_COL) and \
           gpc.is_initialized(ParallelMode.PARALLEL_2P5D_ROW) and \
           gpc.is_initialized(ParallelMode.PARALLEL_2P5D_DEP) and \
           gpc.is_initialized(ParallelMode.PARALLEL_2P5D_XZ), \
        'Both PARALLEL_2P5D_COL, PARALLEL_2P5D_ROW, PARALLEL_2P5D_DEP and PARALLEL_2P5D_XZ must be initialized by the process group initializer'