#!/usr/bin/env python3
"""
Lab 5 - Differentially private query layer using diffprivlib.

Demonstrates:
  1. A non-private baseline (true mean, true count).
  2. DP mean at multiple epsilon values (privacy-utility trade-off table).
  3. A budget-enforcing query layer that refuses queries once
     the total epsilon budget is exhausted (composition as code).

Requires: pip install "diffprivlib==0.6.6"

All data is synthetic (generated inline). Never use real personal data
in a teaching lab.

Usage:
    python dp_query_layer.py
    python dp_query_layer.py --rows 5000 --budget 2.0
"""

import argparse
import random
import sys

# ---------------------------------------------------------------------------
# Synthetic dataset
# ---------------------------------------------------------------------------

def make_synthetic_dataset(n: int, seed: int = 42) -> list[dict]:
    """
    Generate n synthetic records with 'age' (18-80) and 'income' (15000-120000).
    Values are plausible but entirely invented.
    """
    rng = random.Random(seed)
    records = []
    for _ in range(n):
        age = rng.randint(18, 80)
        # Income skewed: most between 25k-70k, tail up to 120k
        income = max(15000.0, min(120000.0, rng.gauss(45000, 18000)))
        records.append({"age": age, "income": income})
    return records


# ---------------------------------------------------------------------------
# diffprivlib import with helpful error
# ---------------------------------------------------------------------------

try:
    import diffprivlib.tools as dp
    import numpy as np
except ImportError:
    print("[!] diffprivlib not found. Install with:")
    print("      pip install 'diffprivlib==0.6.6'")
    sys.exit(1)


# ---------------------------------------------------------------------------
# Non-private baseline
# ---------------------------------------------------------------------------

def baseline_stats(values: list[float]) -> dict:
    arr = np.array(values)
    return {
        "count":        len(arr),
        "true_mean":    float(arr.mean()),
        "true_median":  float(np.median(arr)),
        "true_std":     float(arr.std()),
    }


# ---------------------------------------------------------------------------
# Privacy-utility trade-off sweep
# ---------------------------------------------------------------------------

INCOME_BOUNDS = (0.0, 250_000.0)
EPSILON_VALUES = [0.01, 0.1, 0.5, 1.0, 5.0, 10.0]
SWEEP_RUNS = 5  # average over multiple runs to show jitter


def sweep_epsilon(values: list[float], true_mean: float) -> None:
    """
    For each epsilon value, run SWEEP_RUNS DP mean queries and print the
    average error vs ground truth. This IS the privacy-utility trade-off made
    visible.
    """
    arr = np.array(values)
    print("\n" + "=" * 70)
    print("PRIVACY-UTILITY TRADE-OFF: DP mean income vs epsilon")
    print(f"True mean income: {true_mean:,.2f}")
    print(f"(Each cell is the average absolute error over {SWEEP_RUNS} runs)")
    print("=" * 70)
    print(f"{'epsilon':>10}  {'dp_mean_sample':>16}  {'abs_error':>12}  {'note'}")
    print("-" * 70)
    for eps in EPSILON_VALUES:
        errors = []
        last_dp_mean = None
        for _ in range(SWEEP_RUNS):
            dp_mean = dp.mean(arr, epsilon=eps, bounds=INCOME_BOUNDS)
            errors.append(abs(dp_mean - true_mean))
            last_dp_mean = dp_mean
        avg_err = sum(errors) / len(errors)
        note = "strong privacy, high noise" if eps <= 0.1 else ("weak privacy, low noise" if eps >= 5.0 else "")
        print(f"{eps:>10.2f}  {last_dp_mean:>16,.2f}  {avg_err:>12,.2f}  {note}")
    print()
    print("Observation: run this script twice - the dp_mean column will differ")
    print("because the mechanism is randomised. The error shrinks as epsilon grows.")


# ---------------------------------------------------------------------------
# Budget-enforcing query layer
# ---------------------------------------------------------------------------

