r"""Implement handlers for ``polars.DataFrame``s, ``polars.LazyFrame``s,
and ``polars.Series``es."""

from __future__ import annotations

__all__ = ["PolarsDataFrameEqualHandler", "PolarsLazyFrameEqualHandler", "PolarsSeriesEqualHandler"]

import inspect
import logging
from functools import lru_cache
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_polars_available

if TYPE_CHECKING or is_polars_available():
    import polars as pl
    import polars.selectors as cs
    from polars import testing
else:  # pragma: no cover
    from coola.utils.fallback.polars import polars as pl

if TYPE_CHECKING:
    from coola.equality.config import EqualityConfig

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


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

    This handler returns ``True`` if the two ``polars.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 polars as pl
        >>> from coola.equality.config import EqualityConfig
        >>> from coola.equality.handler import PolarsDataFrameEqualHandler
        >>> config = EqualityConfig()
        >>> handler = PolarsDataFrameEqualHandler()
        >>> handler.handle(
        ...     pl.DataFrame({"col": [1, 2, 3]}),
        ...     pl.DataFrame({"col": [1, 2, 3]}),
        ...     config,
        ... )
        True
        >>> handler.handle(
        ...     pl.DataFrame({"col": [1, 2, 3]}),
        ...     pl.DataFrame({"col": [1, 2, 4]}),
        ...     config,
        ... )
        False

        ```
    """

    def handle(
        self,
        actual: pl.DataFrame,
        expected: pl.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="polars.DataFrames")
            )
        return object_equal


class PolarsLazyFrameEqualHandler(HandlerEqualityMixin, BaseEqualityHandler):
    r"""Check if the two ``polars.LazyFrame``s are equal.

    This handler returns ``True`` if the two ``polars.LazyFrame``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 polars as pl
        >>> from coola.equality.config import EqualityConfig
        >>> from coola.equality.handler import PolarsLazyFrameEqualHandler
        >>> config = EqualityConfig()
        >>> handler = PolarsLazyFrameEqualHandler()
        >>> handler.handle(
        ...     pl.LazyFrame({"col": [1, 2, 3]}),
        ...     pl.LazyFrame({"col": [1, 2, 3]}),
        ...     config,
        ... )
        True
        >>> handler.handle(
        ...     pl.LazyFrame({"col": [1, 2, 3]}),
        ...     pl.LazyFrame({"col": [1, 2, 4]}),
        ...     config,
        ... )
        False

        ```
    """

    def handle(
        self,
        actual: pl.LazyFrame,
        expected: pl.LazyFrame,
        config: EqualityConfig,
    ) -> bool:
        object_equal = frame_equal(actual.collect(), expected.collect(), config)
        if config.show_difference and not object_equal:
            logger.info(
                format_value_difference(actual=actual, expected=expected, name="polars.LazyFrames")
            )
        return object_equal


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

    This handler returns ``True`` if the two ``polars.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 polars as pl
        >>> from coola.equality.config import EqualityConfig
        >>> from coola.equality.handler import PolarsSeriesEqualHandler
        >>> config = EqualityConfig()
        >>> handler = PolarsSeriesEqualHandler()
        >>> handler.handle(pl.Series([1, 2, 3]), pl.Series([1, 2, 3]), config)
        True
        >>> handler.handle(pl.Series([1, 2, 3]), pl.Series([1, 2, 4]), config)
        False

        ```
    """

    def handle(
        self,
        actual: pl.Series,
        expected: pl.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="polars.Series")
            )
        return object_equal


def has_nan(df_or_series: pl.DataFrame | pl.Series) -> bool:
    r"""Indicate if a DataFrame or Series has NaN values.

    Args:
        df_or_series: The DataFrame or series to check.

    Returns:
        ``True`` if the DataFrame or Series has NaN values,
            otherwise ``False``.
    """
    if isinstance(df_or_series, pl.Series):
        return df_or_series.dtype.is_numeric() and df_or_series.is_nan().any()
    frame = df_or_series.select(cs.numeric())
    if frame.is_empty():  # This if is necessary for previous polars version (1.0)
        return False
    return frame.select(pl.any_horizontal(pl.all().is_nan().any())).item()


def frame_equal(df1: pl.DataFrame, df2: pl.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 has_nan(df1):
        return False
    try:
        assert_frame_equal(df1, df2, atol=config.atol, rtol=config.rtol)
    except AssertionError:
        return False
    return True


def series_equal(series1: pl.Series, series2: pl.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 has_nan(series1):
        return False
    try:
        assert_series_equal(
            series1,
            series2,
            atol=config.atol,
            rtol=config.rtol,
        )
    except AssertionError:
        return False
    return True


def assert_frame_equal(
    df1: pl.DataFrame,
    df2: pl.DataFrame,
    rtol: float = 1e-5,
    atol: float = 1e-8,
) -> None:
    r"""Compare two polars DataFrames with tolerance.

    This function supports multiple polars versions.

    Args:
        df1: The first DataFrame to compare.
        df2: The second DataFrame to compare.
        rtol: The relative tolerance parameter.
        atol: The absolute tolerance parameter.
    """
    check_exact = atol == 0 and rtol == 0
    if is_new_naming():
        testing.assert_frame_equal(df1, df2, abs_tol=atol, rel_tol=rtol, check_exact=check_exact)
    else:
        testing.assert_frame_equal(df1, df2, atol=atol, rtol=rtol, check_exact=check_exact)


def assert_series_equal(
    series1: pl.Series,
    series2: pl.Series,
    rtol: float = 1e-5,
    atol: float = 1e-8,
) -> None:
    r"""Compare two polars Series with tolerance.

    This function supports multiple polars versions.

    Args:
        series1: The first series to compare.
        series2: The second series to compare.
        rtol: The relative tolerance parameter.
        atol: The absolute tolerance parameter.
    """
    check_exact = atol == 0 and rtol == 0
    if is_new_naming():
        testing.assert_series_equal(
            series1, series2, abs_tol=atol, rel_tol=rtol, check_exact=check_exact
        )
    else:
        testing.assert_series_equal(series1, series2, atol=atol, rtol=rtol, check_exact=check_exact)


@lru_cache
def is_new_naming() -> bool:
    r"""Indicate if polars uses the new naming for ``abs_tol`` and
    ``rel_tol``.

    Returns:
        ``True`` if ``abs_tol`` and ``rel_tol`` are used,
            otherwise ``atol`` and ``rtol`` are used.
    """
    return "abs_tol" in inspect.signature(pl.testing.assert_frame_equal).parameters
