import pickle

import numpy as np

from numpy.testing import assert_array_almost_equal, assert_array_equal

import pytest

from oddt.scoring.models import classifiers, regressors

@pytest.mark.filterwarnings('ignore:Stochastic Optimizer')

@pytest.mark.parametrize('cls',

                         [classifiers.svm(probability=True),

                          classifiers.neuralnetwork(random_state=42)])

def test_classifiers(cls):

    # toy data

    X = np.concatenate((np.zeros((5, 2)), np.ones((5, 2))))

    Y = np.concatenate((np.ones(5), np.zeros(5)))

    np.random.seed(42)

    cls.fit(X, Y)

    assert_array_equal(cls.predict(X), Y)

    assert cls.score(X, Y) == 1.0

    prob = cls.predict_proba(X)

    assert_array_almost_equal(prob, [[0, 1]] * 5 + [[1, 0]] * 5, decimal=1)

    log_prob = cls.predict_log_proba(X)

    assert_array_almost_equal(np.log(prob), log_prob)

    pickled = pickle.dumps(cls)

    reloaded = pickle.loads(pickled)

    prob_reloaded = reloaded.predict_proba(X)

    assert_array_almost_equal(prob, prob_reloaded)

@pytest.mark.parametrize('reg',

                         [regressors.svm(C=10),

                          regressors.randomforest(random_state=42),

                          regressors.neuralnetwork(solver='lbfgs',

                                                   random_state=42,

                                                   hidden_layer_sizes=(20, 20)),

                          regressors.mlr()])

def test_regressors(reg):

    X = np.vstack((np.arange(30, 10, -2, dtype='float64'),

                   np.arange(100, 90, -1, dtype='float64'))).T

    Y = np.arange(10, dtype='float64')

    np.random.seed(42)

    reg.fit(X, Y)

    pred = reg.predict(X)

    assert (np.abs(pred.flatten() - Y) < 1).all()

    assert reg.score(X, Y) > 0.9

    pickled = pickle.dumps(reg)

    reloaded = pickle.loads(pickled)

    pred_reloaded = reloaded.predict(X)

    assert_array_almost_equal(pred, pred_reloaded)