import json, math
from pathlib import Path
import numpy as np
from scipy.integrate import quad
from scipy.special import gammaln

EPS = 1.0
DELTA_PRIV = 1e-5


def log_norm(p, b):
    return math.log(p) - math.log(2.0) - math.log(b) - gammaln(1.0 / p)


def log_density(x, p, b):
    return log_norm(p, b) - (abs(x) / b) ** p


def hockey_divergence(p, b, sensitivity, epsilon=EPS):
    # Directly integrates [f(x)-exp(eps) f(x-Delta)]+.  Infinite limits
    # are important for p>1, whose likelihood ratio has an unbounded tail.
    d = float(sensitivity)
    ln = log_norm(p, b)
    ee = float(epsilon)

    def integrand(x):
        a = -(abs(x) / b) ** p
        c = ee - (abs(x - d) / b) ** p
        # common normalization factored out; avoids overflow in exp(eps)
        z = math.exp(ln + a) - math.exp(ln + c) if max(ln+a, ln+c) < 700 else 0.0
        return max(0.0, z)

    val, err = quad(integrand, -np.inf, np.inf, epsabs=2e-11,
                    epsrel=2e-8, limit=250, points=None)
    return float(val)


def calibrate_b(p, sensitivity=1.0, epsilon=EPS, delta=DELTA_PRIV):
    # Divergence decreases monotonically with b. Find a conservative bracket.
    lo, hi = 1e-8 * sensitivity, 1.0 * sensitivity
    while hockey_divergence(p, hi, sensitivity, epsilon) > delta:
        hi *= 2.0
        if hi > 1e7 * max(1.0, sensitivity):
            raise RuntimeError("failed to bracket privacy scale")
    for _ in range(55):
        mid = (lo + hi) / 2.0
        if hockey_divergence(p, mid, sensitivity, epsilon) <= delta:
            hi = mid
        else:
            lo = mid
    return hi


def moment(p, b, m=2):
    return math.exp(m * math.log(b) + gammaln((m + 1.0) / p) - gammaln(1.0 / p))


def select_shape(sensitivity=1.0, epsilon=EPS, delta=DELTA_PRIV,
                 grid=None, m=2):
    if grid is None:
        grid = [1.0, 1.25, 1.5, 1.75, 2.0, 2.5, 3.0, 4.0, 6.0, 8.0]
    rows = []
    for p in grid:
        b = calibrate_b(p, sensitivity, epsilon, delta)
        rows.append({"p": p, "b": b, "divergence": hockey_divergence(p, b, sensitivity, epsilon),
                     "moment": moment(p, b, m)})
    best = min(rows, key=lambda r: r["moment"])
    return best, rows


def sample_noise(rng, p, b, size):
    # If T=(|Z|/b)^p, then T ~ Gamma(1/p, 1).
    # The exponential inverse-CDF formula is the special case p=1 only.
    radius = rng.gamma(shape=1.0 / p, scale=1.0, size=size) ** (1.0 / p)
    return b * rng.choice(np.array([-1.0, 1.0]), size=size) * radius


def run_mean_experiment(p_configs, seed=123, n_trials=300, steps=80,
                        n=64, lr=0.25, clip=1.0, true_mu=1.0):
    # Scalar clipped per-example mean-gradient query, with sensitivity 2C/n.
    # Each trial has a fixed dataset; noise is freshly sampled per update.
    out = {name: [] for name in p_configs}
    rng = np.random.default_rng(seed)
    for _ in range(n_trials):
        data = rng.normal(true_mu, 0.7, n)
        for name, cfg in p_configs.items():
            theta = 0.0
            p, b = cfg
            for _step in range(steps):
                per = np.clip(theta - data, -clip, clip)
                g = float(np.mean(per))
                # b was calibrated for scalar sum-query sensitivity convention;
                # this experiment uses the same specified sensitivity in all arms.
                theta -= lr * (g + float(sample_noise(rng, p, b, 1)[0]))
            out[name].append((theta - true_mu) ** 2)
    return {k: {"mse": float(np.mean(v)), "se": float(np.std(v, ddof=1)/math.sqrt(len(v)))}
            for k, v in out.items()}


def main():
    grid = [1.0, 1.25, 1.5, 1.75, 2.0, 2.5, 3.0, 4.0, 6.0, 8.0]
    best, rows = select_shape(grid=grid)
    # Scaling check: normalized geometry should preserve p and scale b.
    scaling = []
    for factor in (0.5, 1.0, 2.0):
        b, rs = select_shape(sensitivity=factor, grid=grid)
        scaling.append({"factor": factor, "best_p": b["p"], "best_b": b["b"],
                        "b_over_factor": b["b"] / factor})
    lap = next(r for r in rows if r["p"] == 1.0)
    gau = next(r for r in rows if r["p"] == 2.0)
    configs = {"laplace_p1": (1.0, lap["b"]), "gaussian_p2": (2.0, gau["b"]),
               "optimized": (best["p"], best["b"])}
    task = run_mean_experiment(configs)
    result = {"settings": {"epsilon": EPS, "delta": DELTA_PRIV, "sensitivity": 1.0,
                            "grid": grid}, "selected": best, "rows": rows,
              "scaling": scaling, "task": task}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()