#!/usr/bin/env python
# -*- coding: utf-8 -*-
import logging
import os

from sklearn.ensemble import AdaBoostClassifier
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.ensemble import VotingClassifier
from sklearn.externals import joblib
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.model_selection import GridSearchCV
from sklearn.model_selection import ShuffleSplit
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import FunctionTransformer
from sklearn.preprocessing import LabelEncoder
from sklearn.svm import LinearSVC
from sklearn.linear_model import LinearRegression
from sklearn.linear_model import SGDClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.naive_bayes import GaussianNB
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import BaggingClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.ensemble import AdaBoostClassifier
from sklearn.gaussian_process.kernels import RBF
from sklearn.gaussian_process import GaussianProcessClassifier
from sklearn.neural_network import MLPClassifier
from sklearn.tree import ExtraTreeClassifier
from sklearn.ensemble import ExtraTreesClassifier
from sklearn.neighbors import RadiusNeighborsClassifier
from sklearn.discriminant_analysis import QuadraticDiscriminantAnalysis
from sklearn.multiclass import OneVsRestClassifier
from sklearn.linear_model import LogisticRegression

import visitor.loader
import visitor.saver


def to_dense(x):
    return x.todense()


class MutliFeatureModels(object):
    def __init__(self,use_behavior= False):
        self._v_clf = None
        self._cv = None
        self._clf = None
        self._label_encoder = None
        self._params = None
        self.use_behavior = use_behavior

    def load(self, directory):
        v_clf_path = os.path.join(directory, "v_clf.dat")
        model_path = os.path.join(directory, "clf.dat")
        encoder_path = os.path.join(directory, "label_encoder.dat")
        param_path = os.path.join(directory, "params.dat")
        self._v_clf = joblib.load(v_clf_path)
        self._clf = joblib.load(model_path)
        self._label_encoder = joblib.load(encoder_path)
        self._params = visitor.loader.JsonLoader.load(param_path)
        return self

    def save(self, directory):
        if not os.path.exists(directory):
            os.makedirs(directory)
        v_clf_path = os.path.join(directory, "v_clf.dat")
        model_path = os.path.join(directory, "clf.dat")
        encoder_path = os.path.join(directory, "label_encoder.dat")
        param_path = os.path.join(directory, "params.dat")
        joblib.dump(self._v_clf, v_clf_path)
        joblib.dump(self._clf, model_path)
        joblib.dump(self._label_encoder, encoder_path)
        # with open(param_path, "w") as param_file:
        #     json.dump(self._params, param_file)
        visitor.saver.JsonSaver.save(self._params, param_path)
        return self

    def __transform(self, tainted_group):
        tmp_str = ""
        concerned = [tainted_group.app_descr,
                     tainted_group.json_string,
                     tainted_group.refer_keywords,
                     tainted_group.api_cls_descr,
                     tainted_group.api_mtd_descr,
                     tainted_group.layout_strings,
                     tainted_group.flow.source, tainted_group.flow.sink]

                     
        for item in concerned:
            if item.is_skip:
                continue
            else:
                if self.use_behavior:
                    if item.behaviors:
                        tmp_str += " ".join(str(x) for x in item.behaviors.pairs)
                        tmp_str +=" " + " ".join(str(x) for x in item.behaviors.actions)
                        tmp_str +=" " + " ".join(str(x) for x in item.behaviors.resources)
                else:
                    tmp_str += item.raw if item.raw else ""
        return tmp_str



    def __transform_groups(self, tainted_groups):

        data = []
        raw_labels = []
        for tainted_group in tainted_groups:
            data.append(self.__transform(tainted_group))
            raw_labels.append(str(tainted_group.label))
        self._label_encoder = LabelEncoder().fit(raw_labels)
        labels = self._label_encoder.transform(raw_labels)

        return data, labels

    def train_single(self, tainted_groups):
        logger = logging.getLogger(__name__)
        data, labels = self.__transform_groups(tainted_groups)
        self._clf.fit(data, labels)

    def train(self, tainted_groups):
        logger = logging.getLogger(__name__)
       
        self._cv = ShuffleSplit(n_splits=3, random_state=5000)

        data, labels = self.__transform_groups(tainted_groups)
        self._params = {}
        self._clf = GridSearchCV(self._v_clf, self._params, cv=self._cv,
                                 n_jobs=8, verbose=10)
        self._clf.fit(data, labels)
        logger.debug("[Best Score]" + str(self._clf.best_score_))


    def predict_one_group(self, tainted_group):
        # self.predict_proba(tainted_group)
        return self._label_encoder.inverse_transform(
            self._clf.predict([self.__transform(tainted_group)]))[0]

    def predict(self,tainted_groups):
        from entity.taint import GroupLabel

        for i, tainted_group in enumerate(tainted_groups):
            feature_result = self._clf.predict([self.__transform(tainted_group)])
            # feature_result = [int(value) for value in feature_result]

            feature_result = self._label_encoder.inverse_transform(feature_result)[0]
            result = GroupLabel[feature_result.split(".")[1]]
            tainted_groups[i].prediction = result
        return tainted_groups
    
    def evaluate(self,tainted_groups):
        import processer.evaluate
        predicted_groups = self.predict(tainted_groups)
        evaluate_core = processer.evaluate.Core()
        evaluate_core.evaluate(predicted_groups)

    def evaluate_models(self,train_groups,test_groups):
        logger = logging.getLogger(__name__)

        linearr_clf = self.linearr_classifier()
        logistir_clf = self.logistir_classifier()
        svc_clf = self.svm_classifier()
        sgd_clf = self.sgd_classifier()
        knn_clf = self.knn_classifier()
        gnb_clf = self.gnb_classifier()
        dt_clf = self.dt_classifier()
        # bknn_clf = self.bknn_classifier()
        # bdt_clf = self.bdt_classifier()
        rdmf_clf = self.rdmf_classifier()
        adaboost_clf = self.adaboost_classifier()
        gboost_clf = self.gboost_classifier()
        gp_clf = self.gp_classifier()
        mlp_clf = self.mlp_classifier()
        # et_clf = self.et_classifier()
        # ensembleet_clf = self.ensembleet_classifier()
        # rn_clf = self.rn_classifier()
        # qda_clf = self.qda_classifier()



        # self._v_clf = VotingClassifier(estimators=[("gboost_clf", gboost_clf),
        #                                        ("svc_clf", svc_clf)],voting="hard")
        # self.train(train_groups)
        # self.evaluate(test_groups)
        # exit()
 

        clf_info_dict = {'linearr_clf':'Linear Regression',
                            'logistir_clf':'Logistic Regression',
                            'svc_clf':'SVM',
                            'sgd_clf':'Stochastic Gradient Descent （随机梯度下降法）',
                            'knn_clf':'k-nearest neighbors',
                            'gnb_clf': 'Gaussian naive bayes',
                            'dt_clf' : 'Decision Tree',
                            'rdmf_clf': 'Random Forest',
                            'adaboost_clf': 'AdaBoost',
                            'gboost_clf': 'Gradient Boosting',
                            'gp_clf': 'Gaussian Process',
                            'mlp_clf': 'Neural Net',
                            }
        clf_obj_dict = {'linearr_clf':linearr_clf,
                            'logistir_clf':logistir_clf,
                            'svc_clf':svc_clf,
                            'sgd_clf':sgd_clf,
                            'knn_clf':knn_clf,
                            'gnb_clf': gnb_clf,
                            'dt_clf' : dt_clf,
                            'rdmf_clf': rdmf_clf,
                            'adaboost_clf': adaboost_clf,
                            'gboost_clf': gboost_clf,
                            'gp_clf': gp_clf,
                            'mlp_clf': mlp_clf,
                            }
        clf_keys = clf_obj_dict.keys()

        for clf_key in clf_keys:
            clf_info = clf_info_dict[clf_key]
            logger.info("********"+clf_info+"********")
            try:
                self._clf = clf_obj_dict[clf_key]
                self.train_single(train_groups)
                self.evaluate(test_groups)
            except Exception as e:
                logger.error(e.message)

        # for idex, clf_key1 in enumerate(clf_keys):
        #     for clf_key2 in clf_keys[idex+1:]:
        #         clf_info1 = clf_info_dict[clf_key1]
        #         clf_info2 = clf_info_dict[clf_key2]
        #         logger.info("********"+clf_info1+" & "+ clf_info2 +"********")

        #         try:
        #             self._v_clf = VotingClassifier(estimators=[(clf_key1, clf_obj_dict[clf_key1]),(clf_key2, clf_obj_dict[clf_key2])],voting="hard")
        #             self.train(train_groups)
        #             self.evaluate(test_groups)
        #         except Error as e:
        #             logger.error(e.message)


    def evaluate_single_classifier(self,train_groups,test_groups,classifier_type):
        logger = logging.getLogger(__name__)


        from sklearn.gaussian_process import GaussianProcessClassifier
        from sklearn.gaussian_process.kernels import RBF
        from sklearn.gaussian_process.kernels import Matern, WhiteKernel, ConstantKernel



        clf_info_dict = {'linearr_clf':'Linear Regression',
                    'logistir_clf':'Logistic Regression',
                    'svc_clf':'SVM',
                    'sgd_clf':'Stochastic Gradient Descent （随机梯度下降法）',
                    'knn_clf':'k-nearest neighbors',
                    'gnb_clf': 'Gaussian naive bayes',
                    'dt_clf' : 'Decision Tree',
                    'rdmf_clf': 'Random Forest',
                    'adaboost_clf': 'AdaBoost',
                    'gboost_clf': 'Gradient Boosting',
                    'gp_clf': 'Gaussian Process',
                    'mlp_clf': 'Neural Net',
                    }

        if classifier_type == "linearr_clf": 
            self._clf = self.linearr_classifier()
        if classifier_type == "logistir_clf": 
            self._clf = self.logistir_classifier()
        if classifier_type == "svc_clf": 
            self._clf = self.svm_classifier()
        if classifier_type == "sgd_clf": 
            self._clf = self.sgd_classifier()
        if classifier_type == "knn_clf": 
            self._clf = self.knn_classifier()
        if classifier_type == "gnb_clf": 
            self._clf = self.gnb_classifier()
        if classifier_type == "dt_clf": 
            self._clf = self.dt_classifier()
        if classifier_type == "rdmf_clf": 
            self._clf = self.rdmf_classifier()
        if classifier_type == "adaboost_clf": 
            self._clf = self.adaboost_classifier()
        if classifier_type == "gboost_clf": 
            self._clf = self.gboost_classifier()
        if classifier_type == "gp_clf": 

            # self._clf = Pipeline([
            #             ("vectorizer", CountVectorizer(stop_words="english")),
            #             ("tfidf", TfidfTransformer()),
            #             ("to_dense", FunctionTransformer(to_dense,
            #                                      accept_sparse=True)),
            #             ("classifier", GaussianProcessClassifier(kernel=0.7* RBF(1.0),max_iter_predict =15, multi_class='one_vs_one',random_state=5000))
            #             ])

            self._clf = Pipeline([
                        ("vectorizer", CountVectorizer(stop_words="english",max_df=1.0)),
                        ("tfidf", TfidfTransformer(norm=None, smooth_idf=False, sublinear_tf=False, use_idf=True)),
                        ("to_dense", FunctionTransformer(to_dense, accept_sparse=True)),
                        ("classifier", GaussianProcessClassifier(kernel=0.5* RBF(1.0), multi_class='one_vs_one',max_iter_predict =15,random_state=0))
                        ])
            # GaussianProcessClassifier(copy_X_train=True, kernel=1* RBF(1.0),
            #                             n_restarts_optimizer=2, max_iter_predict =30, multi_class='one_vs_one',
            #                             optimizer='fmin_l_bfgs_b', random_state=5000)
        if classifier_type == "mlp_clf": 
            self._clf = self.mlp_classifier()

        logger.info("********"+clf_info_dict[classifier_type]+"********")
        logger.info("****************")
        logger.info(" ".join(str(x) for x in self._clf.steps))
        logger.info("****************")

        self.train_single(train_groups)
        self.evaluate(test_groups)
        



    def get_wrapped_clf(self,specific_clf):
        # return specific_clf
        clf = GridSearchCV(specific_clf, {}, cv=ShuffleSplit(random_state=5000),
                                 n_jobs=8, verbose=10)
        return clf 


    def svm_classifier(self):
        # logger = logging.getLogger(__name__)
        svc_clf = Pipeline([
                ("vectorizer", CountVectorizer(stop_words="english",
                                           ngram_range=(1, 2))),
                ("tfidf", TfidfTransformer()),
                ("classifier", LinearSVC(random_state=5000, loss='hinge',multi_class='ovr'))
            ])
        return self.get_wrapped_clf(svc_clf) 

    def mlp_classifier(self):
        from sklearn.neural_network import MLPClassifier

        mlp_clf = Pipeline([
            ("vectorizer", CountVectorizer(stop_words="english",
                                           ngram_range=(1, 2))),
            ("tfidf", TfidfTransformer()),
            ("classifier", MLPClassifier(random_state=5000,solver='lbfgs', activation='relu',max_iter=20,verbose=10,learning_rate_init=.1)) # MLPClassifier()
        ])
        return self.get_wrapped_clf(mlp_clf) 

    def logistir_classifier(self):
        logger = logging.getLogger(__name__)
        logistir_clf = Pipeline([
            ("vectorizer", CountVectorizer(stop_words="english",
                                           ngram_range=(1,2))),
            ("tfidf", TfidfTransformer()),
            ("logistir", OneVsRestClassifier(LogisticRegression(penalty= 'l1', random_state=5000)))
        ])  
        return self.get_wrapped_clf(logistir_clf) 
   
    def sgd_classifier(self):
        from sklearn.linear_model import SGDClassifier
        sgd_clf = Pipeline([
                ("vectorizer", CountVectorizer(stop_words="english",
                                           ngram_range=(1,2))),
                ("tfidf", TfidfTransformer()),
                ("classifier", SGDClassifier(penalty= 'l1', random_state=5000))
            ])
        return self.get_wrapped_clf(sgd_clf) 

    def gboost_classifier(self):
        from sklearn.ensemble import GradientBoostingClassifier

        gboost_clf = Pipeline([
            ("vectorizer", CountVectorizer(stop_words="english",
                                           ngram_range=(1, 2))),
            ("tfidf", TfidfTransformer()),
            ("to_dense", FunctionTransformer(to_dense,
                                             accept_sparse=True)),
            ("gbc", GradientBoostingClassifier(n_estimators=200,random_state=5000))
        ])
        return self.get_wrapped_clf(gboost_clf) 

    def gnb_classifier(self):
        from sklearn.naive_bayes import GaussianNB
        gnb_clf = Pipeline([
                ("vectorizer", CountVectorizer(stop_words="english",
                                           ngram_range=(1, 2))),
                ("tfidf", TfidfTransformer()),
                ("to_dense", FunctionTransformer(to_dense,
                                     accept_sparse=True)),
                ("classifier", GaussianNB(var_smoothing=0.3))
            ])
        return self.get_wrapped_clf(gnb_clf) 


    def adaboost_classifier(self):
        from sklearn.ensemble import AdaBoostClassifier

        adaboost_clf = Pipeline([
                ("vectorizer", CountVectorizer(stop_words="english",
                                           ngram_range=(1, 2))),
                ("tfidf", TfidfTransformer()),
                ("classifier", AdaBoostClassifier(random_state=5000))
            ])
        return self.get_wrapped_clf(adaboost_clf) 


    def gp_classifier(self):
        from sklearn.gaussian_process.kernels import RBF
        from sklearn.gaussian_process import GaussianProcessClassifier

        # gp_clf = Pipeline([
        #     ("vectorizer", CountVectorizer(stop_words="english",
        #                    ngram_range=(1, 2))),
        #     ("tfidf", TfidfTransformer()),
        #     ("to_dense", FunctionTransformer(to_dense,
        #                              accept_sparse=True)),
        #     ("classifier", GaussianProcessClassifier(copy_X_train=True, kernel=1* RBF(1.0),
        #                     n_restarts_optimizer=1, max_iter_predict =50, multi_class='one_vs_one',
        #                     optimizer='fmin_l_bfgs_b', random_state=5000))
        # ])

        gp_clf = Pipeline([
                    ("vectorizer", CountVectorizer(stop_words="english",max_df=1.0)),
                    ("tfidf", TfidfTransformer(norm=None, smooth_idf=False, sublinear_tf=False, use_idf=True)),
                    ("to_dense", FunctionTransformer(to_dense, accept_sparse=True)),
                    ("classifier", GaussianProcessClassifier(kernel=0.7* RBF(1.0), multi_class='one_vs_one',max_iter_predict =15,random_state=0))
                    ])
        return self.get_wrapped_clf(gp_clf) 


    def linearr_classifier(self):
        logger = logging.getLogger(__name__)
        linearr_clf = Pipeline([
            ("vectorizer", CountVectorizer(stop_words="english")),
            ("tfidf", TfidfTransformer()),
            ("linearr", OneVsRestClassifier(LinearRegression()))
        ])    
        return self.get_wrapped_clf(linearr_clf) 
     

    def knn_classifier(self):
        from sklearn.neighbors import KNeighborsClassifier
        knn_clf = Pipeline([
                ("vectorizer", CountVectorizer(stop_words="english")),
                ("tfidf", TfidfTransformer()),
                ("classifier", KNeighborsClassifier(3))
            ])
        # return knn_clf
        return self.get_wrapped_clf(knn_clf) 


    def rdmf_classifier(self):
        from sklearn.ensemble import RandomForestClassifier
        rdmf_clf = Pipeline([
                ("vectorizer", CountVectorizer(stop_words="english")),
                ("tfidf", TfidfTransformer()),
                ("classifier", RandomForestClassifier(n_estimators=10,random_state=5000))
            ])
        # return rdmf_clf
        return self.get_wrapped_clf(rdmf_clf) 


    def dt_classifier(self):
        from sklearn.tree import DecisionTreeClassifier
        dt_clf = Pipeline([
                ("vectorizer", CountVectorizer(stop_words="english")),
                ("tfidf", TfidfTransformer()),
                ("classifier", DecisionTreeClassifier(random_state=5000))
            ])
        # return dt_clf
        return self.get_wrapped_clf(dt_clf) 




    def bknn_classifier(self):
        from sklearn.ensemble import BaggingClassifier
        from sklearn.neighbors import KNeighborsClassifier
        bknn_clf = Pipeline([
                ("vectorizer", CountVectorizer(stop_words="english")),
                ("tfidf", TfidfTransformer()),
                ("classifier", BaggingClassifier(KNeighborsClassifier(), max_samples=0.5,max_features=0.5,random_state=5000))
            ])

        return self.get_wrapped_clf(bknn_clf) 


    def bdt_classifier(self):
        from sklearn.ensemble import BaggingClassifier
        from sklearn.tree import DecisionTreeClassifier
        bdt_clf = Pipeline([
                ("vectorizer", CountVectorizer(stop_words="english")),
                ("tfidf", TfidfTransformer()),
                ("classifier", BaggingClassifier(DecisionTreeClassifier(), max_samples=0.5,max_features=0.5,random_state=5000))
            ])
        return self.get_wrapped_clf(bdt_clf) 


    def et_classifier(self):
        from sklearn.tree import ExtraTreeClassifier

        et_clf = Pipeline([
            ("vectorizer", CountVectorizer(stop_words="english",
                                           ngram_range=(1, 2))),
            ("tfidf", TfidfTransformer()),
            ("classifier", ExtraTreeClassifier())
        ])

        return self.get_wrapped_clf(et_clf) 

    def ensembleet_classifier(self):
        from sklearn.ensemble import ExtraTreesClassifier

        ensembleet_clf = Pipeline([
            ("vectorizer", CountVectorizer(stop_words="english",
                                           ngram_range=(1, 2))),
            ("tfidf", TfidfTransformer()),
            ("classifier", ExtraTreesClassifier(n_estimators=200))
        ])
        return self.get_wrapped_clf(ensembleet_clf) 

    def rn_classifier(self):
        from sklearn.neighbors import RadiusNeighborsClassifier
        rn_clf = Pipeline([
            ("vectorizer", CountVectorizer(stop_words="english",
                                           ngram_range=(1, 2))),
            ("tfidf", TfidfTransformer()),
            ("to_dense", FunctionTransformer(to_dense,
                             accept_sparse=True)),
            ("classifier", RadiusNeighborsClassifier(radius=1.5))
        ])
        
        return self.get_wrapped_clf(rn_clf) 
   

    def qda_classifier(self):
        from sklearn.discriminant_analysis import QuadraticDiscriminantAnalysis
        qda_clf = Pipeline([
            ("vectorizer", CountVectorizer(stop_words="english",
                                           ngram_range=(1, 2))),
            ("tfidf", TfidfTransformer()),
            ("to_dense", FunctionTransformer(to_dense,
                     accept_sparse=True)),
            ("classifier", QuadraticDiscriminantAnalysis())
        ])

        return self.get_wrapped_clf(qda_clf) 



class ChiSquareModel(object):
    pass
