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

from config.constant import Values
import packager.behavior
import packager.taint
import model.feature_model
from semantic.taint import GroupLabel


class PredictionModel(object):
    def __init__(self):
        self.feature_model = model.feature_model.FeatureModel()
        self._hyper_parameters = None

    def train(self, tainted_groups):
        logger = logging.getLogger(__name__)
        logger.info("[Train] start to train feature model")
        self.feature_model.train(tainted_groups)

    def save(self,feature_model_directory=Values.feature_model_directory):
        self.feature_model.save(feature_model_directory)
        return self

    def load(self, feature_model_dir=Values.feature_model_directory):
        self.feature_model = self.feature_model.load(feature_model_dir)

        return self

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

        for i, group in enumerate(tainted_groups):
            feature_result = self.feature_model.predict(group)

            result = GroupLabel[feature_result.split(".")[1]]
            tainted_groups[i].prediction = result
        return tainted_groups
