import unittest

import numpy as np
from sklearn.datasets import load_breast_cancer
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split

from ngboost.distns import Bernoulli
from ngboost.learners import default_tree_learner
from ngboost.ngboost import NGBoost
from ngboost.scores import MLE

np.random.seed(1)


class TestBasic(unittest.TestCase):
    def test_basic(self):
        pass

    def test_classification(self):
        data, target = load_breast_cancer(True)
        x_train, x_test, y_train, y_test = train_test_split(data, target,
                                                            test_size=0.2,
                                                            random_state=42)
        ngb = NGBoost(Base=default_tree_learner, Dist=Bernoulli, Score=MLE(),
                      verbose=False)
        ngb.fit(x_train, y_train)
        preds = ngb.pred_dist(x_test)
        score = roc_auc_score(y_test, preds.prob)
        print(score)
        assert score >= 0.95

    def test_regression(self):
        pass

@@ -76,16 +76,29 @@
 t_cancer
+, load_boston
 %0Afrom sk
@@ -131,16 +131,36 @@
 uc_score
+, mean_squared_error
 %0Afrom sk
@@ -242,16 +242,24 @@
 ernoulli
+, Normal
 %0Afrom ng
@@ -1015,29 +1015,8 @@
 ob)%0A
-        print(score)%0A
     
@@ -1064,29 +1064,566 @@
 gression(self):%0A        
-pass
+data, target = load_boston(True)%0A        x_train, x_test, y_train, y_test = train_test_split(data, target,%0A                                                            test_size=0.2,%0A                                                            random_state=42)%0A        ngb = NGBoost(Base=default_tree_learner, Dist=Normal, Score=MLE(),%0A                      natural_gradient=True, verbose=False)%0A        ngb.fit(x_train, y_train)%0A        preds = ngb.predict(x_test)%0A        score = mean_squared_error(y_test, preds)%0A        assert score %3C= 8.0
 %0A

add test_regression