# !!THIS IS AN OBSOLETE CODES!! NOT USED ANYMORE

import sys
import time
import numpy as np
from scipy.cluster.vq import vq, kmeans, whiten
from scipy.spatial.distance import cdist
import pickle
import os
import analysis_plot as data_plt

__author__ = "Wei-Chung Huang"
__copyright__ = "Copyright 2018, The SJSU MSCS Master project"
__license__ = "MIT"
__version__ = "1.0.0"


working_path = os.path.dirname(os.path.abspath(__file__))
feature_folder_name = 'extracted_feats'
cluster_folder_name = 'clustered_res'
saved_feature_path = working_path + '/../' + feature_folder_name
saved_cluster_path = working_path + '/../' + cluster_folder_name
n_families = 25


def prepare_malware_testing_set(shrink=True):
    """
    Load all malware families as a single malware dataset
    :return: m-by-n ndarray where m is the number of samples and n is the dim
    """
    malware_set = None
    for i in range(n_families):
        f1_fam_name = 'f1_fam_' + str(i) + '_features'
        feature_mat = mask_features(f1_fam_name)
        n_samples = len(feature_mat)
        if shrink:
            if n_samples > 200:  # most training size are around 200
                feature_mat = feature_mat[:200, :]
                n_samples = len(feature_mat)

        feature_mat = feature_mat[:int(n_samples / 5), :]
        if malware_set is None:
            malware_set = np.copy(feature_mat)
        else:
            malware_set = np.append(malware_set, feature_mat, axis=0)
    return malware_set


def prepare_benign_set():
    """
    Load benign samples
    :return: m-by-n ndarray where m is the number of samples and n is the dim
    """
    f1_benign_name = 'benign_features'
    benign_set = mask_features(f1_benign_name)
    return benign_set


def cluster_all_features(feature_mat):
    """
    Run k-means to cluster the input feature vectors
    :param feature_mat: m-by-n ndarray
            M is the number of samples and N is dimensionality
    :return: dictionary<k, (distortion, centroids)>
            This dictionary tells the distortion with what centroids and what's K
            key: k value
            Value: tuple with <distortion, centroids> where centroids are k-by-n ndarray
    """
    n_dims = feature_mat.shape[1]
    whitened = whiten(feature_mat.transpose())
    all_codebooks = dict()
    for k in range(n_dims, 0, -1):
        centroids, distortion = kmeans(whitened, k)
        all_codebooks[k] = (distortion, centroids)

    return all_codebooks


def cluster_feature(feature_mat, k):
    """
    Apply K-means to get the clusters' centroid and distortion
    :param feature_mat: m-by-n ndarray
            M is the number of samples and N is dimensionality
    :param k: int
            Number of centroids
    :return: <centroids, distortions>
            centroids: k-by-n ndarray
            distortion: overall distortion for k centroids
    """
    whitened = whiten(feature_mat.transpose())
    centroid, distortion = kmeans(whitened, k)

    return centroid, distortion


def mask_feature(feature):
    """
    Mask the original feature vector to reduced features with the mask generated from SVM-RFE
    :param feature: 1D List
            The original feature vector.
    :return: 1-by-d ndarray
            Reduced feature matrix.
    """
    with open(saved_feature_path + '/f1_reduced_mask', 'rb') as fp:
        mask = pickle.load(fp)

    feature_vec = list()
    for i in range(len(mask)):
        if mask[i]:
            feature_vec.append(feature[i])

    return np.asarray(feature_vec)


def mask_features(feat_file_name):
    """
    Mask the original feature vector to reduced features with the mask generated from SVM-RFE
    :param feat_file_name: str
            File name of saved feature.
    :return: m-by-n ndarray
            Reduced feature matrix.
            M is the number of samples and N is the reduced length.
    """
    with open(saved_feature_path + '/f1_reduced_mask', 'rb') as fp:
        mask = pickle.load(fp)

    with open(saved_feature_path + '/' + feat_file_name, 'rb') as fp:
        features = pickle.load(fp)

    feature_mat = list()
    for i in range(len(features)):
        feature_vec = list()
        for j in range(len(mask)):
            if mask[j]:
                feature_vec.append(features[i][j])
        feature_mat.append(feature_vec)

    feature_mat = np.asarray(feature_mat)
    return feature_mat


