# -*- coding: utf-8 -*-
from sklearn.linear_model import LogisticRegression, Ridge
from datacloud.ml_utils.common.constants import RANDOM_SEED


class Proba_Logreg(LogisticRegression):
    def predict(self, X):
        return LogisticRegression.predict_proba(self, X)[:, 1]


def get_proba_logistic(random_state=RANDOM_SEED, n_jobs=8, **kwargs):
    return Proba_Logreg(
        random_state=random_state,
        n_jobs=n_jobs,
        **kwargs
    )


def get_ridge(random_state=RANDOM_SEED, copy_X=True, **kwargs):
    return Ridge(
        random_state=random_state,
        copy_X=True,
        **kwargs
    )
