import os

from torch.utils import model_zoo

from classes.fc4.squeezenet.SqueezeNet import SqueezeNet

model_urls = {
    1.0: 'https://download.pytorch.org/models/squeezenet1_0-a815701f.pth',
    1.1: 'https://download.pytorch.org/models/squeezenet1_1-f364aa15.pth',
}


class SqueezeNetLoader:
    def __init__(self, version: float = 1.1):
        self.__version = version
        self.__model = SqueezeNet(self.__version)

    def load(self, pretrained: bool = False) -> SqueezeNet:
        """
        Returns the specified version of SqueezeNet
        @param pretrained: if True, returns a model pre-trained on ImageNet
        """
        if pretrained:
            path_to_local = os.path.join("assets", "pretrained")
            os.environ['TORCH_HOME'] = path_to_local
            self.__model.load_state_dict(model_zoo.load_url(model_urls[self.__version]))
        return self.__model