#!/usr/bin/env python3
"""Warner's randomized response (1965): honest statistics from deniable answers.

Each respondent privately flips a biased coin.  With probability p they answer
"Is it true that you are in group A?", otherwise they answer the complement
question "Is it true that you are NOT in group A?".  The analyst only sees
yes/no, never which question was answered.  If pi is the true fraction in A,

    lambda = P(yes) = p pi + (1 - p)(1 - pi)
    pi_hat = (lambda_hat - (1 - p)) / (2p - 1)                (unbiased, p != 1/2)
    Var(pi_hat) = pi(1-pi)/n + p(1-p) / (n (2p-1)^2)

The second term is the price of privacy.  Each answer is eps-locally
differentially private with eps = ln(p / (1-p)): any single yes/no is at most
e^eps times more likely under one truth than the other.

Usage
  python randomized_response.py [--quick]          simulation + plot
  python randomized_response.py make-example out/example_poll.csv --n 800 --pi 0.3 --p 0.75
  python randomized_response.py poll out/example_poll.csv --p 0.75 [--column answer] [--conf 0.95]
"""
from __future__ import annotations

import argparse
import csv
import os
import sys

import numpy as np
from scipy import stats

OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "out")

YES = {"1", "yes", "y", "true", "t"}
NO = {"0", "no", "n", "false", "f"}


def respond(truth: np.ndarray, p: float, rng) -> np.ndarray:
    """Randomized answers (bool) for true memberships `truth` (bool)."""
    asks_direct = rng.random(truth.shape) < p
    return np.where(asks_direct, truth, ~truth)


def estimate(n_yes, n, p):
    """Unbiased Warner estimator and its plug-in standard error."""
    lam = n_yes / n
    pi_hat = (lam - (1 - p)) / (2 * p - 1)
    se = np.sqrt(lam * (1 - lam) / n) / abs(2 * p - 1)
    return pi_hat, se


def theoretical_sd(pi, n, p):
    return np.sqrt(pi * (1 - pi) / n + p * (1 - p) / (n * (2 * p - 1) ** 2))


def epsilon(p):
    q = max(p, 1 - p)
    return float(np.log(q / (1 - q)))


def confidence_interval(n_yes, n, p, conf=0.95):
    """Exact (Clopper-Pearson) interval for lambda = P(yes), mapped through the
    linear Warner transform and clipped to [0, 1].  Conservative but valid."""
    a = 1 - conf
    lo = stats.beta.ppf(a / 2, n_yes, n - n_yes + 1) if n_yes > 0 else 0.0
    hi = stats.beta.ppf(1 - a / 2, n_yes + 1, n - n_yes) if n_yes < n else 1.0
    ends = sorted(((lo - (1 - p)) / (2 * p - 1), (hi - (1 - p)) / (2 * p - 1)))
    return max(0.0, ends[0]), min(1.0, ends[1])


def simulate_spread(pi, ns, p, reps=2000, seed=0):
    """Empirical mean and sd of pi_hat over reps independent surveys per n."""
    rng = np.random.default_rng(seed)
    means, sds = [], []
    for n in ns:
        # number of yes answers is Binomial(n, lambda); sample it directly
        lam = p * pi + (1 - p) * (1 - pi)
        est, _ = estimate(rng.binomial(n, lam, size=reps), n, p)
        means.append(est.mean())
        sds.append(est.std())
    return np.array(means), np.array(sds)


# --------------------------------------------------------------------------
# CLI: poll analysis
# --------------------------------------------------------------------------


def read_answers(path, column=None):
    with open(path, newline="") as fh:
        rows = list(csv.reader(fh))
    if not rows:
        raise ValueError("empty file")
    header = [h.strip().lower() for h in rows[0]]
    has_header = not all(c.strip().lower() in YES | NO for c in rows[0])
    col = 0
    if has_header:
        if column is not None:
            col = header.index(column.lower())
        rows = rows[1:]
    elif column is not None:
        raise ValueError("--column given but the file has no header row")
    vals, skipped = [], 0
    for r in rows:
        v = r[col].strip().lower() if len(r) > col else ""
        if v in YES:
            vals.append(1)
        elif v in NO:
            vals.append(0)
        else:
            skipped += 1
    return np.array(vals, dtype=int), skipped


def poll(args):
    ans, skipped = read_answers(args.csv, args.column)
    n, n_yes = len(ans), int(ans.sum())
    if n == 0:
        print("no usable answers")
        return 1
    pi_hat, se = estimate(n_yes, n, args.p)
    lo, hi = confidence_interval(n_yes, n, args.p, args.conf)
    z = stats.norm.ppf(0.5 + args.conf / 2)
    print(f"answers used: {n}  (yes {n_yes}, no {n - n_yes}; skipped {skipped} unreadable rows)")
    print(f"design: truthful-question probability p = {args.p}  ->  local DP epsilon = {epsilon(args.p):.3f}")
    print(f"raw 'yes' rate:            {n_yes / n:.4f}  (NOT the answer: it is mixed with the complement question)")
    print(f"debiased estimate pi_hat:  {pi_hat:.4f}   (unbiased; may fall outside [0,1] by chance)")
    print(f"{args.conf:.0%} CI, exact (Clopper-Pearson mapped): [{lo:.4f}, {hi:.4f}]")
    print(f"{args.conf:.0%} CI, normal approx:                  [{max(0, pi_hat - z * se):.4f}, {min(1, pi_hat + z * se):.4f}]")
    # at pi = 1/2: Var_RR = 1/(4 n (2p-1)^2) = Var_honest(n') with n' = n (2p-1)^2
    honest_n = n * (2 * args.p - 1) ** 2
    print(f"privacy cost: this survey is about as precise as ~{honest_n:.0f} direct answers (at pi ~ 1/2)")
    return 0


