import os
import time
import pickle
import datetime
import numpy as np
import pandas as pd

from sklearn.model_selection import KFold
from sklearn.model_selection import LeaveOneGroupOut

from sklearn.metrics import accuracy_score
from sklearn.metrics import confusion_matrix

from sklearn.utils import parallel_backend


########################################################################
class Model(object):
    """
    Class for training and predicting DGA domains.
    """
    
    #----------------------------------------------------------------------
    def __init__(self, logger, config, clf):
        """
        Constructor
        """
        self.logger = logger
        self.config = config
        self.clf = clf
        
        self.labels = self.config.get("model_labels").split(";")
    
    
    #----------------------------------------------------------------------
    def load_pickled_data(self, filepath):
        """
        Self-explaining.
        """
        return pd.read_pickle(filepath)
    
    
    #----------------------------------------------------------------------
    def load_csv(self, filepath):
        """
        Self-explaining.
        """
        return pd.read_csv(filepath, header=0)
    
    
    #----------------------------------------------------------------------
    def load_trained_model(self, filepath):
        """
        Self-explaining.
        """
        with open(filepath, "rb") as f:
            clf = pickle.load(f)
        return clf
    
    
    #----------------------------------------------------------------------
    def save_trained_model(self, filepath):
        """
        Self-explaining.
        """
        with open("{}.bin".format(filepath), "wb") as f:
            pickle.dump(self.clf, f)
        
        
    #----------------------------------------------------------------------
    def predict(self, trained_model, data, subdir, alg, filename):
        """
        Predict data on trained model.
        """
        X_test = data[self.labels]
        y_test = data["dga_flag"]
        
        y_pred = trained_model.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)

        #tn, fp, fn, tp = confusion_matrix(y_test, y_pred).ravel()
        #tpr = tp / (tp + fn)
        #tnr = tn / (tn + fp)
        #fpr = 1 - tnr
        #fnr = 1 - tpr

        predictions_results_filepath = os.path.join(self.config.get("predictions_results_dir"),
                                                    subdir,
                                                    alg,
                                                    "logo_predictions.txt")
        os.makedirs(os.path.dirname(predictions_results_filepath), exist_ok=True)
        self.logger.debug("Saving results of a prediction...")
    
        with open("{}.txt".format(predictions_results_filepath), "a") as out:
            out.write("{},{}\n".format(str(accuracy), str(1 - accuracy)))
            
            
    #----------------------------------------------------------------------
    def predict_real_data(self, trained_model, data, subdir, alg, filename):
        """
        Make predictions on real data.
        """
        X_test = data[self.labels]
        y_pred = trained_model.predict(X_test)
        if alg != "support_vector_classifier":
            y_pred_proba = trained_model.predict_proba(X_test)
        
        
        predictions_filepath = os.path.join(
            self.config.get("predictions_dir"),
            subdir,
            alg,
            "predictions_{}.csv".format(filename)
        )
        os.makedirs(os.path.dirname(predictions_filepath), exist_ok=True)
        self.logger.debug("Saving results of a prediction...")
        
        out = pd.DataFrame()
        out["domain"] = data["raw_domain"]
        out["prediction"] = y_pred
        if alg != "support_vector_classifier":
            out["prediction_proba"] = y_pred_proba[:, 1]
        out.to_csv(predictions_filepath, index=False)
        
        
    #----------------------------------------------------------------------
    def kfold(self, cv_type, k, train_data, subdir, alg, filename):
        """
        """
        def compute_stats(fnc, list_):
            return fnc(list_)
        
        groups = train_data["group"]
        X = train_data[self.labels]
        y = train_data["dga_flag"]
        accuracies = []
        tpr_list = []
        tnr_list = []
        fpr_list = []
        fnr_list = []

        filename = "{}_{}".format(os.getpid(), filename)
        models_filepath = os.path.join(self.config.get("models_dir"), subdir, alg, filename)
        os.makedirs(os.path.dirname(models_filepath), exist_ok=True)

        if cv_type == "mixed":
            kf = KFold(n_splits=k, random_state=42, shuffle=True)
            data_split = kf.split(X)
            self.logger.debug("Performing regular K-Fold test...")
        elif cv_type == "logo":
            logo = LeaveOneGroupOut()
            data_split = logo.split(X, y, groups)
            self.logger.debug("Performing leave one group out test...")
            filename = "logo_{}".format(filename)
        
        for train_index, test_index in data_split:
            if cv_type == "logo":
                group = train_data.iloc[test_index].group.reset_index(drop=True)[0]
                self.logger.debug("Leaving out group {}...".format(group))

                model_filename = "logo_without_group_{}".format(group)
                models_filepath = os.path.join(self.config.get("models_dir"), subdir, alg, model_filename)
            
            data_train = X.iloc[train_index]
            target_train = y.iloc[train_index]
    
            #with parallel_backend("threading", n_jobs=10):
            #    self.clf.fit(data_train, target_train)
            
            self.clf.fit(data_train, target_train)
            
            self.save_trained_model(models_filepath)
    
            data_test = X.iloc[test_index]
            target_test = y.iloc[test_index]

            predictions = self.clf.predict(data_test)
            accuracy = accuracy_score(target_test, predictions)
            accuracies.append(accuracy)
            
            tn, fp, fn, tp = confusion_matrix(target_test, predictions).ravel()
            tpr = tp / (tp + fn)
            tnr = tn / (tn + fp)
            
            tpr_list.append(tpr)
            tnr_list.append(tnr)
            fpr_list.append(1 - tnr)
            fnr_list.append(1 - tpr)
        
        results_filepath = os.path.join(self.config.get("results_dir"), subdir, alg, filename)
        os.makedirs(os.path.dirname(results_filepath), exist_ok=True)
        self.logger.debug("Saving results of a k-fold test...")
        
        with open("{}.txt".format(results_filepath), "w") as out:
            lists = [accuracies, tpr_list, tnr_list, fpr_list, fnr_list]
            functions = [np.mean, np.median, np.min, np.max, np.std]
            
            for list_ in lists:
                out.write(",".join(map(str, list_)))
                out.write("\n")
                
            for list_ in lists:
                temp_list = []
                for fnc in functions:
                    temp_list.append(compute_stats(fnc, list_))
                out.write(",".join(map(str, temp_list)))
                out.write("\n")
        
    
    #----------------------------------------------------------------------
    def main(self):
        """
        Self-explaining.
        """
        pass
    
    
if __name__ == '__main__':
    pass