import unittest
import numpy as np
from modlamp.ml import train_best_model, score_cv, score_testset


class TestTrainModel(unittest.TestCase):
    descriptor = np.array(
        [
            [0.36581538],
            [0.17366842],
            [0.19780249],
            [0.56398772],
            [0.45888646],
            [0.37802301],
            [0.34111542],
            [0.25833683],
            [0.29136177],
            [0.32289102],
            [0.10113402],
            [0.07952263],
            [0.43523497],
            [0.27801969],
            [0.07287307],
            [0.27460903],
            [0.09064536],
            [0.13651603],
            [0.30059792],
            [0.26340581],
        ]
    )
    descriptor_test = np.array([[0.95], [0.91], [0.88], [0.89], [0.02], [0.01], [0.015], [0.05]])

    target = np.array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
    target_test = np.array([1, 1, 1, 1, 0, 0, 0, 0])

    best_svm_model = train_best_model(
        "svm", descriptor, target, cv=3, n_jobs=-1, param_grid=[{"clf__C": [1], "clf__kernel": ["rbf"]}]
    )
    best_rf_model = train_best_model(
        "rf", descriptor, target, cv=3, n_jobs=-1, param_grid=[{"clf__n_estimators": [100]}]
    )

    def test_score_svm(self):
        self.assertAlmostEqual(self.best_svm_model.score(self.descriptor, self.target), 0.7, 3)

    def test_score_rf(self):
        self.assertAlmostEqual(self.best_rf_model.score(self.descriptor, self.target), 1.0, 3)

    def test_score_cv(self):
        score = score_cv(self.best_svm_model, self.descriptor, self.target, cv=2)
        self.assertAlmostEqual(score["mean"].iloc[1], 0.7, 3)

    def test_score_testset(self):
        score = score_testset(self.best_rf_model, self.descriptor_test, self.target_test)
        self.assertEqual(score["Scores"].iloc[0], 1.0)


if __name__ == "__main__":
    unittest.main()
