# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/evasion.ipynb.

# %% auto 0
__all__ = ['Perturber', 'LinfPerturber', 'AdversarialExamplesAttack']

# %% ../nbs/evasion.ipynb 4
from abc import ABC, abstractclassmethod, abstractmethod
from typing import Callable, Type

import torch
from torch import nn
from fastai.vision.all import *


class Perturber(ABC, nn.Module):
    "Perturbs inputs. Perturbations are bounded under some metric (e.g. some norm). Meant to be subclassed and then passed to `AdversarialExamplesAttack`"
    def __init__(self, epsilon, x, rand_init):
        super().__init__()
        self.epsilon = epsilon
        self.p = nn.Parameter(self.rand_init(epsilon, x.shape) if rand_init else torch.zeros(x.shape))
        self.clamp_pixel_values(x)

    def forward(self, x):
        return x + self.p    

# %% ../nbs/evasion.ipynb 5
@patch(cls_method=True)
@abstractmethod
def rand_init(cls: Perturber, epsilon, shape):
    "Initialize a random perturbation"
    ...

# %% ../nbs/evasion.ipynb 6
@patch
@abstractmethod
def clamp_pixel_values(self: Perturber, x):
    "Clamp peturbation such that perturbing `x` maintains valid pixel values"
    ...

# %% ../nbs/evasion.ipynb 7
@patch
@abstractmethod
def callback(self: Perturber) -> Callback:
    "Return a callback in charge of affecting the training loop (e.g. clamping the perturbation after an optimization step, steepest descent, etc.)"
    ...

# %% ../nbs/evasion.ipynb 8
class LinfPerturber(Perturber):
    "Perturbs inputs under a bounded $l_\infty$ norm"
    @classmethod
    def rand_init(cls, epsilon, shape):
        return torch.rand(shape) * epsilon
    
    def clamp_pixel_values(self, x):
        with torch.no_grad():
            x = x.to(self.p.device)
            self.p.data = (x + self.p).clamp(0., 1.) - x

    def callback(self):
        class ProjectionCallback(Callback):
            def before_step(inner_self):
                with torch.no_grad():
                    self.p.grad.sign_()

            def after_batch(inner_self):
                with torch.no_grad():
                    self.p.clamp_(-self.epsilon, self.epsilon)  # keep perturbation small
                    self.clamp_pixel_values(inner_self.x)

        return ProjectionCallback()


# %% ../nbs/evasion.ipynb 9
class AdversarialExamplesAttack(object):
    "Constructs adversarial examples: slightly perturbed inputs that fool classification models"
    def __init__(self,
                 model: Module,
                 loss: Callable = CrossEntropyLossFlat(),
                 perturber_type: Type[Perturber] = LinfPerturber,
                 targeted: bool = False,  # Whether the constructed inputs should be classified as the specified targets or not
                 min_delta: float = 1e-2,  # Minimum loss delta for `ReduceLROnPlateau` and `EarlyStoppingCallback`
                 min_lr: float = 1e-6,  # Minimum lr for `ReduceLROnPlateau`
                 lr: float = None,  # pass `None` to pick `lr` based on other params
                 # defaults taken from advertorch
                 epsilon: float = 0.3,
                 epoch_size: int = 10,  # Affects how often epoch-callbacks are called (e.g. `Recorder`` and `EarlyStoppingCallback`)
                 n_epochs: int = 4,
                 rand_init: bool = True):
        self.loss = loss if targeted else (lambda *args, **kwargs: -loss(*args, **kwargs))
        self.lr = lr or (epsilon / epoch_size)
        store_attr('model, perturber_type, min_delta, min_lr, epsilon, epoch_size, n_epochs, rand_init')

    def perturb(self, dsets):
        x, y = dsets.load()
        x, y = x.detach().clone(), y.detach().clone()  # TODO: can I get rid of this?
        self.model.eval()
        self.model.requires_grad_(False)

        perturber = self.perturber_type(self.epsilon, x, self.rand_init)

        class TrainLoop(TrainEvalCallback):
            def before_train(self):
                super().before_train()
                self.model.eval()

            def before_validate(self):
                raise CancelValidException

        learner = Learner(DataLoaders([(x, y) for _ in range(self.epoch_size)], []),
                          nn.Sequential(perturber, self.model),
                          self.loss,
                          SGD,
                          self.lr,
                          train_bn=False,
                          default_cbs=False,
                          cbs=[TrainLoop, Recorder(valid_metrics=False), ProgressCallback, BnFreeze,
                               perturber.callback(),
                               ReduceLROnPlateau('train_loss', min_delta=self.min_delta, min_lr=self.min_lr),
                               EarlyStoppingCallback('train_loss', min_delta=self.min_delta / 10)
                               ])
        learner.fit(self.n_epochs)
        p = perturber.p.data.detach().cpu()
        return Datasets(tls=[TfmdLists(x.cpu() + p, ToTensor()),  # ToTensor for decoding
                             dsets.tls[1]])