def train_threshold(ref_centroids, feature_mat, k):
    """
    Compute a threshold of a malware by using k-means distortion and its referenced codebook.
    :param ref_centroids: m-by-n ndarray
            The centroids of referenced sample.
    :param feature_mat: m-by-n ndarray
            M is the number of samples and N is dimensionality
    :return: float
            The threshold can distinguish it's family X or not.
    """
    all_dist = list()
    overall_max_dist = 0
    for i in range(len(feature_mat)):
        feature_vec = feature_mat[i]
        centroids_, distortion_ = cluster_feature(feature_vec, k)
        centroids_ = centroids_.reshape(len(centroids_), -1)

        # Compute distortion to decide the threshold
        #   scipy.cdist: Computes distance between each pair of the two collections of inputs
        #   Get the average minimum distance of each i to j pair where i is not equal to j
        dist = cdist(ref_centroids, centroids_, 'euclidean')
        code = dist.argmin(axis=1)
        min_dist_list = dist[np.arange(len(code)), code]
        avg_min_dist = min_dist_list.mean(axis=-1)
        all_dist.append(avg_min_dist)
        overall_max_dist = max(overall_max_dist, avg_min_dist)
    np_all_dist = np.asarray(all_dist)
    mean = np_all_dist.mean()
    std = np_all_dist.std()
    threshold = mean + 2 * std
    return threshold


def score_k(feature_mat, fam_index, k, centroids, codebook=None):
    # get pre-trained threshold
    if codebook is not None:
        subkey = codebook[k]
        threshold = subkey['threshold']
    else:
        threshold = train_threshold(centroids, feature_mat, k)

    # for each families, test inner/outer and get a score
    inner_rate = 0
    outer_rates = list()
    print('testing on other fam ', end='')
    for n in range(n_families):
        print(n, ' ', end='')
        sys.stdout.flush()
        f1_fam_name = 'f1_fam_' + str(n) + '_features'
        feature_mat_outer = mask_features(f1_fam_name)
        n_samples = len(feature_mat_outer)

        # take 80% of samples as training
        feature_mat_outer = feature_mat_outer[int(n_samples / 5):, :]

        # compute accuracy
        hit = 0
        for outer_vec in feature_mat_outer:
            centroids_, distortion_ = cluster_feature(outer_vec, k)
            centroids_ = centroids_.reshape(len(centroids_), -1)
            dist = cdist(centroids, centroids_, 'euclidean')
            code = dist.argmin(axis=1)
            min_dist_list = dist[np.arange(len(code)), code]
            avg_min_dist = min_dist_list.mean(axis=-1)
            if avg_min_dist < threshold:
                hit += 1

        accuracy = hit / len(feature_mat_outer) * 100
        if n == fam_index:
            inner_rate = accuracy
        else:
            outer_rates.append(accuracy)

    score_inner = inner_rate
    score_outer = np.percentile(outer_rates, 25) if len(outer_rates) > 0 else 0
    score = max(1.0 - score_outer / score_inner, 0.0)  # inner is expected larger than outer

    rates = dict()
    rates['inner_rate'] = inner_rate
    rates['outer_rates'] = outer_rates

    return score, rates


