import pytest

from chainer import testing
from chainer.testing import parameterized
import chainerx.testing

from chainerx_tests import cuda_utils


def pytest_configure(config):
    _register_cuda_marker(config)


def pytest_collection(session):
    # Perform pairwise testing.
    # TODO(kataoka): This is a tentative fix. Discuss its public interface.
    pairwise_product_dict = parameterized._pairwise_product_dict
    testing.product_dict = pairwise_product_dict
    parameterized.product_dict = pairwise_product_dict


def pytest_collection_finish(session):
    product_dict = parameterized._product_dict_orig
    testing.product_dict = product_dict
    parameterized.product_dict = product_dict


def pytest_runtest_setup(item):
    _setup_cuda_marker(item)


def pytest_runtest_teardown(item, nextitem):
    current_device = cuda_utils.get_current_device()
    assert current_device is None or current_device == 0


def pytest_generate_tests(metafunc):
    device_fixtures = {'parametrize_device': 'device',
                       'parametrize_device_name': 'device_name'}
    marker = [
        m for m in metafunc.definition.iter_markers()
        if m.name in ('parametrize_device', 'parametrize_device_name')]
    if marker:
        marker, = marker  # asserts len == 1
        device_names, = marker.args
        # fixture name is either device or device_name
        fixture_name = device_fixtures[marker.name]
        metafunc.parametrize(fixture_name, device_names, indirect=True)


def pytest_collection_modifyitems(session, config, items):
    # Make the order of tests deterministic
    items[:] = sorted(items, key=lambda item: item.location)


def _register_cuda_marker(config):
    config.addinivalue_line(
        'markers',
        'cuda(num=1): mark tests needing the specified number of NVIDIA GPUs.')


def _setup_cuda_marker(item):
    """Pytest marker to indicate number of NVIDIA GPUs required to run the test.

    Tests can be annotated with this decorator (e.g., ``@pytest.mark.cuda``) to
    declare that one NVIDIA GPU is required to run.

    Tests can also be annotated as ``@pytest.mark.cuda(2)`` to declare number
    of NVIDIA GPUs required to run. When running tests, if
    ``CHAINERX_TEST_CUDA_DEVICE_LIMIT`` environment variable is set to value
    greater than or equals to 0, test cases that require GPUs more than the
    limit will be skipped.
    """

    cuda_marker = [m for m in item.iter_markers() if m.name == 'cuda']
    if cuda_marker:
        cuda_marker, = cuda_marker  # asserts len == 1
        required_num = cuda_marker.args[0] if cuda_marker.args else 1
        if cuda_utils.get_cuda_limit() < required_num:
            pytest.skip('{} NVIDIA GPUs required'.format(required_num))


def _get_required_cuda_devices_from_device_name(device_name):
    # Returns the number of required CUDA devices to run a test, given a
    # device name. If the device is non-CUDA device, 0 is returned.
    s = device_name.split(':')
    assert len(s) == 2
    if s[0] != 'cuda':
        return 0
    return int(s[1]) + 1


def _skip_if_no_cuda_device_available(device_name):
    # Skip if the device is CUDA device and there's no sufficient CUDA devices.
    cuda_device_count = _get_required_cuda_devices_from_device_name(
        device_name)
    if cuda_device_count > cuda_utils.get_cuda_limit():
        pytest.skip()


@pytest.fixture
def device(request):
    # A fixture to wrap a test with a device scope, given a device name.
    # Device instance is passed to the test.

    device_name = request.param
    _skip_if_no_cuda_device_available(device_name)

    device = chainerx.get_device(device_name)
    device_scope = chainerx.using_device(device)

    def finalize():
        device_scope.__exit__()

    request.addfinalizer(finalize)
    device_scope.__enter__()
    return device


@pytest.fixture
def device_name(request):
    # A fixture to check if a device is available using its name

    device_name = request.param
    _skip_if_no_cuda_device_available(device_name)

    return device_name


@pytest.fixture(params=chainerx.testing.all_dtypes)
def dtype(request):
    return request.param


@pytest.fixture(params=chainerx.testing.float_dtypes)
def float_dtype(request):
    return request.param


@pytest.fixture(params=chainerx.testing.signed_dtypes)
def signed_dtype(request):
    return request.param


@pytest.fixture(params=chainerx.testing.numeric_dtypes)
def numeric_dtype(request):
    return request.param


@pytest.fixture(params=[True, False])
def is_module(request):
    return request.param


_shapes = [
    (),
    (0,),
    (1,),
    (2, 3),
    (1, 1, 1),
    (2, 0, 3),
]

_shapes_as_sequence_or_int = (
    _shapes
    + [[], [0]]  # shape as a list instead of tuple
    + [0, 1, 5])


@pytest.fixture(params=_shapes)
def shape(request):
    return request.param


@pytest.fixture(params=_shapes_as_sequence_or_int)
def shape_as_sequence_or_int(request):
    return request.param
