from collections import namedtuple

from image_prediction.locations import MLRUNS_DIR
from image_prediction.model_loader.loader import ModelLoader
from image_prediction.model_loader.loaders.mlflow import MlflowLoader

ModelClassesPair = namedtuple("ModelClassesPair", ["model", "classes"])


def load_model_and_classes(identifier, model_loader: ModelLoader = None) -> ModelClassesPair:
    if not model_loader:
        model_loader = MlflowLoader(MLRUNS_DIR)

    model = model_loader.load_model(identifier)
    classes = model_loader.load_classes(identifier)

    return ModelClassesPair(model, classes)