import warnings
import numpy as np
import pandas as pd
import pickle
import dill
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.linear_model import LogisticRegression
import os
#Donot modify the method signatures and return dtypes
#you can include additional functions of your choice
import dask as dd
from sklearn import metrics
from sklearn.model_selection import GridSearchCV

usefulFeatures = ['SmartScreen', 
				'RtpStateBitfield',
				'Processor',
				'Platform',
				'OsVer',
				'OsSuite',
				'OsPlatformSubRelease',
				'OsBuildLab',
				'IsProtected',
				'HasTpm',
				'Firewall',
				'Census_TotalPhysicalRAM',
				'Census_ProcessorModelIdentifier',
				'Census_ProcessorManufacturerIdentifier',
				'Census_ProcessorCoreCount',
				'Census_PrimaryDiskTypeName',
				'Census_PowerPlatformRoleName',
				'Census_OSWUAutoUpdateOptionsName',
				'Census_OSVersion',
				'Census_OSSkuName',
				'Census_OSBuildNumber',
				'Census_OSBranch',
				'HasDetections']

usefulFeatures2 = ['SmartScreen', 
				'AVProductsInstalled',
				'AVProductStatesIdentifier',
				'AvSigVersion',
				'AppVersion',
				'Census_TotalPhysicalRAM',
				'EngineVersion',
				'OsBuildLab',
				'Wdft_IsGamer',
				'Census_InternalPrimaryDiagonalDisplaySizeInInches',
				'Wdft_RegionIdentifier',
				'Census_OSInstallTypeName',
				'Census_PrimaryDiskTotalCapacity',
				'CountryIdentifier',
				'Census_ActivationChannel',
				'HasDetections']

class MalwareDetectionModel:
	def __init__(self):
		self.dtypes = {
			'MachineIdentifier':'category',
			'ProductName':'category',
			'EngineVersion':'category',
			'AppVersion':'category',
			'AvSigVersion':'category', 
			'IsBeta':'category', 
			'RtpStateBitfield':'float64',
			'IsSxsPassiveMode':'float64',
			'DefaultBrowsersIdentifier':'float64',
			'AVProductStatesIdentifier':'float64',
			'AVProductsInstalled':'float64',
			'AVProductsEnabled':'float64',
			'HasTpm':'float64',
			'CountryIdentifier':'float64',
			'CityIdentifier':'category',
			'OrganizationIdentifier':'float64',
			'GeoNameIdentifier':'float64',
			'LocaleEnglishNameIdentifier ':'float64',
			'Platform ':'category',
			'Processor':'category',
			'OsVer':'category',
			'OsBuild':'category',
			'OsSuite':'category',
			'OsPlatformSubRelease':'category',
			'OsBuildLab':'category', 
			'SkuEdition':'category',
			'IsProtected':'float64',
			'AutoSampleOptIn':'float64',
			'PuaMode':'category',
			'SMode':'float64',
			'IeVerIdentifier':'float64',
			'SmartScreen':'category',
			'Firewall':'float64',
			'UacLuaenable':'float64',
			'Census_MDC2FormFactor':'category',
			'Census_DeviceFamily':'category',
			'Census_OEMNameIdentifier':'float64',
			'Census_OEMModelIdentifier':'float64',
			'Census_ProcessorCoreCount':'float64', 
			'Census_ProcessorManufacturerIdentifier':'float64',
			'Census_ProcessorModelIdentifier':'float64',
			'Census_ProcessorClass':'category',
			'Census_PrimaryDiskTotalCapacity':'float64',
			'Census_PrimaryDiskTypeName':'category',
			'Census_SystemVolumeTotalCapacity':'float64', 
			'Census_HasOpticalDiskDrive':'float64',
			'Census_TotalPhysicalRAM':'float64',
			'Census_ChassisTypeName':'category',
			'Census_InternalPrimaryDiagonalDisplaySizeInInches':'float64',
			'Census_InternalPrimaryDisplayResolutionHorizontal':'float64',
			'Census_InternalPrimaryDisplayResolutionVertical':'float64',
			'Census_PowerPlatformRoleName':'category',
			'Census_InternalBatteryType':'category',
			'Census_InternalBatteryNumberOfCharges':'float64',
			'Census_OSVersion':'category',
			'Census_OSArchitecture':'category',
			'Census_OSBranch':'category',
			'Census_OSBuildNumber':'float64',
			'Census_OSBuildRevision':'float64',
			'Census_OSEdition':'category',
			'Census_OSSkuName':'category',
			'Census_OSInstallTypeName':'category',
			'Census_OSInstallLanguageIdentifier':'float64',
			'Census_OSUILocaleIdentifier':'float64',
			'Census_OSWUAutoUpdateOptionsName':'category',
			'Census_IsPortableOperatingSystem':'float64',
			'Census_GenuineStateName':'category',
			'Census_ActivationChannel':'category',
			'Census_IsFlightingInternal':'float64',
			'Census_IsFlightsDisabled':'float64', 
			'Census_FlightRing':'category',
			'Census_ThresholdOptIn':'float64',
			'Census_FirmwareManufacturerIdentifier':'float64',
			'Census_FirmwareVersionIdentifier':'float64',
			'Census_IsSecureBootEnabled':'float64',
			'Census_IsWIMBootEnabled':'float64',
			'Census_IsVirtualDevice':'float64',
			'Census_IsTouchEnabled':'float64',
			'Census_IsPenCapable':'float64',
			'Census_IsAlwaysOnAlwaysConnectedCapable':'float64',
			'Wdft_IsGamer':'float64',
			'Wdft_RegionIdentifier':'float64',
			'HasDetections':'float64'               
			}
		
		self.classifier = None

	def preProcessing(self, fileName):
		# remove too much data for now.
		#print("procssinfg the file: {}".format(fileName))
		import dask.dataframe as dd
		data = pd.read_csv(fileName, dtype=self.dtypes, usecols=usefulFeatures2, na_values=[''], skiprows=0, nrows=8500000)   

		print("PREPROCESSING DONE")
		return data


	def transformData(self, inputData):
