import torch

from ..utils import *
from ..attack import Attack



def exp(val):
    return 2**val
def ln(val):
    # return np.log(val)
    # print('change ln')
    return np.log(val)
def power(val, factor=2):
    return val ** factor

def linear(val, factor=1):
    return val * factor

def normalize(fn, len):
    val = np.array([fn(len-i) for i in range(len)])[::-1]
    summation = np.sum(val)
    return 16/255*(val / summation).astype(np.float32)

def pvalue(val,epsion=8/255, order=0.6):
    return epsion/(val+1)**order

def identity(val):
    return 1.

class MIFGSM(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', adapted=False,rho=1.0,**kwargs):
        self.alpha = alpha # epsilon/epoch# alpha
        self.epoch = epoch
        self.decay = decay
        super().__init__(attack, model=model, epsilon=epsilon, targeted=targeted, random_start=random_start, norm=norm, loss=loss, device=device, **kwargs)
        self.alpha = alpha # epsilon/epoch# alpha
        self.epoch = epoch
        self.decay = decay