import torch


def load_checkpoint(filepath, clf_class, device):
    """
    This specific function is used to load a checkpoint that 
    represents a clf_class object that is used in transfer learning.
    
    All params have gradients frozen.
    
    Inputs:
        - Filepath of checkpoint
        - clf_class: The object with which the classifier in the pretrained network was replaced with (i.e. CustomClassifier)
        - device: The device the model should be loaded on to. This is a string
    
    Output:
        - model: Transfer Learning model loaded with a clf_class.
    """

    cuda = torch.cuda.is_available()

    if cuda:
        checkpoint = torch.load(filepath)
    else:
        checkpoint = torch.load(filepath, map_location=lambda storage, location: device) #loads a GPU trained model onto a CPU

    
    new_classifier = clf_class(checkpoint['input_size'],
                             checkpoint['output_size'],
                             checkpoint['hidden_layers'],
                             checkpoint['dropout'])
    
    #loading classifier state dict
    new_classifier.load_state_dict(checkpoint['state_dict'])
    
    # Transfer Architecture type used in creating the classifier
    model_arch = checkpoint['model_architecture']
    
    # Reload that architecture
    model = eval(f"models.{model_arch}(pretrained=True)")
    
    # Replace default classifier with saved classifier
    model.classifier = new_classifier
    
    #model to device
    model.to(device)

    # storing the training data class to idx attribute
    model.class_to_idx = checkpoint['class_to_idx']
    
    # freeze entire Network - no more training to occur.
    for param in model.parameters():
        param.requires_grad = False
    
    return model