"""Definition of load_model function."""

import logging
from pathlib import Path

import torch

from model.src.models import VisualCommunicationModel
from shared.datastore import Datastore

from .get_model_name import get_model_name

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


def load_model(
    datastore: Datastore,
) -> VisualCommunicationModel:
    """Instantiate model with weights loaded from latest model saved in
    MinIO."""
    assert isinstance(datastore, Datastore)
    # instantiate model
    model = VisualCommunicationModel()
    # get model object name
    model_name_path = Path('model_name.txt')
    model_object_name = get_model_name(path=model_name_path)
    logging.info('using model: %s', model_object_name)
    # load model from minio
    model_checkpoint = datastore.get_model(
        object_name=model_object_name,
    )
    model.load_state_dict(model_checkpoint)
    # clean memory
    model_checkpoint.clear()
    # move model to selected device
    model = model.to(DEVICE)
    logging.debug('finished')
    return model