# Copyright 2014-2025 The ODL contributors
#
# This file is part of ODL.
#
# This Source Code Form is subject to the terms of the Mozilla Public License,
# v. 2.0. If a copy of the MPL was not distributed with this file, You can
# obtain one at https://mozilla.org/MPL/2.0/.

"""Test configuration file."""

from __future__ import absolute_import, division, print_function

import operator
import os
from os import path

import numpy as np

import odl
from odl.core.array_API_support import lookup_array_backend, check_device
from odl.core.space.entry_points import tensor_space_impl_names
from odl.trafos.backends import PYFFTW_AVAILABLE, PYWT_AVAILABLE
from odl.core.util.testutils import simple_fixture
from odl.core.util.dtype_utils import INTEGER_DTYPES, FLOAT_DTYPES, COMPLEX_DTYPES

try:
    import pytest
    from pytest import fixture
except ImportError:
    pytest = None

    # Identity fixture
    def fixture(*arg, **kw):
        if arg and callable(arg[0]):
            return arg[0]
        return fixture


# --- Add numpy and ODL to all doctests ---


@fixture(autouse=True)
def _add_doctest_np_odl(doctest_namespace):
    doctest_namespace['np'] = np
    doctest_namespace['odl'] = odl


# --- Ignored tests due to missing modules ---


this_dir = path.dirname(__file__)
odl_root = path.abspath(path.join(this_dir, '..', '..'))
collect_ignore = [path.join(odl_root, 'setup.py'),
                  path.join(odl_root, 'odl', 'contrib')]


# Add example directories to `collect_ignore`
def find_example_dirs():
    dirs = []
    for dirpath, dirnames, _ in os.walk(odl_root):
        if 'examples' in dirnames:
            dirs.append(path.join(dirpath, 'examples'))
    return dirs


collect_ignore.extend(find_example_dirs())


if not PYFFTW_AVAILABLE:
    collect_ignore.append(
        path.join(odl_root, 'odl', 'trafos', 'backends', 'pyfftw_bindings.py')
    )
if not PYWT_AVAILABLE:
    collect_ignore.append(
        path.join(odl_root, 'odl', 'trafos', 'backends', 'pywt_bindings.py')
    )
    # Currently `pywt` is the only implementation
    collect_ignore.append(
        path.join(odl_root, 'odl', 'trafos', 'wavelet.py')
    )


# --- Command-line options --- #


def pytest_addoption(parser):
    """
    This is a pytest template to add options to the CLI when running pytest.
    It is quite handy when we want to run the test on only one backend or device, for instance.
    """
    suite_help = (
        "enable an opt-in test suite NAME. "
        "Available suites: largescale, examples, doc_doctests"
    )
    parser.addoption(
        "-S",
        "--suite",
        nargs='*',
        metavar="NAME",
        type=str,
        default=[],
        help=suite_help,
    )
    parser.addoption(
        "--backend",
        default=None,    # The default if the option is not provided
        help="Select the array backend to run the test on ('numpy'...)",
    )

    parser.addoption(
        "--device",
        default=None,    # The default if the option is not provided
        help="Select the CPU/GPU device to run the test on ('cpu'...)",
    )


def pytest_generate_tests(metafunc):
    """
    Check if the fixture name is used in a test and parametrize it
    based on the command line option.
    """
    # The name of the fixture we want to parametrize
    fixture_name = "odl_impl_device_pairs"

    if fixture_name in metafunc.fixturenames:
        # Get the CLI option
        impl   = metafunc.config.getoption("backend")
        device = metafunc.config.getoption("device")

        # Parse the backend
        if impl is not None:
            array_backend = lookup_array_backend(impl)
            if device is None:
                params=[]
                for device in array_backend.available_devices:
                    params.append((impl, device))
            else:
                device = check_device(impl, device)
                params = [(impl, device)]
        else:
            params = []
            for impl in tensor_space_impl_names():
                array_backend = lookup_array_backend(impl)
                if device is None:
                    for available_device in array_backend.available_devices:
                        params.append((impl, available_device))
                else:
                    device = check_device(impl, device)
                    params = [(impl, device)]
        
        # Parametrize the fixture. 
        # indirect=True is crucial tells pytest to pass these values 
        # to the fixture function (via request.param) rather than directly to # the test.
        metafunc.parametrize(fixture_name, params, indirect=True)


