#!/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 FeatureModel(object):
    def __init__(self,use_behavior=True):
        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,types):
        tmp_str = ""
        if types == 'whole':
            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]
        
        if types == 'appdescr':
            concerned = [tainted_group.app_descr]

        if types == 'jsonstr_clsdescr':
            concerned = [tainted_group.json_string,
                         tainted_group.api_cls_descr]

        if types == 'jsonstr_clsdescr_mtddescr':
            concerned = [tainted_group.json_string,
                         tainted_group.api_cls_descr,
                         tainted_group.api_mtd_descr]

        if types == 'jsonstr_clsdescr_flow':
            concerned = [tainted_group.json_string,
                         tainted_group.api_cls_descr,
                         tainted_group.flow.source, tainted_group.flow.sink]

        if types == 'appdescr_clsdescr':
            if '+*-+*-' in tainted_group.layout_strings.raw:
                tainted_group.layout_strings.raw = tainted_group.layout_strings.raw.split('+*-+*-')[1]
            concerned = [tainted_group.app_descr,
                         tainted_group.api_cls_descr,
                         tainted_group.layout_strings]

        if types == 'appdescr_clsdescr_mtddescr':
            if '+*-+*-' in tainted_group.layout_strings.raw:
                tainted_group.layout_strings.raw = tainted_group.layout_strings.raw.split('+*-+*-')[1]
            concerned = [tainted_group.app_descr,
                         tainted_group.api_cls_descr,
                         tainted_group.api_mtd_descr,
                         tainted_group.layout_strings]

############
        if types == 'appdescr_nolayout_clsdescr':
            concerned = [tainted_group.app_descr,
                         tainted_group.api_cls_descr]
      
        if types == 'appdescr_nolayout_clsdescr_mtddescr':
            concerned = [tainted_group.app_descr,
                         tainted_group.api_cls_descr,
                         tainted_group.api_mtd_descr]

        if types == 'jsonstr_layout_clsdescr_mtddescr_flow':
            concerned = [tainted_group.json_string,
                         tainted_group.layout_strings,
                         tainted_group.api_cls_descr,
                         tainted_group.api_mtd_descr,
                         tainted_group.flow.source, tainted_group.flow.sink]

        if types == 'referkeys_layout_clsdescr_mtddescr_flow':
            concerned = [tainted_group.refer_keywords,
                         tainted_group.layout_strings,
                         tainted_group.api_cls_descr,
                         tainted_group.api_mtd_descr,
                         tainted_group.flow.source, tainted_group.flow.sink]
        
        if types == 'jsonstr_referkeys_layout_clsdescr_mtddescr_flow':
            concerned = [tainted_group.json_string,
                         tainted_group.refer_keywords,
                         tainted_group.layout_strings,
                         tainted_group.api_cls_descr,
                         tainted_group.api_mtd_descr,
                         tainted_group.flow.source, tainted_group.flow.sink]
