# Copyright 2020, The TensorFlow Authors.
#
# 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.
"""Implements DPQuery interface for quantile estimator."""

import collections

import tensorflow as tf
from tensorflow_privacy.privacy.dp_query import dp_query
from tensorflow_privacy.privacy.dp_query import gaussian_query
from tensorflow_privacy.privacy.dp_query import no_privacy_query
from tensorflow_privacy.privacy.dp_query import normalized_query
from tensorflow_privacy.privacy.dp_query import tree_aggregation_query


class QuantileEstimatorQuery(dp_query.SumAggregationDPQuery):
  """DPQuery to estimate target quantile of a univariate distribution.

  Uses the algorithm of Andrew et al. (https://arxiv.org/abs/1905.03871). See
  the paper for details and suggested hyperparameter settings.
  """

  # pylint: disable=invalid-name
  _GlobalState = collections.namedtuple('_GlobalState', [
      'current_estimate', 'target_quantile', 'learning_rate',
      'below_estimate_state'
  ])

  # pylint: disable=invalid-name
  _SampleParams = collections.namedtuple(
      '_SampleParams', ['current_estimate', 'below_estimate_params'])

  # No separate SampleState-- sample state is just below_estimate_query's
  # SampleState.

  def __init__(self,
               initial_estimate,
               target_quantile,
               learning_rate,
               below_estimate_stddev,
               expected_num_records,
               geometric_update=False):
    """Initializes the QuantileEstimatorQuery.

    Args:
      initial_estimate: The initial estimate of the quantile.
      target_quantile: The target quantile. I.e., a value of 0.8 means a value
        should be found for which approximately 80% of updates are less than the
        estimate each round.
      learning_rate: The learning rate. A rate of r means that the estimate will
        change by a maximum of r at each step (for arithmetic updating) or by a
        maximum factor of exp(r) (for geometric updating). Andrew et al.
        recommends that this be set to 0.2 for geometric updating.
      below_estimate_stddev: The stddev of the noise added to the count of
        records currently below the estimate. Andrew et al. recommends that this
        be set to `expected_num_records / 20` for reasonably fast adaptation and
        high privacy.
      expected_num_records: The expected number of records per round.
      geometric_update: If True, use geometric updating of estimate. Geometric
        updating is preferred for non-negative records like vector norms that
        could potentially be very large or very close to zero.
    """

    if target_quantile < 0 or target_quantile > 1:
      raise ValueError(
          f'`target_quantile` must be between 0 and 1, got {target_quantile}.')

    if learning_rate < 0:
      raise ValueError(
          f'`learning_rate` must be non-negative, got {learning_rate}')

    self._initial_estimate = initial_estimate
    self._target_quantile = target_quantile
    self._learning_rate = learning_rate

    self._below_estimate_query = self._construct_below_estimate_query(
        below_estimate_stddev, expected_num_records)
    assert isinstance(self._below_estimate_query,
                      dp_query.SumAggregationDPQuery)

    self._geometric_update = geometric_update

  def _construct_below_estimate_query(self, below_estimate_stddev,
                                      expected_num_records):
    # A DPQuery used to estimate the fraction of records that are less than the
    # current quantile estimate. It accumulates an indicator 0/1 of whether each
    # record is below the estimate, and normalizes by the expected number of
    # records. In practice, we accumulate counts shifted by -0.5 so they are
    # centered at zero. This makes the sensitivity of the below_estimate count
    # query 0.5 instead of 1.0, since the maximum that a single record could
    # affect the count is 0.5. Note that although the l2_norm_clip of the
    # below_estimate query is 0.5, no clipping will ever actually occur
    # because the value of each record is always +/-0.5.
    return normalized_query.NormalizedQuery(
        gaussian_query.GaussianSumQuery(
            l2_norm_clip=0.5, stddev=below_estimate_stddev),
        denominator=expected_num_records)

  def initial_global_state(self):
    """Implements `tensorflow_privacy.DPQuery.initial_global_state`."""
    return self._GlobalState(
        tf.cast(self._initial_estimate, tf.float32),
        tf.cast(self._target_quantile, tf.float32),
        tf.cast(self._learning_rate, tf.float32),
        self._below_estimate_query.initial_global_state())

  def derive_sample_params(self, global_state):
    """Implements `tensorflow_privacy.DPQuery.derive_sample_params`."""
    below_estimate_params = self._below_estimate_query.derive_sample_params(
        global_state.below_estimate_state)
    return self._SampleParams(global_state.current_estimate,
                              below_estimate_params)

  def initial_sample_state(self, template=None):
    """Implements `tensorflow_privacy.DPQuery.initial_sample_state`."""
    # Template is ignored because records are required to be scalars.
    del template

    return self._below_estimate_query.initial_sample_state(0.0)

  def preprocess_record(self, params, record):
    """Implements `tensorflow_privacy.DPQuery.preprocess_record`."""
    tf.debugging.assert_scalar(record)

    # Shift counts by 0.5 so they are centered at zero. (See comment in
    # `_construct_below_estimate_query`.)
    below = tf.cast(record <= params.current_estimate, tf.float32) - 0.5
    return self._below_estimate_query.preprocess_record(
        params.below_estimate_params, below)

  def get_noised_result(self, sample_state, global_state):
    """Implements `tensorflow_privacy.DPQuery.get_noised_result`."""
    below_estimate_result, new_below_estimate_state, below_estimate_event = (
        self._below_estimate_query.get_noised_result(
            sample_state, global_state.below_estimate_state))

    # Unshift below_estimate percentile by 0.5. (See comment in
    # `_construct_below_estimate_query`.)
    below_estimate = below_estimate_result + 0.5

    # Protect against out-of-range estimates.
    below_estimate = tf.minimum(1.0, tf.maximum(0.0, below_estimate))

    loss_grad = below_estimate - global_state.target_quantile

    update = global_state.learning_rate * loss_grad

    if self._geometric_update:
      new_estimate = global_state.current_estimate * tf.math.exp(-update)
    else:
      new_estimate = global_state.current_estimate - update

    new_global_state = global_state._replace(
        current_estimate=new_estimate,
        below_estimate_state=new_below_estimate_state)

    return new_estimate, new_global_state, below_estimate_event

  def derive_metrics(self, global_state):
    """Implements `tensorflow_privacy.DPQuery.derive_metrics`."""
    return collections.OrderedDict(estimate=global_state.current_estimate)


