from os import listdir
from os.path import isfile, join
from sklearn.metrics import confusion_matrix

from features_analyzer import analyzer
from time import sleep
import sys



mypath = '../drebin/feature_vectors'
allfiles = [f for f in listdir(mypath) if isfile(join(mypath, f))]

# Print iterations progress
def printProgressBar (iteration, total, prefix = '', suffix = '', decimals = 1, length = 50, fill = '█'):
    percent = ("{0:." + str(decimals) + "f}").format(100 * (iteration / float(total)))
    filledLength = int(length * iteration // total)
    bar = fill * filledLength + '-' * (length - filledLength)
    print('\r%s |%s| %s%% %s' % (prefix, bar, percent, suffix), end = '\r')
    # Print New Line on Complete
    if iteration == total: 
        print()



lista_malware = list()

def load_malwares():
    with open('../drebin/sha256_family.csv', 'r') as f:
        raw_data = f.read()
        f.close()
    raw_data = raw_data.strip()
    lista_data = raw_data.split('\n')
    for i in range(len(lista_data)):
        if i != 0:
            lista_malware.append((lista_data[i].split(','))[0])
    return lista_malware

def train(lista_file):
    #Training dei dati
    mal,nomal = dict(),dict()
    n_mal,n_nomal = 0,0
    printProgressBar(0, len(lista_file), prefix = 'Training:', suffix = 'Complete')
    for a in range(len(lista_file)):
        printProgressBar(a+1, len(lista_file), prefix = 'Training:', suffix = 'Complete')
        path = '../drebin/feature_vectors/'+lista_file[a]
        with open(path, 'r') as f:
            raw_data = f.read()
            f.close()
        raw_data = raw_data.strip()
        #data.append(raw_data.split('\n'))
        data = raw_data.split('\n')
        for e in range(len(data)):
            if(lista_file[a] in lista_malware):
                n_mal += 1
                if data[e] not in mal.keys():
                    mal[data[e]] = 0
                else:
                    mal[data[e]] +=1
            else:
                n_nomal += 1
                if data[e] not in nomal.keys():
                    nomal[data[e]] = 0
                else:
                    nomal[data[e]] +=1
    return n_mal,n_nomal,mal,nomal

def bernoulli_naive(doc, mal, nomal, n_mal, n_nomal, dic):
    p_mal, p_nomal = 1.0, 1.0
    #print(dic)
    for s in dic:
        if s in mal.keys(): 
            p_m = (mal[s]+1)/(n_mal+2) #probabilita di trovare quella parola in mal
        else:
            p_m = 1/(n_mal+2)
        if s in nomal.keys():
            p_n = (nomal[s]+1)/(n_nomal+2) #probabilita di trovare quella parola in nomal
        else:
            p_n = 1/(n_nomal+2)
        if s in doc:
            p_mal *= p_m #0.09 #andiamo a trovare la probabilita che quella parola s della nostra frase si trovi in mal
            p_nomal *= p_n #0.01
        else:
            p_mal *= (1-p_m)  #0.91 #probabilita di ... e 1 meno la probabilita di trovare quella parola in mal 
            p_nomal *= (1-p_n) #0.99 
    t = max(p_mal, p_nomal)
    if t==p_mal:
        return 'mal'
    elif t==p_nomal:
        return 'nomal'


if __name__ == "__main__":
    lista_malware = load_malwares()
    data = list()
    dic = list()
    tipo = input('What dictionary to be chosen? 1 (All the features) o 2 (top 8 features): ')
    if tipo == '1':
        print("Initializing dictionary")
        printProgressBar(0, len(allfiles), prefix = 'Progress:', suffix = 'Complete')
        for counter,a in enumerate(allfiles):
            path = '../drebin/feature_vectors/'+a
            with open(path, 'r') as f:
                raw_data = f.read()
                f.close()
            raw_data = raw_data.strip()
            data = raw_data.split('\n')
            for i in range(len(data)-1):
                dic.append(data[i]) #senza splittare la frase
            printProgressBar(counter + 1, len(allfiles), prefix = 'Progress:', suffix = 'Complete')
        dic = set(dic)
        print("Dictionary done")
    else:
        dic = analyzer(lista_malware)
    acc = list()

    train_data = allfiles[0:((len(allfiles)*80)//100)]
    test_data = allfiles[len(train_data):len(allfiles)]

    y_pred,y_true = list(),list()
    n_mal,n_nomal,mal,nomal = train(train_data)
    for t in range(len(test_data)):
        printProgressBar(t+1, len(test_data), prefix = 'Testing:', suffix = 'Complete')
        lista_value = list()
        path = '../drebin/feature_vectors/'+test_data[t]
        if test_data[t] in lista_malware:
            truth = 'mal'
        else:
            truth = 'nomal'
        with open(path, 'r') as f:
            raw_data = f.read()
            f.close()
        raw_data = raw_data.strip()
        data = raw_data.split('\n')
        for i in data:
            lista_value.append(i)
        c_map = bernoulli_naive(lista_value,mal,nomal,n_mal,n_nomal,dic)
        y_pred.append(c_map)
        y_true.append(truth)
        if c_map == truth:
            acc.append(1)
        else:
            acc.append(0)
    print('Accuracy is: {}'.format(sum(acc)/len(acc)))
    conf_matrix = confusion_matrix(y_true, y_pred, labels=["mal", "nomal"])
    print('Confusion matrix is: \n {}'.format(conf_matrix))
    TN, FP, FN, TP = confusion_matrix(y_true, y_pred, labels=["mal", "nomal"]).ravel()
    precision = TP/(TP+FP) 
    recall = TP/(TP+FN)
    print("Precision: {}".format(precision))
    print("Recall: {}".format(recall))
    print("False Positive Rate: {}".format(FP/(FP+TN)))
    print("Accuracy: {}".format(((TP+TN)/(TP+FN+TN+FP))))
    print("F-Measure: {}".format(2*(precision*recall)/(precision+recall)))
    

#     Confusion matrix is: 
# [[   33  1092]
#  [   42 24636]]

    # #k-fold cross validation
    # fold_size = len(allfiles)//30
    # acc = list()
    # print('Lunghezza file: {}'.format(len(allfiles)))
    # for i in range(0, len(allfiles), fold_size):
    #     test_data = allfiles[i:i+fold_size]
    #     train_data = allfiles[0:i] + allfiles[(i+fold_size):]
    #     print("Iniziando il training")
    #     n_mal,n_nomal,mal,nomal = train(train_data)
    #     print("Training completato")
    #     f_acc = list()
    #     printProgressBar(0, len(test_data), prefix = 'Testing:', suffix = 'Complete')
    #     for t in range(len(test_data)):
    #         printProgressBar(t+1, len(test_data), prefix = 'Testing:', suffix = 'Complete')
    #         lista_value = list()
    #         path = '../drebin/feature_vectors/'+test_data[t]
    #         if test_data[t] in lista_malware:
    #             truth = 'mal'
    #         else:
    #             truth = 'nomal'
    #         with open(path, 'r') as f:
    #             raw_data = f.read()
    #             f.close()
    #         raw_data = raw_data.strip()
    #         #data.append(raw_data.split('\n'))
    #         data = raw_data.split('\n')
    #         for i in data:
    #             # elemento_test = i.split('::')
    #             # lista_value.append(elemento_test[1])
    #             lista_value.append(i)
    #         c_map = bernoulli_naive(lista_value,mal,nomal,n_mal,n_nomal,dic)
    #         #print("Il programma ha rilevato che il file é: {} / {}".format(c_map,truth))
    #         if c_map == truth:
    #             f_acc.append(1)
    #             acc.append(1)
    #         else:
    #             f_acc.append(0)
    #             acc.append(0)
    #     print('Accuracy for k is: {}'.format(sum(f_acc)/len(f_acc)))
    # print('Final Accuracy for fold size {} is: {}'.format(fold_size, sum(acc)/len(acc)))

