k-anonymity for trajectory datasets

import pandas as pd
import numpy as np


def apply_k_anonymity(
    df: pd.DataFrame,
    quasi_identifiers: list[str],
    entity_col: str = "pseudo_id",
    k: int = 20,
    generalise: dict | None = None,
) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Enforce a k-anonymity floor on distinct entities per quasi-identifier group.

    Parameters
    ----------
    quasi_identifiers : list[str]
        EVERY released attribute an adversary could link on — typically
        [cell, time_bucket] and often more. Applying the floor to a subset
        leaves released tuples that are unique.
    generalise : dict | None
        Optional {column: coarser_column} map. Groups below k are retried at
        the coarser level before being suppressed, which preserves far more
        rural data than suppression alone.

    Returns
    -------
    (released, suppression_report)

    Raises
    ------
    ValueError
        On missing columns, an empty frame, or k < 2.
    """
    missing = set(quasi_identifiers + [entity_col]) - set(df.columns)
    if missing:
        raise ValueError(f"Missing required columns: {sorted(missing)}")
    if df.empty:
        raise ValueError("Input DataFrame is empty.")
    if k < 2:
        raise ValueError("k must be at least 2; k=1 is no control.")

    work = df.copy()
    # nunique, NOT size. This single choice is the difference between a real
    # floor and one that a single chatty device satisfies on its own.
    work["_n_entities"] = work.groupby(quasi_identifiers)[entity_col].transform("nunique")
    passing = work["_n_entities"] >= k

    released = [work[passing]]
    failed = work[~passing]

    # ── optional generalisation pass ──────────────────────────────────
    # Retry the failures at a coarser resolution before giving up. In rural
    # areas this recovers most of what a suppression-only rule destroys.
    if generalise and len(failed):
        coarse_qi = [generalise.get(c, c) for c in quasi_identifiers]
        miss = set(coarse_qi) - set(failed.columns)
        if miss:
            raise ValueError(f"Generalisation targets missing: {sorted(miss)}")
        failed = failed.copy()
        failed["_n_entities"] = (
            failed.groupby(coarse_qi)[entity_col].transform("nunique")
        )
        recovered = failed[failed["_n_entities"] >= k].copy()
        # Overwrite the fine columns with the coarse ones so the released row
        # cannot be mistaken for fine-grained data later.
        for fine, coarse in zip(quasi_identifiers, coarse_qi):
            if fine != coarse:
                recovered[fine] = recovered[coarse]
        released.append(recovered)
        failed = failed[failed["_n_entities"] < k]

    out = pd.concat(released, ignore_index=True).drop(columns=["_n_entities"])

    if len(failed):
        report = (
            failed.groupby(quasi_identifiers, observed=True)
            .size().rename("suppressed_rows").reset_index()
        )
    else:
        report = pd.DataFrame(columns=quasi_identifiers + ["suppressed_rows"])
    out.attrs["suppressed_rows"] = int(len(failed))
    out.attrs["suppressed_frac"] = float(len(failed) / len(df))
    out.attrs["k"] = k
    return out, report

def validate_k_floor(released: pd.DataFrame, quasi_identifiers: list[str],
                     entity_col: str, k: int) -> None:
    """
    Re-derive the floor independently of the code that applied it.

    A shared helper hides the size-versus-nunique bug from itself; a separate
    implementation is the only version of this check worth running.
    """
    sizes = released.groupby(quasi_identifiers)[entity_col].nunique()
    violations = sizes[sizes < k]
    assert violations.empty, (
        f"{len(violations)} released groups below k={k}; worst has "
        f"{violations.min()} entities — the floor was not applied to the "
        "full quasi-identifier"
    )
    # Suppression should not be so heavy that the release is meaningless,
    # nor so light that the floor never bound.
    frac = released.attrs.get("suppressed_frac", 0.0)
    assert frac < 0.5, f"suppressed {frac:.0%} — generalise before suppressing"
    assert frac > 0.0, "nothing suppressed — check that k is binding at all"
    print(f"OK — k={k} holds across {len(sizes)} groups, {frac:.1%} suppressed")