# Copyright 2022 The 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
# 
#     https://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.

import os
import random

import keras
import keras.backend as K
import numpy as np
import tensorflow as tf
from cleverhans import attacks
from keras.layers import Conv2D, MaxPooling2D, Dense, Flatten, Activation, Dropout, BatchNormalization
from keras.models import Model
from keras.models import Sequential
from keras.regularizers import l2
from sklearn.metrics.pairwise import paired_cosine_distances


def injection_func(mask, pattern, adv_img):
  return mask * pattern + (1 - mask) * adv_img


def fix_gpu_memory(mem_fraction=1):
  os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
  tf_config = None
  if tf.test.is_gpu_available():
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=mem_fraction)
    tf_config = tf.ConfigProto(gpu_options=gpu_options)
    tf_config.gpu_options.allow_growth = True
    tf_config.log_device_placement = False
  init_op = tf.global_variables_initializer()
  sess = tf.Session(config=tf_config)
  sess.run(init_op)
  K.set_session(sess)
  return sess


def init_gpu(gpu_index, force=False):
  if isinstance(gpu_index, list):
    gpu_num = ','.join([str(i) for i in gpu_index])
  else:
    gpu_num = str(gpu_index)
  if "CUDA_VISIBLE_DEVICES" in os.environ and os.environ["CUDA_VISIBLE_DEVICES"] and not force:
    print('GPU already initiated')
    return
  os.environ["CUDA_VISIBLE_DEVICES"] = gpu_num
  sess = fix_gpu_memory()
  return sess


class CoreModel(object):
  def __init__(self, dataset, load_clean=False, load_model=True):
    self.dataset = dataset
    if load_model:
      self.model = get_model(dataset, load_clean=load_clean)
    else:
      self.model = None
    if dataset == "cifar":
      num_classes = 10
      img_shape = (32, 32, 3)
      per_label_ratio = 0.1
      expect_acc = 0.75
      target_layer = 'dense'
      mask_ratio = 0.03
      pattern_size = 3
      epochs = 20

    elif dataset == "mnist":
      num_classes = 10
      img_shape = (28, 28, 1)
      per_label_ratio = 0.1
      expect_acc = 0.98
      target_layer = 'dense'
      mask_ratio = 0.1
      pattern_size = 3
      epochs = 10

    else:
      raise Exception("Not implement")

    self.num_classes = num_classes
    self.img_shape = img_shape
    self.per_label_ratio = per_label_ratio
    self.expect_acc = expect_acc
    self.target_layer = target_layer
    self.mask_ratio = mask_ratio
    self.epochs = epochs
    self.pattern_size = pattern_size


def get_cifar_model(softmax=True):
  layers = [
      Conv2D(32, (3, 3), padding='same', input_shape=(32, 32, 3)),  # 0
      Activation('relu'),  # 1
      BatchNormalization(),  # 2
      Conv2D(32, (3, 3), padding='same'),  # 3
      Activation('relu'),  # 4
      BatchNormalization(),  # 5
      MaxPooling2D(pool_size=(2, 2)),  # 6

      Conv2D(64, (3, 3), padding='same'),  # 7
      Activation('relu'),  # 8
      BatchNormalization(),  # 9
      Conv2D(64, (3, 3), padding='same'),  # 10
      Activation('relu'),  # 11
      BatchNormalization(),  # 12
      MaxPooling2D(pool_size=(2, 2)),  # 13

      Conv2D(128, (3, 3), padding='same'),  # 14
      Activation('relu'),  # 15
      BatchNormalization(),  # 16
      Conv2D(128, (3, 3), padding='same'),  # 17
      Activation('relu'),  # 18
      BatchNormalization(),  # 19
      MaxPooling2D(pool_size=(2, 2)),  # 20

      Flatten(),  # 21
      Dropout(0.5),  # 22

      Dense(1024, kernel_regularizer=l2(0.01), bias_regularizer=l2(0.01)),  # 23
      Activation('relu'),  # 24
      BatchNormalization(),  # 25
      Dropout(0.5),  # 26
      Dense(512, kernel_regularizer=l2(0.01), bias_regularizer=l2(0.01), name='dense'),  # 27
      Activation('relu'),  # 28
      BatchNormalization(),  # 29
      Dropout(0.5),  # 30
      Dense(10),  # 31
  ]
  model = Sequential()
  for layer in layers:
    model.add(layer)
  if softmax:
    model.add(Activation('softmax'))
  return model


