Density-Aware Sensitivity Scoring with H3

from __future__ import annotations

from dataclasses import dataclass
from typing import Mapping, Sequence

import math

# Average H3 cell area in square kilometres, derived from the published edge lengths.
H3_AREA_KM2: Mapping[int, float] = {
    5: 189.6, 6: 27.1, 7: 3.87, 8: 0.553, 9: 0.079, 10: 0.0113, 11: 0.00161,
}


@dataclass(frozen=True)
class SensitivityScore:
    """The scorer's output: a tier, the resolution it may publish at, and why."""

    cell: str
    density_per_km2: float
    publish_resolution: int
    expected_k: float
    density_term: float
    facility_term: float
    uniqueness_term: float
    tier: int

    @property
    def composite(self) -> float:
        return self.density_term + self.facility_term + self.uniqueness_term


def finest_resolution_for_k(density_per_km2: float, k_target: float) -> int:
    """Finest H3 resolution whose expected occupancy still clears the floor.

    Expected occupancy in a cell is density * area, so this walks from fine to
    coarse and returns the first resolution that satisfies the floor. Returning the
    coarsest resolution when nothing qualifies is deliberate: the caller must then
    decide between suppression and secure computation, and silently publishing at
    the coarsest grid would hide that decision.
    """
    for res in sorted(H3_AREA_KM2, reverse=True):
        if density_per_km2 * H3_AREA_KM2[res] >= k_target:
            return res
    return min(H3_AREA_KM2)


def score_cell(
    cell: str,
    *,
    density_per_km2: float,
    sensitive_facility_within_m: float | None,
    trajectory_uniqueness: float,
    k_target: float = 5.0,
    facility_weight: float = 0.45,
    uniqueness_weight: float = 0.30,
    facility_buffer_m: float = 500.0,
) -> SensitivityScore:
    """Score one cell and decide the resolution it is allowed to be published at.

    `trajectory_uniqueness` is in [0, 1] — the fraction of visitors whose local
    path is unique within the measurement window. It is the term that makes empty
    countryside score highly despite its low density.
    """
    if not 0.0 <= trajectory_uniqueness <= 1.0:
        raise ValueError("trajectory_uniqueness must be in [0, 1]")

    res = finest_resolution_for_k(density_per_km2, k_target)
    expected_k = density_per_km2 * H3_AREA_KM2[res]

    # Sparse cells are MORE sensitive: the density term is inverted and log-scaled
    # so that the difference between 10 and 100 per km2 matters more than the
    # difference between 10,000 and 100,000.
    density_term = max(0.0, 1.0 - math.log1p(density_per_km2) / math.log1p(20_000.0))

    if sensitive_facility_within_m is None:
        facility_term = 0.0
    else:
        proximity = max(0.0, 1.0 - sensitive_facility_within_m / facility_buffer_m)
        facility_term = facility_weight * proximity

    uniqueness_term = uniqueness_weight * trajectory_uniqueness
    composite = density_term + facility_term + uniqueness_term

    tier = 1 if composite < 0.35 else 2 if composite < 0.60 else 3 if composite < 0.85 else 4
    return SensitivityScore(
        cell=cell,
        density_per_km2=density_per_km2,
        publish_resolution=res,
        expected_k=expected_k,
        density_term=density_term,
        facility_term=facility_term,
        uniqueness_term=uniqueness_term,
        tier=tier,
    )


def smooth_density(
    counts: Mapping[str, float],
    rings: Mapping[str, Sequence[str]],
    area_km2: float,
) -> dict[str, float]:
    """Average raw counts over each cell's k-ring before converting to a density.

    Unsmoothed counts flip tiers between adjacent cells, which produces a mosaic of
    routing decisions that is impossible to explain to an analyst — and which leaks
    the underlying counts through the tier boundaries.
    """
    out: dict[str, float] = {}
    for cell, ring in rings.items():
        members = [cell, *ring]
        total = sum(counts.get(m, 0.0) for m in members)
        out[cell] = total / (len(members) * area_km2)
    return out

def _validate() -> None:
    # 1. The resolution chosen must actually clear the anonymity floor.
    for density in (14.0, 140.0, 1_400.0, 12_000.0):
        res = finest_resolution_for_k(density, 5.0)
        assert density * H3_AREA_KM2[res] >= 5.0

    # 2. Denser areas may publish at a finer resolution — never the reverse.
    assert finest_resolution_for_k(12_000.0, 5.0) >= finest_resolution_for_k(14.0, 5.0)

    # 3. Sparse rural land scores as sensitive even with no facility nearby.
    rural = score_cell("r1", density_per_km2=14.0,
                       sensitive_facility_within_m=None, trajectory_uniqueness=0.9)
    assert rural.tier >= 3, rural.composite

    # 4. A dense street with no facility and low uniqueness is tier 1.
    street = score_cell("u1", density_per_km2=8_000.0,
                        sensitive_facility_within_m=None, trajectory_uniqueness=0.05)
    assert street.tier == 1

    # 5. Facility proximity must be able to lift a dense cell out of tier 1.
    clinic = score_cell("u2", density_per_km2=8_000.0,
                        sensitive_facility_within_m=80.0, trajectory_uniqueness=0.05)
    assert clinic.tier > street.tier

    # 6. Smoothing must not change the total mass, only its distribution.
    counts = {"a": 100.0, "b": 0.0, "c": 20.0}
    rings = {"a": ["b"], "b": ["a", "c"], "c": ["b"]}
    dens = smooth_density(counts, rings, area_km2=0.553)
    assert dens["b"] > 0.0, "an empty cell beside a busy one must not read as empty"

    print("density-aware scoring: all assertions passed")


_validate()