# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Differencing attack and the Laplace mechanism.

Reconstruct one person's value from exact aggregates, then see noise stop it and measure the error."""

import argparse
import json
import numpy as np

def main():
    parser = argparse.ArgumentParser(description=__doc__.split('\n')[0])
    parser.add_argument("--rows", type=int, default=1000, help="Number of salaries in the dataset.")
    parser.add_argument("--eps", type=float, default=0.5, help="Epsilon value for differential privacy.")
    parser.add_argument("--cap", type=int, default=10000, help="Clamp cap for salaries.")
    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": {"laplace_errors": []}}

    print(f"--- 1. Differencing Attack on Exact Aggregates ({args.rows} rows) ---")
    salaries = rng.lognormal(mean=10, sigma=1, size=args.rows)
    
    target_index = 42
    target_salary = salaries[target_index]
    
    sum_all = np.sum(salaries)
    mask = np.ones(args.rows, dtype=bool)
    mask[target_index] = False
    sum_except_target = np.sum(salaries[mask])
    
    differenced_salary = sum_all - sum_except_target
    exact_recovery = bool(abs(target_salary - differenced_salary) < 1e-5)
    
    print(f"Target's actual salary: {target_salary:.2f}")
    print(f"Recovered salary via differencing: {differenced_salary:.2f}")
    print(f"Exact recovery successful: {exact_recovery}")

    results["summary"]["target_salary"] = float(target_salary)
    results["summary"]["differenced_salary"] = float(differenced_salary)
    results["summary"]["exact_recovery"] = exact_recovery

    print(f"\n--- 2. The Laplace Mechanism on Counts (eps = {args.eps}) ---")
    print(f"{'True Count':>12} | {'Mean Abs Error':>15} | {'Relative Error':>15}")
    print("-" * 50)
    
    for true_val in (5, 100, 1_000_000):
        # Add Laplace noise with scale 1/epsilon for 1000 runs
        noisy = true_val + rng.laplace(0, 1/args.eps, 1000)
        err = np.abs(noisy - true_val).mean()
        rel_err = 100 * err / true_val
        
        print(f"{true_val:12} | {round(err, 2):15} | {rel_err:14.5f}%")
        results["tables"]["laplace_errors"].append({
            "true_count": true_val,
            "mean_abs_error": float(err),
            "relative_error_pct": float(rel_err)
        })

    print(f"\n--- 3. Clamping and Noisy SUMs (cap = {args.cap}) ---")
    clamped_salaries = np.clip(salaries, a_min=0, a_max=args.cap)
    
    true_sum = np.sum(salaries)
    clamped_sum = np.sum(clamped_salaries)
    
    noisy_clamped_sum = clamped_sum + rng.laplace(0, args.cap/args.eps)
    distortion = abs(true_sum - clamped_sum)
    distortion_pct = 100 * distortion / true_sum
    
    print(f"True SUM (skewed dataset): {true_sum:,.2f}")
    print(f"Clamped SUM:               {clamped_sum:,.2f}")
    print(f"Noisy Clamped SUM:         {noisy_clamped_sum:,.2f}")
    print(f"Distortion from clamping:  {distortion:,.2f} ({distortion_pct:.2f}%)")

    results["summary"]["true_sum"] = float(true_sum)
    results["summary"]["clamped_sum"] = float(clamped_sum)
    results["summary"]["noisy_clamped_sum"] = float(noisy_clamped_sum)
    results["summary"]["distortion"] = float(distortion)
    results["summary"]["distortion_pct"] = float(distortion_pct)

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

if __name__ == "__main__":
    main()