"""Utilities for testing."""

from functools import partial
import operator
from pathlib import Path
import warnings

import numpy as np
import numpy.testing as npt
from numpy.testing import assert_array_equal

# set path to example data
IO_DATA_PATH = Path(__file__).resolve().parent / ".." / "io" / "tests" / "data"


def assert_operator(value1, value2, msg="", op=operator.eq):
    """Check Boolean statement."""
    try:
        if op == operator.is_:
            value1 = bool(value1)
        assert op(value1, value2)
    except AssertionError as e:
        raise AssertionError(msg.format(str(value2), str(value1))) from e


assert_greater_equal = partial(assert_operator, op=operator.ge, msg="{0} >= {1}")
assert_greater = partial(assert_operator, op=operator.gt, msg="{0} > {1}")
assert_less_equal = partial(assert_operator, op=operator.le, msg="{0} =< {1}")
assert_less = partial(assert_operator, op=operator.lt, msg="{0} < {1}")
assert_true = partial(
    assert_operator, value2=True, op=operator.is_, msg="False is not true"
)
assert_false = partial(
    assert_operator, value2=False, op=operator.is_, msg="True is not false"
)
assert_not_equal = partial(assert_operator, op=operator.ne)


def assert_arrays_equal(arrays1, arrays2):
    for arr1, arr2 in zip(arrays1, arrays2):
        assert_array_equal(arr1, arr2)


def assert_percent_almost_equal(a, b, decimal=7, percent=0.99):
    a = np.asarray(a)
    b = np.asarray(b)
    tol = 1.5 * 10 ** (-decimal)

    diff = np.abs(a - b)
    ok = diff <= tol

    fraction = np.mean(ok)

    if fraction < percent:
        raise AssertionError(
            f"Only {fraction*100:.2f}% of elements match within {decimal} decimals; "
            f"required {percent*100:.2f}%"
        )


class clear_and_catch_warnings(warnings.catch_warnings):
    """Context manager that resets warning registry for catching warnings.

    Warnings can be slippery, because, whenever a warning is triggered, Python
    adds a ``__warningregistry__`` member to the *calling* module.  This makes
    it impossible to retrigger the warning in this module, whatever you put in
    the warnings filters.  This context manager accepts a sequence of `modules`
    as a keyword argument to its constructor and:
    * stores and removes any ``__warningregistry__`` entries in given `modules`
      on entry;
    * resets ``__warningregistry__`` to its previous state on exit.
    This makes it possible to trigger any warning afresh inside the context
    manager without disturbing the state of warnings outside.
    For compatibility with Python 3.0, please consider all arguments to be
    keyword-only.

    Parameters
    ----------
    record : bool, optional
        Specifies whether warnings should be captured by a custom
        implementation of ``warnings.showwarning()`` and be appended to a list
        returned by the context manager. Otherwise None is returned by the
        context manager. The objects appended to the list are arguments whose
        attributes mirror the arguments to ``showwarning()``.
        NOTE: nibabel difference from numpy: default is True
    modules : sequence, optional
        Sequence of modules for which to reset warnings registry on entry and
        restore on exit

    Examples
    --------
    >>> import warnings
    >>> import numpy as np
    >>> with clear_and_catch_warnings(modules=[np.lib.scimath]):
    ...     warnings.simplefilter('always')
    ...     # do something that raises a warning in np.lib.scimath
    ...     _ = np.arccos(90)

    Notes
    -----
    this class is copied (with minor modifications) from Nibabel.
    https://github.com/nipy/nibabel. See COPYING file distributed along with
    the Nibabel package for the copyright and license terms.

    """

    class_modules = ()

    def __init__(self, record=True, modules=()):
        self.modules = set(modules).union(self.class_modules)
        self._warnreg_copies = {}
        super(clear_and_catch_warnings, self).__init__(record=record)

    def __enter__(self):
        for mod in self.modules:
            if hasattr(mod, "__warningregistry__"):
                mod_reg = mod.__warningregistry__
                self._warnreg_copies[mod] = mod_reg.copy()
                mod_reg.clear()
        return super(clear_and_catch_warnings, self).__enter__()

    def __exit__(self, *exc_info):
        super(clear_and_catch_warnings, self).__exit__(*exc_info)
        for mod in self.modules:
            if hasattr(mod, "__warningregistry__"):
                mod.__warningregistry__.clear()
            if mod in self._warnreg_copies:
                mod.__warningregistry__.update(self._warnreg_copies[mod])


def check_for_warnings(warn_printed, w_msg):
    selected_w = [w for w in warn_printed if issubclass(w.category, UserWarning)]
    assert len(selected_w) >= 1
    msg = [str(m.message) for m in selected_w]
    npt.assert_equal(w_msg in msg, True)


def assert_warns(warning_class, *args, **kwargs):
    """Assert that a warning is raised.

    This function can be used as a context manager or as a function wrapper.
    It is a replacement for the deprecated numpy.testing.assert_warns.

    Parameters
    ----------
    warning_class : class
        The class of warning to check for.
    *args : positional arguments
        If provided, the first argument should be a callable, and the
        remaining arguments are passed to that callable.
    **kwargs : keyword arguments
        Keyword arguments passed to the callable.

    Returns
    -------
    ctx : WarningsChecker or result
        If used as a context manager, returns a WarningsChecker instance.
        If used as a function wrapper, returns the result of the callable.

    Examples
    --------
    As a context manager:
    >>> with assert_warns(UserWarning):
    ...     warnings.warn("test", UserWarning)

    As a function wrapper:
    >>> def f():
    ...     warnings.warn("test", UserWarning)
    ...     return 42
    >>> result = assert_warns(UserWarning, f)
    >>> result
    42

    Notes
    -----
    This replaces the deprecated numpy.testing.assert_warns function.
    It uses warnings.catch_warnings internally for better compatibility
    with NumPy 2.x.
    """
    if not args:
        # Used as a context manager
        return _WarningsChecker(warning_class)
    else:
        # Used as a function wrapper
        func = args[0]
        func_args = args[1:]
        with _WarningsChecker(warning_class):
            return func(*func_args, **kwargs)


class _WarningsChecker:
    """Context manager to check that a warning is raised."""

    def __init__(self, warning_class):
        self.warning_class = warning_class
        self.caught_warnings = []

    def __enter__(self):
        self.catch_warnings = warnings.catch_warnings(record=True)
        self.warning_list = self.catch_warnings.__enter__()
        warnings.simplefilter("always")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.catch_warnings.__exit__(exc_type, exc_val, exc_tb)

        # Check if the expected warning was raised
        if exc_type is None:  # No exception was raised
            matching_warnings = [
                w
                for w in self.warning_list
                if issubclass(w.category, self.warning_class)
            ]
            if not matching_warnings:
                raised_warnings = [w.category.__name__ for w in self.warning_list]
                raise AssertionError(
                    f"No warning of class {self.warning_class.__name__} "
                    f"was raised. Warnings raised: {raised_warnings}"
                )
        return False