class DPQueryLayer:
    """
    A simple DP query layer with a fixed total epsilon budget.
    Demonstrates composition: each query spends budget; once exhausted
    further queries are refused.

    In a real system:
      - Budget is tracked server-side per analyst, not per process.
      - The client CANNOT reset the budget.
      - Authentication and authorisation wrap this layer (sections 16.2-16.3).
      - All queries are logged for audit.
    """

    def __init__(self, data: np.ndarray, total_budget: float, bounds: tuple[float, float]) -> None:
        self._data = data
        self._total_budget = total_budget
        self._remaining = total_budget
        self._bounds = bounds
        self._query_log: list[dict] = []

    def mean(self, epsilon: float, label: str = "mean") -> float:
        if epsilon <= 0:
            raise ValueError("epsilon must be positive")
        if epsilon > self._remaining:
            raise RuntimeError(
                f"Budget exhausted: {self._remaining:.4f} remaining, "
                f"query requested {epsilon:.4f}. No further queries allowed."
            )
        result = dp.mean(self._data, epsilon=epsilon, bounds=self._bounds)
        self._remaining -= epsilon
        self._query_log.append({
            "query": label,
            "epsilon_spent": epsilon,
            "remaining": self._remaining,
            "result": result,
        })
        print(f"  [QUERY] {label}: epsilon={epsilon}, remaining_budget={self._remaining:.4f}, result={result:,.2f}")
        return result

    def count(self, epsilon: float, label: str = "count") -> float:
        if epsilon > self._remaining:
            raise RuntimeError(
                f"Budget exhausted: {self._remaining:.4f} remaining, "
                f"query requested {epsilon:.4f}. No further queries allowed."
            )
        result = dp.count_nonzero(self._data, epsilon=epsilon)
        self._remaining -= epsilon
        self._query_log.append({
            "query": label,
            "epsilon_spent": epsilon,
            "remaining": self._remaining,
            "result": result,
        })
        print(f"  [QUERY] {label}: epsilon={epsilon}, remaining_budget={self._remaining:.4f}, result={result}")
        return result

    @property
    def budget_remaining(self) -> float:
        return self._remaining

    def print_log(self) -> None:
        print("\n  Query audit log:")
        for entry in self._query_log:
            print(f"    {entry['query']:30s}  eps_spent={entry['epsilon_spent']:.3f}  remaining={entry['remaining']:.4f}")


# ---------------------------------------------------------------------------
# Demo of budget composition
# ---------------------------------------------------------------------------

def demo_budget_composition(values: list[float], total_budget: float) -> None:
    arr = np.array(values)
    print("\n" + "=" * 70)
    print(f"BUDGET COMPOSITION DEMO  (total budget epsilon={total_budget})")
    print("=" * 70)
    layer = DPQueryLayer(arr, total_budget=total_budget, bounds=INCOME_BOUNDS)

    print("\nQuery 1: mean income (epsilon=0.3)")
    layer.mean(0.3, label="mean_income_all")

    print("\nQuery 2: mean income for age >= 40 subset (epsilon=0.3)")
    subset = np.array([r["income"] for r in make_synthetic_dataset(len(values)) if r["age"] >= 40])
    layer2 = DPQueryLayer(subset, total_budget=layer.budget_remaining, bounds=INCOME_BOUNDS)
    layer2.mean(0.3, label="mean_income_age_40_plus")
    layer._remaining = layer2._remaining  # carry budget across (simplified demo)
    layer._query_log.extend(layer2._query_log)

    print("\nQuery 3: mean income for age < 40 subset (epsilon=0.3)")
    subset2 = np.array([r["income"] for r in make_synthetic_dataset(len(values)) if r["age"] < 40])
    try:
        layer.mean(0.3, label="mean_income_age_under_40")
    except RuntimeError as exc:
        print(f"  [REFUSED] {exc}")

    print(f"\n  Budget used: {total_budget - layer.budget_remaining:.4f} / {total_budget}")
    layer.print_log()

    print()
    print("Composition property: spending epsilon1 + epsilon2 + epsilon3 against the")
    print("same dataset gives (epsilon1+epsilon2+epsilon3)-DP combined. The query layer")
    print("enforces the budget total server-side; once exhausted, no further queries.")

    print()
    print("Authorisation reminder (from the book):")
    print("  - A DP query layer governs *what a statistic reveals*, not *who may ask*.")
    print("  - Put authentication and authorisation (sections 16.2-16.3) in front.")
    print("  - Enforce the per-analyst budget server-side; a budget the client can reset")
    print("    is no budget at all.")
    print("  - Differential privacy protects the output; you still own the front door.")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(description="Lab 5: differentially private query layer.")
    parser.add_argument("--rows",   type=int,   default=2000,  help="Synthetic dataset size (default: 2000)")
    parser.add_argument("--budget", type=float, default=1.0,   help="Total epsilon budget for the query layer demo (default: 1.0)")
    args = parser.parse_args()

    print(f"[+] Generating {args.rows} synthetic records (all invented values)...")
    records = make_synthetic_dataset(args.rows)
    incomes = [r["income"] for r in records]

    print("\n" + "=" * 70)
    print("NON-PRIVATE BASELINE (ground truth)")
    print("=" * 70)
    stats = baseline_stats(incomes)
    for k, v in stats.items():
        print(f"  {k:<16}: {v:,.4f}")

    sweep_epsilon(incomes, stats["true_mean"])
    demo_budget_composition(incomes, total_budget=args.budget)

    print("\n[+] Lab 5 complete.")


if __name__ == "__main__":
    main()
