"""
Test postprocessing 
"""

try:
    import numpy as np
    from sklearn.linear_model import LogisticRegression
    from skdist import postprocessing
    _import_error = None
except Exception as e:
    _import_error = e

def test_import_postprocessing():
    assert _import_error == None

def test_predict():
    X = np.array([[1,2,3], [4,5,6]])
    y = np.array([0,1])

    model1 = LogisticRegression(solver="liblinear")
    model2 = LogisticRegression(solver="lbfgs")
    model1.fit(X,y)
    model2.fit(X,y)
    
    clf = postprocessing.SimpleVoter(
        [("model1", model1), ("model2", model2)],
        voting="soft", classes=model1.classes_
        )
    pred = clf.predict(X)
    probs = clf.predict_proba(X)
    assert pred.shape == y.shape

def test_predict_hard_voting():
    X = np.array([[1,2,3], [4,5,6]])
    y = np.array([0,1])

    model1 = LogisticRegression(solver="liblinear")
    model2 = LogisticRegression(solver="lbfgs")
    model1.fit(X,y)
    model2.fit(X,y)

    clf = postprocessing.SimpleVoter(
        [("model1", model1), ("model2", model2)],
        voting="hard", classes=model1.classes_
        )
    pred = clf.predict(X)
    assert pred.shape == y.shape
