Per-User vs Per-Cell Budget Partitioning

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable, Mapping, Sequence

import numpy as np


class CompositionError(RuntimeError):
    """Raised when a release's structure contradicts its claimed composition rule."""


@dataclass(frozen=True)
class ReleasePlan:
    """One release: which partitions it touches and how a user maps into them."""

    name: str
    epsilon: float
    cell_ids: Sequence[int]                       # the partition being released
    user_cells: Mapping[str, Sequence[int]]       # user -> cells they contribute to
    max_cells_per_user: int                       # the enforced clipping bound

    def validate_parallel(self) -> None:
        """Confirm the release really is a disjoint partition with clipped users.

        Parallel composition is only sound when neighbouring datasets differ in ONE
        cell. Two things break it: duplicated cells in the partition, and a user
        whose contributions were not clipped to the stated bound.
        """
        if len(set(self.cell_ids)) != len(self.cell_ids):
            raise CompositionError(f"{self.name}: partition contains duplicate cells")
        for user, cells in self.user_cells.items():
            if len(set(cells)) > self.max_cells_per_user:
                raise CompositionError(
                    f"{self.name}: user {user} spans {len(set(cells))} cells, "
                    f"bound is {self.max_cells_per_user} — clip before binning"
                )

    def effective_sensitivity(self) -> float:
        """Sensitivity of the whole release given the clipping bound."""
        return float(self.max_cells_per_user)

    def noise_scale(self) -> float:
        """Laplace scale that actually delivers the claimed epsilon."""
        return self.effective_sensitivity() / self.epsilon


def sequential_cost(plans: Iterable[ReleasePlan]) -> float:
    """Total epsilon for releases over overlapping or repeated partitions."""
    return float(sum(p.epsilon for p in plans))


def parallel_cost(plans: Sequence[ReleasePlan]) -> float:
    """Total epsilon for releases over mutually disjoint partitions: the max, not the sum."""
    seen: set[int] = set()
    for p in plans:
        overlap = seen & set(p.cell_ids)
        if overlap:
            raise CompositionError(
                f"{p.name}: partitions overlap on {len(overlap)} cell(s) — "
                "these compose sequentially, not in parallel"
            )
        seen |= set(p.cell_ids)
    return float(max(p.epsilon for p in plans))

def _validate() -> None:
    # A clean disjoint partition with users clipped to one cell each.
    good = ReleasePlan(
        name="hourly-heatmap",
        epsilon=1.0,
        cell_ids=list(range(100)),
        user_cells={"u1": [3], "u2": [17], "u3": [42]},
        max_cells_per_user=1,
    )
    good.validate_parallel()

    # 1. With a clipping bound of 1, the noise scale is the textbook Delta/epsilon.
    assert abs(good.noise_scale() - 1.0) < 1e-9

    # 2. Two disjoint partitions cost the MAX, not the sum.
    north = ReleasePlan("north", 1.0, list(range(0, 50)), {"u1": [3]}, 1)
    south = ReleasePlan("south", 1.0, list(range(50, 100)), {"u9": [77]}, 1)
    assert parallel_cost([north, south]) == 1.0
    assert sequential_cost([north, south]) == 2.0

    # 3. Overlapping partitions must be refused, not silently treated as parallel.
    overlapping = ReleasePlan("res9", 1.0, list(range(40, 90)), {"u1": [42]}, 1)
    try:
        parallel_cost([north, overlapping])
    except CompositionError as exc:
        assert "overlap" in str(exc)
    else:
        raise AssertionError("overlapping partitions must raise")

    # 4. An unclipped user must be caught before the release, not after.
    unclipped = ReleasePlan(
        "od-matrix", 1.0, list(range(100)),
        {"commuter": [1, 2, 3, 4, 5, 6]}, max_cells_per_user=1,
    )
    try:
        unclipped.validate_parallel()
    except CompositionError as exc:
        assert "clip before binning" in str(exc)
    else:
        raise AssertionError("unclipped contributions must raise")

    # 5. Raising the clipping bound raises the noise scale proportionally: the
    #    honest cost of letting a device appear in several cells.
    trajectory = ReleasePlan("od-matrix", 1.0, list(range(100)),
                             {"commuter": [1, 2, 3]}, max_cells_per_user=3)
    trajectory.validate_parallel()
    assert abs(trajectory.noise_scale() - 3.0) < 1e-9

    print("budget partitioning: all assertions passed")


_validate()