# /// script
# requires-python = ">=3.10"
# dependencies = ["pandas", "numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Re-identify a 5-anonymous table and measure the uniqueness of quasi-identifiers.

See that removing names does not remove identity, and that composition breaks k-anonymity.
Two releases are generalized differently so that each is k-anonymous on its own (classes
smaller than k are suppressed); joining them by row id gives classes far smaller than k."""

import argparse
import pandas as pd
import numpy as np
import json

def k_of(df, cols):
    """Size of each person's equivalence class on these columns."""
    return df.groupby(cols)[cols[0]].transform("size")

def main():
    parser = argparse.ArgumentParser(description=__doc__.split('\n')[0])
    parser.add_argument("--people", type=int, default=100_000, help="Number of fake people to generate.")
    parser.add_argument("--k", type=int, default=5, help="Each release suppresses classes smaller than k.")
    parser.add_argument("--seed", type=int, default=1, 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)
    n = args.people

    print(f"--- 1. Generating Fake Population ({n} people) ---")
    df = pd.DataFrame({
        "dob": rng.integers(0, 29_200, n),
        "pc": rng.integers(0, 10_000, n),
        "sex": rng.integers(0, 2, n),
    })
    df["year"] = df.dob // 365          # 80 birth years
    df["decade"] = df.year // 10        # 8 birth decades
    df["pc2"] = df.pc // 100            # 100 postcode districts
    df["pc1"] = df.pc // 1000           # 10 postcode regions

    print("\n--- 2. Uniqueness on Raw Quasi-Identifiers ---")
    unique_raw = int((k_of(df, ["dob", "pc", "sex"]) == 1).sum())
    print(f"unique on (birth date, postcode, sex): {unique_raw} of {n}")

    print(f"\n--- 3. Two Releases, each {args.k}-anonymous ---")
    # Release A is coarse on age and fine on place; release B the other way round.
    releases = {"A (birth decade, district, sex)": ["decade", "pc2", "sex"],
                "B (birth year, region, sex)": ["year", "pc1", "sex"]}
    rows = [{"release": "Raw (birth date, postcode, sex)", "classes": int(df.groupby(["dob", "pc", "sex"]).ngroups),
             "suppressed": 0, "k": int(k_of(df, ["dob", "pc", "sex"]).min()), "unique_count": unique_raw, "below_k": int((k_of(df, ["dob", "pc", "sex"]) < args.k).sum())}]
    keep = pd.Series(True, index=df.index)
    for name, cols in releases.items():
        size = k_of(df, cols)
        kept = size >= args.k                      # suppress classes smaller than k
        keep &= kept
        k_rel = int(size[kept].min())
        rows.append({"release": name, "classes": int(df.groupby(cols).ngroups), "suppressed": int((~kept).sum()),
                     "k": k_rel, "unique_count": int((size[kept] == 1).sum()), "below_k": int((size[kept] < args.k).sum())})
        print(f"k({name}) = {k_rel}, suppressed {int((~kept).sum())} people")

    print("\n--- 4. Composition Attack (Joining Releases by row id) ---")
    both = df[keep]
    size_c = k_of(both, ["decade", "pc2", "year", "pc1", "sex"])
    k_c = int(size_c.min())
    unique_c = int((size_c == 1).sum())
    below_c = int((size_c < args.k).sum())
    rows.append({"release": "A joined with B", "classes": int(both.groupby(["year", "pc2", "sex"]).ngroups), "suppressed": int(n - len(both)),
                 "k": k_c, "unique_count": unique_c, "below_k": below_c})
    print(f"k(A and B) = {k_c}; unique: {unique_c}; in classes smaller than {args.k}: {below_c} of {len(both)}")

    results = {
        "summary": {
            "people": n, "k_target": args.k, "unique_raw": unique_raw,
            "k_A": rows[1]["k"], "k_B": rows[2]["k"], "k_intersection": k_c,
            "unique_intersection": unique_c, "below_k_intersection": below_c,
            "below_k_intersection_pct": 100 * below_c / len(both), "people_in_both": int(len(both)),
        },
        "tables": {"releases": rows},
    }

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

if __name__ == "__main__":
    main()