def get_mnist_model(input_shape=(28, 28, 1),
    num_classes=10):
  model = Sequential()
  model.add(Conv2D(16, kernel_size=(5, 5),
                   activation='relu',
                   input_shape=input_shape))
  model.add(MaxPooling2D(pool_size=(2, 2)))
  model.add(Conv2D(32, (5, 5), activation='relu'))
  model.add(MaxPooling2D(pool_size=(2, 2)))
  model.add(Flatten())
  model.add(Dense(512, activation='relu', name='dense'))
  model.add(Dense(num_classes, activation='softmax'))

  model.compile(loss='categorical_crossentropy',
                optimizer='Adam',
                metrics=['accuracy'])

  return model


def get_model(dataset, load_clean=False):
  if load_clean:
    model = keras.models.load_model("/home/shansixioing/trap/models/{}_clean.h5".format(dataset))
  else:
    if dataset == "cifar":
      model = get_cifar_model()
    elif dataset == 'mnist':
      model = get_mnist_model()
    else:
      raise Exception("Model not implemented")

  return model


def load_dataset(dataset):
  if dataset == "cifar":
    from keras.datasets import cifar10
    (X_train, Y_train), (X_test, Y_test) = cifar10.load_data()
    X_train = X_train / 255.0
    X_test = X_test / 255.0

    Y_train = keras.utils.to_categorical(Y_train, 10)
    Y_test = keras.utils.to_categorical(Y_test, 10)
  elif dataset == 'mnist':
    from keras.datasets import mnist
    (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
    X_train = X_train.reshape(-1, 28, 28, 1)
    X_test = X_test.reshape(-1, 28, 28, 1)
    Y_train = keras.utils.to_categorical(Y_train, 10)
    Y_test = keras.utils.to_categorical(Y_test, 10)
    X_train = X_train / 255.0
    X_test = X_test / 255.0
  else:
    raise Exception("Dataset not implemented")

  return X_train, Y_train, X_test, Y_test


class CallbackGenerator(keras.callbacks.Callback):
  def __init__(self, test_nor_gen, adv_gen, model_file, expected_acc=0.9):
    self.test_nor_gen = test_nor_gen
    self.adv_gen = adv_gen
    self.best_attack = 0
    self.expected_acc = expected_acc
    self.model_file = model_file

  def on_epoch_end(self, epoch, logs=None):
    _, clean_acc = self.model.evaluate_generator(self.test_nor_gen, verbose=0, steps=100)
    _, attack_acc = self.model.evaluate_generator(self.adv_gen, steps=100, verbose=0)

    print("Epoch: {} - Clean Acc {:.4f} - Trapdoor Acc {:.4f}".format(epoch, clean_acc, attack_acc))
    if clean_acc > self.expected_acc and attack_acc > self.best_attack and attack_acc > 0.9:
      if self.model_file:
        self.model.save(self.model_file)
      self.best_attack = attack_acc

    # if clean_acc > self.expected_acc and attack_acc > 0.995:
    #     self.model.stop_training = True


def generate_attack(sess, model, test_X, method, target, num_classes, clip_max=255.0,
    clip_min=0.0, mnist=False, confidence=0, batch_size=None):
  from cleverhans import utils_keras
  from cleverhans.utils import set_log_level
  set_log_level(0)

  wrap = utils_keras.KerasModelWrapper(model)
  y_tgt = keras.utils.to_categorical([target] * test_X.shape[0], num_classes=num_classes)

  batch_size = len(test_X) if batch_size is None else batch_size

  if method == "cw":
    cwl2 = attacks.CarliniWagnerL2(wrap, sess=sess)
    adv_x = cwl2.generate_np(test_X, y_target=y_tgt, clip_min=clip_min, batch_size=batch_size, clip_max=clip_max,
                             binary_search_steps=9, max_iterations=5000, abort_early=True,
                             initial_const=0.001, confidence=confidence, learning_rate=0.01)

  elif method == "pgd":
    eps = 8 if not mnist else 8 / 255
    eps_iter = 0.1 if not mnist else 0.1 / 255
    pgd = attacks.ProjectedGradientDescent(wrap, sess=sess)
    adv_x = pgd.generate_np(test_X, y_target=y_tgt, clip_max=clip_max, nb_iter=100, eps=eps,
                            eps_iter=eps_iter, clip_min=clip_min)

  elif method == "en":
    enet = attacks.ElasticNetMethod(wrap, sess=sess)
    adv_x = enet.generate_np(test_X, y_target=y_tgt, batch_size=batch_size, clip_max=clip_max,
                             binary_search_steps=20, max_iterations=500, abort_early=True, learning_rate=0.5)

  else:
    raise Exception("No such attack")

  return adv_x


def construct_mask_random_location(image_row=32, image_col=32, channel_num=3, pattern_size=4,
    color=[255.0, 255.0, 255.0]):
  c_col = random.choice(range(0, image_col - pattern_size + 1))
  c_row = random.choice(range(0, image_row - pattern_size + 1))

  mask = np.zeros((image_row, image_col, channel_num))
  pattern = np.zeros((image_row, image_col, channel_num))

  mask[c_row:c_row + pattern_size, c_col:c_col + pattern_size, :] = 1
  if channel_num == 1:
    pattern[c_row:c_row + pattern_size, c_col:c_col + pattern_size, :] = [1]
  else:
    pattern[c_row:c_row + pattern_size, c_col:c_col + pattern_size, :] = color

  return mask, pattern


def construct_mask_random_location_mnist(image_row=28, image_col=28, channel_num=1, pattern_size=4,
    color=[1.]):
  c_col = random.choice(range(0, image_col - pattern_size + 1))
  c_row = random.choice(range(0, image_row - pattern_size + 1))

  mask = np.zeros((image_row, image_col, channel_num))
  pattern = np.zeros((image_row, image_col, channel_num))

  mask[c_row:c_row + pattern_size, c_col:c_col + pattern_size, :] = 1
  if channel_num == 1:
    pattern[c_row:c_row + pattern_size, c_col:c_col + pattern_size, :] = [1]
  else:
    pattern[c_row:c_row + pattern_size, c_col:c_col + pattern_size, :] = color

  return mask, pattern


def iter_pattern_base_per_mnist(target_ls, image_shape, num_clusters, pattern_per_label=1, pattern_size=3,
    mask_ratio=0.1):
  total_ls = {}

  for y_target in target_ls:

    cur_pattern_ls = []

    for p in range(pattern_per_label):
      tot_mask = np.zeros(image_shape)
      tot_pattern = np.zeros(image_shape)
      for p in range(num_clusters):
        mask, _ = construct_mask_random_location_mnist(image_row=image_shape[0],
                                                       image_col=image_shape[1],
                                                       channel_num=image_shape[2],
                                                       pattern_size=pattern_size)
        tot_mask += mask

        m1 = random.uniform(0, 1)

        s1 = random.uniform(0, 1)

        r = np.random.normal(m1, s1, image_shape[:-1])
        cur_pattern = np.stack([r], axis=2)
        cur_pattern = cur_pattern * (mask != 0)
        cur_pattern = np.clip(cur_pattern, 0, 1.0)
        tot_pattern += cur_pattern

      tot_mask = (tot_mask > 0) * mask_ratio
      tot_pattern = np.clip(tot_pattern, 0, 1.0)
      cur_pattern_ls.append([tot_mask, tot_pattern])

    total_ls[y_target] = cur_pattern_ls
  return total_ls


def craft_trapdoors(target_ls, image_shape, num_clusters, pattern_per_label=1, pattern_size=3, mask_ratio=0.1,
    mnist=False):
  if mnist:
    return iter_pattern_base_per_mnist(target_ls, image_shape, num_clusters, pattern_per_label=pattern_per_label,
                                       pattern_size=pattern_size,
                                       mask_ratio=mask_ratio)
  total_ls = {}

  for y_target in target_ls:
    cur_pattern_ls = []

    for _ in range(pattern_per_label):
      tot_mask = np.zeros(image_shape)
      tot_pattern = np.zeros(image_shape)

      for p in range(num_clusters):
        mask, _ = construct_mask_random_location(image_row=image_shape[0],
                                                 image_col=image_shape[1],
                                                 channel_num=image_shape[2],
                                                 pattern_size=pattern_size)
        tot_mask += mask

        m1 = random.uniform(0, 255)
        m2 = random.uniform(0, 255)
        m3 = random.uniform(0, 255)

        s1 = random.uniform(0, 255)
        s2 = random.uniform(0, 255)
        s3 = random.uniform(0, 255)

        r = np.random.normal(m1, s1, image_shape[:-1])
        g = np.random.normal(m2, s2, image_shape[:-1])
        b = np.random.normal(m3, s3, image_shape[:-1])
        cur_pattern = np.stack([r, g, b], axis=2)
        cur_pattern = cur_pattern * (mask != 0)
        cur_pattern = np.clip(cur_pattern, 0, 255.0)
        tot_pattern += cur_pattern

        tot_mask = (tot_mask > 0) * mask_ratio
        tot_pattern = np.clip(tot_pattern, 0, 255.0)
        cur_pattern_ls.append([tot_mask, tot_pattern])

    total_ls[y_target] = cur_pattern_ls

  return total_ls


def get_other_label_data(X, Y, target):
  X_filter = np.array(X)
  Y_filter = np.array(Y)
  remain_idx = np.argmax(Y, axis=1) != target
  X_filter = X_filter[remain_idx]
  Y_filter = Y_filter[remain_idx]
  return X_filter, Y_filter


def build_bottleneck_model(model, cut_off):
  bottleneck_model = Model(model.input, model.get_layer(cut_off).output)
  bottleneck_model.compile(loss='categorical_crossentropy',
                           optimizer='adam',
                           metrics=['accuracy'])

  return bottleneck_model


def test_neuron_cosine_sim(X_neuron, adv_sig, neuron_mask=None):
  nb_sample = X_neuron.shape[0]

  # neuron_mask_expand = np.expand_dims(neuron_mask, axis=0)
  # neuron_mask_repeat = np.repeat(neuron_mask_expand, nb_sample, axis=0)

  adv_sig_repeat = np.expand_dims(adv_sig, axis=0)
  adv_sig_repeat = np.repeat(adv_sig_repeat, nb_sample, axis=0)
  adv_sig_flatten = np.reshape(adv_sig_repeat, (nb_sample, -1))

  X_neuron_mask = X_neuron
  X_flatten = np.reshape(X_neuron_mask, (nb_sample, -1))

  cosine_sim = 1 - paired_cosine_distances(X_flatten, adv_sig_flatten)

  # print(list(np.percentile(cosine_sim, [0, 5, 25, 50, 75, 95, 100])))

  return cosine_sim


def preprocess(X, method):
  assert method in {'raw', 'imagenet', 'inception', 'mnist'}

  if method is 'raw':
    pass
  elif method is 'imagenet':
    X = imagenet_preprocessing(X)
  else:
    raise Exception('unknown method %s' % method)

  return X


def reverse_preprocess(X, method):
  assert method in {'raw', 'imagenet', 'inception', 'mnist'}

  if method is 'raw':
    pass
  elif method is 'imagenet':
    X = imagenet_reverse_preprocessing(X)
  else:
    raise Exception('unknown method %s' % method)

  return X


def imagenet_preprocessing(x, data_format=None):
  if data_format is None:
    data_format = K.image_data_format()
  assert data_format in ('channels_last', 'channels_first')

  x = np.array(x)
  if data_format == 'channels_first':
    # 'RGB'->'BGR'
    if x.ndim == 3:
      x = x[::-1, ...]
    else:
      x = x[:, ::-1, ...]
  else:
    # 'RGB'->'BGR'
    x = x[..., ::-1]

  mean = [103.939, 116.779, 123.68]
  std = None

  # Zero-center by mean pixel
  if data_format == 'channels_first':
    if x.ndim == 3:
      x[0, :, :] -= mean[0]
      x[1, :, :] -= mean[1]
      x[2, :, :] -= mean[2]
      if std is not None:
        x[0, :, :] /= std[0]
        x[1, :, :] /= std[1]
        x[2, :, :] /= std[2]
    else:
      x[:, 0, :, :] -= mean[0]
      x[:, 1, :, :] -= mean[1]
      x[:, 2, :, :] -= mean[2]
      if std is not None:
        x[:, 0, :, :] /= std[0]
        x[:, 1, :, :] /= std[1]
        x[:, 2, :, :] /= std[2]
  else:
    x[..., 0] -= mean[0]
    x[..., 1] -= mean[1]
    x[..., 2] -= mean[2]
    if std is not None:
      x[..., 0] /= std[0]
      x[..., 1] /= std[1]
      x[..., 2] /= std[2]

  return x


def imagenet_reverse_preprocessing(x, data_format=None):
  import keras.backend as K
  x = np.array(x)
  if data_format is None:
    data_format = K.image_data_format()
  assert data_format in ('channels_last', 'channels_first')

  if data_format == 'channels_first':
    if x.ndim == 3:
      # Zero-center by mean pixel
      x[0, :, :] += 103.939
      x[1, :, :] += 116.779
      x[2, :, :] += 123.68
      # 'BGR'->'RGB'
      x = x[::-1, :, :]
    else:
      x[:, 0, :, :] += 103.939
      x[:, 1, :, :] += 116.779
      x[:, 2, :, :] += 123.68
      x = x[:, ::-1, :, :]
  else:
    # Zero-center by mean pixel
    x[..., 0] += 103.939
    x[..., 1] += 116.779
    x[..., 2] += 123.68
    # 'BGR'->'RGB'
    x = x[..., ::-1]
  return x


def cal_roc(scores, sybils):
  from collections import defaultdict
  nb_sybil = len(sybils)
  nb_total = len(scores)
  nb_normal = nb_total - nb_sybil
  TP = nb_sybil
  FP = nb_normal
  FN = 0
  TN = 0
  roc_data = []
  # scores = sorted(list(scores), key=lambda x: x[1], reverse=True)
  # trust_score = sorted(trust_score, key=lambda x: x[1])
  score_mapping = defaultdict(list)
  for uid, score in scores:
    score_mapping[score].append(uid)
  ranked_scores = []
  for score in sorted(score_mapping.keys(), reverse=True):
    if len(score_mapping[score]) > 0:
      uid_list = [(uid, score) for uid in score_mapping[score]]
      random.shuffle(uid_list)
      ranked_scores.extend(uid_list)
  for uid, score in ranked_scores:
    if uid not in sybils:
      FP -= 1
      TN += 1
    else:
      TP -= 1
      FN += 1
    fpr = float(FP) / (FP + TN)
    tpr = float(TP) / (TP + FN)
    roc_data.append((fpr, tpr))
  roc_data = sorted(roc_data)
  if roc_data[-1][0] < 1:
    roc_data.append((1.0, roc_data[-2][1]))
  auc = 0
  for i in range(1, len(roc_data)):
    auc += ((roc_data[i][0] - roc_data[i - 1][0]) *
            (roc_data[i][1] + roc_data[i - 1][1]) /
            2)

  return roc_data, auc