from abc import abstractmethod
from secml.array import CArray
from secml_malware.attack.blackbox.c_wrapper_phi import CWrapperPhi
import numpy as np


class CBlackBoxProblem:
	def __init__(
			self,
			model_wrapper: CWrapperPhi,
			latent_space_size: int,
			population_size: int,
			penalty_regularizer: float,
			iterations: int,
			seed: int = None,
			is_debug: bool = False,
			hard_label : bool = False,
			threshold : float = 0,
			loss : str = 'l1'
	):
		self.model_wrapper = model_wrapper
		self.latent_space_size = latent_space_size
		self.population_size = population_size
		self._original_x = None
		self.penalty_regularizer = penalty_regularizer
		self.iterations = iterations
		self.seed = seed
		self.is_debug = is_debug
		self.hard_label = hard_label
		self.threshold = threshold
		self.loss = loss
		self.clear_results()

	def clear_results(self):
		self.confidences_ = []
		self.fitness_ = []
		self.sizes_ = []
		self.advx = []

	def init_starting_point(self, x: CArray):
		self._original_x = x
		padding_positions = x.find(x == 256)
		self.clear_results()
		if padding_positions:
			self._original_x = self._original_x[0,:padding_positions[0]]

	def get_bounds(self):
		return [0] * self.latent_space_size, [1] * self.latent_space_size

	def fitness(self, t):
		candidate = self.apply_feasible_manipulations(t, self._original_x)
		penalty_term = self.compute_penalty_term(self._original_x, candidate, self.penalty_regularizer)
		score = self.score_step(candidate, penalty_term)
		return [score]

	def _compute_loss(self, confidence, penalty):
		if self.loss == 'l1':
			return confidence + penalty
		if self.loss == 'cw':
			return  max(confidence - self.threshold + 0.1, 0) + penalty
		if self.loss == 'log':
			return -np.log(1 - confidence) + penalty
		raise ValueError('NO LOSS')

	def score_step(self, x: CArray, penalty_term: float) -> (float, float):
		_, confidence = self.model_wrapper.predict(x, return_decision_function=True)
		confidence = confidence[0, 1].item()
		if self.hard_label:
			confidence = np.infty if confidence > self.threshold else 0
		fitness_value = self._compute_loss(confidence, penalty_term)
		self.confidences_.append(confidence)
		self.sizes_.append(x.shape[-1])
		self.fitness_.append(fitness_value)
		return fitness_value


	def export_internal_results(self, irregular=None):
		confidence = self._slice_sequence(self.population_size, self.confidences_, irregular)
		fitness = self._slice_sequence(self.population_size, self.fitness_, irregular)
		sizes = self._slice_sequence(self.population_size, self.sizes_, irregular)
		best_idx = [np.argmin(f) for f in fitness]
		fitness = [f[i] for f, i in zip(fitness, best_idx)]
		confidence = [f[i] for f, i in zip(confidence, best_idx)]
		sizes = [f[i] for f, i in zip(sizes, best_idx)]

		return confidence, fitness, sizes

	def _slice_sequence(self, slice_size, sequence, irregular=None):
		if irregular is not None:
			expanded_sequence = [
				sequence[: sum(irregular[:i+1])] for i, _ in enumerate(irregular) #sequence[: sum(irregular[:i])] for i, val in enumerate(irregular)
			]
		else:
			how_many = len(sequence) // slice_size
			expanded_sequence = [
				sequence[: (i + 1) * slice_size]
				for i in range(how_many)
			]
		# expanded_sequence = np.array(expanded_sequence)
		return expanded_sequence

	def compute_penalty_term(self, original_x : CArray, adv_x :CArray, par: float) -> float:
		return par * abs(adv_x.shape[-1] - original_x.shape[-1])

	@abstractmethod
	def apply_feasible_manipulations(self, t, x: CArray) -> CArray:
		raise NotImplementedError("This method is abstract, you should implement it somewhere else!")
