_conditions.py

# Copyright (c) QuantCo 2025-2026
# SPDX-License-Identifier: BSD-3-Clause

import datetime as dt
from collections.abc import Mapping
from typing import cast

import polars as pl
from polars.datatypes import DataType, DataTypeClass

from diffly._utils import (
    ABS_TOL_DEFAULT,
    ABS_TOL_TEMPORAL_DEFAULT,
    REL_TOL_DEFAULT,
    Side,
)


def condition_equal_rows(
    columns: list[str],
    schema_left: pl.Schema,
    schema_right: pl.Schema,
    max_list_lengths_by_column: Mapping[str, int],
    abs_tol_by_column: Mapping[str, float],
    rel_tol_by_column: Mapping[str, float],
    abs_tol_temporal_by_column: Mapping[str, dt.timedelta],
) -> pl.Expr:
    """Build an expression whether two rows are equal, based on all columns' data
    types."""
    if not columns:
        return pl.lit(True)

    return pl.all_horizontal(
        [
            condition_equal_columns(
                column=column,
                dtype_left=schema_left[column],
                dtype_right=schema_right[column],
                max_list_length=max_list_lengths_by_column.get(column),
                abs_tol=abs_tol_by_column[column],
                rel_tol=rel_tol_by_column[column],
                abs_tol_temporal=abs_tol_temporal_by_column[column],
            )
            for column in columns
        ]
    )


def condition_equal_columns(
    column: str,
    dtype_left: pl.DataType,
    dtype_right: pl.DataType,
    max_list_length: int | None,
    abs_tol: float = ABS_TOL_DEFAULT,
    rel_tol: float = REL_TOL_DEFAULT,
    abs_tol_temporal: dt.timedelta = ABS_TOL_TEMPORAL_DEFAULT,
) -> pl.Expr:
    """Build an expression whether two columns are equal, depending on the columns' data
    types."""
    return _compare_columns(
        col_left=pl.col(f"{column}_{Side.LEFT}"),
        col_right=pl.col(f"{column}_{Side.RIGHT}"),
        dtype_left=dtype_left,
        dtype_right=dtype_right,
        max_list_length=max_list_length,
        abs_tol=abs_tol,
        rel_tol=rel_tol,
        abs_tol_temporal=abs_tol_temporal,
    )


# --------------------------------------- UTILS -------------------------------------- #


def _can_compare_dtypes(
    dtype_left: DataType | DataTypeClass,
    dtype_right: DataType | DataTypeClass,
) -> bool:
    return (
        (dtype_left == dtype_right)
        or (dtype_left == pl.Null)
        or (dtype_right == pl.Null)
        or (
            (
                (dtype_left.is_numeric() or dtype_left == pl.Boolean)
                == (dtype_right.is_numeric() or dtype_right == pl.Boolean)
            )
            and (dtype_left.is_temporal() == dtype_right.is_temporal())
            and (dtype_left.is_nested() == dtype_right.is_nested())
            and ((dtype_left == pl.Struct) == (dtype_right == pl.Struct))
        )
    )


