experiment_utils.py

import hashlib
from enum import Enum, auto
from typing import Callable
import numpy as np
import scipy.stats as st
import os, json
from dp_accounting.pld import privacy_loss_distribution
from ks_distance_functions import ks_func_and_dist
from logistic_regression import private_logreg_subset_estimator, make_logreg_scalar_estimator
from plot_functions import plot_interval_metrics, plot_interval_metrics_relative_width, plot_quantile_histograms

ks_distance, KS_FIXED_DIST = ks_func_and_dist('uniform')

from Moshes_experiments.implementations.optimized_quantile_inverse_sensitivity import vectorized_quantile_combined
from private_mechanism import priv_mean_mechanism

BINARY_SEARCH_TOLERANCE = 1e-6
REASONABLE_NUMBER_OF_ITERATION_TO_CHECK = 400
MINIMAL_NUMBER_OF_SUBSETS = 120
NUM_BOOTSTRAPS = 60


####################
#### Enums #########
####################
class ResampleEstimators(Enum):
    MED_INVERSE_SENSITIVITY = auto()
    NONE_PRIVATE_MED = auto()
    NONE_PRIVATE_MEAN = auto()
    NONE_PRIVATE_KS_DIST = auto()
    PRIVATE_MEAN = auto()
    EXP_MECH = auto()
    KS_DIST = auto()
    BLB_MED = auto()
    BLB_MEAN = auto()
    BLB_KS_DIST = auto()
    EMPIRICAL = auto()
    EMPIRICAL_NOISE_ADDITION_MEAN = auto()
    NONE_PRIVATE_LOGREG = auto()
    PRIVATE_LOGREG_OUTPUT_PERT = auto()
    PRIVATE_LOGREG_OBJ_PERT = auto()

class NoiseType(Enum):
    LAPLACE = auto()
    GAUSSIAN = auto()

def default_scaling(n, m): return np.sqrt(m / n)
# def default_scaling(n, m): return np.sqrt(m / (n-m))
####################
### mT func ########
####################

def mT_2(n, params, eps=None, T=100):
    subsample_size = max(50, min(5*np.sqrt(n), 0.1* n))
    return int(subsample_size), int(T)

def mT_mean(n, params, eps=None,  T=200):
    T = max(T, np.log(n)*35)
    subsample_size = max(50, min(5*np.sqrt(n), 0.1* n))
    return int(subsample_size), int(T)

def mT_two_thirds(n, params, eps=None):
    subsample_size = int(n)**(2/3)
    T = max(50, int(n)**(2/3))
    return int(subsample_size), int(T)


