# conduct experiments on androzoo dataset
import os
import sys
from datetime import datetime

import numpy as np
from sklearn.model_selection import train_test_split

from core.feature import feature_type_scope_dict, feature_type_vs_architecture
from core.ensemble import ensemble_method_scope_dict
from tools import utils
from tools.temporal import time_aware_train_test_split
from config import config, logging

logger = logging.getLogger('experiment.androzoo')


# procedure of drebin experiments
# 1. build dataset
# 2. preprocess data
# 3. learn models
# 4. save results for statistical analysis

def run_experiment(feature_type, ensemble_type, n_members=1, proc_numbers=2):
    """
    run this group of experiments
    :param feature_type: the type of features (e.g., drebin, opcode, etc.), feature type associates to the model architecture
    :param ensemble_type: the ensemble method (e.g., vanilla, deep_ensemble, etc.
    :param n_members: the number of base models enclosed by an ensemble
    :param ratio: the ratio of benign files to malware files
    :param proc_numbers: the number of threads
    :return: None
    """
    mal_folder = config.get('androzoo_tesseract', 'malware_dir')
    ben_folder = config.get('androzoo_tesseract', 'benware_dir')
    logger.info('testing:{},{}'.format(feature_type, ensemble_type))

    train_dataset, validation_dataset, iid_test_dataset, iid_test_y, test_data_list, test_y_list, input_dim = \
        data_preprocessing(feature_type, mal_folder, ben_folder, proc_numbers)

    ensemble_obj = get_ensemble_object(ensemble_type)
    # instantiation
    arch_type = feature_type_vs_architecture.get(feature_type)
    saving_dir = config.get('experiments', 'androzoo')
    if ensemble_type in ['vanilla', 'mc_dropout', 'bayesian']:
        ensemble_model = ensemble_obj(arch_type, base_model=None, n_members=1, model_directory=saving_dir)
    else:
        ensemble_model = ensemble_obj(arch_type, base_model=None, n_members=n_members, model_directory=saving_dir)

    ensemble_model.fit(train_dataset, validation_dataset, input_dim=input_dim)
    ensemble_model.evaluate(iid_test_dataset, iid_test_y)

    res_list = [ensemble_model.predict(iid_test_dataset)]
    for i, test_data in enumerate(test_data_list):
        res_list.append(ensemble_model.predict(test_data))
        # In case of unexceptions in a long-time running, we save result in every running
        utils.dump_joblib(res_list, os.path.join(saving_dir, '{}_{}_test.res'.format(feature_type, ensemble_type)))
        logger.info("Test data with ID {} is Done.".format(i))
    return


def run_temperature_scaling(feature_type, ensemble_type):
    """
    Run temperature scaling
    :param feature_type: the type of feature (e.g., drebin, opcode, etc.), feature type associates to the model architecture
    :param ensemble_type: the ensemble method (e.g., vanilla, deep_ensemble, etc.
    """
    from core.post_calibration.temperature_scaling import find_scaling_temperature, \
        apply_temperature_scaling, inverse_sigmoid
    logger.info('run temperature scaling:{},{}'.format(feature_type, ensemble_type))

    # load dataset
    def data_load(feature_type='drebin', malware_dir=None, benware_dir=None):
        assert feature_type in feature_type_scope_dict.keys(), 'Expected {}, but {} are supported.'.format(
            feature_type, list(feature_type_scope_dict.keys()))

        android_features_saving_dir = config.get('metadata', 'naive_data_pool')
        intermediate_data_saving_dir = config.get('androzoo_tesseract', 'intermediate_directory')
        feature_extractor = feature_type_scope_dict[feature_type](android_features_saving_dir,
                                                                  intermediate_data_saving_dir,
                                                                  update=False,
                                                                  proc_number=1)

        save_path = os.path.join(intermediate_data_saving_dir, 'androzoo_database.' + feature_type)
        if os.path.exists(save_path):
            _1, validation_features, iid_test_features, test_features_list, \
            _2, validation_y, iid_test_y, test_y_list, _3, _4 = utils.read_joblib(save_path)
        else:
            raise ValueError

        validation_data, _ = feature_extractor.feature2ipt(validation_features)
        iid_test_data, _ = feature_extractor.feature2ipt(iid_test_features)
        test_data_list = []
        for test_features in test_features_list:
            test_data, _ = feature_extractor.feature2ipt(test_features)
            test_data_list.append(test_data)

        validation_data, _ = feature_extractor.feature2ipt(validation_features)

        return validation_data, iid_test_data, test_data_list, validation_y, iid_test_y, test_y_list

    mal_folder = config.get('androzoo_tesseract', 'malware_dir')
    ben_folder = config.get('androzoo_tesseract', 'benware_dir')
    val_data, iid_test_data, test_data_list, val_label, iid_test_y, test_y_list = \
        data_load(feature_type, mal_folder, ben_folder)

    # load model
    ensemble_obj = get_ensemble_object(ensemble_type)
    arch_type = feature_type_vs_architecture.get(feature_type)
    model_saving_dir = config.get('experiments', 'androzoo')
    if ensemble_type in ['vanilla', 'mc_dropout', 'bayesian']:
        ensemble_model = ensemble_obj(arch_type, base_model=None, model_directory=model_saving_dir)
    else:
        ensemble_model = ensemble_obj(arch_type, base_model=None, model_directory=model_saving_dir)

    temp_save_dir = os.path.join(config.get('androzoo_tesseract', 'intermediate_directory'),
                                 "{}_{}_temp.json".format(feature_type, ensemble_type))
    if not os.path.exists(temp_save_dir):
        probs = np.squeeze(ensemble_model.predict(val_data, use_prob=True))
        logits = inverse_sigmoid(probs)
        temperature = find_scaling_temperature(val_label, logits)
        utils.dump_json({'temperature': temperature}, temp_save_dir)
    temperature = utils.load_json(temp_save_dir)['temperature']

    prob_ = ensemble_model.predict(iid_test_data, use_prob=True)
    res_list = [apply_temperature_scaling(temperature, prob_)]
    for test_data in test_data_list:
        prob_ = ensemble_model.predict(test_data, use_prob=True)
        res_list.append(apply_temperature_scaling(temperature, prob_))

    utils.dump_joblib(res_list,
                      os.path.join(model_saving_dir, '{}_{}_temperature_test.res'.format(feature_type, ensemble_type)))


