Differential Privacy Prod

import threading
from dataclasses import dataclass, field
from datetime import datetime
from typing import Optional
import math


@dataclass
class BudgetAllocation:
    query_id: str
    epsilon_spent: float
    delta_spent: float
    timestamp: datetime
    query_description: str
    analyst_id: str


class PrivacyBudgetManager:
    """
    Track and enforce differential privacy budget across queries.

    Supports both basic and advanced composition theorems.
    """

    def __init__(
        self,
        total_epsilon: float,
        total_delta: float,
        composition_method: str = "advanced"
    ):
        self.total_epsilon = total_epsilon
        self.total_delta = total_delta
        self.composition_method = composition_method
        self.allocations: list[BudgetAllocation] = []
        self._lock = threading.Lock()

    def remaining_budget(self) -> tuple[float, float]:
        """Calculate remaining (epsilon, delta) budget."""
        if self.composition_method == "basic":
            return self._basic_composition_remaining()
        elif self.composition_method == "advanced":
            return self._advanced_composition_remaining()
        elif self.composition_method == "rdp":
            return self._rdp_composition_remaining()
        else:
            raise ValueError(f"Unknown composition method: {self.composition_method}")

    def _basic_composition_remaining(self) -> tuple[float, float]:
        """Basic sequential composition: epsilons and deltas sum."""
        spent_epsilon = sum(a.epsilon_spent for a in self.allocations)
        spent_delta = sum(a.delta_spent for a in self.allocations)
        return (self.total_epsilon - spent_epsilon, self.total_delta - spent_delta)

    def _advanced_composition_remaining(self) -> tuple[float, float]:
        """
        Advanced composition theorem:
        k queries each with epsilon_i satisfy
        (sqrt(2k * ln(1/delta')) * max(epsilon_i) + k * epsilon_i * (e^epsilon_i - 1),
         k * delta_i + delta')-DP
        """
        k = len(self.allocations)
        if k == 0:
            return (self.total_epsilon, self.total_delta)

        epsilons = [a.epsilon_spent for a in self.allocations]
        deltas = [a.delta_spent for a in self.allocations]

        max_eps = max(epsilons)
        sum_delta = sum(deltas)

        # Reserve delta' for composition overhead
        delta_prime = (self.total_delta - sum_delta) / 2
        if delta_prime <= 0:
            return (0.0, 0.0)

        composed_epsilon = (
            math.sqrt(2 * k * math.log(1 / delta_prime)) * max_eps
            + k * max_eps * (math.exp(max_eps) - 1)
        )
        composed_delta = sum_delta + delta_prime

        return (
            max(0.0, self.total_epsilon - composed_epsilon),
            max(0.0, self.total_delta - composed_delta)
        )

    def _rdp_composition_remaining(self) -> tuple[float, float]:
        """Renyi DP composition (simplified)."""
        # RDP provides tighter bounds through Renyi divergence
        spent_epsilon = sum(a.epsilon_spent for a in self.allocations)
        spent_delta = sum(a.delta_spent for a in self.allocations)
        return (self.total_epsilon - spent_epsilon, self.total_delta - spent_delta)

    def request_budget(
        self,
        query_id: str,
        epsilon_requested: float,
        delta_requested: float,
        query_description: str,
        analyst_id: str
    ) -> bool:
        """
        Request budget allocation for a query.

        Returns True if budget is available and allocated, False otherwise.
        """
        with self._lock:
            remaining_eps, remaining_delta = self.remaining_budget()

            if epsilon_requested > remaining_eps or delta_requested > remaining_delta:
                return False

            allocation = BudgetAllocation(
                query_id=query_id,
                epsilon_spent=epsilon_requested,
                delta_spent=delta_requested,
                timestamp=datetime.utcnow(),
                query_description=query_description,
                analyst_id=analyst_id
            )
            self.allocations.append(allocation)
            return True

    def get_usage_report(self) -> dict:
        """Generate budget usage report."""
        remaining_eps, remaining_delta = self.remaining_budget()
        return {
            "total_epsilon": self.total_epsilon,
            "total_delta": self.total_delta,
            "remaining_epsilon": remaining_eps,
            "remaining_delta": remaining_delta,
            "utilization_pct": (1 - remaining_eps / self.total_epsilon) * 100,
            "num_queries": len(self.allocations),
            "composition_method": self.composition_method,
            "allocations": [
                {
                    "query_id": a.query_id,
                    "epsilon": a.epsilon_spent,
                    "delta": a.delta_spent,
                    "timestamp": a.timestamp.isoformat(),
                    "analyst": a.analyst_id,
                }
                for a in self.allocations
            ],
        }

epsilon' = sqrt(2k * ln(1/delta')) * epsilon + k * epsilon * (e^epsilon - 1)

For alpha > 1: epsilon(delta) = RDP_alpha - ln(delta) / (alpha - 1)

Raw Data Store --> Sensitivity Calibration --> DP Mechanism --> Result Cache
                         |                         |
                         v                         v
                  Budget Manager <------------- Audit Log

Client Device --> Local Randomizer --> Aggregation Server --> Estimator
                  (epsilon-LDP)              |                    |
                                             v                    v
                                       Budget Tracker      Utility Monitor

Training Data --> Mini-batch Sampling --> Per-example Gradient
                  (Poisson sampling)            |
                                                v
                                        Gradient Clipping (norm bound C)
                                                |
                                                v
                                        Gaussian Noise Addition (sigma * C)
                                                |
                                                v
                                        Model Update --> Privacy Accountant