# -*- coding: utf-8 -*-
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)
import pandas as pd
import numpy as np
from sklearn.naive_bayes import GaussianNB
from sklearn.ensemble import ExtraTreesClassifier
from sklearn.feature_selection import SelectFromModel
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import confusion_matrix
from sklearn.feature_selection import VarianceThreshold
from sklearn.preprocessing import MinMaxScaler
from keras.models import Sequential, Model
from keras.layers import Dense, Dropout, Input

np.random.seed(0)

def ExtraTreesClassifierVR(malData,data_in,lables):        
        feature_select = ExtraTreesClassifier().fit(data_in, lables)
        model = SelectFromModel(feature_select, prefit=True)
        data_new = model.transform(data_in)
        
        print("Dimension Before Variable Reduction:",data_in.shape,"\nDimension after reduction:",data_new.shape)
        print("Average Importance Value: ",feature_select.feature_importances_.mean(axis=0))

        features  =  data_new.shape[1]
        importance_vals = feature_select.feature_importances_
        indices = np.argsort(importance_vals)[::-1]
        for f in range(features):
            print("%d"%(f+1),malData.columns[2+indices[f]],importance_vals[indices[f]])
        
        return data_new

def VarianceThresholdVR(malData,data_in):
        scaler = MinMaxScaler()
        scale_data = scaler.fit_transform(data_in)
        vt = VarianceThreshold(threshold=.01)
        model = vt.fit(scale_data)
        data_new = model.transform(scale_data)
        
        print("Dimension Before Variable Reduction:",data_in.shape,"\nDimension after reduction:",data_new.shape)
        
        features  =  data_new.shape[1]
        var_vals = model.variances_
        indices = np.argsort(var_vals)[::-1]
        for f in range(features):
            print("%d"%(f+1),malData.columns[2+indices[f]],var_vals[indices[f]])
        
        return data_new
    
class AutoEncoder:
    def __init__(self, encoding_dim, x):
        self.encoding_dim = encoding_dim
        self.x = x

    def _encoder(self):
        inputs = Input(shape=(self.x[0].shape))
        encoded = Dense(self.encoding_dim, activation='relu')(inputs)
        model = Model(inputs, encoded)
        self.encoder = model
        return model

    def _decoder(self):
        inputs = Input(shape=(self.encoding_dim,))
        decoded = Dense(self.x.shape[1])(inputs)
        model = Model(inputs, decoded)
        self.decoder = model
        return model

    def encoder_decoder(self):
        ec = self._encoder()
        dc = self._decoder()
        
        inputs = Input(shape=self.x[0].shape)
        ec_out = ec(inputs)
        dc_out = dc(ec_out)
        model = Model(inputs, dc_out)
        
        self.model = model
        return model

    def fit(self, batch_size=10, epochs=20):
        self.model.compile(optimizer='sgd', loss='mse')
        self.model.fit(self.x, self.x,
                        epochs=epochs,
                        batch_size=batch_size)
        data_red = self.encoder.predict(self.x)
        return data_red 

class AutoEncoderDeep:
    def __init__(self, encoding_dim, x):
        self.encoding_dim = encoding_dim
        self.x = x

    def _encoder(self):
        inputs = Input(shape=(self.x[0].shape))
        encodedOne = Dense(32, activation='relu')(inputs)
        encodedTwo = Dense(self.encoding_dim, activation='relu')(encodedOne)
        model = Model(inputs, encodedOne, encodedTwo)
        self.encoder = model
        return model

    def _decoder(self):
        inputs = Input(shape=(self.encoding_dim,))
        decodedOne = Dense(32)(inputs)
        decodedTwo = Dense(self.x.shape[1])(inputs)
        model = Model(inputs, decodedOne, decodedTwo)
        self.decoder = model
        return model

    def encoder_decoder(self):
        ec = self._encoder()
        dc = self._decoder()
        
        inputs = Input(shape=self.x[0].shape)
        ec_out = ec(inputs)
        dc_out = dc(ec_out)
        model = Model(inputs, dc_out)
        
        self.model = model
        return model

    def fit(self, batch_size=10, epochs=20):
        self.model.compile(optimizer='sgd', loss='mse')
        self.model.fit(self.x, self.x,
                        epochs=epochs,
                        batch_size=batch_size)
        data_red = self.encoder.predict(self.x)
        return data_red       
            
def RandomForestClassify(legit_Train, legit_Test, malware_Train, malware_Test):
        classify = RandomForestClassifier(n_estimators=50)
        classify.fit(legit_Train,malware_Train)
        print("Random Forest Classification Score: ",classify.score(legit_Test, malware_Test)*100)
        result = classify.predict(legit_Test)
        conf_matrix = confusion_matrix(malware_Test,result)
        print(conf_matrix)