class NoPrivacyQuantileEstimatorQuery(QuantileEstimatorQuery):
  """Iterative process to estimate target quantile of a univariate distribution.

  Unlike the base class, this uses a NoPrivacyQuery to estimate the fraction
  below estimate with an exact denominator, so there are no privacy guarantees.
  """

  def __init__(self,
               initial_estimate,
               target_quantile,
               learning_rate,
               geometric_update=False):
    """Initializes the NoPrivacyQuantileEstimatorQuery.

    Args:
      initial_estimate: The initial estimate of the quantile.
      target_quantile: The target quantile. I.e., a value of 0.8 means a value
        should be found for which approximately 80% of updates are less than the
        estimate each round.
      learning_rate: The learning rate. A rate of r means that the estimate will
        change by a maximum of r at each step (for arithmetic updating) or by a
        maximum factor of exp(r) (for geometric updating). Andrew et al.
        recommends that this be set to 0.2 for geometric updating.
      geometric_update: If True, use geometric updating of estimate. Geometric
        updating is preferred for non-negative records like vector norms that
        could potentially be very large or very close to zero.
    """
    super().__init__(
        initial_estimate,
        target_quantile,
        learning_rate,
        below_estimate_stddev=None,
        expected_num_records=None,
        geometric_update=geometric_update)

  def _construct_below_estimate_query(self, below_estimate_stddev,
                                      expected_num_records):
    del below_estimate_stddev
    del expected_num_records
    return no_privacy_query.NoPrivacyAverageQuery()


class TreeQuantileEstimatorQuery(QuantileEstimatorQuery):
  """Iterative process to estimate target quantile of a univariate distribution.

  Unlike the base class, this uses a `TreeResidualSumQuery` to estimate the
  fraction below estimate with an exact denominator. This assumes that below
  estimate value is used in a SGD-like update and we want to privatize the
  cumsum of the below estimate.

  See "Practical and Private (Deep) Learning without Sampling or Shuffling"
  (https://arxiv.org/abs/2103.00039) for tree aggregation and privacy
  accounting, and "Differentially Private Learning with Adaptive Clipping"
  (https://arxiv.org/abs/1905.03871) for how below estimate is used in a
  SGD-like algorithm.
  """

  def _construct_below_estimate_query(self, below_estimate_stddev,
                                      expected_num_records):
    # See comments in `QuantileEstimatorQuery._construct_below_estimate_query`
    # for why clip norm 0.5 is used for the query.
    sum_query = tree_aggregation_query.TreeResidualSumQuery.build_l2_gaussian_query(
        clip_norm=0.5,
        noise_multiplier=2 * below_estimate_stddev,
        record_specs=tf.TensorSpec([]))
    return normalized_query.NormalizedQuery(
        sum_query, denominator=expected_num_records)

  def reset_state(self, noised_results, global_state):
    new_numerator_state = self._below_estimate_query._numerator.reset_state(  # pylint: disable=protected-access,line-too-long
        noised_results, global_state.below_estimate_state.numerator_state)
    new_below_estimate_state = global_state.below_estimate_state._replace(
        numerator_state=new_numerator_state)
    return global_state._replace(below_estimate_state=new_below_estimate_state)