def train_inner_outer_with_hill_climbing(feature_mat, fam_index, init_k):
    # step1: random pick the value k
    # step2: find best score from k to k + t
    # step2: find best score from k to k - t
    # step3: pick the best score from above, named k'
    # step4: if k' is 0 or n, reset k by random
    # step4: else, repeat step2

    # collapse to get an average vector
    features = np.mean(feature_mat, axis=0).reshape((1, feature_mat.shape[1]))
    n_dim = features.shape[1]

    # get all possible centroids with different k from 1 to k
    key_dist_cent = cluster_all_features(features)

    # load the pre-trained codebook if exists
    codebook = None
    codebook_path = saved_cluster_path + '/f1_fam_' + str(fam_index) + '_codebook'
    if os.path.isfile(codebook_path):
        with open(codebook_path, 'rb') as fp:
            codebook = pickle.load(fp)

    # step1: randomly pick the value k
    init_k = init_k
    print('*' * 50, 'current k:', init_k, '*' * 50)

    # step2a: find the best score from k to k + t
    max_score_r = -1
    max_score_r_k = init_k
    max_score_r_rates = dict()
    for k in range(init_k, n_dim):
        print('searching k in right', k)
        distortion, centroids = key_dist_cent[k]
        start_time = time.time()
        score, rates = score_k(feature_mat, fam_index, k, centroids, codebook=codebook)
        print('\nfinal score:', '%.4f' % score)
        elapsed_time = time.time() - start_time
        print('time', '%.2f' % elapsed_time)
        sys.stdout.flush()
        if k == init_k or (abs(score - 1.0) > 10e-7 and score >= max_score_r):
            max_score_r = score
            max_score_r_k = k
            max_score_r_rates = rates
        else:
            break

    # step2b: find best score from k to k - t
    max_score_l = -1
    max_score_l_k = init_k
    max_score_l_rates = dict()
    for k in range(init_k - 1, 0, -1):
        print('searching k in left', k)
        distortion, centroids = key_dist_cent[k]
        start_time = time.time()
        score, rates = score_k(feature_mat, fam_index, k, centroids, codebook=codebook)
        print('\nfinal score:', '%.4f' % score)
        elapsed_time = time.time() - start_time
        print('time', '%.2f' % elapsed_time)
        if k == init_k - 1 or (abs(score - 1.0) > 10e-7 and score >= max_score_l):
            max_score_l = score
            max_score_l_k = k
            max_score_l_rates = rates
        else:
            break

    # step3: pick the best score from above, named k'
    if max_score_r > max_score_l:
        max_score_k = max_score_r_k
        max_score_rates = max_score_r_rates
    else:
        max_score_k = max_score_l_k
        max_score_rates = max_score_l_rates
    print('max score and k for right and left', max_score_r, max_score_l, max_score_r_k, max_score_l_k)

    max_score_rates['k'] = max_score_k
    max_score_rates['score'] = max(max_score_r, max_score_l)

    return max_score_rates


def train_with_inner_outer(feature_mat, fam_index):
    """
    Train the K-means K value by evaluating inner and outer measurement.
    For each k value, compute training samples' centroids, determine a threshold,
    and then compute the accuracy for inner family and outer families. And then
    use the inner/outer accuracy to compute a score which will decide how well is the chosen k.
    :param feature_mat: m-by-n ndarray
            M is the number of samples and N is dimensionality
    :param fam_index: int
            The index of malware family.
    :return: dictionary
            A dictionary storing best_k, inner_rate, outer_rates
    """
    # collapse to get an average vector
    features = np.mean(feature_mat, axis=0).reshape((1, feature_mat.shape[1]))

    # get all possible centroids with different k from 1 to k
    key_dist_cent = cluster_all_features(features)
    max_score = -1
    best_k = 0

    # load the pre-trained codebook if exists
    codebook = None
    codebook_path = saved_cluster_path + '/f1_fam_' + str(fam_index) + '_codebook'
    if os.path.isfile(codebook_path):
        with open(codebook_path, 'rb') as fp:
            codebook = pickle.load(fp)

    result_dict = dict()
    for k in key_dist_cent.keys():
        print('k = ', k)
        distortion, centroids = key_dist_cent[k]

        # get pre-trained threshold
        if codebook is not None:
            subkey = codebook[k]
            threshold = subkey['threshold']
        else:
            threshold = train_threshold(centroids, feature_mat, k)

        # start testing with other families
        inner_rate = 0
        outer_rates = list()
        for n in range(n_families):
            f1_fam_name = 'f1_fam_' + str(n) + '_features'
            feature_mat_outer = mask_features(f1_fam_name)
            n_samples = len(feature_mat_outer)
            if n_samples > 200:  # most training size are around 200
                feature_mat_outer = feature_mat_outer[:, 200, :]
                n_samples = len(feature_mat)

            # take 80% of samples as training
            feature_mat_outer = feature_mat_outer[int(n_samples / 5):, :]

            # compute accuracy
            hit = 0
            for outer_vec in feature_mat_outer:
                centroids_, distortion_ = cluster_feature(outer_vec, k)
                centroids_ = centroids_.reshape(len(centroids_), -1)
                dist = cdist(centroids, centroids_, 'euclidean')
                code = dist.argmin(axis=1)
                min_dist_list = dist[np.arange(len(code)), code]
                avg_min_dist = min_dist_list.mean(axis=-1)
                if avg_min_dist < threshold:
                    hit += 1

            accuracy = hit / len(feature_mat_outer) * 100
            if n == fam_index:
                inner_rate = accuracy
            else:
                outer_rates.append(accuracy)

        score_inner = inner_rate
        score_outer = np.median(outer_rates) if len(outer_rates) > 0 else 0
        score = 1 - score_outer / score_inner
        if score > max_score:
            max_score = score
            best_k = k
            result_dict['k'] = best_k
            result_dict['inner_rate'] = inner_rate
            result_dict['outer_rates'] = outer_rates

    return result_dict


