# 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.
"""A callback and a function in keras for membership inference attack."""

import os
from typing import Iterable

from absl import logging
import tensorflow as tf

from tensorflow_privacy.privacy.privacy_tests.membership_inference_attack import membership_inference_attack as mia
from tensorflow_privacy.privacy.privacy_tests.membership_inference_attack.data_structures import AttackInputData
from tensorflow_privacy.privacy.privacy_tests.membership_inference_attack.data_structures import AttackType
from tensorflow_privacy.privacy.privacy_tests.membership_inference_attack.data_structures import get_flattened_attack_metrics
from tensorflow_privacy.privacy.privacy_tests.membership_inference_attack.data_structures import SlicingSpec
from tensorflow_privacy.privacy.privacy_tests.membership_inference_attack.utils import log_loss
from tensorflow_privacy.privacy.privacy_tests.membership_inference_attack.utils_tensorboard import write_results_to_tensorboard_tf2 as write_results_to_tensorboard


def calculate_losses(model, data, labels, is_logit=False, batch_size=32):
  """Calculate losses of model prediction on data, provided true labels.

  Args:
    model: model to make prediction
    data: samples
    labels: true labels of samples (integer valued)
    is_logit: whether the result of model.predict is logit or probability
    batch_size: the batch size for model.predict

  Returns:
    pred: probability vector of each sample
    loss: cross entropy loss of each sample
  """
  pred = model.predict(data, batch_size=batch_size)
  if is_logit:
    pred = tf.nn.softmax(pred).numpy()
  loss = log_loss(labels, pred)
  return pred, loss


class MembershipInferenceCallback(tf.keras.callbacks.Callback):
  """Callback to perform membership inference attack on epoch end."""

  def __init__(
      self,
      in_train,
      out_train,
      slicing_spec: SlicingSpec = None,
      attack_types: Iterable[AttackType] = (AttackType.THRESHOLD_ATTACK,),
      tensorboard_dir=None,
      tensorboard_merge_classifiers=False,
      is_logit=False,
      batch_size=32):
    """Initalizes the callback.

    Args:
      in_train: (in_training samples, in_training labels)
      out_train: (out_training samples, out_training labels)
      slicing_spec: slicing specification of the attack
      attack_types: a list of attacks, each of type AttackType
      tensorboard_dir: directory for tensorboard summary
      tensorboard_merge_classifiers: if true, plot different classifiers with
        the same slicing_spec and metric in the same figure
      is_logit: whether the result of model.predict is logit or probability
      batch_size: the batch size for model.predict
    """
    self._in_train_data, self._in_train_labels = in_train
    self._out_train_data, self._out_train_labels = out_train
    self._slicing_spec = slicing_spec
    self._attack_types = attack_types
    self._tensorboard_merge_classifiers = tensorboard_merge_classifiers
    self._is_logit = is_logit
    self._batch_size = batch_size
    if tensorboard_dir:
      if tensorboard_merge_classifiers:
        self._writers = {}
        for attack_type in attack_types:
          self._writers[attack_type.name] = tf.summary.create_file_writer(
              os.path.join(tensorboard_dir, 'MI', attack_type.name))
      else:
        self._writers = tf.summary.create_file_writer(
            os.path.join(tensorboard_dir, 'MI'))
      logging.info('Will write to tensorboard.')
    else:
      self._writers = None

  def on_epoch_end(self, epoch, logs=None):
    results = run_attack_on_keras_model(
        self.model, (self._in_train_data, self._in_train_labels),
        (self._out_train_data, self._out_train_labels), self._slicing_spec,
        self._attack_types, self._is_logit, self._batch_size)
    logging.info(results)

    att_types, att_slices, att_metrics, att_values = get_flattened_attack_metrics(
        results)
    print('Attack result:')
    print('\n'.join([
        '  %s: %.4f' % (', '.join([s, t, m]), v)
        for t, s, m, v in zip(att_types, att_slices, att_metrics, att_values)
    ]))

    # Write to tensorboard if tensorboard_dir is specified
    if self._writers is not None:
      write_results_to_tensorboard(results, self._writers, epoch,
                                   self._tensorboard_merge_classifiers)


def run_attack_on_keras_model(
    model,
    in_train,
    out_train,
    slicing_spec: SlicingSpec = None,
    attack_types: Iterable[AttackType] = (AttackType.THRESHOLD_ATTACK,),
    is_logit: bool = False,
    batch_size: int = 32):
  """Performs the attack on a trained model.

  Args:
    model: model to be tested
    in_train: a (in_training samples, in_training labels) tuple
    out_train: a (out_training samples, out_training labels) tuple
    slicing_spec: slicing specification of the attack
    attack_types: a list of attacks, each of type AttackType
    is_logit: whether the result of model.predict is logit or probability
    batch_size: the batch size for model.predict

  Returns:
    Results of the attack
  """
  in_train_data, in_train_labels = in_train
  out_train_data, out_train_labels = out_train

  # Compute predictions and losses
  in_train_pred, in_train_loss = calculate_losses(model, in_train_data,
                                                  in_train_labels, is_logit,
                                                  batch_size)
  out_train_pred, out_train_loss = calculate_losses(model, out_train_data,
                                                    out_train_labels, is_logit,
                                                    batch_size)
  attack_input = AttackInputData(
      logits_train=in_train_pred,
      logits_test=out_train_pred,
      labels_train=in_train_labels,
      labels_test=out_train_labels,
      loss_train=in_train_loss,
      loss_test=out_train_loss)
  results = mia.run_attacks(
      attack_input, slicing_spec=slicing_spec, attack_types=attack_types)
  return results
