r"""Implement handlers for ``pandas.DataFrame``s and
``pandas.Series``s."""

from __future__ import annotations

__all__ = ["PandasDataFrameEqualHandler", "PandasSeriesEqualHandler"]

import logging
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
from coola.utils.imports import is_pandas_available

if TYPE_CHECKING or is_pandas_available():
    import pandas as pd
else:  # pragma: no cover
    from coola.utils.fallback.pandas import pandas as pd

if TYPE_CHECKING:
    from coola.equality.config import EqualityConfig

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


class PandasDataFrameEqualHandler(HandlerEqualityMixin, BaseEqualityHandler):
    r"""Check if the two ``pandas.DataFrame``s are equal.

    This handler returns ``True`` if the two ``pandas.DataFrame``s are
    equal, otherwise ``False``. This handler is designed to be used
    at the end of the chain of responsibility. This handler does
    not call the next handler.

    Example:
        ```pycon
        >>> import pandas
        >>> from coola.equality.config import EqualityConfig
        >>> from coola.equality.handler import PandasDataFrameEqualHandler
        >>> config = EqualityConfig()
        >>> handler = PandasDataFrameEqualHandler()
        >>> handler.handle(
        ...     pandas.DataFrame({"col": [1, 2, 3]}),
        ...     pandas.DataFrame({"col": [1, 2, 3]}),
        ...     config,
        ... )
        True
        >>> handler.handle(
        ...     pandas.DataFrame({"col": [1, 2, 3]}),
        ...     pandas.DataFrame({"col": [1, 2, 4]}),
        ...     config,
        ... )
        False

        ```
    """

    def handle(
        self,
        actual: pd.DataFrame,
        expected: pd.DataFrame,
        config: EqualityConfig,
    ) -> bool:
        object_equal = frame_equal(actual, expected, config)
        if config.show_difference and not object_equal:
            logger.info(
                format_value_difference(actual=actual, expected=expected, name="pandas.DataFrames")
            )
        return object_equal


class PandasSeriesEqualHandler(HandlerEqualityMixin, BaseEqualityHandler):
    r"""Check if the two ``pandas.Series``es are equal.

    This handler returns ``True`` if the two ``pandas.Series``es are
    equal, otherwise ``False``. This handler is designed to be used
    at the end of the chain of responsibility. This handler does
    not call the next handler.

    Example:
        ```pycon
        >>> import pandas
        >>> from coola.equality.config import EqualityConfig
        >>> from coola.equality.handler import PandasSeriesEqualHandler
        >>> config = EqualityConfig()
        >>> handler = PandasSeriesEqualHandler()
        >>> handler.handle(pandas.Series([1, 2, 3]), pandas.Series([1, 2, 3]), config)
        True
        >>> handler.handle(pandas.Series([1, 2, 3]), pandas.Series([1, 2, 4]), config)
        False

        ```
    """

    def handle(
        self,
        actual: pd.Series,
        expected: pd.Series,
        config: EqualityConfig,
    ) -> bool:
        object_equal = series_equal(actual, expected, config)
        if config.show_difference and not object_equal:
            logger.info(
                format_value_difference(actual=actual, expected=expected, name="pandas.Series")
            )
        return object_equal


def frame_equal(df1: pd.DataFrame, df2: pd.DataFrame, config: EqualityConfig) -> bool:
    r"""Indicate if the two DataFrames are equal or not.

    Args:
        df1: The first DataFrame to compare.
        df2: The second DataFrame to compare.
        config: The equality configuration.

    Returns:
        ``True`` if the two DataFrames are equal, otherwise ``False``.
    """
    if not config.equal_nan and df1.isna().any().any():
        return False
    try:
        pd.testing.assert_frame_equal(
            df1,
            df2,
            check_exact=config.atol == 0 and config.rtol == 0,
            check_index_type=True,
            atol=config.atol,
            rtol=config.rtol,
        )
    except AssertionError:
        return False
    return True


def series_equal(series1: pd.Series, series2: pd.Series, config: EqualityConfig) -> bool:
    r"""Indicate if the two series are equal or not.

    Args:
        series1: The first series to compare.
        series2: The second series to compare.
        config: The equality configuration.

    Returns:
        ``True`` if the two series are equal, otherwise ``False``.
    """
    if not config.equal_nan and series1.isna().any():
        return False
    try:
        pd.testing.assert_series_equal(
            series1,
            series2,
            check_exact=config.atol == 0 and config.rtol == 0,
            check_index_type=True,
            atol=config.atol,
            rtol=config.rtol,
        )
    except AssertionError:
        return False
    return True
