import torch
import struct
import numpy as np
from secml.array import CArray
from secml.settings import SECML_PYTORCH_USE_CUDA
from secml_malware.attack.whitebox import CEnd2EndMalwareEvasion

use_cuda = torch.cuda.is_available() and SECML_PYTORCH_USE_CUDA


def gradient_search(
		start_byte: int,
		gradient: torch.Tensor,
		embedding_bytes: torch.Tensor,
		invalid_val=np.infty,
		invalid_pos=-1
):
	"""Given the starting byte, the gradient and the embedding map,
	it returns a list of distances

	Arguments:
		start_byte {int} -- the byte that is used for searching the fittest one.
		gradient {ndarray} -- An array containing the gradient computed in start_byte
		embedding_bytes {ndarray} -- a matrix containing all the embeddings

	Keyword Arguments:
		invalid_val {[type]} -- Value used for indexes that are discrded during the process (default: {np.infty})
		invalid_pos {int} -- Index of invalid value, that is the index 256 (default: {-1})

	Returns:
		list -- matrix of distances
	"""
	if torch.equal(gradient, torch.zeros(gradient.shape)):
		invalid_distances = torch.tensor([invalid_val] * len(embedding_bytes))
		if use_cuda:
			invalid_distances = invalid_distances.cuda()
		return invalid_distances
	distance = torch.zeros(257)
	start_emb_byte = embedding_bytes[start_byte]
	gs = -gradient / torch.norm(gradient)
	if use_cuda:
		distance = distance.cuda()
	for i, b in enumerate(embedding_bytes):
		bts = b - start_emb_byte
		s_i = torch.dot(gs, bts)
		if s_i <= 0:
			distance[i] = invalid_val
		else:
			d_i = torch.norm(b - (start_emb_byte + s_i * gs))
			distance[i] = d_i
	distance[invalid_pos] = invalid_val
	return distance


class CDiscreteBytesEvasion(CEnd2EndMalwareEvasion):
	"""Creates the attack that perturbs the header of a Windows PE malware.
	More details in the reference paper.
	"Explaining Vulnerabilities of Deep Learning to Adversarial Malware Binaries"
	https://arxiv.org/abs/1901.03583
	Demetrio, Biggio, Lagorio, Roli, Armando

	Arguments:
		CMalConvEvasion {[type]} -- [description]
	"""

	def __init__(
			self,
			end2end_model,
			index_to_perturb,
			iterations=100,
			is_debug: bool = False,
			random_init: bool = False,
			threshold: float = 0.5,
			penalty_regularizer: float = 0,
			chunk_hyper_parameter : int = 256
	):
		"""Creates the evasion object

		Arguments:
			malconv {CClassifierMalConvPyTorch} -- The classifier to evade
			surrogate_data {CDataset} -- data used for the surrogate (if any)

		Keyword Arguments:
			how_many {int} -- How many index to perturb. None means use all indexes (default: {None})
			index_to_perturb {list of int} -- Indexes that will be used for the attack (default: [2, ..., 58])
			surrogate_classifier {CClassifierMalConvPyTorch} -- Surrogate classifier, if any (default: {None})
			use_surrogate {bool} -- Specify if should use the surrogate instead of real classifier (default: {False})
			iterations {float} -- How many iterations for the attack (default: {100.0})
			is_debug {bool} -- If true, it will print on console additional information duringthe attack (default: {False})
			random_init {bool} -- Randomize the bytes located at the specified indexes before starting the attack(default: {False})
		"""
		super(CDiscreteBytesEvasion, self).__init__(
			end2end_model=end2end_model,
			indexes_to_perturb=index_to_perturb,
			iterations=iterations,
			is_debug=is_debug,
			random_init=random_init,
			threshold=threshold,
			penalty_regularizer=penalty_regularizer,
		)
		self.chunk_hyper_parameter = chunk_hyper_parameter


	def compute_penalty_term(self, original_x: CArray, adv_x: CArray, par: float):
		emb_original_x = self.classifier.embed(original_x.tondarray())
		emb_adv_x = self.classifier.embed(adv_x.tondarray())
		lambda_variable = torch.autograd.Variable(torch.Tensor([par]), requires_grad=True)
		penalty_term = torch.autograd.Variable(emb_original_x - emb_adv_x, requires_grad=True)
		penalty_term = penalty_term.norm() * lambda_variable
		return penalty_term

	def loss_function_gradient(self, original_x: CArray, adv_x: CArray, penalty_term : torch.Tensor):
		emb_adv_x = self.classifier.embed(adv_x.tondarray())
		if use_cuda:
			penalty_term = penalty_term.cuda()
		y = self.classifier.embedding_predict(emb_adv_x)
		output = y + penalty_term
		if use_cuda:
			output = output.cuda()
		g = torch.autograd.grad(output, emb_adv_x)[0]
		g = torch.transpose(g, 1, 2)[0]
		return g

	def optimization_solver(self, E, gradient_f, index_to_consider, x_init):
		if self.chunk_hyper_parameter:
			best_indexes = gradient_f[index_to_consider].norm(dim=1).argsort(descending=True)[:self.chunk_hyper_parameter]
			best_indexes = [index_to_consider[i] for i in best_indexes]
		else:
			best_indexes = index_to_consider
		results = [self._find_byte_using_gradient(E, gradient_f, i, x_init) for i in best_indexes]
		for res in results:
			if not torch.equal(res[1], self._invalid_value):
				x_init[res[2]] = res[0].item()
		return x_init

	def _find_byte_using_gradient(self, E, gradient_f, i, x_init):
		gradient_f_i = gradient_f[i]
		x_i = x_init[i].tondarray().astype(np.uint16).ravel().item()
		distances = gradient_search(
			x_i,
			gradient_f_i,
			E,
			invalid_val=self._invalid_value,
			invalid_pos=self.invalid_pos,
		)
		min_value, byte_to_choose = torch.min(distances, dim=0, keepdim=True)
		return byte_to_choose, min_value, i

	def infer_step(self, x_init):
		_, confidence = self.classifier.predict(x_init, return_decision_function=True)
		return confidence[1].item()

	def invert_feature_mapping(self, x, x_adv):
		return x_adv

	def apply_feature_mapping(self, x):
		return x