tianxuan/analysis/distance/_geo.py

"""IP distance functions for Polars LazyFrames.

Provides subnet-mask distance and Haversine geographic distance
between pairs of IP columns.  All distances are normalised to [0, 1];
invalid / non-geolocatable IPs produce ``NaN``.
"""

from __future__ import annotations

import math
from typing import Optional

import polars as pl

from analysis.geoip import ip_to_int, lookup

# ---------------------------------------------------------------------------
#  Earth constant
# ---------------------------------------------------------------------------

_HALF_EARTH_CIRCUMFERENCE_KM = 20_037.508  # equatorial / 2


# ---------------------------------------------------------------------------
#  Core distance helpers
# ---------------------------------------------------------------------------

def _subnet_mask_distance(ip1: Optional[str], ip2: Optional[str], mask: int) -> float:
    """Normalised subnet-mask distance between two IP strings.

    Computes ``abs(int(ip1)/2**mask - int(ip2)/2**mask) / 2**(32-mask)``.
    Returns ``NaN`` when either IP is ``None`` or unparseable.
    """
    if ip1 is None or ip2 is None:
        return float("nan")
    if not isinstance(ip1, str):
        ip1 = str(ip1)
    if not isinstance(ip2, str):
        ip2 = str(ip2)

    int1 = ip_to_int(ip1)
    int2 = ip_to_int(ip2)
    if int1 is None or int2 is None:
        return float("nan")

    divisor = 2.0 ** mask
    max_val = 2.0 ** (32 - mask)
    return abs(int1 / divisor - int2 / divisor) / max_val


def _haversine_km(lat1: float, lon1: float, lat2: float, lon2: float) -> float:
    """Great-circle distance (km) via the Haversine formula."""
    dlat = math.radians(lat2 - lat1)
    dlon = math.radians(lon2 - lon1)
    a = (
        math.sin(dlat / 2) ** 2
        + math.cos(math.radians(lat1))
        * math.cos(math.radians(lat2))
        * math.sin(dlon / 2) ** 2
    )
    return 6371.0 * 2.0 * math.asin(math.sqrt(a))


def _geo_distance(ip1: Optional[str], ip2: Optional[str]) -> float:
    """Normalised [0, 1] geographic distance between two IP strings.

    Uses ``analysis.geoip.lookup()`` to obtain lat/lon, then Haversine.
    Returns ``NaN`` when either IP cannot be geolocated.
    """
    if ip1 is None or ip2 is None:
        return float("nan")
    if not isinstance(ip1, str):
        ip1 = str(ip1)
    if not isinstance(ip2, str):
        ip2 = str(ip2)

    geo1 = lookup(ip1)
    geo2 = lookup(ip2)
    if geo1 is None or geo2 is None:
        return float("nan")

    km = _haversine_km(geo1["lat"], geo1["lon"], geo2["lat"], geo2["lon"])
    return km / _HALF_EARTH_CIRCUMFERENCE_KM


# ---------------------------------------------------------------------------
#  Public API
# ---------------------------------------------------------------------------

def ip_distance(
    lf: pl.LazyFrame,
    ip_columns: list[str],
    subnet_masks: Optional[list[int]] = None,
) -> pl.LazyFrame:
    """Append IP distance columns to a LazyFrame.

    For each **consecutive pair** of *ip_columns*
    (``ip_columns[0]`` vs ``ip_columns[1]``,
    ``ip_columns[2]`` vs ``ip_columns[3]``, …) computes:

    * one subnet-mask distance per entry in *subnet_masks*
    * one Haversine geographic distance

    All distances are normalised to [0, 1].  Invalid IPs or IPs not
    found in the GeoIP database produce ``NaN``.

    Parameters
    ----------
    lf:
        LazyFrame containing the IP columns.
    ip_columns:
        List of IP column names (used in consecutive pairs).
    subnet_masks:
        Subnet masks for network distance (default ``[24, 28]``).

    Returns
    -------
    pl.LazyFrame
        Input frame with ``_ip_dist_0``, ``_ip_dist_1``, …
        columns appended.

    Examples
    --------
    >>> lf = pl.LazyFrame({
    ...     "src_ip": ["8.8.8.8", "1.1.1.1"],
    ...     "dst_ip": ["8.8.4.4", "1.0.0.1"],
    ... })
    >>> result = ip_distance(lf, ["src_ip", "dst_ip"], subnet_masks=[24])
    >>> result.collect().columns
    ['src_ip', 'dst_ip', '_ip_dist_0', '_ip_dist_1']
    """
    if subnet_masks is None:
        subnet_masks = [24, 28]

    result = lf
    dist_idx = 0

    # Walk consecutive pairs
    for i in range(0, len(ip_columns) - 1, 2):
        col_a = ip_columns[i]
        col_b = ip_columns[i + 1]

        # -- subnet-mask distances --
        for mask in subnet_masks:
            name = f"_ip_dist_{dist_idx}"
            result = result.with_columns(
                pl.struct([col_a, col_b])
                .map_elements(
                    lambda s, _mask=mask, _ca=col_a, _cb=col_b: (
                        _subnet_mask_distance(s[_ca], s[_cb], _mask)
                    ),
                    return_dtype=pl.Float64,
                )
                .alias(name)
            )
            dist_idx += 1

        # -- geographic distance --
        name = f"_ip_dist_{dist_idx}"
        result = result.with_columns(
            pl.struct([col_a, col_b])
            .map_elements(
                lambda s, _ca=col_a, _cb=col_b: _geo_distance(s[_ca], s[_cb]),
                return_dtype=pl.Float64,
            )
            .alias(name)
        )
        dist_idx += 1

    return result