import abc
from collections import OrderedDict
from pathlib import Path
from typing import Union

import torch


class BaseGenerativeAttack(abc.ABC):

    def __init__(self,
                 device: Union[str, torch.device],
                 epsilon: float = 32 / 255) -> None:
        if isinstance(device, str):
            device = torch.device(device)
        self.device = device
        self.set_adv_gen()
        self.set_mode('eval')
        self.epsilon = epsilon

    @abc.abstractmethod
    def set_adv_gen(self):
        pass

    def load_ckpt(self, ckpt: Union[str, Path, OrderedDict]) -> None:
        if isinstance(ckpt, str):
            ckpt = Path(ckpt)
        if isinstance(ckpt, Path):
            if not ckpt.exists():
                raise FileNotFoundError(f'File not found: {ckpt}')
            ckpt = torch.load(ckpt, map_location=self.device)
        self.adv_gen.load_state_dict(ckpt)
        self.adv_gen.to(self.device)

    def save_ckpt(self, ckpt: Union[str, Path]) -> None:
        if isinstance(ckpt, str):
            ckpt = Path(ckpt)
        _adv_gen_cpu = self.adv_gen.to('cpu')
        torch.save(_adv_gen_cpu.state_dict(), ckpt)

    def get_params(self) -> torch.nn.Parameter:
        return self.adv_gen.parameters()

    def get_model(self) -> torch.nn.Module:
        return self.adv_gen

    def set_mode(self, mode: str) -> None:
        assert mode in ['train', 'eval']
        self.adv_gen.train() if mode == 'train' else self.adv_gen.eval()

    @abc.abstractmethod
    def attack(self, *args) -> torch.Tensor:
        pass

    def __call__(self, x_nat: torch.Tensor, *extra_inputs) -> torch.Tensor:
        x_adv = self.attack(x_nat, *extra_inputs)
        x_adv = torch.min(torch.max(x_adv, x_nat - self.epsilon),
                          x_nat + self.epsilon)
        torch.clamp_(x_adv, 0.0, 1.0)
        return x_adv