def train_k():
    """
    An entrance for training value k in k-means.
    This will dump the result of trained k_value, family_index, and rates
    :return:
    """

    use_hc_train = True
    feature_mat = mask_features('f1_fam_0_features')
    n_dim = feature_mat.shape[1]
    np.random.seed(1)
    next_k = np.random.randint(2, n_dim)
    for i in range(n_families):
        print('[training-k] family:', i, '...')
        f1_fam_name = 'f1_fam_' + str(i) + '_features'
        feature_mat = mask_features(f1_fam_name)
        n_samples = len(feature_mat)
        if n_samples > 200:  # most training size are around 200
            feature_mat = feature_mat[:200, :]
            n_samples = len(feature_mat)

        # use 80% of samples as training set
        feature_mat = feature_mat[int(n_samples / 5):, :]

        if use_hc_train:
            # train the k-value by using hill-climbing approach
            max_score_rates = train_inner_outer_with_hill_climbing(feature_mat, i, next_k)
            print(max_score_rates)
            next_k = np.random.randint(2, n_dim)
        else:
            max_score_rates = train_with_inner_outer(feature_mat, i)

        save_name = '/f1_fam_' + str(i) + '_validation'
        with open(saved_cluster_path + save_name, 'wb') as fp:
            pickle.dump(max_score_rates, fp)


def train_key(shrink=True, inside_test=False):
    """
    Train malware families to get a centroid of each family.
    (Threshold would be computed as mean + 2 * std)
    :param shrink: bool (optional)
            Cap the training sample size to 200
    :param inside_test: bool (optional)
            Inside test on the remaining 20% samples
    :return: dictionary
            <key: family idx, value: dict{}>
                                     <'centroid': family's average sample>
                                     <'threshold': cluster radius>
    """
    keys = dict()
    overall_accuracy = 0
    for i in range(n_families):
        f1_fam_name = 'f1_fam_' + str(i) + '_features'
        feature_mat = mask_features(f1_fam_name)
        n_samples = len(feature_mat)
        n_dim = len(feature_mat[0])
        if shrink:
            if n_samples > 200:  # most training size are around 200
                feature_mat = feature_mat[:200, :]
                n_samples = len(feature_mat)
        feature_mat_old = feature_mat.copy()

        # use 80% of samples as training set
        feature_mat = feature_mat[int(n_samples / 5):, :]
        avg_feature_vec = np.mean(feature_mat, axis=0).reshape((1, feature_mat.shape[1]))

        key = dict()
        key['centroid'] = avg_feature_vec

        distance = list()
        for vec in feature_mat:
            vec = vec.reshape(1, n_dim)
            distance.append(np.linalg.norm(vec - avg_feature_vec))

        distance = np.asarray(distance)
        dist_mean = np.mean(distance)
        dist_std = distance.std()
        threshold = dist_mean + 2 * dist_std

        key['threshold'] = threshold
        keys[i] = key

        if inside_test:
            # testing 20%
            feature_mat = feature_mat_old[:int(n_samples / 5), :]
            accuracy = 0
            for vec in feature_mat:
                h = np.linalg.norm(vec - avg_feature_vec)
                if h <= threshold:
                    accuracy += 1

            accuracy = accuracy / len(feature_mat) * 100
            overall_accuracy += accuracy
            print('inside: accuracy', accuracy)

    if inside_test:
        overall_accuracy = overall_accuracy / n_families
        print('inside: average accuracy:', overall_accuracy)

    return keys


