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

# %% auto 0
__all__ = ['BackdoorAttack', 'DataPoisoningAttack']

# %% ../nbs/Core/backdoor.ipynb 3
from abc import ABC, abstractmethod
from typing import Dict

from fastai.vision.all import *


class BackdoorAttack(ABC, Callback):
    "A `Callback` that affects the training process to install a backdoor. Also allows the measuring of the attack's success"
    def after_create(self):
         test_eq(len(self.dls), 2)
         self.dls.loaders.append(self._asr_dl())

    @abstractmethod
    def _asr_dl(self) -> DataLoader:
         "Returns a `DataLoader` used to measure the ASR (attack success rate)"
         ...

# %% ../nbs/Core/backdoor.ipynb 5
def _poisoned_dataset(clean, poison):
        all = clean + poison
        all.clean, all.poison = clean, poison
        if hasattr(clean, 'loss_func'):
            all.loss_func = clean.loss_func
        return all


class DataPoisoningAttack(BackdoorAttack):
    "A `BackdoorAttack` that installs the backdoor by altering a small portion of the training dataset"
    def __init__(self, test_only=False, poison_fraction=.1):
         super().__init__()
         store_attr('test_only, poison_fraction')

    def after_create(self):
         super().after_create()
         if not self.test_only:
            self._poison_train_dl()

    def _poison_train_dl(self):
        poison_size = int(self.poison_fraction * len(self.dls.train_ds))
        to_be_poisoned = self._subset_to_poison(self.dls.train_ds, poison_size)
        self.dls.train.dataset = _poisoned_dataset(clean=self.dls.train_ds - to_be_poisoned,
                                                   poison=self._poison(to_be_poisoned))

    def _subset_to_poison(self, clean_train_dataset: Datasets, size: int) -> Datasets:
        return clean_train_dataset.random_sub_dsets(size)

    @abstractmethod
    def _poison(self,
                data_to_poison: Datasets  # The prtion of the clean training data that will be replaced by poison. Could be used to construct the poison data
                ) -> Datasets:  # A dataset of poison data
        ...
