import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D

import pandas as pd
import pickle

from sklearn.model_selection import train_test_split

from sklearn.feature_extraction.text import TfidfVectorizer

from sklearn.decomposition import PCA
from sklearn.manifold import Isomap, TSNE, MDS

if __name__ == '__main__':

    with open('Malware-API/all_analysis_data.txt', 'r') as original_data:
        api_calls = original_data.readlines()

    labels = []
    with open('Malware-API/labels.csv', 'r') as data_labels:
        for line in data_labels.readlines():
            labels.append(line.replace('\n', ''))

    train_api_calls, test_api_calls, train_labels, test_labels = train_test_split(api_calls, labels,
                                                                                  test_size=.3, random_state=44)

    tfidf_vect_total = TfidfVectorizer(ngram_range=(1, 3))
    tfidf_apicalls = tfidf_vect_total.fit_transform(api_calls).todense()

    print(tfidf_vect_total.get_feature_names())
    print(len(tfidf_vect_total.get_feature_names()))
    print(tfidf_apicalls)

    tfidf_total_dataframe = pd.DataFrame(tfidf_apicalls, columns=tfidf_vect_total.get_feature_names())
    tfidf_total_dataframe.to_csv('Malware-API/tfidf-all-features-ngrams.csv')

    tfidf_vect_split = TfidfVectorizer(max_features=300)
    tfidf_train_apicalls = tfidf_vect_split.fit_transform(train_api_calls).todense()
    tfidf_test_apicalls = tfidf_vect_split.transform(test_api_calls).todense()

    print(tfidf_vect_split.get_feature_names())
    print(len(tfidf_vect_split.get_feature_names()))
    print(tfidf_train_apicalls)

    iso = Isomap()
    mds = MDS()

    # CREATE 2 SIMPLE EMBEDDING ON GRAPH

    perplexity_array = [5, 10, 15, 20, 25, 30, 35, 40, 45]
    max_iter = 1000
    random_state = 44
    embeddings = []
    for perp in perplexity_array:
        tsne = TSNE(perplexity=perp, n_iter=max_iter, random_state=random_state)
        xy_coordinates = tsne.fit_transform(tfidf_train_apicalls, y=train_labels)
        embeddings.append(xy_coordinates)

    with open('Malware-API/embeddings_iter_1000_stop_300.pickle', 'wb') as out_embeddings:
        pickle.dump(embeddings, out_embeddings)

    # Add labels to the plot
    # I used this as a template for creating the next function.
    # https://matplotlib.org/gallery/subplots_axes_and_figures/demo_constrained_layout.html#sphx-glr-gallery-subplots-axes-and-figures-demo-constrained-layout-py
    def tsne_plot(ax, x, y, perplexity, c=None):
        ax.scatter(x, y, c=c)
        ax.set_title(str(perplexity) + ' perplexity', fontsize=8)
        print('make plots')

    def seperate_x_y(xy_coordinates):
        x = []
        y = []
        for coor in xy_coordinates:
            x.append(coor[0])
            y.append(coor[1])
        return x, y

    with open('Malware-API/embeddings_iter_10000_stop_300.pickle', 'rb') as in_embeddings:
        embeddings = pickle.load(in_embeddings)

    # used this url as a template for the graphs
    # https://stackoverflow.com/questions/33246065/convert-categorical-variable-to-color-with-matplotlib
    def category_to_color(y_classification):
        n = len(y_classification)
        cmap = cm.get_cmap('hsv')
        return [cmap(float(i)/n) for i in range(n)]

    classifier_color = category_to_color(train_labels)

    fig, axs = plt.subplots(nrows=3, ncols=3)
    for embedding, ax, perp, i in zip(embeddings, axs.flat, perplexity_array, range(len(embeddings))):
        x, y = seperate_x_y(embedding)
        if classifier_color is None:
            tsne_plot(ax, x, y, perplexity=perp)
        else:
            tsne_plot(ax, x, y, perplexity=perp, c=classifier_color)
    plt.savefig('2_d_1000_perplexity.png')
    plt.close(fig)

    # CREATE SERIES OF FIGURES IN DIFFERENT PERPLEXITY AND MAX ITER

    # CREATE 3D IN PYPLOT?

    """
    
    tfidf_training_dataframe = pd.DataFrame({'apicalls': tfidf_train_apicalls, 'labels': train_labels},
                                            columns=['api_calls', 'labels'])
    tfidf_test_dataframe = pd.DataFrame({'apicalls': tfidf_test_apicalls, 'labels': test_labels},
                                        columns=['api_calls', 'labels'])
    """

    """
    pd.to_pickle(tfidf_training_dataframe, 'Malware-API/training-dataframe.pickle')
    pd.to_pickle(tfidf_test_dataframe, 'Malware-API/test-dataframe.pickle')

    print(tfidf_training_dataframe.values)

    tfidf_training_dataframe = pd.read_pickle('Malware-API/training-dataframe.pickle')
    tfidf_test_dataframe = pd.read_pickle('Malware-API/test-dataframe.pickle')

    print(tfidf_training_dataframe.values)
    """



