from funcy import rcompose


class PredictionModelHandle:
    """Simplifies usage of ModelHandle instances for prediction purposes."""

    def __init__(self, model_handle):
        self.__model_handle = model_handle
        self.__predict = rcompose(self.__model_handle.prep_images, self.__model_handle.model.predict)
        self.__predict_proba = rcompose(self.__model_handle.prep_images, self.__model_handle.model.predict_proba)

    def predict(self, *args, **kwargs):
        return self.__predict(*args, **kwargs)

    def predict_proba(self, *args, **kwargs):
        return self.__predict_proba(*args, **kwargs)

    def __call__(self, *args, **kwargs):
        return self.predict_proba(*args, **kwargs)