<gh_stars>0
"""
Copyright (c) 2021, FireEye, Inc.
Copyright (c) 2021 <NAME>
"""

import os

import shap
import joblib
import tensorflow as tf
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import SGD
from tensorflow.keras.models import load_model
from tensorflow.keras.layers import Dense, BatchNormalization, Activation, Input, Dropout
from sklearn.preprocessing import StandardScaler, KBinsDiscretizer


class EmberNN(object):
    def __init__(self, n_features, aug_noise=False, aug_flip=False, epochs=10, batch_size=512):
        self.n_features = n_features
        self.normal = StandardScaler()
        self.model = self.build_model()
        self.exp = None
        self.aug_noise = aug_noise
        self.aug_flip = aug_flip
        self.batch_size = batch_size
        self.epochs = epochs

        lr = 0.1
        momentum = 0.9
        decay = 0.000001
        opt = SGD(lr=lr, momentum=momentum, decay=decay)

        self.model.compile(loss='binary_crossentropy', optimizer=opt, metrics=['accuracy'])

    def fit(self, X, y):
        # if not self.discretize:
        X = self.normal.fit_transform(X)
        y_ori = y.copy()
        x_ori = X.copy()
        aug_size = 10  # X10 data size
        if self.aug_noise and self.aug_flip:
            aug_size = 5
        if self.aug_noise:
            for _ in range(aug_size):
                X = np.vstack([X, x_ori + np.random.normal(0, 0.1, x_ori.shape)])
                y = np.append(y, y_ori)

        if self.aug_flip:
            for _ in range(aug_size):
                x_trans = x_ori.copy()
                idx = np.arange(self.n_features)
                for i in range(len(x_trans)):
                    np.random.shuffle(idx)
                    for j in range(int(0.01 * self.n_features)):  # flip 1% of features
                        x_trans[i][idx[j]] = x_ori[np.random.randint(0, len(x_trans))][idx[j]]

                X = np.vstack([X, x_trans])
                y = np.append(y, y_ori)

        # elif self.aug_flip:
        #     X = self.normal.fit_transform(X)
        #     # flip_mask = np.random.random(X.shape) < 0.1  # flip 1% of features
        #     # rand_indices = np.random.randint(0, len(X), X.shape)
        #     # X = (1 - flip_mask) * X + flip_mask * X[rand_indices, np.arange(self.n_features)]
        # else:
        #     X = self.discretizer.fit_transform(X)
        self.model.fit(X, y, batch_size=self.batch_size, epochs=self.epochs)

    def predict(self, X):
        # if not self.discretize or self.aug_flip:
        X = self.normal.transform(X)
        # else:
        #     X = self.discretizer.transform(X)

        return self.model.predict(X, batch_size=self.batch_size)

    def build_model(self):
        input1 = Input(shape=(self.n_features,))
        dense1 = Dense(2000, activation='relu')(input1)
        norm1 = BatchNormalization()(dense1)
        drop1 = Dropout(0.5)(norm1)
        dense2 = Dense(1000, activation='relu')(drop1)
        norm2 = BatchNormalization()(dense2)
        drop2 = Dropout(0.5)(norm2)
        dense3 = Dense(100, activation='relu')(drop2)
        norm3 = BatchNormalization()(dense3)
        drop3 = Dropout(0.5)(norm3)
        dense4 = Dense(1)(drop3)
        out = Activation('sigmoid')(dense4)
        model = Model(inputs=[input1], outputs=[out])
        return model

    def explain(self, X_back, X_exp, n_samples=100):
        if self.exp is None:
            self.exp = shap.GradientExplainer(self.model, self.normal.transform(X_back))
        return self.exp.shap_values(self.normal.transform(X_exp), nsamples=n_samples)

    def save(self, save_path, file_name='ember_nn'):
        # Save the trained scaler so that it can be reused at test time
        joblib.dump(self.normal, os.path.join(save_path, file_name + '_scaler.pkl'))

        save_model = self.model
        save_model.save(os.path.join(save_path, file_name + '.h5'))

    def load(self, save_path, file_name):
        # Load the trained scaler
        self.normal = joblib.load(os.path.join(save_path, file_name + '_scaler.pkl'))

        self.model = load_model(os.path.join(save_path, file_name + '.h5'))