def _compare_columns(
    col_left: pl.Expr,
    col_right: pl.Expr,
    dtype_left: DataType | DataTypeClass,
    dtype_right: DataType | DataTypeClass,
    max_list_length: int | None,
    abs_tol: float,
    rel_tol: float,
    abs_tol_temporal: dt.timedelta,
) -> pl.Expr:
    """Build an expression whether two expressions yield the same value.

    This method is more generic than :meth:`condition_equal_columns` as it accepts two
    arbitrary expressions rather than a "base column name".
    """
    if not _can_compare_dtypes(dtype_left, dtype_right):
        return pl.repeat(pl.lit(False), pl.len())

    # If we encounter nested dtypes, we have to treat them specially
    if dtype_left.is_nested():
        if isinstance(dtype_left, pl.Struct):
            assert isinstance(dtype_right, pl.Struct)
            # For two structs, we necessarily need to have matching field names (the
            # order does not matter). If that isn't the case, we cannot observe equality
            fields_left = {f.name: f.dtype for f in dtype_left.fields}
            fields_right = {f.name: f.dtype for f in dtype_right.fields}
            if fields_left.keys() != fields_right.keys():
                return pl.repeat(pl.lit(False), pl.len())

            # Otherwise, we simply compare all fields independently
            return pl.all_horizontal(
                [
                    _compare_columns(
                        col_left=col_left.struct[field],
                        col_right=col_right.struct[field],
                        dtype_left=fields_left[field],
                        dtype_right=fields_right[field],
                        max_list_length=max_list_length,
                        abs_tol=abs_tol,
                        rel_tol=rel_tol,
                        abs_tol_temporal=abs_tol_temporal,
                    )
                    for field in fields_left
                ]
            )
        elif isinstance(dtype_left, pl.List | pl.Array) and isinstance(
            dtype_right, pl.List | pl.Array
        ):
            if _needs_element_wise_comparison(dtype_left.inner, dtype_right.inner):
                return _compare_sequence_columns(
                    col_left=col_left,
                    col_right=col_right,
                    dtype_left=dtype_left,
                    dtype_right=dtype_right,
                    max_list_length=max_list_length,
                    abs_tol=abs_tol,
                    rel_tol=rel_tol,
                    abs_tol_temporal=abs_tol_temporal,
                )
            return col_left.eq_missing(col_right)

    if _different_enums(dtype_left, dtype_right) or _enum_and_categorical(
        dtype_left, dtype_right
    ):
        # Enums with different categories as well as enums and categoricals
        # can't be compared directly.
        # Fall back to comparison of strings.
        return _compare_columns(
            col_left=col_left.cast(pl.String),
            col_right=col_right.cast(pl.String),
            dtype_left=pl.String,
            dtype_right=pl.String,
            abs_tol=abs_tol,
            rel_tol=rel_tol,
            abs_tol_temporal=abs_tol_temporal,
            max_list_length=max_list_length,
        )

    return _compare_primitive_columns(
        col_left=col_left,
        col_right=col_right,
        dtype_left=dtype_left,
        dtype_right=dtype_right,
        abs_tol=abs_tol,
        rel_tol=rel_tol,
        abs_tol_temporal=abs_tol_temporal,
    )


def _compare_sequence_columns(
    col_left: pl.Expr,
    col_right: pl.Expr,
    dtype_left: pl.List | pl.Array,
    dtype_right: pl.List | pl.Array,
    max_list_length: int | None,
    abs_tol: float,
    rel_tol: float,
    abs_tol_temporal: dt.timedelta,
) -> pl.Expr:
    """Compare Array/List columns element-wise with tolerance."""
    n_elements: int
    has_same_length: pl.Expr

    if isinstance(dtype_left, pl.Array) and isinstance(dtype_right, pl.Array):
        if dtype_left.shape != dtype_right.shape:
            return pl.repeat(pl.lit(False), pl.len())
        n_elements = dtype_left.shape[0]
        has_same_length = pl.repeat(pl.lit(True), pl.len())
    elif isinstance(dtype_left, pl.Array) and isinstance(dtype_right, pl.List):
        n_elements = dtype_left.shape[0]
        has_same_length = col_right.list.len().eq(pl.lit(n_elements))
    elif isinstance(dtype_left, pl.List) and isinstance(dtype_right, pl.Array):
        n_elements = dtype_right.shape[0]
        has_same_length = col_left.list.len().eq(pl.lit(n_elements))
    else:  # pl.List vs pl.List
        n_elements = cast(int, max_list_length)
        has_same_length = col_left.list.len().eq_missing(col_right.list.len())

    if n_elements == 0:
        return _eq_missing(pl.lit(True), col_left, col_right)

    def _get_element(col: pl.Expr, dtype: DataType | DataTypeClass, i: int) -> pl.Expr:
        if isinstance(dtype, pl.Array):
            return col.arr.get(i)
        return col.list.get(i, null_on_oob=True)

    elements_match = pl.all_horizontal(
        [
            _compare_columns(
                col_left=_get_element(col_left, dtype_left, i),
                col_right=_get_element(col_right, dtype_right, i),
                dtype_left=dtype_left.inner,
                dtype_right=dtype_right.inner,
                abs_tol=abs_tol,
                rel_tol=rel_tol,
                abs_tol_temporal=abs_tol_temporal,
                max_list_length=max_list_length,
            )
            for i in range(n_elements)
        ]
    )

    return _eq_missing(has_same_length & elements_match, col_left, col_right)