def make_example(args):
    rng = np.random.default_rng(args.seed)
    truth = rng.random(args.n) < args.pi
    answers = respond(truth, args.p, rng)
    os.makedirs(os.path.dirname(os.path.abspath(args.csv)), exist_ok=True)
    with open(args.csv, "w", newline="") as fh:
        w = csv.writer(fh)
        w.writerow(["respondent", "answer"])
        for i, a in enumerate(answers):
            w.writerow([i, "yes" if a else "no"])
    print(f"wrote {args.n} randomized answers to {args.csv} (true pi = {args.pi}, true sample "
          f"fraction = {truth.mean():.4f}, p = {args.p})")
    return 0


# --------------------------------------------------------------------------
# demo
# --------------------------------------------------------------------------


def demo(quick=False):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    os.makedirs(OUT, exist_ok=True)
    pi = 0.3
    ns = np.array([100, 200, 500, 1000, 2000, 5000, 10000])
    reps = 1000 if quick else 5000
    designs = [0.6, 0.75, 0.9]
    print(f"true pi = {pi}; {reps} simulated surveys per n")
    print(f"{'design':<20}{'eps':>6}{'n':>7}{'mean est':>10}{'sd (sim)':>10}{'sd (theory)':>12}{'n-inflation':>12}")
    fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))
    honest_sd = np.sqrt(pi * (1 - pi) / ns)
    axes[0].loglog(ns, honest_sd, "k-o", ms=3, label="honest answers (eps = ∞)")
    for p in designs:
        means, sds = simulate_spread(pi, ns, p, reps, seed=int(100 * p))
        infl = theoretical_sd(pi, 1, p) ** 2 / (pi * (1 - pi))
        for n, m, s in zip(ns, means, sds):
            if n in (100, 1000, 10000):
                print(f"{'Warner p=' + str(p):<20}{epsilon(p):>6.2f}{n:>7}{m:>10.4f}{s:>10.4f}"
                      f"{theoretical_sd(pi, n, p):>12.4f}{infl:>11.1f}x")
        line, = axes[0].loglog(ns, sds, "o", ms=4, label=f"p={p} (ε={epsilon(p):.2f}) simulated")
        axes[0].loglog(ns, theoretical_sd(pi, ns, p), "-", color=line.get_color(), lw=1)
    axes[0].set_xlabel("respondents n")
    axes[0].set_ylabel("sd of estimate of π")
    axes[0].set_title("Precision vs privacy (dots: simulation, lines: formula)")
    axes[0].legend(fontsize=8)
    axes[0].grid(alpha=0.3, which="both")

    ps = np.linspace(0.52, 0.99, 200)
    infl = np.array([theoretical_sd(pi, 1, p) ** 2 / (pi * (1 - pi)) for p in ps])
    eps = np.log(ps / (1 - ps))
    axes[1].semilogy(eps, infl)
    axes[1].set_xlabel("privacy loss ε = ln(p/(1-p))  (smaller = more private)")
    axes[1].set_ylabel("respondents needed / honest survey")
    axes[1].set_title(f"Sample-size price of deniability (π = {pi})")
    axes[1].grid(alpha=0.3, which="both")
    fig.tight_layout()
    path = os.path.join(OUT, "randomized_response.png")
    fig.savefig(path, dpi=130)
    plt.close(fig)
    print(f"saved {path}")


def main(argv=None):
    argv = sys.argv[1:] if argv is None else argv
    if argv and argv[0] in ("poll", "make-example"):
        ap = argparse.ArgumentParser(prog=f"randomized_response.py {argv[0]}")
        ap.add_argument("csv")
        ap.add_argument("--p", type=float, default=0.75,
                        help="probability the respondent answered the direct question")
        if argv[0] == "poll":
            ap.add_argument("--column", default=None, help="column name holding yes/no answers")
            ap.add_argument("--conf", type=float, default=0.95)
        else:
            ap.add_argument("--n", type=int, default=800)
            ap.add_argument("--pi", type=float, default=0.3)
            ap.add_argument("--seed", type=int, default=0)
        args = ap.parse_args(argv[1:])
        if abs(args.p - 0.5) < 1e-9:
            ap.error("p = 0.5 carries no information")
        return poll(args) if argv[0] == "poll" else make_example(args)
    ap = argparse.ArgumentParser(description="randomized response simulation")
    ap.add_argument("--quick", action="store_true")
    demo(ap.parse_args(argv).quick)
    return 0


if __name__ == "__main__":
    sys.exit(main())