def pytest_configure(config):
    # Register an additional marker
    config.addinivalue_line(
        'markers', "suite(name): mark test to belong to an opt-in suite"
    )


def pytest_runtest_setup(item):
    suites = [mark.args[0] for mark in item.iter_markers(name='suite')]
    if suites:
        if not any(val in suites for val in item.config.getoption('-S')):
            pytest.skip(f"test not in suites {suites}")


# Remove duplicates
collect_ignore = list(set(collect_ignore))
collect_ignore = [path.normcase(ignored) for ignored in collect_ignore]


# NB: magical `path` param name is needed
def pytest_ignore_collect(path):
    """This is to ignore paths during test collection"""
    normalized = os.path.normcase(str(path))
    return any(normalized.startswith(ignored) for ignored in collect_ignore)


# --- Reusable fixtures --- #

# NOTE: All global fixtures need to be prefixed with `odl_` to make them
# non-conflicting with other packages' fixture names, since these fixtures
# are exposed globally across packages by setuptools.

# Simple ones, use helper
odl_tspace_impl = simple_fixture(name='tspace_impl',
                                 params=tensor_space_impl_names())

real_floating_dtypes = FLOAT_DTYPES
odl_real_floating_dtype = simple_fixture(name='dtype',
                                    params=real_floating_dtypes)

complex_floating_dtypes = COMPLEX_DTYPES
odl_complex_floating_dtype = simple_fixture(name='dtype',
                                    params=complex_floating_dtypes)

floating_dtypes = FLOAT_DTYPES + COMPLEX_DTYPES
odl_floating_dtype = simple_fixture(name='dtype',
                                    params=floating_dtypes)

scalar_dtypes = INTEGER_DTYPES + FLOAT_DTYPES + COMPLEX_DTYPES
odl_scalar_dtype = simple_fixture(name='dtype',
                                  params=scalar_dtypes)


@pytest.fixture(scope='module')
def odl_impl_device_pairs(request):
    """Fixture for testing accross device and backends. See
    pytest_generate_tests function to see how to modify it through the CLI
    """
    return request.param

IMPL_DEVICE_PAIRS = []
    
for impl in tensor_space_impl_names():
    array_backend = lookup_array_backend(impl)
    for device in array_backend.available_devices:
        IMPL_DEVICE_PAIRS.append((impl, device))

odl_alternative_impl_device_pair = simple_fixture(name='alternative_impl_device', params=IMPL_DEVICE_PAIRS)

if 'pytorch' in tensor_space_impl_names():
    CUDA_DEVICES = []
    for device in lookup_array_backend('pytorch').available_devices:
        CUDA_DEVICES.append(device)

    cuda_device = simple_fixture(name='cuda_device', params=CUDA_DEVICES)

odl_elem_order = simple_fixture(name='order', params=['C'])

odl_reduction = simple_fixture('reduction', ['sum', 'prod', 'min', 'max'])

# More complicated ones with non-trivial documentation
arithmetic_op_par = [operator.add,
                     operator.truediv,
                     operator.mul,
                     operator.sub,
                     operator.iadd,
                     operator.itruediv,
                     operator.imul,
                     operator.isub]
arithmetic_op_ids = [
    f" op = '{op}' " for op in ['+', '/', '*', '-', '+=', '/=', '*=', '-=']]


@fixture(ids=arithmetic_op_ids, params=arithmetic_op_par)
def odl_arithmetic_op(request):
    """An arithmetic operator, e.g. +, -, // etc."""
    return request.param
