import os
import sys
import logging
import configparser
import numpy as np
import pandas as pd

from sklearn.ensemble import RandomForestClassifier
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.svm import LinearSVC
from sklearn.naive_bayes import GaussianNB
from sklearn.linear_model import LogisticRegression

from .extract_features import FeatureExtraction
from .modify_data_class import DataModification
from .model import Model
from .feature_assessment import FeatureAssessment

CONFIG_PATH = "_config.ini"


########################################################################
class MainWrapper(object):
    """
    Main for processing datasets and saving results.
    """
    
    #----------------------------------------------------------------------
    def __init__(self):
        """
        Constructor
        """
        self.name = sys.argv[1]
        self.config_parsed = parse_config_file(path=CONFIG_PATH)
        self.config = self.config_parsed[self.name]
        self.logger = logging.getLogger("logger")
        
        self.subdir = sys.argv[3] if len(sys.argv) == 4 else ""
        self.alg = self.config.get("name")
        
        log_dir = self.config.get("log_dir")
        os.makedirs(log_dir, exist_ok=True)
        
        debug_file_path = os.path.join(log_dir, self.subdir, self.alg, "{}_debug.log".format(os.getpid()))
        info_file_path = os.path.join(log_dir, self.subdir, self.alg, "{}_info.log".format(os.getpid()))
        os.makedirs(os.path.dirname(debug_file_path), exist_ok=True)
        
        initialize_logging(
            debug_file_name=debug_file_path,
            info_file_name=info_file_path,
            logger=self.logger,
            uncaught_exception=True,
            reinit=False
        )


    #-----------------------------------------------------------------------
    def __enter__(self):
        self.logger.info("Started ...")
        return self


    #-----------------------------------------------------------------------
    def __exit__(self, *ignore):
        self.logger.info("Done.\n{}".format("-" * 80))


    #-----------------------------------------------------------------------
    def load_pickled_data(self, filepath):
        """
        Self-explaining.
        """
        return pd.read_pickle(filepath)
    
    
    #----------------------------------------------------------------------
    def parse_config_file(path=None, cfg=None):
        """
        Parse configuration file.
        """
        cfg = configparser.ConfigParser(allow_no_value=True, interpolation=configparser.ExtendedInterpolation())

        if path:
            if os.path.isfile(path):
                with open(path) as file:
                    cfg.read_file(file)

        return cfg
    
    
    #----------------------------------------------------------------------
    def initialize_logging(debug_file_name=None, info_file_name=None, errors_to_stderr=True,
                           logger="", reinit=True, uncaught_exception=False, debug_log_size_mb=30):
        """
        Basic initialisation with log rotation.
        """
    
        def __create_handler(file_name, level, formatter, max_log_size_mb=30, backup_count=1):
            handler = logging.handlers.RotatingFileHandler(
                    file_name, maxBytes=1024 * 1024 * max_log_size_mb, backupCount=backup_count, encoding="utf8")
        
            handler.setFormatter(formatter)
            handler.setLevel(level)
            return handler
    
        if isinstance(logger, str):
            logger = logging.getLogger(logger)
    
        if not reinit and logger in initialized_loggers:
            return  # Skips already initialized logger.
    
        log_dir = "."
        name = None
        log_levels = [("DEBUG", debug_file_name), ("INFO", info_file_name)]
        for (i, (level, file_name)) in enumerate(log_levels):
            if file_name:
                if isinstance(file_name, bool):
                    if name is None:
                        import __main__
                        # Get file name of file with main and omit ".py" extension.
                        name = os.path.basename(__main__.__file__)[:-3]
                        name = os.path.join(log_dir, name)
                
                    file_name = "{}.{}".format(name, level.lower())
                    log_levels[i] = (level, file_name)
            
                LOGGING_FILENAMES.add(file_name)
    
        formatter = logging.Formatter(fmt="%(levelname)-5s %(asctime)s %(name)s %(thread)-5d %(message)s",
                                      datefmt="%Y-%m-%d %H:%M:%S")
        logger.setLevel("DEBUG")
        initialized_loggers.add(logger)
    
        unassigned_level = None
        for (level_name, file) in log_levels:
            if file:
                if unassigned_level:
                    level_name = unassigned_level
            
                handler = __create_handler(file, level_name, formatter, debug_log_size_mb, 2)
                logger.addHandler(handler)
                unassigned_level = None
            else:
                unassigned_level = level_name
    
        if errors_to_stderr:
            error_handler = logging.StreamHandler(sys.stderr)
            error_handler.setFormatter(formatter)
            error_handler.setLevel('WARNING')
            logger.addHandler(error_handler)
    
        # UNCAUGHT EXCEPTIONS
        if uncaught_exception:
            def excepthook(type_, value, traceback_):
                logger.critical("Uncaught exception:", exc_info=(type_, value, traceback_))
                if not errors_to_stderr:
                    sys.__excepthook__(type_, value, traceback_)
        
            sys.excepthook = excepthook


    #----------------------------------------------------------------------
    def extract_features(self):
        """
        Extract features for all sets of domains.
        """
        domains_dir = self.config.get("domains_dir")
        self.logger.info("Extracting features...")
        for root, dirs, files in os.walk(domains_dir):
            for file in files:
                filepath_full = os.path.join(root, file)
                self.logger.debug("Extracting features of domains in file {}...".format(filepath_full))
                
                filename = file.split(".")[0]
                filepath = filepath_full.split(".")[0]
                filetype = filepath_full.split(".")[1]

                if "clean" in file:
                    fe = FeatureExtraction(self.logger, self.config, filetype, filepath, filename, 0)
                    fe.main(self.subdir)
                else:
                    fe = FeatureExtraction(self.logger, self.config, filetype, filepath, filename, 1)
                    fe.main(self.subdir)
                    
        self.logger.info("Feature extraction is done.".format(filepath_full))


    #----------------------------------------------------------------------
    def generate_datasets(self):
        """
        Generate datasets for experiments.
        """
        features_dir = os.path.join(self.config.get("features_dir"), self.subdir)
        clean_filepath = os.path.join(features_dir, "clean.bin")
        self.logger.info("Generating datasets...")

        for root, dirs, files in os.walk(features_dir):
            for file in files:
                if "clean" in file:
                    continue

                filepath_full = os.path.join(root, file)
                filename = file.split(".")[0]
                filepath = filepath_full.split(".")[0]
                filetype = filepath_full.split(".")[1]
                
                if filetype == "bin":
                    self.logger.debug("Generating dataset for {}...".format(file))
                    clean = self.load_pickled_data(clean_filepath)
                    dm = DataModification(self.logger, self.config, filetype, filepath, filename)
                    dm.concat_dataframes(clean, self.subdir)
                    
        self.logger.info("Datasets generation is done.")


    #----------------------------------------------------------------------
    def perform_kfold(self):
        """
        Perform K-fold cross validation on each dataset and save results.
        """
        model = Model(self.logger, self.config, self.clf)
        self.logger.info("Performing K-Fold tests...")
        
        filepath_full = os.path.join(self.config.get("datasets_dir"), "dataset_mixed.bin")
        filename = "dataset"
        filepath = filepath_full.split(".")[0]
        filetype = filepath_full.split(".")[1]
        
        self.logger.debug("Performing K-Fold (algorithm {})...".format(self.alg))
        train_data = self.load_pickled_data(filepath_full)
        model.kfold("mixed", 5, train_data, self.subdir, self.alg, filename)
        model.kfold("logo", 5, train_data, self.subdir, self.alg, filename)
        
        self.logger.info("K-Fold tests are done.")
        
        
    #----------------------------------------------------------------------
    def make_predictions(self):
        """
        Make predictions for LOGO tests.
        """
        model = Model(self.logger, self.config, self.clf)
        self.logger.info("Making predictions...")
        
        for id in range(1, 74):
            filename = "logo_without_group_{}.bin".format(id)
            
            trained_model_path = os.path.join(
                self.config.get("models_dir"),
                self.subdir,
                self.alg,
                filename
            )
            trained_model = self.load_pickled_data(trained_model_path)
            
            data_path = os.path.join(self.config.get("families_dir"), "{}.bin".format(id))
            data = self.load_pickled_data(data_path)
            
            self.logger.debug("Making predictions for family {}...".format(id))
            model.predict(trained_model, data, self.subdir, self.alg, filename)
            
        self.logger.info("Predictions are done.")
        
        
    #----------------------------------------------------------------------
    def make_predictions_on_real_data(self):
        """
        Make predictions on real data - domains collected from URLDB.
        """
        model = Model(self.logger, self.config, self.clf)
        self.logger.info("Making predictions on real data...")
        
        trained_model_path = os.path.join(
            self.config.get("models_dir"),
            self.subdir,
            self.alg,
            "model.bin"
        )
        trained_model = self.load_pickled_data(trained_model_path)
        
        for root, dirs, files in os.walk(self.config.get("features_dir")):
            for file in files:
                data_path = os.path.join(root, file)
                filename = file.split(".")[0]
                filetype = file.split(".")[1]
                
                if filetype == "bin":
                    data = self.load_pickled_data(data_path)
                    model.predict_real_data(trained_model, data, self.subdir, self.alg, filename)
                
        self.logger.info("Predictions are done.")
        
        
    #----------------------------------------------------------------------
    def process_predictions_results(self):
        """
        Process results for all tests after LOGO predictions.
        """
        def print_row(fnc_name, fnc, file):
            print(template.format(
                fnc_name,
                "{0:.18f}".format(fnc(tprs)),
                "{0:.18f}".format(fnc(fnrs)),
            ), file=file)

        results_dir = os.path.join(self.config.get("predictions_results_dir"), self.subdir, self.alg)
        results_file = os.path.join(results_dir, "logo_predictions.txt")
        summary_file = os.path.join(results_dir, "summary.txt")
        
        tprs = []
        fnrs = []
        
        with open(results_file, "r") as f:
            lines = f.readlines()
            for line in lines:
                line = line.strip().split(",")
                tprs.append(float(line[0]))
                fnrs.append(float(line[1]))

        with open(summary_file, "w") as summary:
            template = "{0:10}|{1:22}|{2:22}"
    
            print("=" * 60, file=summary)
            print(template.format("Function",
                                  "True positive rate",
                                  "False negitive rate",
                                  ), file=summary)
            print("=" * 60, file=summary)
    
            print_row("Mean", np.mean, summary)
            print_row("Median", np.median, summary)
            print_row("Min", np.min, summary)
            print_row("Max", np.max, summary)
            print_row("Std", np.std, summary)
        
        
    #----------------------------------------------------------------------
    def process_results(self):
        """
        Process results for all datasets after KFold.
        """
        results_dir = os.path.join(self.config.get("results_dir"), self.subdir, self.alg)

        accuracies = []
        tprs = []
        tnrs = []
        fprs = []
        fnrs = []
        
        logo_accuracies = []
        logo_tprs = []
        logo_tnrs = []
        logo_fprs = []
        logo_fnrs = []
        
        for root, dirs, files in os.walk(results_dir):
            for file in files:
                if not "ratio" in file:
                    filepath = os.path.join(root, file)
                    
                    if file.startswith("logo"):
                        logo_accuracies, logo_tprs, logo_tnrs, logo_fprs, logo_fnrs = self.process_results_file(filepath)
                        continue
                    
                    accuracies, tprs, tnrs, fprs, fnrs = self.process_results_file(filepath)

        summary_filepath = os.path.join(results_dir, "results_summary.txt")
        logo_summary_filepath = os.path.join(results_dir, "logo_results_summary.txt")
        self.save_results_summary(summary_filepath, accuracies, tprs, tnrs, fprs, fnrs)
        self.save_results_summary(logo_summary_filepath, logo_accuracies, logo_tprs, logo_tnrs, logo_fprs, logo_fnrs)
        self.append_logo_summary(logo_summary_filepath, logo_accuracies, logo_tprs, logo_tnrs, logo_fprs, logo_fnrs)
        
        
    #----------------------------------------------------------------------
    def process_results_file(self, filepath):
        """
        """
        accuracies = []
        tprs = []
        tnrs = []
        fprs = []
        fnrs = []
        
        with open(filepath, "r") as f:
            accuracies.extend(list(map(float, f.readline().strip().split(","))))
            tprs.extend(list(map(float, f.readline().strip().split(","))))
            tnrs.extend(list(map(float, f.readline().strip().split(","))))
            fprs.extend(list(map(float, f.readline().strip().split(","))))
            fnrs.extend(list(map(float, f.readline().strip().split(","))))
        
        return accuracies, tprs, tnrs, fprs, fnrs
    
    
    #----------------------------------------------------------------------
    def save_results_summary(self, filepath, accuracies, tprs, tnrs, fprs, fnrs):
        """
        """
        def print_row(fnc_name, fnc, file):
            print(template.format(
                fnc_name,
                "{0:.18f}".format(fnc(accuracies)),
                "{0:.18f}".format(fnc(tprs)),
                "{0:.18f}".format(fnc(fnrs)),
                "{0:.18f}".format(fnc(tnrs)),
                "{0:.18f}".format(fnc(fprs)),
            ), file=file)

        with open(filepath, "w") as summary:
            template = "{0:10}|{1:22}|{2:22}|{3:22}|{4:22}|{5:22}"
    
            print("=" * 130, file=summary)
            print(template.format("Function",
                                  "Accuracy",
                                  "True positive rate",
                                  "False negitive rate",
                                  "True negative rate",
                                  "False positive rate",
                                  ), file=summary)
            print("=" * 130, file=summary)
    
            print_row("Mean", np.mean, summary)
            print_row("Median", np.median, summary)
            print_row("Min", np.min, summary)
            print_row("Max", np.max, summary)
            print_row("Std", np.std, summary)
            
            
    #----------------------------------------------------------------------
    def append_logo_summary(self, filepath, accuracies, tprs, tnrs, fprs, fnrs):
        """
        """
        def print_row(group, file, accuracy, tpr, tnr, fpr, fnr):
            print(template.format(
                group,
                "{0:.18f}".format(accuracy),
                "{0:.18f}".format(tpr),
                "{0:.18f}".format(fnr),
                "{0:.18f}".format(tnr),
                "{0:.18f}".format(fpr),
            ), file=file)

        with open(filepath, "a") as summary:
            template = "{0:10}|{1:22}|{2:22}|{3:22}|{4:22}|{5:22}"
    
            print("=" * 130, file=summary)
            print(template.format("Groups",
                                  "Accuracy",
                                  "True positive rate",
                                  "False negitive rate",
                                  "True negative rate",
                                  "False positive rate",
                                  ), file=summary)
            print("=" * 130, file=summary)
            
            for i in range(1, 74):
                print_row("Group {}".format(i), summary, accuracies[i-1], tprs[i-1], tnrs[i-1], fprs[i-1], fnrs[i-1])
                print("-" * 130, file=summary)
                
                
    #----------------------------------------------------------------------
    def combine_logo_results(self):
        """
        """
        pids = []
        
        accuracies = []
        tprs = []
        tnrs = []
        fprs = []
        fnrs = []
        accuracy_stats = []
        tpr_stats = []
        tnr_stats = []
        fpr_stats = []
        fnr_stats = []
        
        dir = os.path.join(self.config.get("results_dir"), self.subdir, self.alg)
        
        for pid in pids:
            filepath = os.path.join(dir, "logo_{}_dataset.txt".format(pid))
            with open(filepath, "r") as f:
                accuracies.extend(list(map(float, f.readline().strip().split(","))))
                tprs.extend(f.readline().strip().split(","))
                tnrs.extend(f.readline().strip().split(","))
                fprs.extend(f.readline().strip().split(","))
                fnrs.extend(f.readline().strip().split(","))
                accuracy_stats.extend(f.readline().strip().split(","))
                tpr_stats.extend(f.readline().strip().split(","))
                tnr_stats.extend(f.readline().strip().split(","))
                fpr_stats.extend(f.readline().strip().split(","))
                fnr_stats.extend(f.readline().strip().split(","))
                
        with open(os.path.join(dir, "logo_dataset.txt"), "w") as out:
            lists = [accuracies, tprs, tnrs, fprs, fnrs, accuracy_stats, tpr_stats, tnr_stats, fpr_stats, fnr_stats]
            for l in lists:
                out.write(",".join(l) + "\n")
    
        
    #----------------------------------------------------------------------
    def feature_assessment(self):
        """
        Compute statistical tests on features.
        """
        self.logger.info("Performing feature assessment...")
        
        dataset = os.path.join(self.config.get("datasets_dir"), "dataset_mixed.bin")
        fa = FeatureAssessment(self.logger, self.config)
        fa.main(dataset, self.subdir)
        
        self.logger.info("Feature assessment is done.")


    #----------------------------------------------------------------------
    def main(self):
        """
        Main function.
        """
        print(os.getpid())
        
        if sys.argv[1] == "gnb":
            self.clf = GaussianNB()
        elif sys.argv[1] == "rf":
            self.clf = RandomForestClassifier()
        elif sys.argv[1] == "lr":
            self.clf = LogisticRegression(solver="saga", max_iter=4000, tol=0.002)
        elif sys.argv[1] == "gbc":
            self.clf = GradientBoostingClassifier()
        elif sys.argv[1] == "svc":
            self.clf = LinearSVC(dual=False)
        else:
            print("Unknown classifier!")
            sys.exit(0)

        allowed_functions = set("fdkracpsx")
        if not set(sys.argv[2]).issubset(allowed_functions):
            print("Unknown function!")
            sys.exit(0)
            
        if "a" in sys.argv[2]:
            self.feature_assessment()
            return
        
        if "f" in sys.argv[2]:
            self.extract_features()
            
        if "d" in sys.argv[2]:
            self.generate_datasets()
            
        if "k" in sys.argv[2]:
            self.perform_kfold()
            
        if "r" in sys.argv[2]:
            self.process_results()
            
        if "c" in sys.argv[2]:
            self.combine_logo_results()
            
        if "p" in sys.argv[2]:
            self.make_predictions()
            
        if "s" in sys.argv[2]:
            self.process_predictions_results()
            
        if "x" in sys.argv[2]:
            self.make_predictions_on_real_data()



if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: main_wrapper.py %classifier% %function% [%subdir%]")
        sys.exit(0)
    
    with MainWrapper() as m:
        m.main()
    