# Copyright 2020 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Implementing Privacy Loss Distribution.

This file implements the privacy loss distribution (PLD) and its basic
functionalities. The main feature of PLD is that it allows for accurate
computation of privacy parameters under composition. Please refer to the
supplementary material below for more details:
../../common_docs/Privacy_Loss_Distributions.pdf
"""

import collections
import logging
import math
from typing import Any, Callable, Mapping, Optional, Sequence, Tuple, Union

import numpy as np

from dp_accounting import privacy_accountant
from dp_accounting.pld import common
from dp_accounting.pld import pld_pmf
from dp_accounting.pld import privacy_loss_mechanism


NeighborRel = privacy_accountant.NeighboringRelation


class PrivacyLossDistribution:
  """Class for privacy loss distributions and computation involving them.

  The privacy loss distribution (PLD) of two discrete distributions, the upper
  distribution mu_upper and the lower distribution mu_lower, is defined as a
  distribution on real numbers generated by first sampling an outcome o
  according to mu_upper and then outputting the privacy loss
  ln(mu_upper(o) / mu_lower(o)) where mu_lower(o) and mu_upper(o) are the
  probability masses of o in mu_lower and mu_upper respectively. This class
  allows one to create and manipulate privacy loss distributions.

  PLD allows one to (approximately) compute the epsilon-hockey stick divergence
  between mu_upper and mu_lower, which is defined as
  sum_{o} [mu_upper(o) - e^{epsilon} * mu_lower(o)]_+. This quantity in turn
  governs the parameter delta of (eps, delta)-differential privacy of the
  corresponding protocol. (See Observation 1 in the supplementary material.)

  The above definitions extend to continuous distributions. The PLD of two
  continuous distributions mu_upper and mu_lower is defined as a distribution on
  real numbers generated by first sampling an outcome o according to mu_upper
  and then outputting the privacy loss ln(f_{mu_upper}(o) / f_{mu_lower}(o))
  where f_{mu_lower}(o) and f_{mu_upper}(o) are the probability density
  functions at o in mu_lower and mu_upper respectively. Moreover, for continuous
  distributions the epsilon-hockey stick divergence is defined as
  integral [f_{mu_upper}(o) - e^{epsilon} * f_{mu_lower}(o)]_+ do.

  A single privacy loss distribution is represented as an object of the class
  BasicPrivacyLossDistribution. This class, on the other hand, holds the higher
  level logic.

  Namely, this class maintains up to two BasicPrivacyLossDistribution objects.
  One for the 'add' adjacency type, which specifies the privacy loss
  distribution for a mechanism M with mu_upper = M(D) and mu_lower = M(D'),
  where D' contains one more datapoint than D.
  And one for the 'remove' adjacency type, which specifies the privacy loss
  distribution for a mechanism M, with mu_upper = M(D) and mu_lower = M(D'),
  where D' contains one less datapoint than D.
  In the case where both privacy loss distributions are the same, only one copy
  is maintained.

  While this class offers additional support with respect to the ADD/REMOVE
  adjacency, this is not an inherent limitation; this class can also be used
  in the case of other adjacencies such as Substitution.

  Factory methods in this module provide a convenient way to generate objects of
  this class associated to various mechanisms.

  Attributes:
    _pmf_remove: the privacy loss distribution probability mass function with
      respect to REMOVE adjacency.
    _pmf_add: the privacy loss distribution probability mass function with
      respect to ADD adjacency.
    _symmetric: When True, _pmf_add is assumed to be the same as _pmf_remove.
  """

  def __init__(self,
               pmf_remove: pld_pmf.PLDPmf,
               pmf_add: Optional[pld_pmf.PLDPmf] = None):
    """Initialization method for PrivacyLossDistribution."""
    self._pmf_remove = pmf_remove
    self._symmetric = pmf_add is None
    self._pmf_add = pmf_remove if pmf_add is None else pmf_add

  @classmethod
  def create_from_rounded_probability(
      cls,
      rounded_probability_mass_function: Mapping[int, float],
      infinity_mass: float,
      value_discretization_interval: float,
      pessimistic_estimate: bool = True,
      rounded_probability_mass_function_add: Optional[Mapping[int,
                                                              float]] = None,
      infinity_mass_add: Optional[float] = None,
      symmetric: bool = True) -> 'PrivacyLossDistribution':
    """Create PrivacyLossDistribution from rounded probability mass function(s).

    Args:
      rounded_probability_mass_function: rounded probability mass function of
        the basic privacy loss distribution, with respect to REMOVE adjacency.
      infinity_mass: infinity_mass for basic privacy loss distribution with
        respect to the REMOVE adjacency.
      value_discretization_interval: the interval length for which the values of
        the privacy loss distribution are discretized. In particular, the values
        are always integer multiples of value_discretization_interval. Smaller
        value results in more accurate estimates of the privacy loss, at the
        cost of increased run-time / memory usage. Smaller value results in more
        accurate estimates of the privacy loss, at the cost of increased
        run-time / memory usage.
      pessimistic_estimate: whether the rounding is done in such a way that the
        resulting epsilon-hockey stick divergence computation gives an upper
        estimate to the real value.
      rounded_probability_mass_function_add: rounded probability mass function
        of the basic privacy loss distribution,with respect to ADD adjacency.
      infinity_mass_add: infinity_mass for basic privacy loss distribution with
        respect to the ADD adjacency.
      symmetric: When True, the basic privacy loss distribution with respect to
        ADD adjacency is assumed to be the same as that for REMOVE adjacency.
        Arguments rounded_probability_mass_function_add, infinity_mass_add are
        ignored in this case.

    Returns:
      Privacy Loss Distribution object.
    """
    pmf_remove = pld_pmf.create_pmf(rounded_probability_mass_function,
                                    value_discretization_interval,
                                    infinity_mass, pessimistic_estimate)
    pmf_add = None
    if symmetric:
      if (rounded_probability_mass_function_add is not None or
          infinity_mass_add is not None):
        raise ValueError('Details about privacy loss distribution with respect '
                         'to ADD adjacency cannot be specified when symmetric')
    else:
      if (rounded_probability_mass_function_add is None or
          infinity_mass_add is None):
        raise ValueError('Details about privacy loss distribution with respect'
                         'to ADD adjacency should be specified when not '
                         'symmetric')
      pmf_add = pld_pmf.create_pmf(rounded_probability_mass_function_add,
                                   value_discretization_interval,
                                   infinity_mass_add, pessimistic_estimate)
    return cls(pmf_remove, pmf_add)

  def get_delta_for_epsilon(
      self, epsilon: Union[float, Sequence[float]]) -> Union[float, np.ndarray]:
    """Computes the epsilon-hockey stick divergence between mu_upper, mu_lower.

    When this privacy loss distribution corresponds to a mechanism, the
    epsilon-hockey stick divergence gives the value of delta for which the
    mechanism is (epsilon, delta)-differentially private. (See Observation 1 in
    the supplementary material.)

    Args:
      epsilon: the epsilon in epsilon-hockey stick divergence.

    Returns:
      A non-negative real number which is the epsilon-hockey stick divergence
      between the upper (mu_upper) and the lower (mu_lower) distributions
      corresponding to this privacy loss distribution.
    """
    delta_remove = self._pmf_remove.get_delta_for_epsilon(epsilon)
    if self._symmetric:
      return delta_remove
    delta_add = self._pmf_add.get_delta_for_epsilon(epsilon)
    return np.maximum(delta_remove, delta_add)

  def get_epsilon_for_delta(self, delta: float) -> float:
    """Computes epsilon for which hockey stick divergence is at most delta.

    This function computes the smallest non-negative epsilon for which the
    epsilon-hockey stick divergence between mu_upper, mu_lower is at most delta.

    When this privacy loss distribution corresponds to a mechanism and the
    rounding is pessimistic, the returned value corresponds to an epsilon for
    which the mechanism is (epsilon, delta)-differentially private. (See
    Observation 1 in the supplementary material.)

    Args:
      delta: the target epsilon-hockey stick divergence.

    Returns:
      A non-negative real number which is the smallest epsilon such that the
      epsilon-hockey stick divergence between the upper (mu_upper) and the
      lower (mu_lower) distributions is at most delta. When no such finite
      epsilon exists, return math.inf.
    """
    epsilon_remove = self._pmf_remove.get_epsilon_for_delta(delta)
    if self._symmetric:
      return epsilon_remove
    epsilon_add = self._pmf_add.get_epsilon_for_delta(delta)
    return max(epsilon_remove, epsilon_add)

  def validate_composable(self,
                          privacy_loss_distribution: 'PrivacyLossDistribution'):
    """Verifies that a given PLD can be composed with this PLD.

    The two privacy loss distributions must have the same discretization
    interval and estimate type for the composition to be allowed.

    Args:
      privacy_loss_distribution: the privacy loss distribution to be composed
        with the current privacy loss distribution.

    Raises:
      ValueError if the value_discretization_interval or estimate_type of the
      two PLDs are different.
    """
    self._pmf_remove.validate_composable(privacy_loss_distribution._pmf_remove)  # pylint:disable=protected-access

  def compose(
      self,
      privacy_loss_distribution: 'PrivacyLossDistribution',
      tail_mass_truncation: float = 1e-15,
  ) -> 'PrivacyLossDistribution':
    """Computes a privacy loss distribution resulting from composing two PLDs.

    Args:
      privacy_loss_distribution: the privacy loss distribution to be composed
        with the current privacy loss distribution. The two must have the same
        value_discretization_interval.
      tail_mass_truncation: an upper bound on the tails of the probability mass
        of the PLD that might be truncated.

    Returns:
      A privacy loss distribution which is the result of composing the two.
    """
    # pylint:disable=protected-access
    pld_pmf_remove = pld_pmf.compose_pmfs(self._pmf_remove,
                                          privacy_loss_distribution._pmf_remove,
                                          tail_mass_truncation)
    if self._symmetric and privacy_loss_distribution._symmetric:
      return PrivacyLossDistribution(pld_pmf_remove)
    pld_pmf_add = pld_pmf.compose_pmfs(self._pmf_add,
                                       privacy_loss_distribution._pmf_add,
                                       tail_mass_truncation)
    # pylint:enable=protected-access
    return PrivacyLossDistribution(pld_pmf_remove, pld_pmf_add)

  def get_delta_for_epsilon_for_composed_pld(
      self, privacy_loss_distribution: 'PrivacyLossDistribution',
      epsilon: float) -> float:
    """Computes delta for given epsilon for the result of composing this PLD and a given PLD.

    The output of this function should be the same as first composing this PLD
    and privacy_loss_distribution, and then call get_delta_for_epsilon on the
    resulting PLD. The main advantage is that this function is faster.

    Args:
      privacy_loss_distribution: the privacy loss distribution to be composed
        with the current privacy loss distribution. The two must have the same
        value_discretization_interval.
      epsilon: the epsilon in epsilon-hockey stick divergence.

    Returns:
      A non-negative real number which is the epsilon-hockey stick divergence
      of the privacy loss distribution which is the result of composing this PLD
      with privacy_loss_distribution.
    """
    # pylint:disable=protected-access
    delta_remove = self._pmf_remove.get_delta_for_epsilon_for_composed_pld(
        privacy_loss_distribution._pmf_remove, epsilon)
    if self._symmetric and privacy_loss_distribution._symmetric:
      return delta_remove
    delta_add = self._pmf_add.get_delta_for_epsilon_for_composed_pld(
        privacy_loss_distribution._pmf_add, epsilon)
    # pylint:enable=protected-access
    return max(delta_remove, delta_add)

  def self_compose(
      self,
      num_times: int,
      tail_mass_truncation: float = 1e-15) -> 'PrivacyLossDistribution':
    """Computes PLD resulting from repeated composing the PLD with itself.

    Args:
      num_times: the number of times to compose this PLD with itself.
      tail_mass_truncation: an upper bound on the tails of the probability mass
        of the PLD that might be truncated. Currently only supports for
        pessimistic estimates.

    Returns:
      A privacy loss distribution which is the result of the composition.
    """
    pmf_remove = self._pmf_remove.self_compose(num_times, tail_mass_truncation)
    if self._symmetric:
      return PrivacyLossDistribution(pmf_remove)
    pmf_add = self._pmf_add.self_compose(num_times, tail_mass_truncation)
    return PrivacyLossDistribution(pmf_remove, pmf_add)


def identity(
    value_discretization_interval: float = 1e-4) -> PrivacyLossDistribution:
  """Constructs an identity privacy loss distribution.

  Args:
    value_discretization_interval: the dicretization interval for the privacy
      loss distribution. The values will be rounded up/down to be integer
      multiples of this number. Smaller value results in more accurate estimates
      of the privacy loss, at the cost of increased run-time / memory usage.

  Returns:
    The privacy loss distribution corresponding to an algorithm with no
    privacy leak (i.e. output is independent of input).
  """
  return PrivacyLossDistribution.create_from_rounded_probability(
      {0: 1}, 0, value_discretization_interval)


def from_two_probability_mass_functions(
    log_probability_mass_function_lower: Mapping[Any, float],
    log_probability_mass_function_upper: Mapping[Any, float],
    pessimistic_estimate: bool = True,
    value_discretization_interval: float = 1e-4,
    log_mass_truncation_bound: float = -math.inf,
    symmetric: bool = True) -> PrivacyLossDistribution:
  """Constructs a privacy loss distribution from mu_lower and mu_upper.

  Args:
    log_probability_mass_function_lower: the probability mass function of
      mu_lower represented as a dictionary where each key is an outcome o of
      mu_lower and the corresponding value is the natural log of the
      probability mass of mu_lower at o.
    log_probability_mass_function_upper: the probability mass function of
      mu_upper represented as a dictionary where each key is an outcome o of
      mu_upper and the corresponding value is the natural log of the
      probability mass of mu_upper at o.
    pessimistic_estimate: whether the rounding is done in such a way that the
      resulting epsilon-hockey stick divergence computation gives an upper
      estimate to the real value.
    value_discretization_interval: the dicretization interval for the privacy
      loss distribution. The values will be rounded up/down to be integer
      multiples of this number. Smaller value results in more accurate estimates
      of the privacy loss, at the cost of increased run-time / memory usage.
    log_mass_truncation_bound: when the log of the probability mass of the upper
      distribution is below this bound, it is either (i) included in
      infinity_mass in the case of pessimistic estimate or (ii) discarded
      completely in the case of optimistic estimate. The larger
      log_mass_truncation_bound is, the more error it may introduce in
      divergence calculations.
    symmetric: if True it creates a symmetric PrivacyLossDistribution.

  Returns:
    The privacy loss distribution constructed as specified.
  """

  def _create_rounded_probability_mass_function(
      log_pmf_lower: Mapping[Any, float],
      log_pmf_upper: Mapping[Any, float]
  ) -> Tuple[float, Mapping[int, float]]:
    """Helper function for creating rounded pmf."""
    infinity_mass = 0
    for outcome in log_pmf_upper:
      if log_pmf_lower.get(outcome, -math.inf) == -math.inf:
        # When an outcome only appears in the upper distribution but not in the
        # lower distribution, then it must be counted in infinity_mass as such
        # an outcome contributes to the hockey stick divergence.
        infinity_mass += math.exp(log_pmf_upper[outcome])
    # Compute the (non-discretized) probability mass function for the privacy
    # loss distribution.
    pmf = {}
    for outcome, log_prob_lower in log_pmf_lower.items():
      if log_prob_lower == -math.inf:
        # This outcome never occurs in mu_lower. This case was already included
        # as infinity_mass above.
        continue
      elif log_pmf_upper.get(outcome, -math.inf) > log_mass_truncation_bound:
        # When the probability mass of mu_upper at the outcome is greater than
        # the threshold, add it to the distribution.
        privacy_loss_value = log_pmf_upper[outcome] - log_prob_lower
        pmf[privacy_loss_value] = (
            pmf.get(privacy_loss_value, 0) + math.exp(log_pmf_upper[outcome])
        )
      else:
        if pessimistic_estimate:
          # When the probability mass of mu_upper at the outcome is no more than
          # the threshold and we would like to get a pessimistic estimate,
          # account for this in infinity_mass.
          infinity_mass += math.exp(log_pmf_upper.get(outcome, -math.inf))
    # Discretize the probability mass so that the values are integer multiples
    # of value_discretization_interval
    rounded_pmf = collections.defaultdict(lambda: 0)
    round_fn = math.ceil if pessimistic_estimate else math.floor
    for val in pmf:
      rounded_pmf[round_fn(val / value_discretization_interval)] += pmf[val]
    return infinity_mass, rounded_pmf

  infinity_mass, rounded_probability_mass_function = (
      _create_rounded_probability_mass_function(
          log_pmf_lower=log_probability_mass_function_lower,
          log_pmf_upper=log_probability_mass_function_upper
      )
  )

  if symmetric:
    return PrivacyLossDistribution.create_from_rounded_probability(
        rounded_probability_mass_function, infinity_mass,
        value_discretization_interval, pessimistic_estimate,
    )

  infinity_mass_add, rounded_probability_mass_function_add = (
      _create_rounded_probability_mass_function(
          log_pmf_lower=log_probability_mass_function_upper,
          log_pmf_upper=log_probability_mass_function_lower,
      )
  )
  return PrivacyLossDistribution.create_from_rounded_probability(
      rounded_probability_mass_function, infinity_mass,
      value_discretization_interval, pessimistic_estimate,
      rounded_probability_mass_function_add, infinity_mass_add,
      symmetric=False,
  )


def _create_pld_pmf_from_monotone_privacy_loss(
    monotone_privacy_loss:
    'privacy_loss_mechanism.MonotonePrivacyLoss',
    pessimistic_estimate: bool = True,
    value_discretization_interval: float = 1e-4,
    use_connect_dots: bool = False) -> pld_pmf.PLDPmf:
  """Constructs the privacy loss distribution of an additive noise mechanism.

  An additive noise mechanism for computing a scalar-valued function f is a
  mechanism that outputs the sum of the true value of the function and a noise
  drawn from a certain distribution mu. This function calculates the privacy
  loss distribution for such an additive noise mechanism.

  This method supports two algorithms for constructing the privacy loss
  distribution. One given by the "Privacy Buckets" algorithm and other given by
  "Connect the Dots" algorithm. See Sections 2.1 and 2.2 of supplementary
  material for more details.

  Args:
    monotone_privacy_loss: the privacy loss representation of the mechanism.
    pessimistic_estimate: a value indicating whether the rounding is done in
      such a way that the resulting epsilon-hockey stick divergence computation
      gives an upper estimate to the real value.
    value_discretization_interval: the length of the dicretization interval for
      the privacy loss distribution. The values will be rounded up/down to be
      integer multiples of this number. Smaller value results in more accurate
      estimates of the privacy loss, at the cost of increased run-time / memory
      usage.
    use_connect_dots: when False (default), the privacy buckets algorithm will
      be used to construct the privacy loss distribution. When True, the
      connect-the-dots algorithm would be used.

  Returns:
    The privacy loss distribution constructed as specified.
  """
  if use_connect_dots:
    if not pessimistic_estimate:
      logging.warning('Connect-the-Dots implementation does not support '
                      'pessimistic_estimate=False. Using Privacy Buckets '
                      'algorithm instead. Set use_connect_dots=False to avoid '
                      'this warning.')
    else:
      connect_dots_bounds = monotone_privacy_loss.connect_dots_bounds()

      if monotone_privacy_loss.is_discrete:
        if (connect_dots_bounds.lower_x is None or
            connect_dots_bounds.upper_x is None):
          raise ValueError('Connect dots bounds does not contain lower_x and '
                           'upper_x for discrete additive noise mechanism.')
        rounded_epsilons = []
        for x in range(connect_dots_bounds.upper_x,
                       connect_dots_bounds.lower_x - 1,
                       -1):
          scaled_epsilon = (monotone_privacy_loss.privacy_loss(x)
                            / value_discretization_interval)
          if (not rounded_epsilons or
              math.floor(scaled_epsilon) > rounded_epsilons[-1]):
            rounded_epsilons.append(math.floor(scaled_epsilon))
          if math.ceil(scaled_epsilon) > rounded_epsilons[-1]:
            rounded_epsilons.append(math.ceil(scaled_epsilon))
        rounded_epsilons = np.array(rounded_epsilons)

        deltas = monotone_privacy_loss.get_delta_for_epsilon(
            rounded_epsilons * value_discretization_interval)

        return pld_pmf.create_pmf_pessimistic_connect_dots(
            value_discretization_interval,
            rounded_epsilons,
            deltas)
      else:
        if (connect_dots_bounds.epsilon_upper is None or
            connect_dots_bounds.epsilon_lower is None):
          raise ValueError('Connect dots bounds does not contain epsilon_upper '
                           'and epsilon_lower for continuous additive noise '
                           'mechanism.')
        rounded_epsilon_upper = math.ceil(connect_dots_bounds.epsilon_upper /
                                          value_discretization_interval)
        rounded_epsilon_lower = math.floor(connect_dots_bounds.epsilon_lower /
                                           value_discretization_interval)
        rounded_epsilons = np.arange(rounded_epsilon_lower,
                                     rounded_epsilon_upper + 1)

        deltas = monotone_privacy_loss.get_delta_for_epsilon(
            rounded_epsilons * value_discretization_interval)

        # Use a specialized numerically stable approach for continuous noise
        return pld_pmf.create_pmf_pessimistic_connect_dots_fixed_gap(
            value_discretization_interval,
            rounded_epsilon_lower,
            rounded_epsilon_upper,
            deltas)

  round_fn = math.ceil if pessimistic_estimate else math.floor

  tail_pld = monotone_privacy_loss.privacy_loss_tail()
  lower_x, upper_x = tail_pld.lower_x_truncation, tail_pld.upper_x_truncation

  rounded_probability_mass_function = collections.defaultdict(lambda: 0)
  infinity_mass = tail_pld.tail_probability_mass_function.get(math.inf, 0)
  for privacy_loss in tail_pld.tail_probability_mass_function:
    if privacy_loss != math.inf:
      rounded_probability_mass_function[round_fn(
          privacy_loss / value_discretization_interval
      )] += tail_pld.tail_probability_mass_function[privacy_loss]

  if monotone_privacy_loss.is_discrete:
    xs = list(range(math.ceil(lower_x) - 1, math.floor(upper_x) + 1))

    # Compute PMF for the x's. Note that a vectorized call to mu_upper_cdf can
    # be much faster than many scalar calls.
    cdf_values = monotone_privacy_loss.mu_upper_cdf(xs)
    probability_mass = cdf_values[1:] - cdf_values[:-1]

    for x, prob in zip(xs[1:], probability_mass):
      rounded_probability_mass_function[round_fn(
          monotone_privacy_loss.privacy_loss(x) /
          value_discretization_interval)] += prob
  else:
    rounded_down_value = math.floor(
        monotone_privacy_loss.privacy_loss(lower_x) /
        value_discretization_interval)
    upper_x_privacy_loss = monotone_privacy_loss.privacy_loss(upper_x)

    # Compute discretization intervals for PLD approximation.
    xs, rounded_values = [lower_x], []
    x = lower_x
    while x < upper_x:
      if (value_discretization_interval * rounded_down_value <=
          upper_x_privacy_loss):
        x = upper_x
      else:
        x = monotone_privacy_loss.inverse_privacy_loss(
            value_discretization_interval * rounded_down_value)

      xs.append(x)
      rounded_values.append(round_fn(rounded_down_value + 0.5))
      rounded_down_value -= 1

    # Compute PLD for discretization intervals. Note that a vectorized call to
    # mu_upper_cdf is much faster than many scalar calls.
    cdf_values = monotone_privacy_loss.mu_upper_cdf(xs)
    probability_mass = cdf_values[1:] - cdf_values[:-1]

    # Each x in [lower_x, upper_x] results in privacy loss that lies in
    # [value_discretization_interval * rounded_down_value,
    #  value_discretization_interval * (rounded_down_value + 1)]
    for rounded_value, prob in zip(rounded_values, probability_mass):
      rounded_probability_mass_function[rounded_value] += prob

  return pld_pmf.create_pmf(
      dict(rounded_probability_mass_function),
      value_discretization_interval,
      infinity_mass,
      pessimistic_estimate=pessimistic_estimate)


def create_from_cdf(
    cdf: Callable[[float], float],
    pessimistic_estimate: bool = True,
    value_discretization_interval: float = 1e-4,
    tail_mass_truncation: float = 1e-15) -> PrivacyLossDistribution:
  """Constructs the privacy loss distribution from its cumulative density function.

  Args:
    cdf: the cumulative density function of the privacy loss distribution.
    pessimistic_estimate: a value indicating whether the rounding is done in
      such a way that the resulting epsilon-hockey stick divergence computation
      gives an upper estimate to the real value.
    value_discretization_interval: the length of the dicretization interval for
      the privacy loss distribution. The values will be rounded up/down to be
      integer multiples of this number. Smaller value results in more accurate
      estimates of the privacy loss, at the cost of increased run-time / memory
      usage.
    tail_mass_truncation: an upper bound on the tails of the probability mass of
      the PLD that might be truncated.

  Returns:
    The privacy loss distribution constructed as specified.
  """
  rounded_probability_mass_function = {}

  # Construct the distribution for value greater than or equal to zero.
  rounded_value = 1 if pessimistic_estimate else 0
  value = 0
  while cdf(value) < 1 - tail_mass_truncation / 2:
    rounded_probability_mass_function[rounded_value] = (
        cdf(value + value_discretization_interval) - cdf(value))
    rounded_value += 1
    value += value_discretization_interval

  # Construct the distribution for value less than zero.
  rounded_value = 0 if pessimistic_estimate else -1
  value = 0
  while cdf(value) > tail_mass_truncation / 2:
    rounded_probability_mass_function[rounded_value] = (
        cdf(value) - cdf(value - value_discretization_interval))
    rounded_value -= 1
    value -= value_discretization_interval

  return PrivacyLossDistribution.create_from_rounded_probability(
      rounded_probability_mass_function,
      tail_mass_truncation if pessimistic_estimate else 0,
      value_discretization_interval,
      pessimistic_estimate=pessimistic_estimate)


def from_randomized_response(
    noise_parameter: float,
    num_buckets: int,
    pessimistic_estimate: bool = True,
    value_discretization_interval: float = 1e-4,
    neighboring_relation: NeighborRel = NeighborRel.REPLACE_ONE,
) -> PrivacyLossDistribution:
  """Constructs the privacy loss distribution of Randomized Response.

  The Randomized Response over k buckets with noise parameter p takes in an
  input which is one of the k buckets. With probability 1 - p, it simply
  outputs the input bucket. Otherwise, with probability p, it outputs a bucket
  drawn uniformly at random from the k buckets.

  This function calculates the privacy loss distribution for the
  aforementioned Randomized Response with a given number of buckets, and a
  given noise parameter, for REPLACE_ONE and REPLACE_SPECIAL neighboring
  relations. For defining the privacy loss distribution for the REPLACE_SPECIAL
  neighboring relation, we augment the mechanism to also handle a "special"
  input bucket on which the output is uniformly random over all k buckets.

  Case REPLACE_ONE:
    Suppose that the original input is x is changed to x'. The privacy loss
    distribution of the Randomized Response mechanism is generated as follows:
      first pick o according to R(x), where R(x) denote the output distribution
      of the Randomized Response mechanism on input x. Then, the privacy loss is
      ln(Pr[R(x) = o] / Pr[R(x') = o]).
    There are three cases here:
      - When o = x, ln(Pr[R(x) = o] / Pr[R(x') = o]) =
        ln(Pr[R(x) = x] / Pr[R(x') = x]). Here Pr[R(x) = x] = 1 - p + p / k
        and Pr[R(x') = x] = p / k.
      - When o = x', ln(Pr[R(x) = o] / Pr[R(x') = o]) =
        ln(Pr[R(x') = x'] / Pr[R(x) = x']), which is just the negation of the
        previous privacy loss.
      - When o != x, x', the privacy loss is zero.

  Case REPLACE_SPECIAL:
    Suppose that the original input is x is changed to the special bucket 'bot'.
    The privacy loss distribution of the Randomized Response mechanism in this
    case is asymmetric corresponding to the following two distributions:
    R(x): that outputs x with probability 1 - p + p / k, and outputs any other
      o != x with probability p/k.
    R('bot'): that outputs o with probability 1/k.
    In particular, the privacy loss distribution for the REMOVE adjacency
    corresponds to sampling o according to R(x) and returning the privacy loss
    ln(Pr[R(x) = o] / Pr[R('bot') = o]), and the privacy loss distribution for
    the ADD adjacency corresponds to sampling o according to R('bot') and
    returning the privacy loss ln(Pr[R('bot') = o] / Pr[R(x) = o]).
    Note: While REMOVE and ADD adjacency types usually refer to the
    ADD_OR_REMOVE_ONE neighboring relation, we overload the same terminology
    to correspond to "x -> 'bot'" and "'bot' -> x" adjacencies respectively.

  Args:
    noise_parameter: the probability that the Randomized Response outputs a
      completely random bucket.
    num_buckets: the total number of possible input values (which is equal to
      the total number of possible output values).
    pessimistic_estimate: a value indicating whether the rounding is done in
      such a way that the resulting epsilon-hockey stick divergence computation
      gives an upper estimate to the real value.
    value_discretization_interval: the length of the dicretization interval for
      the privacy loss distribution. The values will be rounded up/down to be
      integer multiples of this number. Smaller value results in more accurate
      estimates of the privacy loss, at the cost of increased run-time / memory
      usage.
    neighboring_relation: the neighboring relation for analyzing the mechanism.
      The default value is REPLACE_ONE, which corresponds to the case where
      the input is replaced by a different bucket. The other option is
      REPLACE_SPECIAL, which corresponds to the case, where the input is
      replaced by a uniformly random bucket.

  Returns:
    The privacy loss distribution constructed as specified.
  """

  if noise_parameter <= 0 or noise_parameter >= 1:
    raise ValueError(f'Noise parameter must be strictly between 0 and 1: '
                     f'Found {noise_parameter}')

  if num_buckets <= 1:
    raise ValueError(
        f'Number of buckets must be strictly greater than 1: {num_buckets}')

  if neighboring_relation not in [
      NeighborRel.REPLACE_ONE, NeighborRel.REPLACE_SPECIAL,
  ]:
    raise ValueError(
        'Neighboring relation must be either REPLACE_ONE or REPLACE_SPECIAL: '
        f'Found {neighboring_relation}.')

  if neighboring_relation == NeighborRel.REPLACE_ONE:
    log_pmf_upper = {
        0: math.log(1 - noise_parameter * (num_buckets - 1) / num_buckets),
        1: math.log(noise_parameter / num_buckets),
    }
    log_pmf_lower = {0: log_pmf_upper[1], 1: log_pmf_upper[0]}
    if num_buckets > 2:
      # Since all other buckets correspond to the same privacy loss, we can
      # combine them into a single bucket.
      log_pmf_upper[2] = math.log(
          (num_buckets - 2) * noise_parameter / num_buckets
      )
      log_pmf_lower[2] = log_pmf_upper[2]
    return from_two_probability_mass_functions(
        log_pmf_lower,
        log_pmf_upper,
        pessimistic_estimate,
        value_discretization_interval,
        symmetric=True,
    )
  else:  # Case of REPLACE_SPECIAL
    # Since all buckets other than the input bucket correspond to the same
    # privacy loss, we can combine them into a single bucket `1`.
    log_pmf_upper = {
        0: math.log(1 - noise_parameter * (num_buckets - 1) / num_buckets),
        1: math.log((num_buckets - 1) * noise_parameter / num_buckets),
    }
    log_pmf_lower = {
        0: - math.log(num_buckets),  # equals log(1 / num_buckets)
        1: math.log1p(- 1 / num_buckets),  # equals log(1 - 1 / num_buckets)
    }
    return from_two_probability_mass_functions(
        log_pmf_lower,
        log_pmf_upper,
        pessimistic_estimate,
        value_discretization_interval,
        symmetric=False,
    )


def _pld_for_subsampled_mechanism(
    single_pld_pmf: Callable[[privacy_loss_mechanism.AdjacencyType],
                             pld_pmf.PLDPmf],
    sampling_prob: float = 1.0) -> PrivacyLossDistribution:
  """Computes the privacy loss distribution for subsampled mechanisms.

  It is assumed that when sub-sampling probability is 1, the privacy loss
  distributions corresponding to ADD and REMOVE adjacencies are identical.

  Args:
    single_pld_pmf: method for computing the privacy loss distributions with
    respect to ADD and REMOVE adjacency types.
    sampling_prob: sub-sampling probability, a value in (0,1].

  Returns:
    A symmetric privacy loss distribution when sampling_prob = 1; An
    asymmetric privacy loss distribution corresponding to ADD and REMOVE
    adjacency types when sampling_prob < 1.
  """
  pmf_remove = single_pld_pmf(privacy_loss_mechanism.AdjacencyType.REMOVE)
  if sampling_prob == 1.0:
    return PrivacyLossDistribution(pmf_remove)

  pmf_add = single_pld_pmf(privacy_loss_mechanism.AdjacencyType.ADD)
  return PrivacyLossDistribution(pmf_remove, pmf_add)


def from_laplace_mechanism(
    parameter: float,
    sensitivity: float = 1,
    pessimistic_estimate: bool = True,
    value_discretization_interval: float = 1e-4,
    sampling_prob: float = 1.0,
    use_connect_dots: bool = True) -> PrivacyLossDistribution:
  """Computes the privacy loss distribution of the Laplace mechanism.

  This method supports two algorithms for constructing the privacy loss
  distribution. One given by the "Privacy Buckets" algorithm and other given by
  "Connect the Dots" algorithm. See Sections 2.1 and 2.2 of supplementary
  material for more details.

  Args:
    parameter: the parameter of the Laplace distribution.
    sensitivity: the sensitivity of function f. (i.e. the maximum absolute
      change in f when an input to a single user changes.)
    pessimistic_estimate: a value indicating whether the rounding is done in
      such a way that the resulting epsilon-hockey stick divergence computation
      gives an upper estimate to the real value.
    value_discretization_interval: the length of the dicretization interval for
      the privacy loss distribution. The values will be rounded up/down to be
      integer multiples of this number. Smaller value results in more accurate
      estimates of the privacy loss, at the cost of increased run-time / memory
      usage.
    sampling_prob: sub-sampling probability, a value in (0,1].
    use_connect_dots: When True (default), the connect-the-dots algorithm will
      be used to construct the privacy loss distribution. When False, the
      privacy buckets algorithm will be used.

  Returns:
    The privacy loss distribution corresponding to the Laplace mechanism with
    given parameters.
  """

  def single_laplace_pld(
      adjacency_type: privacy_loss_mechanism.AdjacencyType) -> pld_pmf.PLDPmf:
    return _create_pld_pmf_from_monotone_privacy_loss(
        privacy_loss_mechanism.LaplacePrivacyLoss(
            parameter,
            sensitivity=sensitivity,
            sampling_prob=sampling_prob,
            adjacency_type=adjacency_type),
        pessimistic_estimate=pessimistic_estimate,
        value_discretization_interval=value_discretization_interval,
        use_connect_dots=use_connect_dots)

  return _pld_for_subsampled_mechanism(single_laplace_pld, sampling_prob)


def from_gaussian_mechanism(
    standard_deviation: float,
    sensitivity: float = 1,
    pessimistic_estimate: bool = True,
    value_discretization_interval: float = 1e-4,
    log_mass_truncation_bound: float = -50,
    sampling_prob: float = 1.0,
    use_connect_dots: bool = True,
    neighboring_relation: NeighborRel = NeighborRel.ADD_OR_REMOVE_ONE,
) -> PrivacyLossDistribution:
  """Creates the privacy loss distribution of the Gaussian mechanism.

  This method supports two algorithms for constructing the privacy loss
  distribution. One given by the "Privacy Buckets" algorithm and other given by
  "Connect the Dots" algorithm. See Sections 2.1 and 2.2 of supplementary
  material for more details.

  Args:
    standard_deviation: the standard_deviation of the Gaussian distribution.
    sensitivity: the sensitivity of function f. (i.e. the maximum absolute
      change in f when an input to a single user changes.)
    pessimistic_estimate: a value indicating whether the rounding is done in
      such a way that the resulting epsilon-hockey stick divergence computation
      gives an upper estimate to the real value.
    value_discretization_interval: the length of the dicretization interval for
      the privacy loss distribution. The values will be rounded up/down to be
      integer multiples of this number. Smaller value results in more accurate
      estimates of the privacy loss, at the cost of increased run-time / memory
      usage.
    log_mass_truncation_bound: the ln of the probability mass that might be
      discarded from the noise distribution. The larger this number, the more
      error it may introduce in divergence calculations.
    sampling_prob: sub-sampling probability, a value in (0,1].
    use_connect_dots: When True (default), the connect-the-dots algorithm will
      be used to construct the privacy loss distribution. When False, the
      privacy buckets algorithm will be used.
    neighboring_relation: the neighboring relation for analyzing the mechanism.
      The default value is ADD_OR_REMOVE_ONE, which corresponds to the case
      where an input is added or removed from the dataset. The value of
      REPLACE_SPECIAL, corresponding to replacing an element by a special
      element that corresponds to zero value, results in the same privacy loss
      distribution. The other option is REPLACE_ONE, which corresponds to the
      case, where an input is replaced by a different input.

  Returns:
    The privacy loss distribution corresponding to the Gaussian mechanism with
    given parameters.
  """

  def single_gaussian_pld(
      adjacency_type: privacy_loss_mechanism.AdjacencyType) -> pld_pmf.PLDPmf:
    return _create_pld_pmf_from_monotone_privacy_loss(
        privacy_loss_mechanism.GaussianPrivacyLoss(
            standard_deviation,
            sensitivity=sensitivity,
            pessimistic_estimate=pessimistic_estimate,
            log_mass_truncation_bound=log_mass_truncation_bound,
            sampling_prob=sampling_prob,
            adjacency_type=adjacency_type),
        pessimistic_estimate=pessimistic_estimate,
        value_discretization_interval=value_discretization_interval,
        use_connect_dots=use_connect_dots)

  if neighboring_relation in [
      NeighborRel.ADD_OR_REMOVE_ONE, NeighborRel.REPLACE_SPECIAL,
  ]:
    return _pld_for_subsampled_mechanism(single_gaussian_pld, sampling_prob)
  if neighboring_relation == NeighborRel.REPLACE_ONE:
    return PrivacyLossDistribution(
        single_gaussian_pld(privacy_loss_mechanism.AdjacencyType.REPLACE)
    )
  raise ValueError(
      'Unsupported neighboring relation: %s' % neighboring_relation
  )


def from_discrete_laplace_mechanism(
    parameter: float,
    sensitivity: int = 1,
    pessimistic_estimate: bool = True,
    value_discretization_interval: float = 1e-4,
    sampling_prob: float = 1.0,
    use_connect_dots: bool = False) -> PrivacyLossDistribution:
  """Computes the privacy loss distribution of the Discrete Laplace mechanism.

  This method supports two algorithms for constructing the privacy loss
  distribution. One given by the "Privacy Buckets" algorithm and other given by
  "Connect the Dots" algorithm. See Sections 2.1 and 2.2 of supplementary
  material for more details.

  Args:
    parameter: the parameter of the discrete Laplace distribution.
    sensitivity: the sensitivity of function f. (i.e. the maximum absolute
      change in f when an input to a single user changes.)
    pessimistic_estimate: a value indicating whether the rounding is done in
      such a way that the resulting epsilon-hockey stick divergence computation
      gives an upper estimate to the real value.
    value_discretization_interval: the length of the dicretization interval for
      the privacy loss distribution. The values will be rounded up/down to be
      integer multiples of this number. Smaller value results in more accurate
      estimates of the privacy loss, at the cost of increased run-time / memory
      usage.
    sampling_prob: sub-sampling probability, a value in (0,1].
    use_connect_dots: when False (default), the privacy buckets algorithm will
      be used to construct the privacy loss distribution. When True, the
      connect-the-dots algorithm would be used.

  Returns:
    The privacy loss distribution corresponding to the Discrete Laplace
    mechanism with given parameters.
  """

  def single_discrete_laplace_pld(
      adjacency_type: privacy_loss_mechanism.AdjacencyType) -> pld_pmf.PLDPmf:
    return _create_pld_pmf_from_monotone_privacy_loss(
        privacy_loss_mechanism.DiscreteLaplacePrivacyLoss(
            parameter,
            sensitivity=sensitivity,
            sampling_prob=sampling_prob,
            adjacency_type=adjacency_type),
        pessimistic_estimate=pessimistic_estimate,
        value_discretization_interval=value_discretization_interval,
        use_connect_dots=use_connect_dots)

  return _pld_for_subsampled_mechanism(single_discrete_laplace_pld,
                                       sampling_prob)


def from_discrete_gaussian_mechanism(
    sigma: float,
    sensitivity: int = 1,
    truncation_bound: Optional[int] = None,
    pessimistic_estimate: bool = True,
    value_discretization_interval: float = 1e-4,
    sampling_prob: float = 1.0,
    use_connect_dots: bool = False) -> PrivacyLossDistribution:
  """Creates the privacy loss distribution of the discrete Gaussian mechanism.

  This method supports two algorithms for constructing the privacy loss
  distribution. One given by the "Privacy Buckets" algorithm and other given by
  "Connect the Dots" algorithm. See Sections 2.1 and 2.2 of supplementary
  material for more details.

  Args:
    sigma: the parameter of the discrete Gaussian distribution. Note that
      unlike the (continuous) Gaussian distribution this is not equal to the
      standard deviation of the noise.
    sensitivity: the sensitivity of function f. (i.e. the maximum absolute
      change in f when an input to a single user changes.)
    truncation_bound: bound for truncating the noise, i.e. the noise will only
      have a support in [-truncation_bound, truncation_bound]. When not
      specified, truncation_bound will be chosen in such a way that the mass
      of the noise outside of this range is at most 1e-30.
    pessimistic_estimate: a value indicating whether the rounding is done in
      such a way that the resulting epsilon-hockey stick divergence computation
      gives an upper estimate to the real value.
    value_discretization_interval: the length of the dicretization interval for
      the privacy loss distribution. The values will be rounded up/down to be
      integer multiples of this number. Smaller value results in more accurate
      estimates of the privacy loss, at the cost of increased run-time / memory
      usage.
    sampling_prob: sub-sampling probability, a value in (0,1].
    use_connect_dots: when False (default), the privacy buckets algorithm will
      be used to construct the privacy loss distribution. When True, the
      connect-the-dots algorithm would be used.

  Returns:
    The privacy loss distribution corresponding to the discrete Gaussian
    mechanism with given parameters.
  """

  def single_discrete_gaussian_pld(
      adjacency_type: privacy_loss_mechanism.AdjacencyType) -> pld_pmf.PLDPmf:
    return _create_pld_pmf_from_monotone_privacy_loss(
        privacy_loss_mechanism.DiscreteGaussianPrivacyLoss(
            sigma,
            sensitivity=sensitivity,
            truncation_bound=truncation_bound,
            sampling_prob=sampling_prob,
            adjacency_type=adjacency_type),
        pessimistic_estimate=pessimistic_estimate,
        value_discretization_interval=value_discretization_interval,
        use_connect_dots=use_connect_dots)

  return _pld_for_subsampled_mechanism(single_discrete_gaussian_pld,
                                       sampling_prob)


def from_mixture_gaussian_mechanism(
    standard_deviation: float,
    sensitivities: Sequence[float],
    sampling_probs: Sequence[float],
    pessimistic_estimate: bool = True,
    value_discretization_interval: float = 1e-4,
    log_mass_truncation_bound: float = -50,
    use_connect_dots: bool = True,
) -> PrivacyLossDistribution:
  """Creates the privacy loss distribution of a Mixture of Gaussians mechanism.

  This method supports two algorithms for constructing the privacy loss
  distribution. One given by the "Privacy Buckets" algorithm and other given by
  "Connect the Dots" algorithm. See Sections 2.1 and 2.2 of supplementary
  material for more details.

  Args:
    standard_deviation: the standard_deviation of the Gaussian distribution.
    sensitivities: the support of the sensitivity distribution. Must be the same
      length as sampling_probs, and both should be 1D.
    sampling_probs: the probabilities associated with the sensitivities.
    pessimistic_estimate: a value indicating whether the rounding is done in
      such a way that the resulting epsilon-hockey stick divergence computation
      gives an upper estimate to the real value.
    value_discretization_interval: the length of the dicretization interval for
      the privacy loss distribution. The values will be rounded up/down to be
      integer multiples of this number. Smaller value results in more accurate
      estimates of the privacy loss, at the cost of increased run-time / memory
      usage.
    log_mass_truncation_bound: the ln of the probability mass that might be
      discarded from the noise distribution. The larger this number, the more
      error it may introduce in divergence calculations.
    use_connect_dots: When True (default), the connect-the-dots algorithm will
      be used to construct the privacy loss distribution. When False, the
      privacy buckets algorithm will be used.

  Returns:
    The privacy loss distribution corresponding to the Mixture of Gaussians
    mechanism with given parameters.
  """

  def single_pld_pmf(
      adjacency_type: privacy_loss_mechanism.AdjacencyType,
  ) -> pld_pmf.PLDPmf:
    return _create_pld_pmf_from_monotone_privacy_loss(
        privacy_loss_mechanism.MixtureGaussianPrivacyLoss(
            standard_deviation,
            sensitivities,
            sampling_probs,
            pessimistic_estimate=pessimistic_estimate,
            log_mass_truncation_bound=log_mass_truncation_bound,
            adjacency_type=adjacency_type,
        ),
        pessimistic_estimate=pessimistic_estimate,
        value_discretization_interval=value_discretization_interval,
        use_connect_dots=use_connect_dots,
    )

  pmf_remove = single_pld_pmf(privacy_loss_mechanism.AdjacencyType.REMOVE)
  pmf_add = single_pld_pmf(privacy_loss_mechanism.AdjacencyType.ADD)
  return PrivacyLossDistribution(pmf_remove, pmf_add)


def from_privacy_parameters(
    privacy_parameters: common.DifferentialPrivacyParameters,
    value_discretization_interval: float = 1e-4) -> PrivacyLossDistribution:
  """Constructs pessimistic PLD from epsilon and delta parameters.

  When the mechanism is (epsilon, delta)-differentially private, the following
  is a pessimistic estimate of its privacy loss distribution (see Section 3.5
  of the supplementary material for more explanation):
    - infinity with probability delta.
    - epsilon with probability (1 - delta) / (1 + exp(-eps))
    - -epsilon with probability (1 - delta) / (1 + exp(eps))

  Args:
    privacy_parameters: the privacy guarantee of the mechanism.
    value_discretization_interval: the length of the dicretization interval for
      the privacy loss distribution. The values will be rounded up/down to be
      integer multiples of this number. Smaller value results in more accurate
      estimates of the privacy loss, at the cost of increased run-time / memory
      usage.

  Returns:
    The privacy loss distribution constructed as specified.
  """
  delta = privacy_parameters.delta
  epsilon = privacy_parameters.epsilon

  rounded_probability_mass_function = collections.defaultdict(lambda: 0)

  rounded_probability_mass_function[math.ceil(
      epsilon /
      value_discretization_interval)] = (1 - delta) / (1 + math.exp(-epsilon))
  rounded_probability_mass_function[math.ceil(
      -epsilon /
      value_discretization_interval)] += (1 - delta) / (1 + math.exp(epsilon))

  return PrivacyLossDistribution.create_from_rounded_probability(
      rounded_probability_mass_function, privacy_parameters.delta,
      value_discretization_interval)
