r"""Implement handlers to check if the objects are equal."""

from __future__ import annotations

__all__ = ["NanEqualHandler", "ScalarEqualHandler"]

import logging
import math
from typing import TYPE_CHECKING

from coola.equality.handler.base import BaseEqualityHandler
from coola.equality.handler.format import format_value_difference
from coola.equality.handler.mixin import HandlerEqualityMixin

if TYPE_CHECKING:
    from coola.equality.config import EqualityConfig


logger: logging.Logger = logging.getLogger(__name__)


class NanEqualHandler(HandlerEqualityMixin, BaseEqualityHandler):
    r"""Check if the two NaNs are equal.

    This handler returns ``True`` if the two numbers are NaNs,
    otherwise it passes the inputs to the next handler.

    Example:
        ```pycon
        >>> from coola.equality.config import EqualityConfig
        >>> from coola.equality.handler import NanEqualHandler, FalseHandler
        >>> config = EqualityConfig()
        >>> handler = NanEqualHandler(next_handler=FalseHandler())
        >>> handler.handle(float("nan"), float("nan"), config)
        False
        >>> config.equal_nan = True
        >>> handler.handle(float("nan"), float("nan"), config)
        True

        ```
    """

    def handle(
        self,
        actual: float,
        expected: float,
        config: EqualityConfig,
    ) -> bool:
        if config.equal_nan and math.isnan(actual) and math.isnan(expected):
            return True
        return self._handle_next(actual, expected, config=config)


class ScalarEqualHandler(HandlerEqualityMixin, BaseEqualityHandler):
    r"""Check if the two numbers are equal or not.

    This handler returns ``False`` if the two numbers are
    different, otherwise it returns ``True``. It is possible to
    control the tolerance by using ``atol`` and ``rtol``.
    By default, the tolerances are set to 0.

    Example:
        ```pycon
        >>> from coola.equality.config import EqualityConfig
        >>> from coola.equality.handler import ScalarEqualHandler
        >>> config = EqualityConfig()
        >>> handler = ScalarEqualHandler()
        >>> handler.handle(42.0, 42.0, config)
        True
        >>> config.atol = 1e-3
        >>> handler.handle(42.0, 42.0001, config)
        True
        >>> handler.handle(float("nan"), float("nan"), config)
        False

        ```
    """

    def handle(self, actual: float, expected: float, config: EqualityConfig) -> bool:
        object_equal = number_equal(actual, expected, config)
        if not object_equal and config.show_difference:
            logger.info(format_value_difference(actual=actual, expected=expected, name="numbers"))
        return object_equal


def number_equal(number1: float, number2: float, config: EqualityConfig) -> bool:
    r"""Indicate if the two numbers are equal within a tolerance.

    Args:
        number1: The first number to compare.
        number2: The second number to compare.
        config: The equality configuration.

    Returns:
        ``True`` if the two numbers are equal within a tolerance,
            otherwise ``False``.
    """
    if config.atol > 0.0 or config.rtol > 0.0:
        return math.isclose(number1, number2, abs_tol=config.atol, rel_tol=config.rtol)
    return number1 == number2