def NaiveBayesClassify(legit_Train, legit_Test, malware_Train, malware_Test):
        model=GaussianNB()
        model.fit(legit_Train,malware_Train)
        print("Naive Bayes Score: ",model.score(legit_Test, malware_Test)*100)
        
        conf_matrix = confusion_matrix(malware_Test,model.predict(legit_Test))
        print(conf_matrix)

def DnnTwoL(legit_Train, legit_Test, malware_Train, malware_Test):
    numInputs = legit_Train.shape[1]
    model = Sequential()
    model.add(Dense(1024, input_dim=numInputs, activation='relu')) 
    model.add(Dense(32, activation='relu'))
    model.add(Dropout(.1))
    model.add(Dense(1, activation='sigmoid')) 

    model.compile(loss="binary_crossentropy", optimizer="adam", metrics=['accuracy'])
    model.fit(legit_Train, malware_Train, epochs = 50, batch_size=200, validation_data=(legit_Test, malware_Test))
    
def DnnSevenL(legit_Train, legit_Test, malware_Train, malware_Test):
    numInputs = legit_Train.shape[1]
    model = Sequential()
    model.add(Dense(1024, input_dim=numInputs, activation='relu')) 
    model.add(Dense(512, activation='relu'))
    model.add(Dense(256, activation='relu'))
    model.add(Dense(128, activation='relu'))
    model.add(Dense(64, activation='relu'))
    model.add(Dense(32, activation='relu'))
    model.add(Dense(16, activation='relu'))
    model.add(Dropout(.1))
    model.add(Dense(1, activation='sigmoid')) 

    model.compile(loss="binary_crossentropy", optimizer="adam", metrics=['accuracy'])
    model.fit(legit_Train, malware_Train, epochs = 50, batch_size=200, validation_data=(legit_Test, malware_Test))
    

if __name__ == '__main__':
    
    malData = pd.read_csv("MalwareData.csv", sep = "|")
    #print(malData.columns)
    
    legit = malData[0:41323].drop(['legitimate'], axis = 1)
    malware = malData[41323::].drop(['legitimate'], axis = 1)
    
    data_in = malData.drop(['Name', 'md5', 'legitimate'], axis=1).values
    lables = malData['legitimate'].values
    
    print("Legit Dimensions: %s X %s"%(legit.shape[0],legit.shape[1]))
    print("Malware Dimensions: %s X %s"%(malware.shape[0],malware.shape[1]))
    scaler = MinMaxScaler()
    
    print("Reduce variables:\n1. Extra Trees Classifier\n2. Variance Threshold\n3. Auto Encoder\n4. Deep Auto Encoder")
    vrInput = input()
    
    if vrInput == "1":
        print("\nReducing variables using Extra Trees Classifier: \n")
        data_red = ExtraTreesClassifierVR(malData,data_in,lables)
        data_new = scaler.fit_transform(data_red)
    
    elif vrInput == "2":
        print("\nReducing variables using Variance Threshold: \n")
        data_new = VarianceThresholdVR(malData,data_in)
    
    elif vrInput == "3":
        print("\nReducing variables using Auto Encoder: \n")
        data_scl = scaler.fit_transform(data_in)
        ae = AutoEncoder(encoding_dim=12, x=data_scl)
        ae.encoder_decoder()
        data_new = ae.fit(batch_size=200, epochs=50)

    elif vrInput == "4":
        print("\nReducing variables using Deep Auto Encoder: \n")
        data_scl = scaler.fit_transform(data_in)
        ae = AutoEncoderDeep(encoding_dim=12, x=data_scl)
        ae.encoder_decoder()
        data_new = ae.fit(batch_size=200, epochs=50)

    else:
        print("Invalid Input")
    
    legit_Train, legit_Test, malware_Train, malware_Test = train_test_split(data_new, lables ,test_size=0.2)
    
    print("\n\nClassify using\n1.Machine Learning\n2.7 Layer DNN\n3.2 Layer DNN\n4.All DNN\n5.All Methods")
    clsInp = input()
    if clsInp == "1" or "5":        
        print("\nClassifying Using Random Forest\n")
        RandomForestClassify(legit_Train, legit_Test, malware_Train, malware_Test)
        
        print("\nClassifying Using Naive Bayes\n")
        NaiveBayesClassify(legit_Train, legit_Test, malware_Train, malware_Test)
    
    elif clsInp == "2" or "4"  or "5":
        print("\nClassifying Using DNN with 7 Layers\n")
        DnnSevenL(legit_Train, legit_Test, malware_Train, malware_Test)
    elif clsInp == "3" or "4"  or "5":    
        print("\nClassifying Using DNN with 2 Layers\n")
        DnnTwoL(legit_Train, legit_Test, malware_Train, malware_Test)
    
    else:
        print("Invalid Input")