############



       
        if types == 'appdescr_clsdescr_mtddescr_flow':
            if '+*-+*-' in tainted_group.layout_strings.raw:
                tainted_group.layout_strings.raw = tainted_group.layout_strings.raw.split('+*-+*-')[1]
            tainted_group.layout_strings.raw = tainted_group.layout_strings.raw.replace(tainted_group.json_string.raw,'')  
            concerned = [tainted_group.app_descr,
                         tainted_group.api_cls_descr,
                         tainted_group.api_mtd_descr,
                         tainted_group.layout_strings,
                         tainted_group.flow.source, tainted_group.flow.sink]

        if types == 'jsonstr_clsdescr_mtddescr_flow':
            concerned = [tainted_group.json_string,
                         tainted_group.api_cls_descr,
                         tainted_group.api_mtd_descr,
                         tainted_group.flow.source, tainted_group.flow.sink]


        if types == 'referkeys_clsdescr_mtddescr_flow':
            concerned = [tainted_group.refer_keywords,
                         tainted_group.api_cls_descr,
                         tainted_group.api_mtd_descr,
                         tainted_group.flow.source, tainted_group.flow.sink]

        if types == 'appdescr_jsonstr_clsdescr_mtddescr':
            concerned = [tainted_group.app_descr,
                         tainted_group.json_string,
                         tainted_group.api_cls_descr,
                         tainted_group.api_mtd_descr,
                         tainted_group.layout_strings]

        if types == 'appdescr_jsonstr_referkeys_mtddescr_flow':
            concerned = [tainted_group.app_descr,
                         tainted_group.json_string,
                         tainted_group.refer_keywords,
                         tainted_group.api_mtd_descr,
                         tainted_group.layout_strings,
                         tainted_group.flow.source, tainted_group.flow.sink]

        if types == 'appdescr_jsonstr_referkeys_clsdescr_flow':
            concerned = [tainted_group.app_descr,
                         tainted_group.json_string,
                         tainted_group.refer_keywords,
                         tainted_group.api_cls_descr,
                         tainted_group.layout_strings,
                         tainted_group.flow.source, tainted_group.flow.sink]


        if types == 'appdescr_referkeys_clsdescr_mtddescr_flow':

            concerned = [tainted_group.app_descr,
                         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]

        if types == 'appdescr_jsonstr_referkeys_clsdescr_mtddescr':

            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]
 

        if types == 'jsonstr_referkeys_clsdescr_mtddescr_flow':
            concerned = [tainted_group.json_string,
                         tainted_group.refer_keywords,
                         tainted_group.api_cls_descr,
                         tainted_group.api_mtd_descr,
                         tainted_group.flow.source, tainted_group.flow.sink]
        
        if types == 'appdescr_jsonstr_clsdescr_mtddescr_flow':
            if '+*-+*-' in tainted_group.layout_strings.raw:
                tainted_group.layout_strings.raw = tainted_group.layout_strings.raw.split('+*-+*-')[1]
            concerned = [tainted_group.app_descr,
                tainted_group.json_string,
                tainted_group.api_cls_descr,
                tainted_group.api_mtd_descr,                         
                tainted_group.layout_strings,
                tainted_group.flow.source, tainted_group.flow.sink]
        
        if types == 'layout_clsdescr_mtddescr_flow_false':
            tainted_group.layout_strings.is_skip = False
            concerned = [tainted_group.layout_strings,
                tainted_group.api_cls_descr,
                tainted_group.api_mtd_descr,                         
                tainted_group.flow.source, tainted_group.flow.sink]

        if types == 'layout_clsdescr_mtddescr_flow_true':
            concerned = [tainted_group.layout_strings,
                tainted_group.api_cls_descr,
                tainted_group.api_mtd_descr,                         
                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 predict(self,tainted_groups,types):
        from entity.taint import GroupLabel

        for i, tainted_group in enumerate(tainted_groups):
            feature_result = self._clf.predict([self.__transform(tainted_group,types)])
            # 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,types):
        import processer.evaluate
        predicted_groups = self.predict(tainted_groups,types)
        evaluate_core = processer.evaluate.Core()
        evaluate_core.evaluate(predicted_groups)


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

        logger.info("********Gradient Boosting********")
        self._clf = self.gboost_classifier()
        self.train_single(train_groups,'whole')
        logger.info("++++++++whole++++++++")
        self.evaluate(test_groups,'whole')

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

        logger.info("********train sub semantics and test with the same semantics********")
        logger.info("********Gradient Boosting********")
        self._clf = self.gboost_classifier()
        self.train_single(train_groups,'whole')
        logger.info("++++++++whole++++++++")
        self.evaluate(test_groups,'whole')


        self._clf = self.gboost_classifier()     
        logger.info("++++++++appdescr++++++++")
        self.train_single(train_groups,'appdescr')
        self.evaluate(test_groups,'appdescr')

        self._clf = self.gboost_classifier()
        logger.info("++++++++jsonstr_referkeys_layout_clsdescr_mtddescr_flow++++++++")
        self.train_single(train_groups,'jsonstr_referkeys_layout_clsdescr_mtddescr_flow')
        self.evaluate(test_groups,'jsonstr_referkeys_layout_clsdescr_mtddescr_flow')


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


    def __transform_groups(self, tainted_groups,types):

        data = []
        raw_labels = []
        for tainted_group in tainted_groups:
            data.append(self.__transform(tainted_group,types))
            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 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))
            ])
        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")),
            ("tfidf", TfidfTransformer()),
            ("logistir", OneVsRestClassifier(LogisticRegression(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")),
                ("tfidf", TfidfTransformer()),
                ("classifier", SGDClassifier(loss="hinge", penalty="l2", random_state=1))
            ])
        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")),
                ("tfidf", TfidfTransformer()),
                ("to_dense", FunctionTransformer(to_dense,
                                     accept_sparse=True)),
                ("classifier", GaussianNB())
            ])
        return self.get_wrapped_clf(gnb_clf) 


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

        adaboost_clf = Pipeline([
                ("vectorizer", CountVectorizer(stop_words="english")),
                ("tfidf", TfidfTransformer()),
                ("classifier", AdaBoostClassifier(n_estimators=200, 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(1.0 * RBF(1.0), random_state=1)) #GaussianProcessClassifier(1.0 * RBF(1.0), warm_start=True)
        ])
        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=200,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
