import pytest
import cuml
from cuml.test.utils import ClassEnumerator
import numpy as np

from sklearn.datasets import make_classification


def func_positional_arg(func):
    if hasattr(func, '__code__'):
        all_args = func.__code__.co_argcount
        if func.__defaults__ is not None:
            kwargs = len(func.__defaults__)
        else:
            kwargs = 0
        return all_args - kwargs
    return 2


@pytest.fixture(scope="session")
def dataset():
    X, y = make_classification(100, 5, random_state=42)
    X = X.astype(np.float64)
    y = y.astype(np.float64)
    return X, y


models_config = ClassEnumerator(module=cuml)
models = models_config.get_models()


@pytest.mark.parametrize('model_name', list(models.keys()))
def test_fit_function(dataset, model_name):
    if model_name in ['SparseRandomProjection', 'TSNE', 'TruncatedSVD']:
        pytest.xfail("These models are not tested yet")

    n_pos_args_constr = func_positional_arg(models[model_name].__init__)

    if model_name in ['SparseRandomProjection', 'GaussianRandomProjection']:
        model = models[model_name](n_components=2)
    elif model_name == 'ExponentialSmoothing':
        model = models[model_name](endog=np.zeros((10,)))
    else:
        if n_pos_args_constr == 1:
            model = models[model_name]()
        elif n_pos_args_constr == 2:
            model = models[model_name](5)
        else:
            model = models[model_name](5, 5)

    if hasattr(model, 'fit'):
        n_pos_args_fit = func_positional_arg(models[model_name].fit)

        X, y = dataset

        if model_name == 'RandomForestClassifier':
            y = y.astype(np.int32)
            assert model.fit(X, y) is model
        else:
            if n_pos_args_fit == 1:
                assert model.fit() is model
            elif n_pos_args_fit == 2:
                assert model.fit(X) is model
            else:
                assert model.fit(X, y) is model
