import torch

from ..utils import *
from ..attack import Attack
import torch.nn.functional as F
class RGIMIFGSM(Attack):
    """
    MI-FGSM Attack
    'Boosting Adversarial Attacks with Momentum (CVPR 2018)'(https://arxiv.org/abs/1710.06081)
    Arguments:
        model (torch.nn.Module): the surrogate model for attack.
        epsilon (float): the perturbation budget.
        alpha (float): the step size.
        epoch (int): the number of iterations.
        decay (float): the decay factor for momentum calculation.
        targeted (bool): targeted/untargeted attack.
        random_start (bool): whether using random initialization for delta.
        norm (str): the norm of perturbation, l2/linfty.
        loss (str): the loss function.
        device (torch.device): the device for data. If it is None, the device would be same as model
        
    Official arguments:
        epsilon=16/255, alpha=epsilon/epoch=1.6/255, epoch=10, decay=1.
    """
    
    def __init__(self, model, epsilon=16/255, alpha=1.6/255, epoch=10, decay=1., targeted=False, random_start=False, 
                norm='linfty', loss='crossentropy', device=None, attack='MI-FGSM', pre_epoch=5, s=10, num_directions=5, **kwargs):
        super().__init__(attack, model, epsilon, targeted, random_start, norm, loss, device, **kwargs)
        self.alpha = alpha
        self.epoch = epoch
        self.decay = decay
        self.pre_epoch = pre_epoch
        self.s = s
        self.num_directions = num_directions

    def forward(self, data, label, **kwargs):
        """
        The general attack procedure
        Arguments:
            data: (N, C, H, W) tensor for input images
            labels: (N,) tensor for ground-truth labels if untargetd, otherwise targeted labels
        """
        data = data.clone().detach().to(self.device)
        label = label.clone().detach().to(self.device)
        momentum = 0.
        self.random_start = True
        for _ in range(self.num_directions):
            direction_momentum = 0.
            delta = self.init_delta(data).to(self.device)
            for _ in range(self.pre_epoch):
                # Obtain the output
                logits = self.get_logits(self.transform(data+delta, momentum=direction_momentum))
                # Calculate the loss
                loss = self.get_loss(logits, label)
                # Calculate the gradients
                grad = self.get_grad(loss, delta)
                # Calculate the momentum
                direction_momentum = self.get_momentum(grad, direction_momentum)
                # Update adversarial perturbation
                delta = self.update_delta(delta, data, direction_momentum, self.alpha*self.s)
            momentum += direction_momentum
        momentum /= self.num_directions
        self.random_start = False
        delta = self.init_delta(data).to(self.device)
        for _ in range(self.epoch):
            # Obtain the output
            logits = self.get_logits(self.transform(data+delta, momentum=momentum))
            # Calculate the loss
            loss = self.get_loss(logits, label)
            # Calculate the gradients
            grad = self.get_grad(loss, delta)
            # Calculate the momentum
            momentum = self.get_momentum(grad, momentum)
            # Update adversarial perturbation
            delta = self.update_delta(delta, data, momentum, self.alpha)
            # print('iter: ', _,'mean: ', torch.mean(delta),'std: ', torch.std(delta),'grad l2 norm: ', torch.norm(grad,p=2))
        # exit()
        # exit()
        return delta.detach()
    
    

class RGINIFGSM(RGIMIFGSM):
    """
    NI-FGSM Attack
    'Nesterov Accelerated Gradient and Scale Invariance for Adversarial Attacks (ICLR 2020)'(https://arxiv.org/abs/1908.06281)
    Arguments:
        model (torch.nn.Module): the surrogate model for attack.
        epsilon (float): the perturbation budget.
        alpha (float): the step size.
        epoch (int): the number of iterations.
        decay (float): the decay factor for momentum calculation.
        targeted (bool): targeted/untargeted attack.
        random_start (bool): whether using random initialization for delta.
        norm (str): the norm of perturbation, l2/linfty.
        loss (str): the loss function.
        device (torch.device): the device for data. If it is None, the device would be same as model
        
    Official arguments:
        epsilon=16/255, alpha=epsilon/epoch=1.6/255, epoch=10, decay=1.
    """
    
    def __init__(self, model, epsilon=16/255, alpha=1.6/255, epoch=10, decay=1., targeted=False, random_start=False, 
                norm='linfty', loss='crossentropy', device=None, attack='NI-FGSM', num_directions=5, **kwargs):
        super().__init__(model, epsilon, alpha, epoch, decay, targeted, random_start, norm, loss, device, attack, **kwargs)
    
        self.num_directions = num_directions
    def transform(self, x, momentum, **kwargs):
        """
        look ahead for NI-FGSM
        """
        return x + self.alpha*self.decay*momentum

