import os
import pickle
import pandas as pd

from sklearn.feature_selection import chi2
from sklearn.feature_selection import f_classif
from sklearn.feature_selection import mutual_info_classif

from sklearn.utils import parallel_backend


########################################################################
class FeatureAssessment(object):
    """
    Class for statistical tests of extracted features.
    """
    
    #----------------------------------------------------------------------
    def __init__(self, logger, config):
        """
        Constructor
        """
        self.logger = logger
        self.config = config
        
        self.labels = self.config.get("feature_assessment_labels").split(";")
        
        
    #----------------------------------------------------------------------
    def load_pickled_data(self, filepath):
        """
        Self-explaining.
        """
        return pd.read_pickle(filepath)
    
    
    #----------------------------------------------------------------------
    def test_wrapper(self, test_name, input, output):
        """
        """
        train_data = self.load_pickled_data(input)
        X = train_data[self.labels]
        y = train_data["dga_flag"]
        
        if test_name == mutual_info_classif:
            values = test_name(X, y)
        else:
            values = test_name(X, y)[0]
        
        values_zipped = zip(self.labels, values)
        values_sorted = sorted(values_zipped, key=lambda x:x[1], reverse=True)
        
        with open(output, "w") as out:
            for value in values_sorted:
                out.write("{}\n".format(value))
    
    
    #----------------------------------------------------------------------
    def main(self, dataset, subdir):
        
        output = os.path.join(self.config.get("results_dir"), subdir, "{}.txt")
        os.makedirs(os.path.dirname(output), exist_ok=True)
        
        self.logger.debug("Performing chi2 test...")
        self.test_wrapper(chi2, dataset, output.format("no_ngrams_chi2"))

        self.logger.debug("Performing f_classif test...")
        self.test_wrapper(f_classif, dataset, output.format("no_ngrams_f_classif"))

        self.logger.debug("Performing mutual_info_classif test...")
        self.test_wrapper(mutual_info_classif, dataset, output.format("no_ngrams_mutual_info_classif"))


if __name__ == '__main__':
    pass