Randomized Response for Grid-Cell Reporting

from __future__ import annotations

import math
import secrets
from dataclasses import dataclass

import numpy as np


@dataclass(frozen=True)
class RRParams:
    """Generalised randomised response over a grid-cell domain."""

    epsilon: float
    domain: int

    def __post_init__(self) -> None:
        if self.domain < 2:
            raise ValueError("randomised response needs at least two cells")
        if self.epsilon <= 0:
            raise ValueError("epsilon must be positive")

    @property
    def p_truth(self) -> float:
        """Probability the client reports its real cell."""
        e = math.exp(self.epsilon)
        return e / (e + self.domain - 1)

    @property
    def p_other(self) -> float:
        """Probability of reporting any one specific other cell."""
        return (1.0 - self.p_truth) / (self.domain - 1)

    def estimator_variance(self, n: int) -> float:
        """Variance of one cell's estimated count over n reporters."""
        p, q = self.p_truth, self.p_other
        return n * q * (1.0 - q) / ((p - q) ** 2)


def report_cell(true_cell: int, params: RRParams) -> int:
    """Randomise one grid-cell report on the device.

    `secrets` is used rather than `random` because the guarantee rests entirely on
    the unpredictability of this draw: an adversary who can reproduce the RNG state
    recovers the true cell exactly, and there is no partial degradation.
    """
    if not 0 <= true_cell < params.domain:
        raise ValueError("cell outside the declared domain")
    if secrets.randbelow(10 ** 9) < int(params.p_truth * 10 ** 9):
        return true_cell
    # Pick uniformly among the OTHER cells — including the true cell here would
    # skew p_truth upward and silently break the claimed epsilon.
    other = secrets.randbelow(params.domain - 1)
    return other if other < true_cell else other + 1


def estimate_counts(reports: np.ndarray, params: RRParams) -> np.ndarray:
    """Invert the randomisation into unbiased per-cell counts.

    Raw report tallies are biased upward for every cell, because each cell collects
    a share of everyone else's random reports. Publishing tallies without this
    inversion is the classic implementation error: it produces a map on which every
    cell looks moderately busy.
    """
    n = int(reports.size)
    tally = np.bincount(reports, minlength=params.domain).astype(float)
    p, q = params.p_truth, params.p_other
    return (tally - n * q) / (p - q)


def confidence_interval(params: RRParams, n: int, z: float = 1.96) -> float:
    """Half-width of the interval around each estimated count."""
    return z * math.sqrt(params.estimator_variance(n))

def _validate() -> None:
    rng = np.random.default_rng(23)
    params = RRParams(epsilon=2.0, domain=16)

    # A known population: 40% in cell 3, 25% in cell 9, the rest spread out.
    truth = np.concatenate([
        np.full(40_000, 3), np.full(25_000, 9),
        rng.integers(0, 16, size=35_000),
    ])
    reports = np.fromiter((report_cell(int(c), params) for c in truth),
                          dtype=int, count=truth.size)
    estimates = estimate_counts(reports, params)
    true_counts = np.bincount(truth, minlength=16).astype(float)
    half_width = confidence_interval(params, truth.size)

    # 1. The estimator is unbiased: every cell lands inside its interval.
    assert np.all(np.abs(estimates - true_counts) < half_width * 1.5), (
        np.max(np.abs(estimates - true_counts)), half_width
    )

    # 2. Raw tallies are NOT unbiased — this is the bug the inversion prevents.
    raw = np.bincount(reports, minlength=16).astype(float)
    empty_cells = np.where(true_counts < 100)[0]
    if empty_cells.size:
        assert raw[empty_cells].mean() > 500, "raw tallies must show the noise floor"
        assert abs(estimates[empty_cells].mean()) < half_width

    # 3. Truth-telling probability must match the closed form.
    matches = float(np.mean(reports == truth))
    assert abs(matches - params.p_truth) < 0.01, (matches, params.p_truth)

    # 4. Widening the domain at fixed epsilon must increase variance.
    wide = RRParams(epsilon=2.0, domain=256)
    assert wide.estimator_variance(100_000) > params.estimator_variance(100_000)

    # 5. The randomiser must never be able to return an out-of-domain cell.
    assert all(0 <= report_cell(0, params) < params.domain for _ in range(2_000))

    print("randomised response: all assertions passed")


_validate()