def predict(X, keys):
    """
    Classify a single vector on all malware family models
    :param X: 1-by-n ndarray
            Testing sample with dimension is equal to keys[]['centroid']
    :param keys: dictionary from train_key
    :return: int
            Classified family index
    """
    min_dist = np.finfo(np.float32).max
    fit_index = -1
    for index in keys.keys():
        avg_feature_vec = keys[index]['centroid']
        threshold = keys[index]['threshold']
        d = np.linalg.norm(X - avg_feature_vec)
        if d < min_dist:
            min_dist = d
            fit_index = index

    return fit_index


def predict_malware_all(plot_cm=False, shrink=True):
    """
    Classify all malware samples by using its cluster centroid
    :param plot_cm: bool (optional)
            Plot a confusion matrix is set
    :param shrink: bool (optional)
            Cap the sample size
    :return: float
            Overall accuracy
    """
    keys = train_key()
    overall_accuracy = 0
    confusion_matrix = np.zeros([n_families, n_families], dtype=np.int32)
    for i in range(n_families):
        f1_fam_name = 'f1_fam_' + str(i) + '_features'
        feature_mat = mask_features(f1_fam_name)
        n_samples = len(feature_mat)
        if shrink:
            if n_samples > 200:  # most training size are around 200
                feature_mat = feature_mat[:200, :]
                n_samples = len(feature_mat)

        feature_mat = feature_mat[:int(n_samples / 5), :]
        accuracy = 0
        for vec in feature_mat:
            fit_index = predict(vec, keys)
            if fit_index == i:
                accuracy += 1

            confusion_matrix[i, fit_index] += 1

        accuracy = accuracy / len(feature_mat) * 100
        overall_accuracy += accuracy

    overall_accuracy = overall_accuracy / n_families
    print('outside test: average accuracy', overall_accuracy)

    if plot_cm:
        classes_label = list()
        for i in range(n_families):
            classes_label.append(str(i + 1))
        data_plt.plot_confusion_matrix(confusion_matrix, classes_label, normalize=True)

    return overall_accuracy


def predict_malware_benign_roc():
    """
    1. Take the min distance as a threshold.
    2. Generate possible threshold(distance)
    3. Plot a ROC curve
    :return:
    """
    malware_set = prepare_malware_testing_set()
    benign_set = prepare_benign_set()
    x = list()
    y = list()
    keys = train_key()
    for scale in range(0, 101):
        threshold = 5 * scale / 100
        tp = 0
        for vec in malware_set:
            min_dist = np.finfo(np.float32).max
            for index in keys.keys():
                avg_feature_vec = keys[index]['centroid']
                d = np.linalg.norm(vec - avg_feature_vec)
                if d < min_dist:
                    min_dist = d
            if min_dist <= threshold:
                tp += 1
        tpr = tp / len(malware_set)

        tn = 0
        for vec in benign_set:
            min_dist = np.finfo(np.float32).max
            for index in keys.keys():
                avg_feature_vec = keys[index]['centroid']
                d = np.linalg.norm(vec - avg_feature_vec)
                if d < min_dist:
                    min_dist = d
            if min_dist > threshold:
                tn += 1

        tnr = tn / len(benign_set)
        fpr = 1.0 - tnr
        x.append(fpr)
        y.append(tpr)

    data_plt.plot_roc(x, y)


