# Copyright (c) OpenMMLab. All rights reserved.
import multiprocessing as mp
import os
import platform
import sys

import cv2
import mmcv
import pytest
from mmcv import Config

import mmflow
from mmflow.utils import setup_multi_processes


@pytest.mark.parametrize('workers_per_gpu', (0, 2))
@pytest.mark.parametrize(('valid', 'env_cfg'), [(True,
                                                 dict(
                                                     mp_start_method='fork',
                                                     opencv_num_threads=0,
                                                     omp_num_threads=1,
                                                     mkl_num_threads=1)),
                                                (False,
                                                 dict(
                                                     mp_start_method=1,
                                                     opencv_num_threads=0.1,
                                                     omp_num_threads='s',
                                                     mkl_num_threads='1'))])
def test_setup_multi_processes(workers_per_gpu, valid, env_cfg):
    # temp save system setting
    sys_start_mehod = mp.get_start_method(allow_none=True)
    sys_cv_threads = cv2.getNumThreads()
    # pop and temp save system env vars
    sys_omp_threads = os.environ.pop('OMP_NUM_THREADS', default=None)
    sys_mkl_threads = os.environ.pop('MKL_NUM_THREADS', default=None)

    config = dict(
        data=dict(
            train_dataloader=dict(dict(workers_per_gpu=workers_per_gpu))))
    config.update(env_cfg)
    cfg = Config(config)
    setup_multi_processes(cfg)

    # test when cfg is valid and workers_per_gpu > 0
    # setup_multi_processes will work
    if valid and workers_per_gpu > 0:
        # test config without setting env

        assert os.getenv('OMP_NUM_THREADS') == str(env_cfg['omp_num_threads'])
        assert os.getenv('MKL_NUM_THREADS') == str(env_cfg['mkl_num_threads'])
        # when set to 0, the num threads will be 1
        assert cv2.getNumThreads() == env_cfg[
            'opencv_num_threads'] if env_cfg['opencv_num_threads'] > 0 else 1
        if platform.system() != 'Windows':
            assert mp.get_start_method() == env_cfg['mp_start_method']

        # revert setting to avoid affecting other programs
        if sys_start_mehod:
            mp.set_start_method(sys_start_mehod, force=True)
        cv2.setNumThreads(sys_cv_threads)
        if sys_omp_threads:
            os.environ['OMP_NUM_THREADS'] = sys_omp_threads
        else:
            os.environ.pop('OMP_NUM_THREADS')
        if sys_mkl_threads:
            os.environ['MKL_NUM_THREADS'] = sys_mkl_threads
        else:
            os.environ.pop('MKL_NUM_THREADS')

    elif valid and workers_per_gpu == 0:

        if platform.system() != 'Windows':
            assert mp.get_start_method() == env_cfg['mp_start_method']
        assert cv2.getNumThreads() == env_cfg[
            'opencv_num_threads'] if env_cfg['opencv_num_threads'] > 0 else 1
        assert 'OMP_NUM_THREADS' not in os.environ
        assert 'MKL_NUM_THREADS' not in os.environ
        if sys_start_mehod:
            mp.set_start_method(sys_start_mehod, force=True)
        cv2.setNumThreads(sys_cv_threads)
        if sys_omp_threads:
            os.environ['OMP_NUM_THREADS'] = sys_omp_threads
        if sys_mkl_threads:
            os.environ['MKL_NUM_THREADS'] = sys_mkl_threads

    else:
        assert mp.get_start_method(allow_none=True) == sys_start_mehod
        assert cv2.getNumThreads() == sys_cv_threads
        assert 'OMP_NUM_THREADS' not in os.environ
        assert 'MKL_NUM_THREADS' not in os.environ


def test_collect_env():
    try:
        import torch  # noqa: F401
    except ModuleNotFoundError:
        pytest.skip('skipping tests that require PyTorch')

    from mmflow.utils import collect_env
    env_info = collect_env()
    expected_keys = [
        'sys.platform', 'Python', 'CUDA available', 'PyTorch',
        'PyTorch compiling details', 'OpenCV', 'MMCV', 'MMCV Compiler',
        'MMCV CUDA Compiler', 'MMFlow', 'GCC'
    ]
    for key in expected_keys:
        assert key in env_info

    if env_info['CUDA available']:
        for key in ['CUDA_HOME', 'NVCC']:
            assert key in env_info

    assert env_info['sys.platform'] == sys.platform
    assert env_info['Python'] == sys.version.replace('\n', '')
    assert env_info['MMCV'] == mmcv.__version__
    assert mmflow.__version__ in env_info['MMFlow']
