# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""A privacy budget ledger.

Feel that a privacy budget is depleted by reading by simulating queries against a fixed budget."""

import argparse
import json
import numpy as np

class PrivacyLedger:
    def __init__(self, initial_budget, rng):
        self.budget = initial_budget
        self.rng = rng

    def dp_count(self, true_count, eps):
        # Check if budget is depleted (with small tolerance for float math)
        if self.budget - eps < -1e-9:
            return "REFUSED: Privacy budget depleted"
        
        self.budget -= eps
        noisy_count = true_count + self.rng.laplace(0, 1/eps)
        return noisy_count

def main():
    parser = argparse.ArgumentParser(description=__doc__.split('\n')[0])
    parser.add_argument("--budget", type=float, default=1.0, help="Initial global privacy budget.")
    parser.add_argument("--eps", type=float, default=0.1, help="Privacy budget consumed per query.")
    parser.add_argument("--queries", type=int, default=15, help="Number of exploratory queries to run.")
    parser.add_argument("--repeats", type=int, default=10, help="Number of repeated queries to average.")
    parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility.")
    parser.add_argument("--json", type=str, help="Also write the results to this JSON file.")
    args = parser.parse_args()

    rng = np.random.default_rng(args.seed)
    results = {"summary": {}, "tables": {"exploratory_queries": [], "repeated_queries": []}}

    print(f"--- 1. Exploratory Queries (eps = {args.eps}, budget = {args.budget}) ---")
    ledger1 = PrivacyLedger(args.budget, rng)
    true_answer = 100
    answered = 0
    refused = 0
    
    for i in range(1, args.queries + 1):
        res = ledger1.dp_count(true_answer, args.eps)
        if isinstance(res, str) and res.startswith("REFUSED"):
            refused += 1
            status = "refused"
            val = None
            print(f"Query {i:2d}: {res} (Budget remaining: {max(0, ledger1.budget):.2f})")
        else:
            answered += 1
            status = "answered"
            val = float(res)
            print(f"Query {i:2d}: Noisy Result = {val:.2f} (Budget remaining: {ledger1.budget:.2f})")
            
        results["tables"]["exploratory_queries"].append({
            "query_id": i,
            "status": status,
            "result": val,
            "budget_remaining": float(max(0, ledger1.budget))
        })

    print(f"\nSummary: {answered} queries answered, {refused} refused.")
    results["summary"]["answered_count"] = answered
    results["summary"]["refused_count"] = refused

    print(f"\n--- 2. Averaging Repeated Queries (eps = {args.eps}, repeats = {args.repeats}) ---")
    # Reset ledger for the averaging experiment
    ledger2 = PrivacyLedger(args.budget, rng)
    true_answer_2 = 500
    repeated_results = []
    
    for i in range(1, args.repeats + 1):
        res = ledger2.dp_count(true_answer_2, args.eps)
        if not isinstance(res, str):
            repeated_results.append(float(res))
            results["tables"]["repeated_queries"].append({
                "query_id": i,
                "result": float(res)
            })

    if repeated_results:
        average_result = np.mean(repeated_results)
        print(f"True Answer: {true_answer_2}")
        print(f"Individual noisy results: {[round(x, 1) for x in repeated_results]}")
        print(f"Averaged result ({len(repeated_results)} queries): {average_result:.2f}")
        print(f"Budget remaining: {ledger2.budget:.2f}")
        
        results["summary"]["true_answer_repeated"] = true_answer_2
        results["summary"]["average_repeated_result"] = float(average_result)
    else:
        print("Not enough budget to run repeated queries.")

    if args.json:
        with open(args.json, "w") as f:
            json.dump(results, f, indent=2)

if __name__ == "__main__":
    main()