import torchvision
import torch
import numpy as np
from advertorch.attacks import PGDAttack

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class MNISTdata():
    def __init__(self):
        self.num_channel = 1
        self.img_size = 28
        self.num_lables = 10
        data_train = torchvision.datasets.MNIST('./', train=True, download=True,
                                                transform=torchvision.transforms.ToTensor())
        data_test = torchvision.datasets.MNIST('./', train=False, download=True,
                                               transform=torchvision.transforms.ToTensor())

        self.train_data = torch.mul(data_train.data, 1 / 255).reshape((-1, 1, 28, 28)).type(torch.float32)
        self.train_labels = data_train.targets
        self.test_data = torch.mul(data_test.data, 1 / 255).reshape((-1, 1, 28, 28)).type(torch.float32)
        self.test_labels = data_test.targets


class dataSet(torch.utils.data.Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        data = self.data[index]
        labels = self.labels[index]
        return data, labels


class normalMnist():
    def __init__(self, data_type="test", loader_batch=128):
        data = MNISTdata()
        if data_type == 'train':
            self.data = data.train_data
            self.labels = data.train_labels
        else:
            self.data = data.test_data
            self.labels = data.test_labels

        self.loader = torch.utils.data.DataLoader(dataSet(self.data, self.labels), batch_size=loader_batch)


class attackMnist():
    def __init__(self, attack_model, attack_method="FGSM", eps=0.3, data_type="test", rand_seed=0, rand_min=0,
                 rand_max=1, loader_batch=128, for_trainning=False, atk_loss=None, quantize=False):

        normal_data = normalMnist(data_type=data_type, loader_batch=loader_batch)
        self.noarmal_data = normal_data.data
        self.labels = normal_data.labels

        x_atk = torch.tensor([]).to(device)
        bs = loader_batch
        for batch_data, batch_labels in normal_data.loader:
            if (attack_method == "FGSM"):
                if isinstance(eps, str):
                    batch_pn = FGSM(attack_model, loss_fn=atk_loss, getAtkpn=True).perturb(batch_data.to(device),
                                                                                           batch_labels.to(device))
                    eps_temp = (1) * torch.rand((len(batch_pn), 1, 1, 1))
                    eps_temp = eps_temp.to(device)
                    batch_atk = torch.clamp(batch_data.to(device) + eps_temp * batch_pn, min=0, max=1)
                else:
                    batch_atk = FGSM(attack_model, loss_fn=atk_loss, eps=eps).perturb(batch_data.to(device),
                                                                                      batch_labels.to(device))
            if (attack_method == "PGD"):
                batch_atk = PGDAttack(attack_model, loss_fn=atk_loss, eps=eps).perturb(batch_data.to(device),
                                                                                       batch_labels.to(device))
            x_atk = torch.cat((x_atk, batch_atk))
        # x_atk = torch.tensor(x_atk)
        self.data = x_atk.cpu()
        if quantize:
            self.data = (self.data * 255).type(torch.int) / 255.

        if for_trainning:
            self.loader = torch.utils.data.DataLoader(train_dataSet(self.noarmal_data, self.labels, self.data),
                                                      batch_size=loader_batch)
        else:
            self.loader = torch.utils.data.DataLoader(dataSet(self.data, self.labels), batch_size=loader_batch)


class train_dataSet(torch.utils.data.Dataset):
    def __init__(self, data, labels, data_pn):
        self.data = data
        self.labels = labels
        self.data_pn = data_pn

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        data = self.data[index]
        labels = self.labels[index]
        data_pn = self.data_pn[index]
        return data, labels, data_pn


from advertorch.attacks.base import Attack, LabelMixin
from advertorch.utils import clamp
from advertorch.utils import normalize_by_pnorm


# modify from advertorch.attacks.FGSM
class FGSM(Attack, LabelMixin):

    def __init__(self, predict, loss_fn=None, eps=0.3, clip_min=0.,
                 clip_max=1., targeted=False, getAtkpn=False):
        """
        Create an instance of the GradientSignAttack.
        """
        super(FGSM, self).__init__(
            predict, loss_fn, clip_min, clip_max)

        self.eps = eps
        self.targeted = targeted
        if self.loss_fn is None:
            self.loss_fn = torch.nn.CrossEntropyLoss(reduction="sum")
        self.getAtkpn = getAtkpn

    def perturb(self, x, y=None):
        """
        Given examples (x, y), returns their adversarial counterparts with
        an attack length of eps.
        :param x: input tensor.
        :param y: label tensor.
                  - if None and self.targeted=False, compute y as predicted
                    labels.
                  - if self.targeted=True, then y must be the targeted labels.
        :return: tensor containing perturbed inputs.
        """

        x, y = self._verify_and_process_inputs(x, y)
        xadv = x.requires_grad_()
        outputs = self.predict(xadv)

        loss = self.loss_fn(outputs, y)
        if self.targeted:
            loss = -loss
        loss.backward()
        grad_sign = xadv.grad.detach().sign()

        if self.getAtkpn:
            xadv = grad_sign
        else:
            xadv = xadv + self.eps * grad_sign
            xadv = clamp(xadv, self.clip_min, self.clip_max)

        return xadv.detach()