import mlflow

from rarelens_ml.features import RAW_COLUMNS, build


class PathogenicityModel(mlflow.pyfunc.PythonModel):
    """Serving contract: raw VEP columns in, P(pathogenic) out.

    The stock LightGBM pyfunc flavour calls `predict`, which returns class labels; wrapping the
    classifier keeps feature engineering and `predict_proba` inside the registered artifact.
    """

    def __init__(self, classifier):
        self.classifier = classifier

    def predict(self, context, model_input, params=None):
        return self.classifier.predict_proba(build(model_input[RAW_COLUMNS]))[:, 1]