import optuna
import pandas as pd
from ngboost import NGBRegressor
from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split
from sklearn.utils.estimator_checks import check_estimator

from yikit.models import NNRegressor, Objective

SEED = 334


def test_models():
    diabetes = load_diabetes()
    X = pd.DataFrame(diabetes["data"], columns=diabetes["feature_names"])
    y = pd.Series(diabetes["target"])

    # objective = Objective(RandomForestRegressor(), X, y, scoring = 'neg_mean_squared_error')
    # trial = optuna.trial.FixedTrial({
    #     'min_samples_split': 2,
    #     'max_depth': 10,
    #     'n_estimators': 100,
    # })
    # print(objective(trial))

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=SEED
    )

    objective = Objective(
        NGBRegressor(random_state=SEED, verbose=False),
        X,
        y,
        scoring="neg_mean_squared_error",
    )
    study = optuna.create_study(sampler=objective.sampler)
    # trial = optuna.trial.FixedTrial({
    #     'Base__max_depth': 3,
    #     'Base__criterion': 'mse',
    #     'n_estimators': 100,
    #     'minibatch_frac': 1.0,
    # })
    study.optimize(objective, n_trials=1)
    print(objective.get_best_estimator(study))

    # estimator = EnsembleRegressor(scoring = ['r2', 'neg_mean_squared_error'], random_state = 334, verbose = 1, boruta = True, opt = False, method = 'stacking')
    # estimator = EnsembleRegressor(estimators = [RandomForestRegressor(), LinearRegression()], scoring = ['r2', 'neg_mean_squared_error'], random_state = 334, boruta = False, opt = False, verbose = 1, method = 'stacking')
    # estimator = EnsembleRegressor(estimators = [RandomForestRegressor(), LGBMRegressor()], scoring = None, random_state = 334, boruta = False, opt = False)
    # estimator.fit(X_train, y_train)
    # print(mean_squared_error(estimator.predict(X_test), y_test, squared = False))
