from opacus.accountants import create_accountant
from scipy import optimize


EPSILON_TOL = 0.01
DEFAULT_ALPHAS = [1 + x / 100.0 for x in range(1, 1000)] + list(range(12, 99)) + list(range(100, 999, 10))


def get_sampling_rate(
        target_epsilon: float,
        target_delta: float,
        noise_multiplier: float,
        n_iters: int
    ) -> float:
    r"""
    Calculates the proper sampling rate that achieves :math:`(\epsilon, \delta)`-DP in `n_iters` iterations given `noise_multiplier`, using a bisection method.

    Args:
        target_epsilon: The target :math:`\epsilon` to satisfy.
        target_delta: The target :math:`\delta` to satisfy.
        noise_multiplier: The noise scale divided by clipping threshold.
        n_iters: Total number of iterations.
    """
    eps_high = float('inf')
    accountant = create_accountant(mechanism='rdp')

    q_low, q_high = 0.0, 1.0

    accountant.history = [(noise_multiplier, q_high, n_iters)]
    eps_high = accountant.get_epsilon(delta=target_delta, alphas=DEFAULT_ALPHAS)
    if eps_high < target_epsilon:
        return 1.0
    
    def bisect_objective(q):
        accountant.history = [(noise_multiplier, q, n_iters)]
        epsilon = 0 if q == 0 else accountant.get_epsilon(delta=target_delta, alphas=DEFAULT_ALPHAS)   
        return epsilon - target_epsilon
    
    return optimize.bisect(bisect_objective, q_low, q_high, maxiter=100, disp=False)
