Truncating home locations with spatial cloaking

import numpy as np
import pandas as pd


def truncate_trip_ends(
    df: pd.DataFrame,
    density: pd.Series | None = None,
    trip_col: str = "trip_id",
    dist_col: str = "cum_dist_m",
    base_m: float = 500.0,
    min_retained_m: float = 500.0,
    max_m: float = 3000.0,
) -> pd.DataFrame:
    """
    Remove the first and last portion of every trip, sized by local density.

    Parameters
    ----------
    df : pd.DataFrame
        Fix-level frame with a trip id and cumulative along-trip distance in
        METRES. Computing cum_dist in degrees produces a truncation whose
        real length varies with latitude.
    density : pd.Series | None
        Optional entities-per-km² at each trip's origin cell, indexed by
        trip id. Where it is low the cloaking radius is widened, because a
        fixed 500 m that anonymises a city block names a farmhouse.
    min_retained_m : float
        A trip must retain at least this much after truncation or it is
        dropped entirely. Partial truncation leaves endpoints that are still
        effectively the origin and destination.

    Returns
    -------
    pd.DataFrame
        The retained middles, with 'cloak_m' recording what was applied.

    Raises
    ------
    ValueError
        On missing columns, an empty frame, or non-monotone distance.
    """
    required = {trip_col, dist_col}
    missing = required - set(df.columns)
    if missing:
        raise ValueError(f"Missing required columns: {missing}")
    if df.empty:
        raise ValueError("Input DataFrame is empty.")

    out = df.sort_values([trip_col, dist_col]).copy()
    if (out.groupby(trip_col)[dist_col].diff().fillna(0) < 0).any():
        raise ValueError(
            "cum_dist_m decreases within a trip — recompute it from the "
            "sorted fixes; a non-monotone distance makes truncation arbitrary."
        )

    trip_len = out.groupby(trip_col)[dist_col].transform("max")

    # ── adaptive cloaking radius ──────────────────────────────────────
    # Scale inversely with the square root of density, so the cloaked area
    # contains a roughly constant number of candidate addresses.
    if density is not None:
        d = density.reindex(out[trip_col]).to_numpy()
        d = np.where(np.isfinite(d) & (d > 0), d, np.nan)
        ref = np.nanmedian(d)
        scale = np.sqrt(ref / d)
        cloak = np.clip(base_m * np.nan_to_num(scale, nan=1.0), base_m, max_m)
    else:
        cloak = np.full(len(out), base_m)
    out["cloak_m"] = cloak

    keep = (out[dist_col] > out["cloak_m"]) & (
        out[dist_col] < trip_len - out["cloak_m"]
    )
    long_enough = trip_len - 2 * out["cloak_m"] >= min_retained_m

    kept = out[keep & long_enough].copy()
    dropped_trips = out.loc[~long_enough, trip_col].nunique()
    kept.attrs["dropped_trips"] = int(dropped_trips)
    kept.attrs["retained_frac"] = float(len(kept) / len(df))
    return kept

def validate_cloaking(original: pd.DataFrame, cloaked: pd.DataFrame,
                      home_truth: pd.Series, trip_col: str = "trip_id",
                      entity_col: str = "pseudo_id", radius_m: float = 200.0) -> dict:
    """
    Attack the cloaked data with the detector the cloaking is meant to defeat.

    home_truth : known home coordinates per entity, held back for this test.
    A match rate near zero is the only evidence the control worked; the
    truncation distance alone proves nothing.
    """
    from scipy.spatial import cKDTree

    # Naive home detection: the most common first-fix location per entity.
    first_fixes = (cloaked.sort_values("cum_dist_m")
                          .groupby([entity_col, trip_col]).first()
                          .groupby(level=0)[["x", "y"]].median())
    truth = home_truth.reindex(first_fixes.index).dropna()
    if truth.empty:
        raise ValueError("No overlap between the attack set and the release.")

    tree = cKDTree(np.column_stack(truth.tolist()))
    d, _ = tree.query(first_fixes.loc[truth.index].to_numpy(), k=1)
    hit_rate = float((d < radius_m).mean())

    assert hit_rate < 0.05, (
        f"{hit_rate:.0%} of homes still recoverable within {radius_m} m — "
        "increase the cloaking radius or check that truncation ran before "
        "any downstream re-derivation of trip ends"
    )
    return {"home_hit_rate": hit_rate,
            "retained_frac": cloaked.attrs.get("retained_frac"),
            "dropped_trips": cloaked.attrs.get("dropped_trips")}