def data_preprocessing(feature_type='drebin', malware_dir=None, benware_dir=None, proc_numbers=2):
    assert feature_type in feature_type_scope_dict.keys(), 'Expected {}, but {} are supported.'.format(
        feature_type, list(feature_type_scope_dict.keys()))

    android_features_saving_dir = config.get('metadata', 'naive_data_pool')
    intermediate_data_saving_dir = config.get('androzoo_tesseract', 'intermediate_directory')
    feature_extractor = feature_type_scope_dict[feature_type](android_features_saving_dir,
                                                              intermediate_data_saving_dir,
                                                              update=False,
                                                              proc_number=proc_numbers)
    save_path = os.path.join(intermediate_data_saving_dir, 'androzoo_database.' + feature_type)
    if os.path.exists(save_path):
        train_filenames, validation_filenames, iid_test_filenames, test_filenames_list, \
        train_y, validation_y, iid_test_y, test_y_list, _1, _2 = utils.read_joblib(save_path)

        train_features = [os.path.join(android_features_saving_dir, filename) for filename in
                          train_filenames]
        validation_features = [os.path.join(android_features_saving_dir, filename) for filename in
                               validation_filenames]
        iid_test_features = [os.path.join(android_features_saving_dir, filename) for filename in
                             iid_test_filenames]
        test_features_list = [np.array([os.path.join(android_features_saving_dir, filename) for filename in
                                        test_filenames]) for test_filenames in test_filenames_list]
    else:
        def merge_mal_ben(mal, ben):
            mal_feature_list = feature_extractor.feature_extraction(mal)
            n_malware = len(mal_feature_list)
            ben_feature_list = feature_extractor.feature_extraction(ben)
            n_benware = len(ben_feature_list)
            feature_list = mal_feature_list + ben_feature_list
            gt_labels = np.zeros((n_malware + n_benware,), dtype=np.int32)
            gt_labels[:n_malware] = 1
            import random
            random.seed(0)
            random.shuffle(feature_list)
            random.seed(0)
            random.shuffle(gt_labels)
            return feature_list, gt_labels

        malware_path_list = utils.retrive_files_set(malware_dir, "", ".apk|")
        benware_path_list = utils.retrive_files_set(benware_dir, "", ".apk|")
        feature_paths, gt_labels = merge_mal_ben(malware_path_list, benware_path_list)

        date_dict = utils.load_json(config.get('androzoo_tesseract', 'date_stamp'))
        associated_date = np.array([datetime.strptime(date_dict[utils.get_file_name(path)],
                                                      '%Y-%m-%d') for path in feature_paths])

        train_features, test_features_list, train_y, test_y_list, train_date, test_date_list = \
            time_aware_train_test_split(np.array(feature_paths, dtype=np.object),
                                        gt_labels,
                                        associated_date,
                                        train_size=12,
                                        test_size=1,
                                        granularity='month',
                                        start_date=datetime.strptime('2014-01-01', '%Y-%m-%d')
                                        )

        train_features, iid_test_features, train_y, iid_test_y = train_test_split(train_features,
                                                                                  train_y,
                                                                                  test_size=0.0833,
                                                                                  random_state=0)

        train_features, validation_features, train_y, validation_y = train_test_split(train_features,
                                                                                      train_y,
                                                                                      test_size=0.091,
                                                                                      random_state=0)

        train_filenames = [os.path.basename(path) for path in train_features]
        validation_filenames = [os.path.basename(path) for path in validation_features]
        iid_test_filenames = [os.path.basename(path) for path in iid_test_features]
        test_filenames_list = [[os.path.basename(path) for path in test_filenames] for test_filenames in test_features_list]
        utils.dump_joblib(
            (train_filenames, validation_filenames, iid_test_filenames, test_filenames_list,
             train_y, validation_y, iid_test_y, test_y_list, train_date, test_date_list),
            save_path)

    # obtain data in a format for ML algorithms
    feature_extractor.feature_preprocess(train_features, train_y)  # produce datasets products
    train_dataset, input_dim = feature_extractor.feature2ipt(train_features, train_y, is_training_set=True)
    validation_dataset, _ = feature_extractor.feature2ipt(validation_features, validation_y)
    iid_test_dataset, _ = feature_extractor.feature2ipt(iid_test_features)
    test_data_list = []
    for test_features in test_features_list:
        test_data, _ = feature_extractor.feature2ipt(test_features)
        test_data_list.append(test_data)

    return train_dataset, validation_dataset, iid_test_dataset, iid_test_y, test_data_list, test_y_list, input_dim


def get_ensemble_object(ensemble_type):
    assert ensemble_type in ensemble_method_scope_dict.keys(), '{} expected, but {} are supported'.format(
        ensemble_type,
        ','.join(ensemble_method_scope_dict.keys())
    )
    return ensemble_method_scope_dict[ensemble_type]