if __name__ == '__main__':
    # features = None
    # with open('f1_fam_0_features', 'rb') as fp:
    #     features = pickle.load(fp)
    #
    # features = np.asarray(features)
    # print(features.shape)
    # features = np.mean(features, axis=0).reshape((1, features.shape[1]))
    #
    # print(features.shape)
    # quantize_feature(features)

    # train_k()
    # predict_malware_all()
    # test_benign()
    predict_malware_benign_roc()

    # for n in range(25):
    #     f1_fam_name = 'f1_fam_' + str(n) + '_features'
    #     feature_mat_fam1 = mask_features(f1_fam_name)
    #     n_fold = 5
    #     n_samples = len(feature_mat_fam1)
    #     accuracy_avg = 0
    #     for i in range(n_fold):
    #         fold_start = i * int(n_samples / n_fold)
    #         fold_end = fold_start + int(n_samples / n_fold) - 1
    #         testing_feat_mat = feature_mat_fam1[fold_start:fold_end + 1, :]
    #         training_feat_mat = np.append(feature_mat_fam1[0:fold_start, :], feature_mat_fam1[fold_end + 1:, :], axis=0)
    #         # collapse all samples to get an average features
    #         features = np.mean(training_feat_mat, axis=0).reshape((1, training_feat_mat.shape[1]))
    #         codebooks = cluster_all_features(features)
    #
    #         # Temp. assuming k = 50 is the best k (by trail-n-error)
    #         # This k should from 1 to k
    #         best_k = 50
    #         distortion, centroids = codebooks[best_k]
    #         threshold = train_threshold(distortion, centroids, training_feat_mat)
    #         print('final threshold', threshold)
    #
    #         # testing-----
    #         hit = 0
    #         for testing_vec in testing_feat_mat:
    #             centroids_, distortion_ = cluster_feature(testing_vec, best_k)
    #             centroids_ = centroids_.reshape(len(centroids_), -1)
    #             dist = cdist(centroids, centroids_, 'euclidean')
    #             code = dist.argmin(axis=1)
    #             min_dist_list = dist[np.arange(len(code)), code]
    #             avg_min_dist = min_dist_list.mean(axis=-1)
    #             if avg_min_dist < threshold:
    #                 hit += 1
    #
    #         accuracy = hit / len(testing_feat_mat) * 100
    #         accuracy_avg += accuracy
    #
    #     accuracy_avg = accuracy_avg / n_fold
    #     print(n, 'average accuracy', accuracy_avg)

    # approach1
    # my key would be: <k + centroids + threshold>
    # how to decide key?
    #   run every k and its centroids for every samples, and then get a distortion
    #   find the 'maximum' distortion as threshold
    # next time when I got a new input, cluster it with key.k, and compute the distortion by
    # key.centroids and key.threshold.


    # key: k + centroids + threshold
    # 1. Trained a threshold by:
    #   Get maximum distortion for all pairs of cluster result of training samples to average sample
    #   Set the maximum distortion as the threshold
    #   Set the key to (k + averaged centroids + threshold)
    #   (Think of the radius of a circle)
    # 2. Test the rest samples by:
    #   Use the same k to do a k-means
    #   Get centroids
    #   Compute the distortion from current centroids to referenced centroids
    #   If the distortion is less than the threshold, then hit
    #
    # Tested for all families with a specific k = 50, got 99% accuracy.
    # But it is for inner-distortion test
    # Need to consider outer-distortion by:
    #  Train a threshold:
    #    Get inner accuracy
    #    Get outer accuracy
    #    Find a best k to maximize the inner and minimize the outer
    #
    #  LDA: as referenced
    #
    # How to train k?
    #   distortion is declining as k
    #

