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

# %% auto 0
__all__ = ['BadNetsAttack']

# %% ../nbs/badnets.ipynb 3
from fastai.vision.all import *

from .backdoor import *

@delegates()
class BadNetsAttack(DataPoisoningAttack):
    "https://arxiv.org/abs/1708.06733"
    def __init__(self,
                 trigger: np.array,  # Expected to be the same shape as the input
                 target_class: str,  # The target class for inputs with the trigger
                 *args, **kwargs):
        super().__init__(*args, **kwargs)
        store_attr('trigger, target_class')

    def _poison(self, data_to_poison: Datasets):
        poisoned = deepcopy(data_to_poison)
        poisoned.tls[0].tfms.add(mk_transform(self._insert_trigger))
        poisoned.tls[1].tfms = Pipeline([lambda _: self.target_class, data_to_poison.categorize])

        return poisoned

    def _insert_trigger(self, img):
        patched_np = (np.array(img)+self.trigger) % 256
        return type(img).create(patched_np)

    def _asr_dl(self):
        return self._poison(self.dls.valid_ds).dl()
