from abc import abstractmethod

from secml.array import CArray
from secml.ml.classifiers import CClassifier
from secml_malware.models import End2EndModel


class CWrapperPhi():

	def __init__(self, model: CClassifier):
		self.classifier = model

	@abstractmethod
	def extract_features(self, x):
		raise NotImplementedError("This method is abstract, you should implement it somewhere else!")

	@abstractmethod
	def predict(self, x, return_decision_function=True):
		padding_position = x.find(x == 256)
		if padding_position:
			x = x[0, :padding_position[0]]
		feature_vector = self.extract_features(x)
		return self.classifier.predict(feature_vector, return_decision_function=return_decision_function)


class CEmberWrapperPhi(CWrapperPhi):
	def extract_features(self, x):
		padding_positions = x.find(x == 256)
		if padding_positions:
			return self.classifier.extract_features(x[0, :padding_positions[0]])
		else:
			return self.classifier.extract_features(x)


class CEnd2EndWrapperPhi(CWrapperPhi):
	def extract_features(self, x):
		padded_x = CArray.zeros((1, self.classifier.get_input_max_length())) + self.classifier.get_embedding_value()
		length = min(x.shape[-1], self.classifier.get_input_max_length())
		padded_x[0, :length] = x[0, :length] + self.classifier.get_is_shifting_values()
		return padded_x
