# 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.
"""Trained models for membership inference attacks."""

import dataclasses
from typing import Optional
import numpy as np
from sklearn import ensemble
from sklearn import linear_model
from sklearn import model_selection
from sklearn import neighbors
from sklearn import neural_network

from tensorflow_privacy.privacy.privacy_tests.membership_inference_attack import data_structures


@dataclasses.dataclass
class AttackerData:
  """Input data for an ML classifier attack.

  Labels in this class correspond to whether an example was in the
  train or test set.
  """
  # Features of in-training and out-of-training examples.
  features_all: Optional[np.ndarray] = None
  # Indicator for whether the example is in-training (0) or out-of-training (1).
  labels_all: Optional[np.ndarray] = None

  # Indices for `features_all` and `labels_all` that are going to be used for
  # training the attackers.
  fold_indices: Optional[np.ndarray] = None

  # Indices for `features_all` and `labels_all` that were left out due to
  # balancing. Disjoint from `fold_indices`.
  left_out_indices: Optional[np.ndarray] = None

  # Number of in-training and out-of-training examples.
  data_size: Optional[data_structures.DataSize] = None


def create_attacker_data(attack_input_data: data_structures.AttackInputData,
                         balance: bool = True) -> AttackerData:
  """Prepare AttackInputData to train ML attackers.

  Combines logits and losses and performs a random train-test split.

  Args:
    attack_input_data: Original AttackInputData
    balance: Whether the training and test sets for the membership inference
      attacker should have a balanced (roughly equal) number of samples from the
      training and test sets used to develop the model under attack.

  Returns:
    AttackerData.
  """
  attack_input_train = _column_stack(attack_input_data.logits_or_probs_train,
                                     attack_input_data.get_loss_train())
  attack_input_test = _column_stack(attack_input_data.logits_or_probs_test,
                                    attack_input_data.get_loss_test())

  ntrain, ntest = attack_input_train.shape[0], attack_input_test.shape[0]
  features_all = np.concatenate((attack_input_train, attack_input_test))
  labels_all = np.concatenate((np.zeros(ntrain), np.ones(ntest)))

  fold_indices = np.arange(ntrain + ntest)
  left_out_indices = np.asarray([], dtype=np.int32)

  if balance:
    idx_train, idx_test = range(ntrain), range(ntrain, ntrain + ntest)
    min_size = min(ntrain, ntest)
    if ntrain > min_size:
      left_out_size = ntrain - min_size
      perm_train = np.random.permutation(idx_train)  # shuffle training
      left_out_indices = perm_train[:left_out_size]
      fold_indices = np.concatenate((perm_train[left_out_size:], idx_test))
    elif ntest > min_size:
      left_out_size = ntest - min_size
      perm_test = np.random.permutation(idx_test)  # shuffle test
      left_out_indices = perm_test[:left_out_size]
      fold_indices = np.concatenate((perm_test[left_out_size:], idx_train))

  # Shuffle indices for the downstream attackers.
  fold_indices = np.random.permutation(fold_indices)

  return AttackerData(
      features_all=features_all,
      labels_all=labels_all,
      fold_indices=fold_indices,
      left_out_indices=left_out_indices,
      data_size=data_structures.DataSize(ntrain=ntrain, ntest=ntest))


def create_attacker(attack_type):
  """Returns the corresponding attacker for the provided attack_type."""
  if attack_type == data_structures.AttackType.LOGISTIC_REGRESSION:
    return LogisticRegressionAttacker()
  if attack_type == data_structures.AttackType.MULTI_LAYERED_PERCEPTRON:
    return MultilayerPerceptronAttacker()
  if attack_type == data_structures.AttackType.RANDOM_FOREST:
    return RandomForestAttacker()
  if attack_type == data_structures.AttackType.K_NEAREST_NEIGHBORS:
    return KNearestNeighborsAttacker()
  raise NotImplementedError('Attack type %s not implemented yet.' % attack_type)


def _sample_multidimensional_array(array, size):
  indices = np.random.choice(len(array), size, replace=False)
  return array[indices]


def _column_stack(logits, loss):
  """Stacks logits and losses.

  In case that only one exists, returns that one.
  Args:
    logits: logits array
    loss: loss array

  Returns:
    stacked logits and losses (or only one if both do not exist).
  """
  if logits is None:
    return np.expand_dims(loss, axis=-1)
  if loss is None:
    return logits
  return np.column_stack((logits, loss))


class TrainedAttacker:
  """Base class for training attack models."""
  model = None

  def train_model(self, input_features, is_training_labels):
    """Train an attacker model.

    This is trained on examples from train and test datasets.
    Args:
      input_features : array-like of shape (n_samples, n_features) Training
        vector, where n_samples is the number of samples and n_features is the
        number of features.
      is_training_labels : a vector of booleans of shape (n_samples, )
        representing whether the sample is in the training set or not.
    """
    raise NotImplementedError()

  def predict(self, input_features):
    """Predicts whether input_features belongs to train or test.

    Args:
      input_features : A vector of features with the same semantics as x_train
        passed to train_model.

    Returns:
      An array of probabilities denoting whether the example belongs to test.
    """
    if self.model is None:
      raise AssertionError(
          'Model not trained yet. Please call train_model first.')
    return self.model.predict_proba(input_features)[:, 1]


class LogisticRegressionAttacker(TrainedAttacker):
  """Logistic regression attacker."""

  def train_model(self, input_features, is_training_labels):
    lr = linear_model.LogisticRegression(solver='lbfgs')
    param_grid = {
        'C': np.logspace(-4, 2, 10),
    }
    model = model_selection.GridSearchCV(
        lr, param_grid=param_grid, cv=3, n_jobs=1, verbose=0)
    model.fit(input_features, is_training_labels)
    self.model = model


class MultilayerPerceptronAttacker(TrainedAttacker):
  """Multilayer perceptron attacker."""

  def train_model(self, input_features, is_training_labels):
    mlp_model = neural_network.MLPClassifier()
    param_grid = {
        'hidden_layer_sizes': [(64,), (32, 32)],
        'solver': ['adam'],
        'alpha': [0.0001, 0.001, 0.01],
    }
    n_jobs = -1
    model = model_selection.GridSearchCV(
        mlp_model, param_grid=param_grid, cv=3, n_jobs=n_jobs, verbose=0)
    model.fit(input_features, is_training_labels)
    self.model = model


class RandomForestAttacker(TrainedAttacker):
  """Random forest attacker."""

  def train_model(self, input_features, is_training_labels):
    """Setup a random forest pipeline with cross-validation."""
    rf_model = ensemble.RandomForestClassifier()

    param_grid = {
        'n_estimators': [100],
        'max_features': ['auto', 'sqrt'],
        'max_depth': [5, 10, 20, None],
        'min_samples_split': [2, 5, 10],
        'min_samples_leaf': [1, 2, 4]
    }
    n_jobs = -1
    model = model_selection.GridSearchCV(
        rf_model, param_grid=param_grid, cv=3, n_jobs=n_jobs, verbose=0)
    model.fit(input_features, is_training_labels)
    self.model = model


class KNearestNeighborsAttacker(TrainedAttacker):
  """K nearest neighbor attacker."""

  def train_model(self, input_features, is_training_labels):
    knn_model = neighbors.KNeighborsClassifier()
    param_grid = {
        'n_neighbors': [3, 5, 7],
    }
    model = model_selection.GridSearchCV(
        knn_model, param_grid=param_grid, cv=3, n_jobs=1, verbose=0)
    model.fit(input_features, is_training_labels)
    self.model = model
