tianxuan/analysis/distance/_text.py

"""String distance, enum detection, and boolean distance functions."""

import polars as pl


# ---------------------------------------------------------------------------
#  Enum-like detection helper
# ---------------------------------------------------------------------------

def _is_enum_like(s: pl.Series) -> bool:
    """Check if ≥50% of rows have values appearing ≥2 times."""
    total = len(s)
    if total == 0:
        return False
    non_null = s.drop_nulls()
    if len(non_null) == 0:
        return False
    vc = non_null.value_counts()
    if 'count' not in vc.columns:
        return False
    duplicated = vc.filter(pl.col('count') >= 2)
    if len(duplicated) == 0:
        return False
    duplicated_count: int = duplicated['count'].sum()  # type: ignore[union-attr]
    return duplicated_count >= 0.5 * total


def _levenshtein_normalized(s1: str | None, s2: str) -> float:
    """Normalised Levenshtein distance between *s1* and *s2*.

    Empty or ``None`` values return 1.0 (maximum distance).
    """
    if s1 is None or s1 == '':
        return 1.0
    if s2 is None or s2 == '':
        return 1.0
    s1_str = str(s1)
    s2_str = str(s2)
    max_len = max(len(s1_str), len(s2_str))
    if max_len == 0:
        return 0.0
    import Levenshtein
    return Levenshtein.distance(s1_str, s2_str) / max_len


# ---------------------------------------------------------------------------
#  String distance
# ---------------------------------------------------------------------------

def string_distance(lf: pl.LazyFrame, str_columns: list[str]) -> pl.LazyFrame:
    """Compute string distance columns.

    For each column in *str_columns*:

    * If **≥50%** of rows have values appearing ≥2 times, the column is
      treated as an enum: 0 when the value equals the mode, 1 otherwise.
    * Otherwise, the normalised Levenshtein distance to the mode is
      computed for each row.

    Empty or ``None`` values always produce the maximum distance (1.0).

    Parameters
    ----------
    lf:
        Input LazyFrame.
    str_columns:
        List of string column names to compute distances for.

    Returns
    -------
    pl.LazyFrame
        Input LazyFrame extended with ``_str_dist_0``, ``_str_dist_1``,
        ... columns (one per entry in *str_columns*).
    """
    df = lf.collect(streaming=True)
    dist_exprs: list[pl.Expr] = []

    for idx, col in enumerate(str_columns):
        if col not in df.columns:
            continue

        series = df[col].cast(pl.Utf8)
        enum_like = _is_enum_like(series)

        non_null = series.drop_nulls()
        if len(non_null) > 0:
            mode_val = non_null.value_counts().sort('count', descending=True).row(0)[0]
            mode_str = str(mode_val)
        else:
            mode_str = ''

        col_name = f'_str_dist_{idx}'

        if enum_like:
            dist_exprs.append(
                pl.when(pl.col(col).is_null())
                .then(pl.lit(1.0))
                .when(pl.col(col).cast(pl.Utf8) == pl.lit(mode_str))
                .then(pl.lit(0.0))
                .otherwise(pl.lit(1.0))
                .alias(col_name)
            )
        else:
            # Capture in closure for map_elements
            _mode = mode_str
            dist_exprs.append(
                pl.col(col).map_elements(
                    lambda v: _levenshtein_normalized(v, _mode),
                    return_dtype=pl.Float64,
                ).alias(col_name)
            )

    return lf.with_columns(dist_exprs)


# ---------------------------------------------------------------------------
#  Boolean distance
# ---------------------------------------------------------------------------

def bool_distance(lf: pl.LazyFrame, bool_columns: list[str]) -> pl.LazyFrame:
    """Compute boolean distance columns.

    For each column in *bool_columns*, each row gets:
    0 if its value matches the column mode, 1 otherwise.
    ``None`` values always produce 1.0 (maximum distance).

    Parameters
    ----------
    lf:
        Input LazyFrame.
    bool_columns:
        List of boolean column names to compute distances for.

    Returns
    -------
    pl.LazyFrame
        Input LazyFrame extended with ``_bool_dist_0``, ``_bool_dist_1``,
        ... columns (one per entry in *bool_columns*).
    """
    df = lf.collect(streaming=True)
    dist_exprs: list[pl.Expr] = []

    for idx, col in enumerate(bool_columns):
        if col not in df.columns:
            continue

        series = df[col]
        non_null = series.drop_nulls()
        col_name = f'_bool_dist_{idx}'

        if len(non_null) == 0:
            dist_exprs.append(pl.lit(0.0).alias(col_name))
        else:
            mode_val = non_null.value_counts().sort('count', descending=True).row(0)[0]
            dist_exprs.append(
                pl.when(pl.col(col).is_null())
                .then(pl.lit(1.0))
                .when(pl.col(col) == pl.lit(mode_val))
                .then(pl.lit(0.0))
                .otherwise(pl.lit(1.0))
                .alias(col_name)
            )

    return lf.with_columns(dist_exprs)