class RGIPIFGSM(RGINIFGSM):
    """
    PI-FGSM Attack
    'Patch-wise Attack for Fooling Deep Neural Network (ECCV 2020)'(https://arxiv.org/abs/2007.06765)
    Arguments:
        model (torch.nn.Module): the surrogate model for attack.
        epsilon (float): the perturbation budget.
        alpha (float): the step size.
        epoch (int): the number of iterations.
        decay (float): the decay factor for momentum calculation.
        targeted (bool): targeted/untargeted attack.
        random_start (bool): whether using random initialization for delta.
        norm (str): the norm of perturbation, l2/linfty.
        loss (str): the loss function.
        device (torch.device): the device for data. If it is None, the device would be same as model
        
    """
    
    def __init__(self, model, epsilon=16/255, alpha=1.6/255, epoch=10, decay=1., targeted=False, random_start=False, 
                norm='linfty', loss='crossentropy', device=None, attack='PI-FGSM', num_directions=5, **kwargs):
        super().__init__(model, epsilon, alpha, epoch, decay, targeted, random_start, norm, loss, device, attack, **kwargs)
        self.kern_size = 3
        self.alpha_beta = self.alpha * epoch
        self.gamma = self.alpha_beta
        P_kern = self.project_kern(self.kern_size)
        self.device = "cuda:0"# next(model.parameters()).device if device is None else device
        self.w = torch.from_numpy(P_kern).to(self.device) # project kernel
        self.num_directions = num_directions
        self.s = 10

    
    def project_kern(self, kern_size):
        kern = np.ones((kern_size, kern_size), dtype=np.float32) / (kern_size ** 2 - 1)
        kern[kern_size // 2, kern_size // 2] = 0.0
        kern = kern.astype(np.float32)
        stack_kern = np.stack([kern, kern, kern])
        stack_kern = np.expand_dims(stack_kern, 1)
        return stack_kern
    
    def project_noise(self, x, stack_kern, padding_size):
        # x = tf.pad(x, [[0,0],[kern_size,kern_size],[kern_size,kern_size],[0,0]], "CONSTANT")
        x = F.conv2d(x, stack_kern, padding=(padding_size, padding_size), groups=3)
        return x
    
    
    def update_delta(self, delta, data, grad,  **kwargs):
        if self.norm == 'linfty':
            delta = torch.clamp(delta +  grad, -self.epsilon, self.epsilon)
        else:
            grad_norm = torch.norm(grad.view(grad.size(0), -1), dim=1).view(-1, 1, 1, 1)
            scaled_grad = grad / (grad_norm + 1e-20)
            delta = (delta + scaled_grad ).view(delta.size(0), -1).renorm(p=2, dim=0, maxnorm=self.epsilon).view_as(delta)
        delta = clamp(delta, img_min-data, img_max-data)
        return delta

    def forward(self, data, label, **kwargs):
        """
        Overriden for PI-FGSM
        Arguments:
            data: (N, C, H, W) tensor for input images
            labels: (N,) tensor for ground-truth labels if untargetd, otherwise targeted labels
        """
        data = data.clone().detach().to(self.device)
        label = label.clone().detach().to(self.device)

        # Initialize adversarial perturbation
        delta = self.init_delta(data)
        delta.requires_grad = True
        grad = 0.

        amplification = 0.0 
        momentum = 0.
        self.random_start = True
        for _ in range(self.num_directions):
            direction_momentum = 0.
            delta = self.init_delta(data).to(self.device)
            for _ in range(self.pre_epoch):
                # Obtain the output
                logits = self.get_logits(self.transform(data+delta, momentum=direction_momentum))
                # Calculate the loss
                loss = self.get_loss(logits, label)
                # Calculate the gradients
                grad = self.get_grad(loss, delta)
                # Calculate the momentum
                direction_momentum = self.get_momentum(grad, direction_momentum)
                # Update adversarial perturbation
                delta = self.update_delta(delta, data, direction_momentum*self.alpha*self.s)
            momentum += direction_momentum
        momentum /= self.num_directions
        self.random_start = False
        delta = self.init_delta(data).to(self.device)

        for _ in range(self.epoch):
            # Obtain the output
            logits = self.get_logits(self.transform(data+delta, momentum=momentum))

            # Calculate the loss
            loss = self.get_loss(logits, label)
            penalty = self.penalty(delta)
            loss = loss + self.penalty_factor * penalty

            # Calculate the gradients
            grad = self.get_grad(loss, delta)
            momentum = self.get_momentum(grad, momentum)
            # Calculate the cut noise
            amplification += self.alpha_beta * momentum.sign()
            cut_noise = torch.clamp(abs(amplification) - self.epsilon, 0., 10000.0) * torch.sign(amplification)
            projection = self.gamma * torch.sign(self.project_noise(cut_noise, self.w, self.kern_size//2))
            amplification += projection
            # sum_grad = torch.sum(grad)
            # second_grad = torch.autograd.grad(sum_grad, delta, create_graph=True)[0]
            delta = self.update_delta(delta, data, self.alpha_beta * momentum.sign() + projection)
            # print('iter: ', _,'mean: ', torch.mean(delta),'std: ', torch.std(delta),'grad l2 norm: ', torch.norm(grad,p=2),'2-th order grad: mean:', torch.mean(second_grad),'std: ', torch.std(second_grad))
        # exit()

        return delta.detach()
    
class RGIEMIFGSM(RGIPIFGSM):
    """
    EMI-FGSM Attack
    'Boosting Adversarial Transferability through Enhanced Momentum (BMVC 2021)'(https://arxiv.org/abs/2103.10609)
    Arguments:
        model (torch.nn.Module): the surrogate model for attack.
        epsilon (float): the perturbation budget.
        alpha (float): the step size.
        epoch (int): the number of iterations.
        decay (float): the decay factor for momentum calculation.
        num_sample (int): the number of samples to enhance the momentum.
        radius (float): the relative radius for sampling.
        sample_method (str): the sampling method (linear/uniform/gaussian).
        targeted (bool): targeted/untargeted attack.
        random_start (bool): whether using random initialization for delta.
        norm (str): the norm of perturbation, l2/linfty.
        loss (str): the loss function.
        device (torch.device): the device for data. If it is None, the device would be same as model
        
    Official arguments:
        epsilon=16/255, alpha=epsilon/epoch=1.6/255, epoch=10, decay=1., num_sample=11, radius=7, sample_method='linear'
    """
    def __init__(self, model, epsilon=16/255, alpha=1.6/255, epoch=10, decay=1., num_sample=11, radius=7, sample_method='linear', 
                targeted=False, random_start=False, norm='linfty', loss='crossentropy', device=None, attack='EMI-FGSM', num_directions=5, **kwargs):
        super().__init__(model, epsilon, alpha, epoch, decay, targeted, random_start, norm, loss, device, attack)
        self.num_sample = num_sample
        self.radius = radius
        self.sample_method = sample_method.lower()
        self.num_directions = num_directions

    def get_factors(self):
        """
        Generate the sampling factors
        """
        if self.sample_method == 'linear':
            return np.linspace(-self.radius, self.radius, num=self.num_sample)
        elif self.sample_method == 'uniform':
            return np.random.uniform(-self.radius, self.radius, size=self.num_sample)
        elif self.sample_method == 'gaussian':
            return np.clip(np.random.normal(size=self.num_sample)/3, -1, 1)*self.radius
        else:
            raise Exception('Unsupported sampling method {}!'.format(self.sample_method)) 

    def transform(self, x, momentum, **kwargs):
        """
        Admix the input for Admix Attack
        """
        factors = np.linspace(-self.radius, self.radius, num=self.num_sample)
        return torch.concat([x+factor*self.alpha*momentum for factor in factors])

    def get_loss(self, logits, label):
        """
        Calculate the loss
        """
        return self.loss(logits, label.repeat(self.num_sample))


class RGIVMIFGSM(Attack):
    """
    VMI-FGSM Attack
    'Enhancing the transferability of adversarial attacks through variance tuning (CVPR 2021)'(https://arxiv.org/abs/2103.15571)
    Arguments:
        model (torch.nn.Module): the surrogate model for attack.
        epsilon (float): the perturbation budget.
        alpha (float): the step size.
        beta (float): the relative value for the neighborhood.
        num_neighbor (int): the number of samples for estimating the gradient variance.
        epoch (int): the number of iterations.
        decay (float): the decay factor for momentum calculation.
        targeted (bool): targeted/untargeted attack.
        random_start (bool): whether using random initialization for delta.
        norm (str): the norm of perturbation, l2/linfty.
        loss (str): the loss function.
        device (torch.device): the device for data. If it is None, the device would be same as model
        
    Official arguments:
        epsilon=16/255, alpha=epsilon/epoch=1.6/255, beta=1.5, num_neighbor=20, epoch=10, decay=1.
    """
    
    def __init__(self, model, epsilon=16/255, alpha=1.6/255, beta=1.5, num_neighbor=20, epoch=10, decay=1., targeted=False, 
                random_start=False, norm='linfty', loss='crossentropy', device=None, attack='VMI-FGSM', num_directions=5, **kwargs):
        super().__init__(attack, model, epsilon, targeted, random_start, norm, loss, device, **kwargs)
        self.alpha = alpha
        self.radius = beta * epsilon
        self.epoch = epoch
        self.decay = decay
        self.num_neighbor = num_neighbor
        self.num_directions = num_directions
        self.pre_epoch = 5
        self.s = 10

    def get_variance(self, data, delta, label, cur_grad, momentum, **kwargs):
        """
        Calculate the gradient variance    
        """
        grad = 0
        for _ in range(self.num_neighbor):
            # Obtain the output
            # This is inconsistent for transform!
            logits = self.get_logits(self.transform(data+delta+torch.zeros_like(delta).uniform_(-self.radius, self.radius).to(self.device), momentum=momentum))

            # Calculate the loss
            loss = self.get_loss(logits, label)

            # Calculate the gradients
            grad += self.get_grad(loss, delta)

        return grad / self.num_neighbor - cur_grad

    def forward(self, data, label, **kwargs):
        """
        The attack procedure for VMI-FGSM
        Arguments:
            data: (N, C, H, W) tensor for input images
            labels: (N,) tensor for ground-truth labels if untargetd, otherwise targeted labels
        """
        data = data.clone().detach().to(self.device)
        label = label.clone().detach().to(self.device)

        # Initialize adversarial perturbation
        momentum, variance = 0, 0
        self.random_start = True
        for _ in range(self.num_directions):
            direction_momentum = 0.
            direction_variance = 0.
            delta = self.init_delta(data).to(self.device)
            for _ in range(self.pre_epoch):
                # Obtain the output
                logits = self.get_logits(self.transform(data+delta, momentum=direction_momentum))
                # Calculate the loss
                loss = self.get_loss(logits, label)
                # Calculate the gradients
                grad = self.get_grad(loss, delta)
                # Calculate the momentum
                direction_momentum = self.get_momentum(grad+direction_variance, direction_momentum)
                direction_variance = self.get_variance(data, delta, label, grad, direction_momentum)
                # Update adversarial perturbation
                delta = self.update_delta(delta, data, direction_momentum, self.alpha*self.s)
            momentum += direction_momentum
            variance += direction_variance
        momentum /= self.num_directions
        variance /= self.num_directions
        self.random_start = False
        delta = self.init_delta(data).to(self.device)

        for _ in range(self.epoch):
            # Obtain the output
            logits = self.get_logits(self.transform(data+delta, momentum=momentum))

            # Calculate the loss
            loss = self.get_loss(logits, label)
            penalty = self.penalty(delta)
            loss = loss + self.penalty_factor * penalty

            # Calculate the gradients
            grad = self.get_grad(loss, delta)

            # Calculate the momentum
            momentum = self.get_momentum(grad+variance, momentum)

            # Calculate the variance
            variance = self.get_variance(data, delta, label, grad, momentum)
            # sum_grad = torch.sum(grad)
            # second_grad = torch.autograd.grad(sum_grad, delta, create_graph=True)[0]
            # Update adversarial perturbation
            delta = self.update_delta(delta, data, momentum, self.alpha)
            # print('iter: ', _,'mean: ', torch.mean(delta),'std: ', torch.std(delta),'grad l2 norm: ', torch.norm(grad,p=2),'2-th order grad: mean:', torch.mean(second_grad),'std: ', torch.std(second_grad))
        # exit()

        return delta.detach()