import lightgbm
import numpy as np
from ember import PEFeatureExtractor
from secml.array import CArray
from secml.ml.classifiers import CClassifier


class CClassifierEmber(CClassifier):
	def __init__(self, tree_path=None):
		super(CClassifierEmber, self).__init__()
		self._lightgbm_model = self.load_tree(tree_path)

	def extract_features(self, x : CArray):
		extractor = PEFeatureExtractor(2)
		x_bytes = bytes(x.astype(np.int).tolist()[0])
		features = CArray([np.array(extractor.feature_vector(x_bytes), dtype=np.float32)])
		return features


	def load_tree(self, tree_path):
		booster = lightgbm.Booster(model_file=tree_path)
		self._classes = 2
		self._n_features = booster.num_feature()

		return booster

	def _fit(self, dataset):
		raise NotImplementedError("Fit is not implemented.")

	def _forward(self, x):
		confidence = self._lightgbm_model.predict(x.tondarray())[0]
		confidence = CArray([1 - confidence, confidence])
		return confidence
