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

# %% auto 0
__all__ = ['PGDCallback', 'LinfPGD', 'L2PGD']

# %% ../nbs/evasion.ipynb 5
from abc import abstractmethod

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

from .input_optimization import *


class PGDCallback(PerturbationCallback):
    "Implementes Projected Gradient Descent by bounding some $l_p$ norm of the perturbation"
    def __init__(self, epsilon=.3, rand_init=True):
        super().__init__()
        self.epsilon = epsilon
        self.p = nn.Parameter()

    def init_pert(self, x):
        self.p.data = self.rand_init(x.shape) if self.rand_init else torch.zeros(x.shape)
        self.clamp_pixel_values(x)
    
    def clamp_pixel_values(self, x):
        "Clamp peturbation such that perturbing `x` maintains valid pixel values"
        with torch.no_grad():
            x = x.to(self.p.device)
            self.p.data = (x + self.p).clamp(0., 1.) - x 

    def suggest_lr(self, input_opt: InputOptimizer):
        return self.epsilon / input_opt.epoch_size

    def before_step(self):
        with torch.no_grad():
            self.steepest_descent()

    def after_batch(self):
        with torch.no_grad():
            self.project_pert()
            self.clamp_pixel_values(self.x)

# %% ../nbs/evasion.ipynb 6
@patch()
@abstractmethod
def rand_init(self: PGDCallback, shape) -> Tensor:
    "Initialize a random perturbation in the $\epsilon$-ball"
    ...

# %% ../nbs/evasion.ipynb 7
@patch
@abstractmethod
def steepest_descent(self: PGDCallback) -> None:
    "Edit the perturbation's gradient to implement steepest descent"
    ...

# %% ../nbs/evasion.ipynb 8
@patch
@abstractmethod
def project_pert(self: PGDCallback) -> None:
    "Project the perturbation to the $\epsilon$-ball"
    ...

# %% ../nbs/evasion.ipynb 20
class LinfPGD(PGDCallback):
    "Implements PGD by bounding the $l_\infty$ norm"
    def rand_init(self, shape):
        return torch.rand(shape) * self.epsilon

    def steepest_descent(self):
        self.p.grad.sign_()

    def project_pert(self):
        self.p.clamp_(-self.epsilon, self.epsilon)

# %% ../nbs/evasion.ipynb 30
import torch.nn.functional as F
from torch.linalg import vector_norm


@delegates(F.normalize)
def _batch_normalize(t, **kwargs):
    return F.normalize(t.view(t.shape[0], -1), **kwargs).view(*t.shape)


@delegates(_batch_normalize)
def _random_unit_vector(shape, **kwargs):
    return _batch_normalize(torch.randn(shape), **kwargs)


def _batch_norm(t, **kwargs):
    return vector_norm(t.view(t.size(0), -1), **kwargs)


class L2PGD(PGDCallback):
    "Implements PGD by bounding the $l_2$ norm"
    def rand_init(self, shape):
        rand_dir = _random_unit_vector(shape)
        rand_len = torch.rand(1) * self.epsilon
        return rand_dir * rand_len

    def steepest_descent(self):
        self.p.grad.data = _batch_normalize(self.p.grad)

    def project_pert(self):
        norm = _batch_norm(self.p)
        self.p.mul_(torch.min(self.epsilon/norm, torch.ones_like(norm)))
