#!/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

import visitor.loader
import visitor.saver


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


class FeatureModel(object):
    def __init__(self):
        self._v_clf = None
        self._cv = None
        self._clf = None
        self._label_encoder = None
        self._params = None

    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 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)
                # 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)
        # logger = logging.getLogger(__name__)
        # logger.debug("[Transform Data]" + " ".join(str(x) for x in data))
        # logger.debug("[Transform Labels]" + " ".join(str(x) for x in labels))
        return data, labels

    def train(self, tainted_groups):
        logger = logging.getLogger(__name__)
        

        data, labels = self.__transform_groups(tainted_groups)
        self._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))
        ])
        self._clf.fit(data, labels)
        # logger.debug("[Best Score]" + str(self._clf.best_score_))
       

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


    def score(self, tainted_groups):
        data, labels = self.__transform_groups(tainted_groups)
        try:
            print("score", self._clf.score(data, labels))
        except Exception as e:
            print(e)


class ChiSquareModel(object):
    pass
