import mlflow
import mlflow.sklearn
from loguru import logger


class ModelLoader(object):
    def __init__(self, model_info):
        self.model_info = model_info

    def load_model(self):
        logger.info("Loading the Model")
        if self.model_info["type"] == "mlflow.sklearn":
            return self._load_mlflow_sklearn_model()
        else:
            logger.info("Unsupported Model Type")

    def _load_mlflow_sklearn_model(self):
        try:
            _model = mlflow.sklearn.load_model(self.model_info["path"])
            logger.debug("Model loaded successfully!")
            return _model
        except Exception as e:
            logger.error("Error while loading mlflow.sklearn model : {}".format(str(e)))
