import logging
import torch
import torch.nn as nn
from abc import ABC, abstractmethod

log = logging.getLogger(__name__)


class Detector(ABC):
    def __init__(self, lr=0.001):
        self.lr = lr
        self.__model = None
        self.__optimizer = None
        self.__criterion = nn.CrossEntropyLoss()

    @property
    def model(self):
        return self.__model

    @property
    def optimizer(self):
        return self.__optimizer

    @property
    def criterion(self):
        return self.__criterion

    @abstractmethod
    def init_model(self, char_vocab, hidden_size, n_domain_type, n_layers):
        pass

    @abstractmethod
    def train_model(self, epoch, train_dataset):
        pass

    @abstractmethod
    def predict(self, epoch, train_dataset):
        pass

    def load_model(self, file_path):
        """ This function load already saved model and sets cuda parameters.

        :param file_path: File path of a model to loaded.
        :type file_path: string
        """
        model = torch.load(file_path)
        model.eval()
        self.__model = model
        self.__set_model2cuda()
        self.__set_optimizer()

    def save_model(self, file_path):
        """ This function saves model to given location.

        :param file_path: File path to save model.
        :type file_path: string
        """
        torch.save(self.model, file_path)

    def __set_parallelism(self):
        gpu_count = torch.cuda.device_count()
        if gpu_count > 1:
            log.info("%s GPUs!" % (gpu_count))
            self.__model = nn.DataParallel(self.model)
            self.__set_model2cuda()
        else:
            self.__set_model2cuda()

    def __set_optimizer(self):
        self.__optimizer = torch.optim.RMSprop(
            self.model.parameters(), self.lr, weight_decay=0.0
        )

    def __set_model2cuda(self):
        if torch.cuda.is_available():
            log.info("Setting cuda")
            self.model.cuda()

    def leverage_model(self, model):
        """This function leverages model by setting parallelism parameters.

        :param model: Model instance.
        :type model: RNNClassifier
        """
        self.__model = model
        self.__set_parallelism()
        self.__set_optimizer()

[(6, 'logging.getLogger', 'logging.getLogger', False, 'import logging\n'), (14, 'torch.nn.CrossEntropyLoss', 'nn.CrossEntropyLoss', True, 'import torch.nn as nn\n'), (46, 'torch.load', 'torch.load', False, 'import torch\n'), (58, 'torch.save', 'torch.save', False, 'import torch\n'), (61, 'torch.cuda.device_count', 'torch.cuda.device_count', False, 'import torch\n'), (75, 'torch.cuda.is_available', 'torch.cuda.is_available', False, 'import torch\n'), (64, 'torch.nn.DataParallel', 'nn.DataParallel', True, 'import torch.nn as nn\n')]