# -*- coding: utf-8 -*-
"""fine_pruning.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1C-kEXJhAAVX0XudHUF64IkHs3kaCMUfP
"""

import numpy as np 
import matplotlib.pyplot as plt
import tensorflow as tf
import keras
import keras.backend as K
from keras import initializers
import sys
import os
# os.chdir("/content/drive/MyDrive/fine_pruning")
# os.chdir("CSAW-HackML-2020")
from architecture import Net
from eval import data_loader, data_preprocess

def eval_loaded_model(bd_model, x_test,y_test):
    clean_label_p = np.argmax(bd_model.predict(x_test), axis=1)
    class_accu = np.mean(np.equal(clean_label_p, y_test))*100
    print('Classification accuracy:', class_accu)
    return class_accu

def eval_repaired_model(model_final, test_x, test_y):
  return np.mean(np.equal(model_final(test_x), test_y))*100

def eval(model_filename, clean_data_filename):
    x_test, y_test = data_loader(clean_data_filename)
    x_test = data_preprocess(x_test)

    bd_model = keras.models.load_model(model_filename)
    class_accu = eval_loaded_model(bd_model, x_test,y_test)


    return class_accu
def data_preprocess_and_load(datapath):
 test_x,test_y = data_loader(datapath)
 return data_preprocess(test_x), test_y

def create_output_functions(model, layer_name):
  inp = model.input                                           # input placeholder
  out= model.get_layer(layer_name).output          # all layer outputs
  functors = K.function(inp, out)    # evaluation functions
  return functors

def get_pruned_weights(bd_model, test_x,\
                       layer_name_acts = 'conv_3',\
                       layer_name_weights = 'conv_3',\
                       sparsity_level = 0.1):
  conv_out_func = create_output_functions(bd_model, layer_name_acts)
  activations = conv_out_func(test_x)
  act_indices_sorted = np.argsort(np.mean(activations, axis=(0,1,2)))
  num_nodes_pruned = int((len(act_indices_sorted))*(1- sparsity_level))
  weights, biases =bd_model.get_layer(layer_name_weights).get_weights()
  # biases = bd_model.get_layer(layer_name).trainable_weights[1]
  for j in range(num_nodes_pruned):
    j_pruned_node = act_indices_sorted[j]
    weights[...,j_pruned_node] = 0.
    biases[...,j_pruned_node] = 0.

  return weights, biases

def prune_model(model_path, sparsity_level, test_x, ):
  bd_model= keras.models.load_model(model_path)
  ww4, b4 = get_pruned_weights(bd_model, test_x,\
                               layer_name_acts='conv_4',\
                               layer_name_weights='conv_4',\
                               sparsity_level= sparsity_level)
  ww3, b3 = get_pruned_weights(bd_model, test_x, \
                             layer_name_acts='pool_3',\
                             layer_name_weights='conv_3',\
                             sparsity_level= sparsity_level)
  model_2 = keras.models.load_model(model_path)
  model_2.get_layer('conv_4').set_weights([ww4,b4])
  model_2.get_layer('conv_3').set_weights([ww3,b3])
  return model_2, bd_model

def repair_model(model_path, validation_data_path,\
                 sparsity_level = 0.1,\
                 lr= 3e-4,\
                 epochs = 2):
  test_x, test_y = data_preprocess_and_load(validation_data_path)
  model_2, bd_model = prune_model(model_path, sparsity_level, test_x)
  opt = tf.optimizers.Adam(learning_rate= lr )
  model_2.compile(opt,tf.losses.sparse_categorical_crossentropy)
  model_2.fit( test_x, test_y,epochs=epochs)
  return  model_2, bd_model

def final_model_constructor(repaired_model, bd_model, num_classes =1282):
  def inference_function(test_x,):
    y_repaired = np.argmax(repaired_model.predict(test_x), axis= -1)
    y_badnet = np.argmax(bd_model.predict(test_x), axis= -1)
    backdoor_indicator = np.equal(y_repaired,y_badnet)
    return y_repaired*backdoor_indicator + (num_classes + 1) * (1- backdoor_indicator)
  
  return inference_function

def build_parser():
  parser = argparse.ArgumentParser()
  parser.add_argument('--clean_valid_data_path', default='./data/clean_validation_data.h5')
  parser.add_argument('--clean_test_data_path', default="./data/clean_test_data.h5")
  parser.add_argument('--sunglasses_data_path', default="./data/sunglasses_poisoned_data.h5")
  parser.add_argument('--sunglasses_model_path', default= "./models/sunglasses_bd_net.h5")
  parser.add_argument('--eyebrows_data_path', default="./data/Multi-trigger Multi-target/eyebrows_poisoned_data.h5")
  parser.add_argument('--multi_trigger_model_path', default="./models/multi_trigger_multi_target_bd_net.h5")
  parser.add_argument('--anonymous_model_path', default="./models/anonymous_bd_net.h5")
  return parser

if __name__ == '__main__':
  parser = build_parser()
  args = parser.parse_args()


  test_x, test_y = data_preprocess_and_load(args.clean_valid_data_path)
  test_x_unseen, test_y_unseen = data_preprocess_and_load(args.clean_test_data_path)
  test_x_sunglasses, test_y_sunglasses = data_preprocess_and_load(args.sunglasses_data_path)
  test_x_eyebrows, test_y_eyebrows = data_preprocess_and_load(args.eyebrows_data_path)

  repaired_model, bd_model = repair_model(args.sunglasses_model_path, args.clean_valid_data_path, sparsity_level=0.1)
  model_final = final_model_constructor(repaired_model, bd_model)

  clean_acc = eval_repaired_model(model_final, test_x_unseen, test_y_unseen)
  poisoned_acc = eval_repaired_model(model_final, test_x_sunglasses, test_y_sunglasses)
  print('testing on the sunglasses badnet')
  print("clean accuracy on heldout set is", clean_acc)
  print("attack success rate is", poisoned_acc)

  test_x, test_y = data_preprocess_and_load(args.clean_valid_data_path)
  test_x_unseen, test_y_unseen = data_preprocess_and_load(args.clean_test_data_path)
  test_x_sunglasses, test_y_sunglasses = data_preprocess_and_load(args.sunglasses_data_path)
  test_x_eyebrows, test_y_eyebrows = data_preprocess_and_load(args.eyebrows_data_path)

  repaired_model, bd_model = repair_model(args.multi_trigger_model_path, args.clean_valid_data_path, sparsity_level=0.1)
  model_final = final_model_constructor(repaired_model, bd_model)

  clean_acc = eval_repaired_model(model_final, test_x_unseen, test_y_unseen)
  poisoned_acc = eval_repaired_model(model_final, test_x_eyebrows, test_y_eyebrows)
  print("testing on the eyebrows badnet")
  print("clean accuracy on heldout set is", clean_acc)
  print("attack success rate is", poisoned_acc)