def _is_float_numeric_pair(
    dtype_left: DataType | DataTypeClass,
    dtype_right: DataType | DataTypeClass,
) -> bool:
    return (dtype_left.is_float() or dtype_right.is_float()) and (
        dtype_left.is_numeric() and dtype_right.is_numeric()
    )


def _is_temporal_pair(
    dtype_left: DataType | DataTypeClass,
    dtype_right: DataType | DataTypeClass,
) -> bool:
    return dtype_left.is_temporal() and dtype_right.is_temporal()


def _needs_element_wise_comparison(
    dtype_left: DataType | DataTypeClass,
    dtype_right: DataType | DataTypeClass,
) -> bool:
    """Check if two dtypes require element-wise comparison (tolerances or special
    handling).

    Returns False when eq_missing() on the whole column would produce identical results,
    allowing us to skip the expensive element-wise iteration for list/array columns.
    """
    if (
        _is_float_numeric_pair(dtype_left, dtype_right)
        or _is_temporal_pair(dtype_left, dtype_right)
        or _different_enums(dtype_left, dtype_right)
        or _enum_and_categorical(dtype_left, dtype_right)
    ):
        return True
    if isinstance(dtype_left, pl.Struct) and isinstance(dtype_right, pl.Struct):
        fields_left = {f.name: f.dtype for f in dtype_left.fields}
        fields_right = {f.name: f.dtype for f in dtype_right.fields}
        return any(
            _needs_element_wise_comparison(fields_left[name], fields_right[name])
            for name in fields_left
            if name in fields_right
        )
    if isinstance(dtype_left, pl.List | pl.Array) and isinstance(
        dtype_right, pl.List | pl.Array
    ):
        return _needs_element_wise_comparison(dtype_left.inner, dtype_right.inner)
    return False


def _compare_primitive_columns(
    col_left: pl.Expr,
    col_right: pl.Expr,
    dtype_left: DataType | DataTypeClass,
    dtype_right: DataType | DataTypeClass,
    abs_tol: float,
    rel_tol: float,
    abs_tol_temporal: dt.timedelta,
) -> pl.Expr:
    if _is_float_numeric_pair(dtype_left, dtype_right):
        return col_left.is_close(col_right, abs_tol=abs_tol, rel_tol=rel_tol).pipe(
            _eq_missing_with_nan, lhs=col_left, rhs=col_right
        )
    elif _is_temporal_pair(dtype_left, dtype_right):
        diff_less_than_tolerance = (col_left - col_right).abs() <= abs_tol_temporal
        return diff_less_than_tolerance.pipe(_eq_missing, lhs=col_left, rhs=col_right)

    return col_left.eq_missing(col_right)


def _eq_missing(expr: pl.Expr, lhs: pl.Expr, rhs: pl.Expr) -> pl.Expr:
    both_null = lhs.is_null() & rhs.is_null()
    both_not_null = lhs.is_not_null() & rhs.is_not_null()
    return (expr & both_not_null) | both_null


def _eq_missing_with_nan(expr: pl.Expr, lhs: pl.Expr, rhs: pl.Expr) -> pl.Expr:
    both_nan = lhs.is_nan() & rhs.is_nan()
    return _eq_missing(expr, lhs, rhs) | both_nan


def _different_enums(
    left: DataType | DataTypeClass, right: DataType | DataTypeClass
) -> bool:
    return isinstance(left, pl.Enum) and isinstance(right, pl.Enum) and left != right


def _enum_and_categorical(
    left: DataType | DataTypeClass, right: DataType | DataTypeClass
) -> bool:
    return (isinstance(left, pl.Enum) and isinstance(right, pl.Categorical)) or (
        isinstance(left, pl.Categorical) and isinstance(right, pl.Enum)
    )