#		print("get input data")

		# change ProductName
		
		#inputData['ProductName'] = inputData['ProductName'].astype('category')#pd.Categorical(inputData['ProductName'])
		inputData['EngineVersion'] = inputData['EngineVersion'].astype('category')
		inputData['AppVersion'] = inputData['AppVersion'].astype('category')        
		inputData['AvSigVersion'] = inputData['AvSigVersion'].astype('category')
		#inputData['Platform'] = inputData['Platform'].astype('category')
		#inputData['Processor'] = inputData['Processor'].astype('category')
		#inputData['OsVer'] = inputData['OsVer'].astype('category')
		#inputData['OsPlatformSubRelease'] = inputData['OsPlatformSubRelease'].astype('category')
		#inputData['OsBuildLab'] = inputData['OsBuildLab'].astype('category')
		#inputData['SkuEdition'] = inputData['SkuEdition'].astype('category')
		inputData['SmartScreen'] = inputData['SmartScreen'].astype('category')
		#inputData['Census_MDC2FormFactor'] = inputData['Census_MDC2FormFactor'].astype('category')
		#inputData['Census_DeviceFamily'] = inputData['Census_DeviceFamily'].astype('category')
		#inputData['Census_PrimaryDiskTypeName'] = inputData['Census_PrimaryDiskTypeName'].astype('category')
		#inputData['Census_ChassisTypeName'] = inputData['Census_ChassisTypeName'].astype('category')
		#inputData['Census_PowerPlatformRoleName'] = inputData['Census_PowerPlatformRoleName'].astype('category')
		#inputData['Census_ChassisTypeName'] = inputData['Census_ChassisTypeName'].astype('category')
		#inputData['Census_InternalBatteryType'] = inputData['Census_InternalBatteryType'].astype('category')
		#inputData['Census_OSVersion'] = inputData['Census_OSVersion'].astype('category')
		#inputData['Census_OSArchitecture'] = inputData['Census_ChassisTypeName'].astype('category')
		#inputData['Census_OSBranch'] = inputData['Census_OSBranch'].astype('category')
		#inputData['Census_OSEdition'] = inputData['Census_ChassisTypeName'].astype('category')
		#inputData['Census_OSSkuName'] = inputData['Census_OSSkuName'].astype('category')
		#inputData['Census_OSSkuName'] = inputData['Census_ChassisTypeName'].astype('category')
		#inputData['Census_OSWUAutoUpdateOptionsName'] = inputData['Census_OSWUAutoUpdateOptionsName'].astype('category')
		#inputData['Census_GenuineStateName'] = inputData['Census_GenuineStateName'].astype('category')
		inputData['Census_ActivationChannel'] = inputData['Census_ActivationChannel'].astype('category')
		inputData['Census_OSInstallTypeName'] = inputData['Census_OSInstallTypeName'].astype('category')
		#inputData['Census_FlightRing'] = inputData['Census_FlightRing'].astype('category')
		#inputData['Census_ProcessorClass'] = inputData['Census_ProcessorClass'].astype('category')
		
		cat_columns = inputData.select_dtypes(['category']).columns

		inputData[cat_columns] = inputData[cat_columns].apply(lambda x: x.cat.codes)

		#inputData = inputData.fillna(inputData.mean())

		
		print("DATA TRANSFORMATION DONE")
		return inputData

	def returnFeatureAndOutput(self, inputData):
		columnCount = len(self.dtypes)
		cols = [col for col in inputData.columns if col != 'HasDetections']

		all_inputs = inputData.ix[:, cols].values
		print("returning data for {} columns".format(len(cols)))
		
		all_labels = inputData['HasDetections'].values
		return all_inputs, all_labels

		#trains a model, you can access the malware_train.csv file for training your model
	def train(self):
		currentDirectory = os.getcwd()
		fileName = os.path.join(currentDirectory, "malware_train.csv")

		print('Train Model')

		inputData = self.preProcessing(fileName)
		inputData = self.transformData(inputData)

		all_inputs, all_labels = self.returnFeatureAndOutput(inputData)

		print("input data count = {}".format(len(all_inputs)))
		print("labels count = {} ".format(len(all_labels)))
		(training_inputs,
		testing_inputs,
		training_classes,
		testing_classes) = train_test_split(all_inputs, all_labels, test_size=0.1, random_state=1)
		
		# Create the classifier
		'''
		decision_tree_classifier = DecisionTreeClassifier()
		decision_tree_classifier.fit(training_inputs, training_classes)
		print("Score = {}".format(decision_tree_classifier.score(testing_inputs, testing_classes)))
		self.classifier = decision_tree_classifier
		'''
		'''
		log_reg = LogisticRegression()
		log_reg.fit(training_inputs, training_classes)
		#print("TRAINING DONE")
		print("Score = {}".format(log_reg.score(testing_inputs, testing_classes)))
		self.classifier = log_reg
		'''
		'''
		# Neural Network
		from keras.models import Sequential
		from keras.layers import Dense

		model = Sequential()
		model.add(Dense(5, input_dim=22, activation='sigmoid'))
		model.add(Dense(1, activation='sigmoid'))
		model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])
		model.fit(X_train,y_train,epochs=100)
		y_pred_neural_network = model.predict(X_test)
		fpr, tpr, thresholds = roc_curve(y_test, y_pred_neural_network[:, 0], pos_label=1)
		roc_auc = auc(fpr, tpr)
		print("AUC for Neural Network: %.2f" % roc_auc)
		'''
		'''
		from sklearn.svm import SVC # "Support Vector Classifier" 
		svc = SVC(kernel='linear') 
		svc.fit(training_inputs, training_classes)
		print("Score = {}".format(svc.score(testing_inputs, testing_classes)))
		self.classifier = svc
		'''
		'''
		lgb_model = lgb.LGBMClassifier(max_depth=7,
									 num_leaves=400,
                                   objective='binary',
                                   )
		import lightgbm as lgb
		lgb_model = lgb.LGBMClassifier(
                                   n_estimators=50000,
                                   learning_rate=0.05,
                                   num_leaves=2 ** 12 - 1,
                                   colsample_bytree=0.28,
                                   objective='binary',
                                   n_jobs=-1)'''
		import lightgbm as lgb
		lgb_model = lgb.LGBMClassifier()
		lgb_model.fit(training_inputs, training_classes, eval_metric='auc',
                  eval_set=[(testing_inputs, testing_classes)],
                  verbose=100, early_stopping_rounds=100)

		#pred_classes = lgb_model.predict(testing_inputs)

		print("Score = {}".format(lgb_model.score(testing_inputs, testing_classes)))
		
		#print('Accuracy Score: {}'.format(metrics.accuracy_score(pred_classes, testing_classes)))
		
		#print('AUC: {}'.format(metrics.roc_auc_score(pred_classes, testing_classes)))
		
		self.classifier = lgb_model
		
		'''
		from sklearn.ensemble import AdaBoostClassifier

		abc = AdaBoostClassifier(n_estimators=70, learning_rate=1)
		model = abc.fit(training_inputs, training_classes)
		print("Score = {}".format(abc.score(testing_inputs, testing_classes)))		
		self.classifier = abc
		'''

		'''
		from xgboost import XGBClassifier
		import xgboost as xgb
		import matplotlib.pyplot as plt
		ee
		xgbc = XGBClassifier(booster='gbtree',
							eta=0.1,
							max_depth=9,
							objective='binary:logistic',
							eval_metric='auc',
							min_child_weight=6,
							gamma=0.0,
							subsample=1)
		xgbc.fit(training_inputs, training_classes)
		pred_classes = xgbc.predict(testing_inputs)
		print("Score = {}".format(xgbc.score(testing_inputs, testing_classes)))
		print('AUC = {}'.format(metrics.roc_auc_score(testing_classes,pred_classes)))
		xgb.plot_importance(xgbc)
		plt.rcParams['figure.figsize'] = [5, 5]
		plt.show()
		self.classifier = xgbc'''
		
		'''param_test1 = {'max_depth':range(3,11,2),
						'min_child_weight':range(1,7,2)}

		gsearch1 = GridSearchCV(estimator = XGBClassifier( learning_rate =0.1, n_estimators=100, max_depth=5,
		min_child_weight=1, gamma=0, subsample=0.8, colsample_bytree=0.8,
		objective= 'binary:logistic', seed=27), 
		param_grid = param_test1, scoring='roc_auc',n_jobs=4,iid=False, cv=5)
		gsearch1.fit(training_inputs,training_classes)
		gsearch1.best_params_, gsearch1.best_score_'''
		
				
		



		print("Perform prediction..")
		firstRow = inputData[0:100]
		firstRow.to_csv("testingData.csv", index = False)
		firstRow = firstRow.drop('HasDetections', axis=1)

		self.predict_probablities(firstRow)

	#predicts the probablity of a machine affected by malware. The input test_data is a pandas dataframe.
	#This function must return the predicted probabilities for the test machines
	def predict_probablities(self, test_data):
		print('Predict malware')
		copy = dict(test_data)
		test_data = self.transformData(test_data)
		cols = [col for col in test_data.columns if col != 'HasDetections']
		all_inputs = test_data.ix[:, cols].values
		predictions = self.classifier.predict(all_inputs)
		return predictions



#calls the train function and saves the model_file
#donot modify the contents
if __name__ == "__main__":
	detectionModel = MalwareDetectionModel()
	currentDirectory = os.getcwd()

	fileName = os.path.join(currentDirectory, "malware_train.csv")

	#detectionModel.preProcessing(fileName)
	import time
	import datetime

	currentDT = datetime.datetime.now()

	start_time = time.time()
	print ('\nProcess started at {}\n'.format(currentDT.strftime("%I:%M:%S %p")))
	
	detectionModel.train()
	
	
	elapsed_time = time.time() - start_time
	print(time.strftime("%H:%M:%S", time.gmtime(elapsed_time)))
	
	filename = 'malwaremodel.pkl'
	f = open(filename, 'wb')
	dill.dump(detectionModel, f)
	f.close()
	