# Copyright 2017 The 'Scalable Private Learning with PATE' Authors All Rights Reserved.
#
# 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.
# ==============================================================================

"""Functions for smooth sensitivity analysis for PATE mechanisms.

This library implements functionality for doing smooth sensitivity analysis
for Gaussian Noise Max (GNMax), Threshold with Gaussian noise, and Gaussian
Noise with Smooth Sensitivity (GNSS) mechanisms.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math
from absl import app
import numpy as np
import scipy
import sympy as sp

import core as pate

################################
# SMOOTH SENSITIVITY FOR GNMAX #
################################

# Global dictionary for storing cached q0 values keyed by (sigma, order).
_logq0_cache = {}


def _compute_logq0(sigma, order):
  key = (sigma, order)
  if key in _logq0_cache:
    return _logq0_cache[key]

  logq0 = compute_logq0_gnmax(sigma, order)

  _logq0_cache[key] = logq0  # Update the global variable.
  return logq0


def _compute_logq1(sigma, order, num_classes):
  logq0 = _compute_logq0(sigma, order)  # Most likely already cached.
  logq1 = math.log(_compute_bl_gnmax(math.exp(logq0), sigma, num_classes))
  assert logq1 <= logq0
  return logq1


def _compute_mu1_mu2_gnmax(sigma, logq):
  # Computes mu1, mu2 according to Proposition 10.
  mu2 = sigma * math.sqrt(-logq)
  mu1 = mu2 + 1
  return mu1, mu2


def _compute_data_dep_bound_gnmax(sigma, logq, order):
  # Applies Theorem 6 in Appendix without checking that logq satisfies necessary
  # constraints. The pre-conditions must be assured by comparing logq against
  # logq0 by the caller.
  variance = sigma**2
  mu1, mu2 = _compute_mu1_mu2_gnmax(sigma, logq)
  eps1 = mu1 / variance
  eps2 = mu2 / variance

  log1q = np.log1p(-math.exp(logq))  # log1q = log(1-q)
  log_a = (order - 1) * (
      log1q - (np.log1p(-math.exp((logq + eps2) * (1 - 1 / mu2)))))
  log_b = (order - 1) * (eps1 - logq / (mu1 - 1))

  return np.logaddexp(log1q + log_a, logq + log_b) / (order - 1)


def _compute_rdp_gnmax(sigma, logq, order):
  logq0 = _compute_logq0(sigma, order)
  if logq >= logq0:
    return pate.rdp_data_independent_gaussian(sigma, order)
  else:
    return _compute_data_dep_bound_gnmax(sigma, logq, order)


def compute_logq0_gnmax(sigma, order):
  """Computes the point where we start using data-independent bounds.

  Args:
    sigma: std of the Gaussian noise
    order: Renyi order lambda

  Returns:
    logq0: the point above which the data-ind bound overtakes data-dependent
    bound.
  """

  def _check_validity_conditions(logq):
    # Function returns true iff logq is in the range where data-dependent bound
    # is valid. (Theorem 6 in Appendix.)
    mu1, mu2 = _compute_mu1_mu2_gnmax(sigma, logq)
    if mu1 < order:
      return False
    eps2 = mu2 / sigma**2
    # Do computation in the log space. The condition below comes from Lemma 9
    # from Appendix.
    return (logq <= (mu2 - 1) * eps2 - mu2 * math.log(mu1 / (mu1 - 1) * mu2 /
                                                      (mu2 - 1)))

  def _compare_dep_vs_ind(logq):
    return (_compute_data_dep_bound_gnmax(sigma, logq, order) -
            pate.rdp_data_independent_gaussian(sigma, order))

  # Natural upper bounds on q0.
  logub = min(-(1 + 1. / sigma)**2, -((order - .99) / sigma)**2, -1 / sigma**2)
  assert _check_validity_conditions(logub)

  # If data-dependent bound is already better, we are done already.
  if _compare_dep_vs_ind(logub) < 0:
    return logub

  # Identifying a reasonable lower bound to bracket logq0.
  loglb = 2 * logub  # logub is negative, and thus loglb < logub.
  while _compare_dep_vs_ind(loglb) > 0:
    assert loglb > -10000, "The lower bound on q0 is way too low."
    loglb *= 1.5

  logq0, r = scipy.optimize.brentq(
      _compare_dep_vs_ind, loglb, logub, full_output=True)
  assert r.converged, "The root finding procedure failed to converge."
  assert _check_validity_conditions(logq0)  # just in case.

  return logq0


def _compute_bl_gnmax(q, sigma, num_classes):
  return ((num_classes - 1) / 2 * scipy.special.erfc(
      1 / sigma + scipy.special.erfcinv(2 * q / (num_classes - 1))))


def _compute_bu_gnmax(q, sigma, num_classes):
  return min(1, (num_classes - 1) / 2 * scipy.special.erfc(
      -1 / sigma + scipy.special.erfcinv(2 * q / (num_classes - 1))))


def _compute_local_sens_gnmax(logq, sigma, num_classes, order):
  """Implements Algorithm 3 (computes an upper bound on local sensitivity).

  (See Proposition 13 for proof of correctness.)
  """
  logq0 = _compute_logq0(sigma, order)