def mT_sqrt(n): return int(np.sqrt(n)), int(np.sqrt(n))
def mT_split_min_T50(n): return (int(n // MINIMAL_NUMBER_OF_SUBSETS), MINIMAL_NUMBER_OF_SUBSETS)
def choose_mT_2_3(n): return int(n ** (2 / 3)), int(n ** (2 / 3))
def mT_for_bootstrap_m_out_of_n(n):
    m = n
    T = max(min(5 * np.sqrt(n), 400), 200)
    return int(m), int(T)

def mT_bootstrap(n): return int(n), int(NUM_BOOTSTRAPS)

def scaling(n, m, with_replacement, eps=1):
    if with_replacement:return np.sqrt(m / n)
    return np.sqrt(m / (n-m))

def true_value_for_stats(dist, stats, n_values, quant):
    if stats == 'median': return  {n: dist.median() for n in n_values}
    elif stats == 'mean': return  {n: dist.mean() for n in n_values}
    elif stats == 'ks':  return {n: np.sqrt(np.pi/2) * np.log(2)/np.sqrt(n) for n in n_values}
    elif stats == 'quant': return {n: dist.ppf(quant) for n in n_values}
    else: raise ValueError(f"Unsupported stats: {stats}")

def build_priv_nonpriv(stats):
    if stats == 'median' or stats == 'quant':
        non_priv_mech = ResampleEstimators.NONE_PRIVATE_MED
        priv_mech = ResampleEstimators.MED_INVERSE_SENSITIVITY

    elif stats == 'mean':
        non_priv_mech = ResampleEstimators.NONE_PRIVATE_MEAN
        priv_mech = ResampleEstimators.PRIVATE_MEAN

    elif stats =='ks':
        non_priv_mech = ResampleEstimators.NONE_PRIVATE_KS_DIST
        priv_mech = ResampleEstimators.KS_DIST

    elif stats == 'logreg out':
        non_priv_mech = ResampleEstimators.NONE_PRIVATE_LOGREG
        priv_mech = ResampleEstimators.PRIVATE_LOGREG_OUTPUT_PERT

    elif stats == 'logreg obj':
        non_priv_mech = ResampleEstimators.NONE_PRIVATE_LOGREG
        priv_mech = ResampleEstimators.PRIVATE_LOGREG_OBJ_PERT

    else:
        raise ValueError(f"Unsupported stats: {stats}")
    return non_priv_mech, priv_mech

####################
### CI Functions ###
####################
def calc_new_normal_ci(significance, estimated_std, cmp):
    z_score = st.norm.ppf(1 - significance / 2)
    L_new = cmp - z_score * estimated_std
    U_new = cmp + z_score * estimated_std
    return np.array([L_new, U_new])

def CI_normal_estimation(estimations, significance, scaling: float):
    mean = np.mean(estimations)
    estimated_std = np.std(estimations, ddof=1) * scaling
    return calc_new_normal_ci(significance, estimated_std, mean)


def CI_normal_with_n_estimation(estimations, significance, scaling: float, n_est_priv):
    estimated_std = np.sqrt(1/(len(estimations)-1) * np.sum((estimations - n_est_priv)**2)) * scaling
    return calc_new_normal_ci(significance, estimated_std, n_est_priv)


def CI_quantile_estimation(estimations, significance: float, scaling: float) -> np.ndarray:
    mid = np.mean(estimations)
    return calc_quantile_ci(estimations, significance, scaling, mid)

def empirical_var_ci(distribution, significance, n, num_experiments):
    large_sample = distribution.rvs(size=(2000, n))
    estimates = np.median(large_sample, axis=1)
    lower_quantile = np.quantile(estimates, significance / 2, method="lower")
    upper_quantile = np.quantile(estimates, 1 - significance / 2, method="higher")
    quantile_approx = np.array([[lower_quantile, upper_quantile]] * num_experiments)
    return quantile_approx

def CI_width(intervals):
    if len(intervals) == 0:
        return 0.0, 0.0
    width_arr = intervals[:, 1] - intervals[:, 0]
    return np.mean(width_arr), np.std(width_arr) / np.sqrt(intervals.shape[0])


def coverage_probability(intervals, true_value: float):
    if len(intervals) == 0:
        return 0.0, 0.0
    coverage = (intervals[:, 0] <= true_value) * (intervals[:, 1] >= true_value)
    p = np.mean(coverage)
    return p, np.sqrt(p * (1 - p) / intervals.shape[0])


def calc_quantile_ci(estimations,significance, scaling, mid_pnt):
    L = np.quantile(estimations, significance / 2, method='lower')
    U = np.quantile(estimations, 1 - (significance / 2), method='higher')
    L_new = mid_pnt - scaling * (mid_pnt - L)
    U_new = mid_pnt + scaling * (U - mid_pnt)
    return np.array([L_new, U_new])


def ks_approx(center, est, params):
    scaling, significance, z, T= params.scaling_factor, params.significance, params.z, params.T

    diffs = est - center[:, None]  # (B,T)
    sd = np.sqrt(np.sum(diffs * diffs, axis=1) / (T - 1)) * scaling  # with cent var as well
    Lb = center - z * sd
    Ub = center + z * sd
    normal_approx = np.column_stack([Lb, Ub])

    Lq = np.quantile(est, significance / 2, axis=1, method="lower")
    Uq = np.quantile(est, 1 - significance / 2, axis=1, method="higher")

    c_m = est.mean(axis=1)
    Lq = Lq - (c_m - center)
    Uq = Uq - (c_m - center)
    Lb = center - scaling * (center - Lq)
    Ub = center + scaling * (Uq - center)
    quantile_approx = np.column_stack([Lb, Ub])
    return normal_approx, quantile_approx


def debias_approx(center, est, params):
    """ former version (debiasing)"""
    """noise addition debiasing: dependes on the mechanism (noise addition of what kind for noise
     - scale like the variance, that is different between lap and N, for example)"""
    T, z, significance = params.T, params.z, params.significance
    scaling_factor, debias_func = params.scaling_factor, params.debias_func

    # diffs = est - center[:, None]  # (B,T)
    diffs = est - est.mean(axis=1, keepdims=True)  # (B,T)
    var_hat = np.sum(diffs * diffs, axis=1) / (T - 1)

    var_hat = debias_func(var_hat, params) if debias_func else var_hat*(scaling_factor**2) # depends on the mechanism
    sd = np.maximum(var_hat, 0)  # avoid negative variance after debias
    sd = np.sqrt(sd)

    Lb = center - z * sd
    Ub = center + z * sd
    normal_approx = np.column_stack([Lb, Ub])
    # normal_approx =[]
    #########

    #--- quantile approx ---
    # c=center
    # # center =  4.355794424605847
    # Lq = np.quantile(est, significance / 2, axis=1, method="lower") # og
    # Uq = np.quantile(est, 1 - significance / 2, axis=1, method="higher") # og
    # Lb = center - scaling_factor * (center - Lq) # og
    # Ub = center + scaling_factor * (Uq - center) #og

    # plot_quantile_histograms(Lq, Uq, Ub, Lb, c, begin_with='OG')

    # Lq = np.quantile(est, significance / 2, axis=1, method="lower")
    # Uq = np.quantile(est, 1 - significance / 2, axis=1, method="higher")
    # Lb = center + scaling_factor * (center - Uq)
    # Ub = center - scaling_factor * (Lq - center)

    quantile_approx = np.column_stack([Lb, Ub])
    return normal_approx, quantile_approx

def new_approx_method(center, est, params):
    T, z, significance = params.T, params.z, params.significance
    scaling_factor, debias_func = params.scaling_factor, params.debias_func

    Lq = np.quantile(est, significance / 2, axis=1, method="lower")
    Uq = np.quantile(est, 1 - significance / 2, axis=1, method="higher")
    Lb = center + scaling_factor * (center - Uq)
    Ub = center - scaling_factor * (Lq - center)

    # plot_quantile_histograms(Lq, Uq, Ub, Lb, center, begin_with='new')

    quantile_approx = np.column_stack([Lb, Ub])
    return [], quantile_approx

def approx_quant_wo_center(center, est, params):
    sig = params.significance
    Lq = np.quantile(est, sig / 2, axis=1, method="lower")  # og
    Uq = np.quantile(est, 1 - sig / 2, axis=1, method="higher") # og
    return [], np.column_stack([Lq, Uq])

def mechanism(samples, lcl_est, center_est, params, sort=False):
    """smallest component of the mechanism - given samples (B,n), return intervals (B,2)"""
    n, m, T = params.n, params.m, params.T
    subsets = params.resampling_fn(data=samples, n=n, m=m, T=T, sort=sort)  # create a subset for bootstrap / subsample / data-splitting
    approx_method = params.approx_func

    est = lcl_est(subsets)  # (B, T, m) -> (B, T)
    center = center_est(samples) if center_est else est.mean(axis=1)  # (B,)

    normal_approx, quantile_approx = approx_method(center, est, params)
    # normal_approx, quantile_approx = debias_approx(center, est, params)
    # normal_approx, quantile_approx = ks_approx(center, est, params)

    return normal_approx, quantile_approx

############################
### Resampling Functions ###
############################

def _stable_argsort(keys):
    return np.argsort(keys, kind="mergesort")

def _stable_perm(n, tag):
    keys = np.fromiter(
        (int.from_bytes(hashlib.sha256(f"{tag}|{i}".encode()).digest()[:8], "big")
         for i in range(n)),
        dtype=np.uint64, count=n)
    return _stable_argsort(keys)

def _stable_group_order(T, tag):
    keys = np.fromiter(
        (int.from_bytes(hashlib.sha256(f"{tag}|{g}".encode()).digest()[:8], "big")
         for g in range(T)),
        dtype=np.uint64, count=T)
    return _stable_argsort(keys)

def permute_and_split_rows_numpy_nd(X: np.ndarray, T: int, *, shared_extra_groups: bool = False,
       shared_permutation: bool = False, return_indices: bool = True, pad_to_array=False):
    """
    X: shape (R, n). (Values are ignored if return_indices=True.)
    Returns: list over rows -> list of T numpy arrays of indices (ragged).
    Deterministic; no RNG.
    """
    if T <= 0:
        raise ValueError("T must be positive")
    R, n = X.shape
    base_perm = _stable_perm(n, "perm|base") if shared_permutation else None
    base_order = _stable_group_order(T, "extra|base") if shared_extra_groups else None

    out = []
    for r in range(R):
        perm = base_perm if base_perm is not None else _stable_perm(n, f"perm|row={r}")
        q, rem = divmod(n, T)
        sizes = np.full(T, q, dtype=int)
        order = base_order if base_order is not None else _stable_group_order(T, f"extra|row={r}")
        sizes[order[:rem]] += 1
        cuts = sizes.cumsum()[:-1]
        idx_groups = np.split(perm, cuts)
        out.append(idx_groups if return_indices else [X[r, g] for g in idx_groups])
    if not pad_to_array:
        return out

def small_subsampling_indices(sample_size: int, num_subsamples: int, subsample_size: int, sort:bool=True) -> np.ndarray:
    # print(num_subsamples, type(num_subsamples), subsample_size, type(subsample_size))
    assert subsample_size <= sample_size and subsample_size > 0
    assert isinstance(subsample_size, int) and isinstance(num_subsamples, int)

    random_keys = np.random.random((num_subsamples, sample_size))
    sorted_indices = np.argsort(random_keys, axis=1)
    subsamples = sorted_indices[:, :subsample_size]
    if sort:
        subsamples.sort(axis=1)

    return subsamples

def bootstrap_subsets(data, n, m, T, sort:bool=True):
    subsets_idx = bootstrap_indices(T, data, m, n, sort)
    return np.take_along_axis(data[:, None, :], subsets_idx, axis=2)


def bootstrap_indices(T, data, m, n, sort):
    B = data.shape[0]
    subsets_idx = np.random.randint(0, n, size=(B, T, m))
    if sort:
        subsets_idx.sort(axis=2)
    return subsets_idx


def subsample_subsets(data, n, m, T, sort:bool=True):
    subsets_idx = subsample_indices(T, data, m, n, sort)
    return np.take_along_axis(data[:, None, :], subsets_idx, axis=2)


def subsample_indices(T, data, m, n, sort):
    B = data.shape[0]
    subsets_idx = np.empty((B, T, m), dtype=int)
    for b in range(B):
        subsets_idx[b] = small_subsampling_indices(n, T, m, sort=sort)  # probably can be way more efficient
    return subsets_idx


def data_splitting_subsets(data, n, m, T, sort=True):
    groups_per_row = permute_and_split_rows_numpy_nd(data, T, return_indices=False)
    if not sort:
        return groups_per_row
    return [[np.sort(g) for g in row_groups] for row_groups in groups_per_row]  # NEW: sort each subset

def wrap_logreg_estimator_to_3d(X_full, y_full, est_on_index_set):
    def est3d(subset_indices):
        if subset_indices.ndim == 2:
            T, m = subset_indices.shape
            out = np.empty(T, dtype=float)
            for t in range(T):
                idx = subset_indices[t]
                out[t] = est_on_index_set(X_full[idx], y_full[idx])
            return out

        B, T, m = subset_indices.shape
        out = np.empty((B, T), dtype=float)
        for b in range(B):
            for t in range(T):
                idx = subset_indices[b, t]
                out[b, t] = est_on_index_set(X_full[b, idx], y_full[b, idx])
        return out
    return est3d

def _wrap_estimator_to_3d(est2d: Callable[[np.ndarray], np.ndarray]) -> Callable[[np.ndarray], np.ndarray]:
    """Lift an estimator mapping (T, m)->(T) to handle (B, T, m)->(B, T) without Python loops."""
    def est3d(subsets) -> np.ndarray:
        if type(subsets) is list:
            B = len(subsets)
            T = len(subsets[0])
            est = np.empty((B, T), dtype=float)
            for b in range(len(subsets)):
                for t in range(T):
                    est[b, t] = est2d(subsets[b][t])
            return est

        if subsets.ndim == 2:
            return est2d(subsets)
        try:
            B, T, m = subsets.shape
            flat = subsets.reshape(B*T, m)
            out = est2d(flat)  # (B*T,)
        except Exception as e:
            print("large 3D shpae estimation, falling back to loops")
            out = np.empty((B, T), dtype=float)
            for b in range(B):
                out[b] = est2d(subsets[b])

        return out.reshape(B, T)
    return est3d

def epsilon_splitter(fraction=0.5):
    """return eps_cent, eps_lcl"""
    return lambda n, m, eps: (eps * fraction, eps * (1 - fraction))

def epsilon_splitter_m_over_nplusm():
    return lambda n, m, eps: (eps * m / (n + m), eps * (1 - m / (n + m)))

def build_estimator_from_enum(lcl_mech_enum, eps, L, U, wrap_3d=True, axis=-1, is_sorted=False, noise_func=None, quantile=0.5):
    if lcl_mech_enum == ResampleEstimators.NONE_PRIVATE_MED:
        if quantile ==0.5: fn = lambda X: np.median(X, axis=axis)
        fn = lambda X: np.quantile(X, quantile, axis=axis)
        # fn = lambda X: np.median(X, axis)

    elif lcl_mech_enum == ResampleEstimators.MED_INVERSE_SENSITIVITY:
        if eps==np.inf:
            if quantile==0.5: return lambda X: np.median(X, axis)
            return lambda X: np.quantile(X, quantile, axis=axis)
        # if eps==np.inf: return lambda X: np.median(X, axis)
        sigma = 2.0 / max(eps, 1e-12)
        fn = lambda X: vectorized_quantile_combined(X, L=L, U=U, sigma=sigma, quantile=quantile, is_sorted=is_sorted)

    elif lcl_mech_enum == ResampleEstimators.KS_DIST:
        fn = lambda X: ks_distance(X, lcl_eps=eps, priv=True, noise_func=noise_func)

    elif lcl_mech_enum == ResampleEstimators.NONE_PRIVATE_KS_DIST:
        fn = lambda X: ks_distance(X, lcl_eps=eps, priv=False, noise_func=noise_func)

    elif lcl_mech_enum == ResampleEstimators.NONE_PRIVATE_MEAN:
        fn = lambda X: np.mean(X, axis)

    elif lcl_mech_enum == ResampleEstimators.PRIVATE_MEAN:
        sigma = 2.0 / max(eps, 1e-12)
        fn = lambda X: priv_mean_mechanism(X, L=L, U=U, eps=eps)

    elif lcl_mech_enum == ResampleEstimators.EMPIRICAL:
        fn = lambda X: np.var(np.median(X, axis=axis))

    # elif lcl_mech_enum == ResampleEstimators.NONE_PRIVATE_LOGREG:
    #     base_fn = lambda Xs, ys: private_logreg_subset_estimator(Xs, ys, noise_func=noise_func, coef_index=coef_index,
    #                                 reg_lambda=reg_lambda, fit_intercept=fit_intercept)
    #     fn = wrap_logreg_estimator_to_3d(X_full, y_full, base_fn)
    #     return fn
    #
    # elif lcl_mech_enum == ResampleEstimators.PRIVATE_LOGREG_OUTPUT_PERT:
    #     base_fn = lambda Xs, ys: private_logreg_subset_estimator(Xs, ys, coef_index=coef_index,
    #                                 reg_lambda=reg_lambda, fit_intercept=fit_intercept, noise_func=noise_func)
    #     fn = wrap_logreg_estimator_to_3d(X_full, y_full, base_fn)
    #     return fn
    else:
        raise ValueError(f"Unknown mechanism {lcl_mech_enum}")

    if wrap_3d:
        fn = _wrap_estimator_to_3d(fn)

    return fn


def sampling_amplification_inv(epsilon: float, delta: float, sampling_prob: float):
    """ returns amplified (epsilon', delta') """
    return np.log1p(np.expm1(epsilon)) / sampling_prob, delta / sampling_prob


def sampling_amplification(epsilon: float, delta: float, sampling_prob: float):
    """ returns amplified (epsilon', delta') """
    return np.log(1 + sampling_prob * (np.exp(epsilon) - 1)), sampling_prob * delta


def compose_epsilon(epsilon, delta, num_subsamples, print_choice=False):
    """returns composed epsilon with num_subsamples iterations with (epsilon, delta)-DP each """
    eps_total_basic = num_subsamples * epsilon
    if delta == 0:
        return eps_total_basic
    eps_total_adv = np.sqrt(2 * num_subsamples * np.log(1 / delta)) * epsilon + num_subsamples * epsilon * (np.exp(epsilon) - 1) / (
                np.exp(epsilon) + 1)

    if print_choice:
        if eps_total_basic < eps_total_adv:
            print("basic wins")
        else:
            print("advanced wins")

    return min(eps_total_basic, eps_total_adv)


def calc_local_delta(delta, sample_size, subsample_size, num_subsamples, is_local_pure):
    if delta == 0:
        assert is_local_pure, "Local delta cannot be zero for approximate DP"
        return 0, 0

    elif is_local_pure: # all delta goes to comp, per-step delta is 0
        return delta, 0

    else:
        delta_add = delta / (num_subsamples + 1)  # compute delta'' only here
        _, delta_local = sampling_amplification_inv(1, delta_add, subsample_size / sample_size)
        return delta_add, delta_local

# def calc_lcl_eps(sample_size, subsample_size, num_subsamples, epsilon, delta, is_local_pure):
#     lcl_epsilon, _ = calc_local_epsilon_delta(sample_size=sample_size, subsample_size=subsample_size,
#                                               num_subsamples=num_subsamples, epsilon=epsilon, delta=delta,
#                                               is_local_pure=is_local_pure)
#     return lcl_epsilon, _

def no_split_or_without_composition_eps(sample_size, subsample_size, num_subsamples, epsilon, delta, is_local_pure):
    return epsilon, delta


def epsilon_bin_search(epsilon_target, delta_tag, sample_size, subsample_size, num_subsamples):
    lower, upper = 0.0, max(4, epsilon_target * 10)
    p = subsample_size / sample_size

    amp_epsilon, _ = sampling_amplification(epsilon_target, 0, p)
    comp_amp_epsilon = compose_epsilon(amp_epsilon, delta_tag, num_subsamples)

    i = 0
    mid = upper
    while abs(epsilon_target - comp_amp_epsilon) > BINARY_SEARCH_TOLERANCE and (REASONABLE_NUMBER_OF_ITERATION_TO_CHECK > i):
        i+= 1
        mid = (lower + upper) / 2.0

        amp_epsilon, _ = sampling_amplification(mid, 0, p)
        comp_amp_epsilon = compose_epsilon(amp_epsilon, delta_tag, num_subsamples)

        if comp_amp_epsilon < epsilon_target:
            lower = mid
        else:
            upper = mid
    if i == REASONABLE_NUMBER_OF_ITERATION_TO_CHECK:
        print(f"IT MIGHT BE THE WRONG DIRECTION.\n Comp&Amp local eps is: {comp_amp_epsilon}, target eps is: {epsilon_target}")
    return mid


def bin_search_function(func, y_target, bounds, x_tol, y_tol):
    """
    Binary search to find x where func(x) satisfies domination constraint.
    Assumes func is monotonic decreasing.

    Returns:
        x such that there exists x' for which:
        1. |x - x'| < x_tol
        2. y_target - y_tol <= func(x') <= y_target
    """
    left, right = bounds
    f_right = func(right)
    if f_right > y_target:
        return right

    while right - left > x_tol:
        mid = (left + right) / 2
        f_mid = func(mid)
        if f_mid <= y_target and f_mid + y_tol > y_target:
            return mid
        if f_mid < y_target:
            right = mid
        else:
            left = mid
    return right


def calc_local_epsilon_delta(sample_size, subsample_size, num_subsamples, epsilon, delta, is_local_pure,
                             bootstrap=True):
    if type(delta) is Callable:
        delta = delta(sample_size)
    if epsilon == np.inf or epsilon==0:
        return epsilon, delta
    delta_add, delta_local = calc_local_delta(delta, sample_size, subsample_size, num_subsamples, is_local_pure)
    # preform binary search of eps s.t. the epsilon will be close to target_epsilon
    epsilon_local = epsilon_bin_search(epsilon, delta_add, sample_size, subsample_size, num_subsamples)
    # todo de-comment this assertion
    # assert (epsilon_local >= 1/np.sqrt(subsample_size)), f"Local epsilon is too small ({epsilon_local}), local mechanism cannot be applied (1/sqrt(m))={1/np.sqrt(subsample_size)}, subsample_size={subsample_size}"
    return epsilon_local, delta_local


def eps_from_sigma(T: int, n: int, m: int, delta: float, sigma: float, sensitivity_factor):
  return privacy_loss_distribution.from_gaussian_mechanism(
      standard_deviation=sigma,
      sensitivity=sensitivity_factor, # was (1.0/n) hardcode, should check with Moshe
      sampling_prob=float(m)/n,
      value_discretization_interval=1e-3,
      pessimistic_estimate=True,
  ).self_compose(T).get_epsilon_for_delta(delta)

def sigma_from_eps(T: int, n: int, m: int, delta: float, eps_global: float, sensitivity_factor):
  func = lambda sigma: eps_from_sigma(T, n, m, delta, sigma, sensitivity_factor)
  return bin_search_function(func=func, y_target=eps_global, bounds=(0, 100), x_tol=1e-5, y_tol=eps_global/100)

def calc_epsilons(alg, n, m, T, eps, d):
    center_eps, sbsmp_eps = alg.eps_split_func(n, m, eps) if alg.eps_split_func else (0, eps)
    lcl_epsilon, _ = alg.calc_lcl_eps_func(sample_size=n, subsample_size=m, num_subsamples=T, epsilon=sbsmp_eps,
                                           delta=d, is_local_pure=alg.is_lcl_pure)

    # for privsubGDP only:
    if alg.noise_type == NoiseType.GAUSSIAN and 'PrivSub' in alg.name:
        lcl_priv_param = sigma_from_eps(T, n, m, d, eps-center_eps, sensitivity_factor=(1 / m))
        cent_priv_param = sigma_from_eps(1, n, n, d, center_eps, sensitivity_factor=(1 / n))
    else:
        lcl_priv_param = lcl_epsilon
        cent_priv_param = center_eps

    noise_gen_lcl = build_noise_generator(alg, lcl_priv_param)
    noise_gen_center = build_noise_generator(alg, cent_priv_param)

    return center_eps, lcl_epsilon, noise_gen_center, noise_gen_lcl

def build_noise_generator(alg, priv_param):
    if alg.noise_type == NoiseType.LAPLACE:
        return lambda c, size: np.random.laplace(loc=0.0, scale=(c / priv_param), size=size)
    elif alg.noise_type == NoiseType.GAUSSIAN:
        return lambda c, size: np.random.normal(loc=0.0, scale=priv_param, size=size)

def rows_from_single_n_results(results, settings, n):
    rows = []
    true_val = settings.true_value(n)
    eps = settings.epsilon

    for r in results:
        alg = r["alg"]
        params = r["params"]

        for app in alg.approx_method:
            approximations = r[f"{app}"]
            rows.append(metrics_row(alg, app, eps, true_val, approximations, params))
    return rows



def _append_jsonl(log_path, row):
    if log_path is None:
        return
    os.makedirs(os.path.dirname(log_path) or ".", exist_ok=True)
    with open(log_path, "a") as f:
        f.write(json.dumps(row, default=lambda x: float(x) if hasattr(x, "item") else x) + "\n")

def metrics_row(alg, approx, eps, true_val, intervals, params):
    coverage_mean, coverage_sd = coverage_probability(intervals, true_val)
    width_mean, width_sd = CI_width(intervals)
    return {'method': alg.name, 'approx': approx, 'n': params.n, 'm': params.m, 'T': params.T, 'lcl_eps': params.lcl_eps,
            'eps': eps, 'coverage_mean': coverage_mean, 'coverage_sd': coverage_sd, 'CI_width_mean': width_mean,
            'CI_width_sd': width_sd}

def append_rows(df, rows, log_path=None):
    for row in rows:
        df.append(row)
        _append_jsonl(log_path, row)

def init_log_file(log_path, overwrite=False):
    if not log_path: return
    os.makedirs(os.path.dirname(log_path) or ".", exist_ok=True)

    if overwrite and os.path.exists(log_path):
        with open(log_path, "w", encoding="utf-8"):
            pass

def log_and_plot(df_res, dist, eps, json_path, plot, save_plot, sbs_description, significance, stats):
    df_res.to_json(json_path, orient='records', lines=False)
    df_sbs = df_res[df_res['method'] == "PrivSub (ours)" + sbs_description]
    df_sbs = df_sbs[['n', 'm', 'T', 'lcl_eps']].drop_duplicates().sort_values(by='n').reset_index(drop=True)
    df_sbs['epsilon'] = eps
    if plot:
        plot_interval_metrics(df_res, title=f"{dist.name} {stats}, eps={eps}", significance=significance,save=save_plot)
        plot_interval_metrics_relative_width(df_res, reference_method="Non-private Bootstrap", x="n", title=f" Relative Width {dist.name} {stats}, eps={eps}",)



### 3D
def create_logreg_est(enum, params, n, m, T, eps_to_check, log_exp_setting):
    reg = log_exp_setting.reg_lambda

    if enum == ResampleEstimators.NONE_PRIVATE_LOGREG or eps_to_check == np.inf:
        sigma_res = 0
        noise_fn = lambda c, size: np.random.normal(loc=0.0, scale=0, size=size)

    elif enum == ResampleEstimators.PRIVATE_LOGREG_OUTPUT_PERT:
        sigma_res = sigma_from_eps(T, n, m, params.delta, eps_to_check, sensitivity_factor=(1 / (2 * reg * m)))
        noise_fn= lambda c, size: np.random.normal(loc=0.0, scale=sigma_res, size=size)

    # elif enum == ResampleEstimators.PRIVATE_LOGREG_OBJ_PERT:

    est = make_logreg_scalar_estimator(log_exp_setting, noise_fn)
    return est, sigma_res