import torch

from ..utils import *
from ..attack import Attack
import os
import sys

class NIT_IFGSM(Attack):
    """
    I-FGSM Attack
    'Adversarial Examples in the Physical World (ICLR 2017)'(https://arxiv.org/abs/1607.02533)

    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.
        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
    """

    def __init__(self, model_name, epsilon=16/255, alpha=1.6/255, epoch=10, targeted=False, random_start=False, 
                norm='linfty', loss='crossentropy', device=None, mode='ensemble', num_ens=5, **kwargs):
        super().__init__('I-FGSM', model=model_name, epsilon=epsilon, targeted=targeted, random_start=random_start, norm=norm, loss=loss, device=device, **kwargs)
        self.alpha = alpha
        self.epoch = epoch
        self.decay = 0
        self.alpha = epsilon/self.epoch
        self.noise_coeff = self.load_noise(model_name)
        self.mode = mode
        self.num_ens = num_ens
        
        
    def load_noise(self, model_name):
        noise_name = 'noise_' + model_name + '.pt'
        noise = torch.load('/home/cxu-serve/p1/zzh136/nips2023/tricks/transferAttack/bayes/models/' + noise_name,map_location=torch.device('cuda'))
        # print(noise)
        # exit()
        return noise
    
    def get_logits(self, x, **kwargs):
        """
        The inference stage, which should be overridden when the attack need to change the models (e.g., ensemble-model attack, ghost, etc.) or the input (e.g. DIM, SIM, etc.)
        """
        if self.mode == 'random':
            logit =  self.model(x) 
            output = logit + 1e-3*torch.randn_like(logit)
            # print('------------------------------')
            # print(logit)
            # print('------------------------------')
            # print(output)
            # print('------------------------------')
            # exit()
            return output
        elif self.mode == 'ensemble':
            outputs = torch.zeros((x.shape[0], 1000))
            for i in range(self.num_ens):
                logit =  self.model(x) 
                outputs = logit + 1e-4 * torch.randn_like(logit)
            return outputs / self.num_ens
    
    def load_model(self, model_name):
        nit_model = 'nit_' + model_name + '.pt'
        # whether it is in the models dic
        models = os.listdir('/home/cxu-serve/p1/zzh136/nips2023/tricks/transferAttack/bayes/models')
        if nit_model in models:
            # model_stat = torch.load('/home/cxu-serve/p1/zzh136/nips2023/tricks/transferAttack/bayes/models/' + nit_model, map_location=torch.device('cuda'))
            model = model_list[model_name](weights='DEFAULT').eval().cuda()
            # model.load_state_dict(model_stat)
            model = wrap_model(model)
            return model
        else:
            # throw an exception
            raise Exception('No such model')
        
    
        