import pytest
import numpy as np
import lightgbm as lgb
from arfs.allrelevant import Leshy, BoostAGroota, GrootCV
from arfs.utils import generated_corr_dataset_regr, generated_corr_dataset_classification
from arfs.utils import LightForestClassifier, LightForestRegressor


class TestLeshy:
    """
    Test suite for all-relevant FS boruta-like method: Leshy
    """

    def test_borutaPy_vs_leshy_with_rfc_and_native_feature_importance(self):
        # too slow for circleci to run them in a reasonable time
        # takes 2 min on laptop, 1h or more on circleci
        # sklearn random forest implementation
        # X, y, w = _generated_corr_dataset_classification()
        # rfc = RandomForestClassifier(max_features='sqrt', max_samples=0.632, n_estimators=100)
        # bt = BorutaPy(rfc)
        # bt.fit(X.values, y)
        # borutapy_rfc_list = sorted(list(X.columns[bt.support_]))

        # lightGBM random forest implementation
        baseline_list = ['var0', 'var1', 'var2', 'var3', 'var4']
        X, y, w = generated_corr_dataset_classification(size=100)
        n_feat = X.shape[1]
        rfc = LightForestClassifier(n_feat)
        # RandomForestClassifier(max_features='sqrt', max_samples=0.632, n_estimators=100) # --> too slow
        arfs = Leshy(rfc, verbose=0, max_iter=10, random_state=42, importance='native')
        arfs.fit(X, y)
        leshy_rfc_list = sorted(arfs.support_names_)

        # assert borutapy_rfc_list == leshy_rfc_list, "same selected features are expected"
        assert bool(set(baseline_list) & set(leshy_rfc_list)), "expect non-empty intersection"

    def test_borutaPy_vs_leshy_with_rfr_and_native_feature_importance(self):
        # too slow for circleci to run them in a reasonable time
        # takes 2 min on laptop, 1h or more on circleci
        # # sklearn random forest implementation
        # X, y, w = _generated_corr_dataset_regr()
        # rfr = RandomForestRegressor(max_features=0.3, max_samples=0.632, n_estimators=100)
        # bt = BorutaPy(rfr)
        # bt.fit(X.values, y)
        # borutapy_rfc_list = sorted(list(X.columns[bt.support_]))

        # lightGBM random forest implementation
        baseline_list = ['var0', 'var1', 'var2', 'var3', 'var4']
        X, y, w = generated_corr_dataset_regr(size=100)
        n_feat = X.shape[1]
        rfr = LightForestRegressor(n_feat)
        # rfr = RandomForestRegressor(max_features=0.3, max_samples=0.632, n_estimators=10)
        arfs = Leshy(rfr, verbose=0, max_iter=10, random_state=42, importance='native')
        arfs.fit(X, y)
        leshy_rfc_list = sorted(arfs.support_names_)

        # assert borutapy_rfc_list == leshy_rfc_list, "same selected features are expected"
        assert bool(set(baseline_list) & set(leshy_rfc_list)), "expect non-empty intersection"

    def test_borutaPy_vs_leshy_with_rfc_and_shap_feature_importance(self):
        # too slow for circleci to run them in a reasonable time
        # takes 2 min on laptop, 1h or more on circleci
        # # sklearn random forest implementation
        # X, y, w = _generated_corr_dataset_classification()
        # rfc = RandomForestClassifier(max_features='sqrt', max_samples=0.632, n_estimators=100)
        # bt = BorutaPy(rfc)
        # bt.fit(X.values, y)
        # borutapy_rfc_list = sorted(list(X.columns[bt.support_]))

        # lightGBM random forest implementation
        baseline_list = ['var0', 'var1', 'var2', 'var3', 'var4']
        X, y, w = generated_corr_dataset_classification(size=100)
        n_feat = X.shape[1]
        model = LightForestClassifier(n_feat)
        arfs = Leshy(model, verbose=0, max_iter=10, random_state=42, importance='shap')
        arfs.fit(X, y)
        leshy_rfc_list = sorted(arfs.support_names_)

        # assert borutapy_rfc_list == leshy_rfc_list, "same selected features are expected"
        assert bool(set(baseline_list) & set(leshy_rfc_list)), "expect non-empty intersection"

    def test_borutaPy_vs_leshy_with_rfr_and_shap_feature_importance(self):
        # too slow for circleci to run them in a reasonable time
        # takes 2 min on laptop, 1h or more on circleci
        # # sklearn random forest implementation
        # X, y, w = _generated_corr_dataset_regr()
        # rfr = RandomForestRegressor(max_features=0.3, max_samples=0.632, n_estimators=100)
        # bt = BorutaPy(rfr)
        # bt.fit(X.values, y)
        # borutapy_rfc_list = sorted(list(X.columns[bt.support_]))

        # lightGBM random forest implementation
        baseline_list = ['var0', 'var1', 'var2', 'var3', 'var4']
        X, y, w = generated_corr_dataset_regr(size=500)
        n_feat = X.shape[1]
        model = LightForestRegressor(n_feat)
        arfs = Leshy(model, verbose=0, max_iter=10, random_state=42, importance='shap')
        arfs.fit(X, y)
        leshy_rfc_list = sorted(arfs.support_names_)

        # assert borutapy_rfc_list == leshy_rfc_list, "same selected features are expected"
        assert bool(set(baseline_list) & set(leshy_rfc_list)), "expect non-empty intersection"

    def test_leshy_clf_with_lgb_and_shap_feature_importance_and_sample_weight(self):
        baseline_list = ['var0', 'var1', 'var2', 'var3', 'var4']

        X, y, w = generated_corr_dataset_classification(size=500)
        model = lgb.LGBMClassifier(verbose=-1, force_col_wise=True, n_estimators=10)
        arfs = Leshy(model, verbose=0, max_iter=10, random_state=42, importance='shap')
        arfs.fit(X, y, w)
        leshy_list = sorted(arfs.support_names_)

        assert bool(set(baseline_list) & set(leshy_list)), "expect non-empty intersection"

    def test_leshy_regr_with_lgb_and_shap_feature_importance_and_sample_weight(self):
        baseline_list = ['var0', 'var1', 'var2', 'var3', 'var4', 'var5']

        X, y, w = generated_corr_dataset_classification(size=500)
        model = lgb.LGBMRegressor(verbose=-1, force_col_wise=True, n_estimators=10)
        arfs = Leshy(model, verbose=0, max_iter=10, random_state=42, importance='shap')
        arfs.fit(X, y, w)
        leshy_list = sorted(arfs.support_names_)

        assert bool(set(baseline_list) & set(leshy_list)), "expect non-empty intersection"


class TestBoostAGroota:
    """
    Test suite for all-relevant FS boruta-like method: Leshy
    """

    def test_boostagroota_clf_with_lgb_and_shap_feature_importance_and_sample_weight(self):
        baseline_list = ['var0', 'var1', 'var2', 'var3', 'var4']

        X, y, w = generated_corr_dataset_classification(size=500)
        model = lgb.LGBMClassifier(verbose=-1, force_col_wise=True, n_estimators=10)
        arfs = BoostAGroota(est=model, cutoff=1, iters=3, max_rounds=3, delta=0.1, silent=False, importance='shap')
        arfs.fit(X, y, w)
        leshy_list = sorted(arfs.support_names_)

        assert bool(set(baseline_list) & set(leshy_list)), "expect non-empty intersection"

    def test_boostagroota_clf_with_lgb_and_pimp_feature_importance_and_sample_weight(self):
        baseline_list = ['var0', 'var1', 'var2', 'var3', 'var4']

        X, y, w = generated_corr_dataset_classification(size=500)
        model = lgb.LGBMClassifier(verbose=-1, force_col_wise=True, n_estimators=10)
        arfs = BoostAGroota(est=model, cutoff=1, iters=3, max_rounds=3, delta=0.1, silent=False, importance='pimp')
        arfs.fit(X, y, w)
        leshy_list = sorted(arfs.support_names_)

        assert bool(set(baseline_list) & set(leshy_list)), "expect non-empty intersection"

    def test_boostagroota_rgr_with_lgb_and_shap_feature_importance_and_sample_weight(self):
        baseline_list = ['var0', 'var1', 'var2', 'var3', 'var4', 'var5']

        X, y, w = generated_corr_dataset_regr(size=500)
        model = lgb.LGBMRegressor(verbose=-1, force_col_wise=True, n_estimators=10)
        arfs = BoostAGroota(est=model, cutoff=1, iters=3, max_rounds=3, delta=0.1, silent=False, importance='shap')
        arfs.fit(X, y, w)
        leshy_list = sorted(arfs.support_names_)

        assert bool(set(baseline_list) & set(leshy_list)), "expect non-empty intersection"

    def test_boostagroota_regr_with_lgb_and_pimp_feature_importance_and_sample_weight(self):
        baseline_list = ['var0', 'var1', 'var2', 'var3', 'var4', 'var5']

        X, y, w = generated_corr_dataset_regr(size=500)
        model = lgb.LGBMRegressor(verbose=-1, force_col_wise=True, n_estimators=10)
        arfs = BoostAGroota(est=model, cutoff=1, iters=3, max_rounds=3, delta=0.1, silent=False, importance='pimp')
        arfs.fit(X, y, w)
        leshy_list = sorted(arfs.support_names_)

        assert bool(set(baseline_list) & set(leshy_list)), "expect non-empty intersection"


class TestGrootCV:
    """
    Test suite for all-relevant FS boruta-like method: Leshy
    """

    def test_grootcv_classification_with_and_sample_weight(self):
        baseline_list = ['var0', 'var1', 'var2', 'var3', 'var4']

        X, y, w = generated_corr_dataset_classification(size=100)
        arfs = GrootCV(objective='binary', cutoff=1, n_folds=3, n_iter=3, silent=False)
        arfs.fit(X, y, w)
        grootcv_list = sorted(arfs.support_names_)

        assert bool(set(baseline_list) & set(grootcv_list)), "expect non-empty intersection"

    def test_grootcv_regression_with_and_sample_weight(self):
        baseline_list = ['var0', 'var1', 'var2', 'var3', 'var4', 'var5']

        X, y, w = generated_corr_dataset_regr(size=100)
        arfs = GrootCV(objective='l2', cutoff=1, n_folds=3, n_iter=3, silent=False)
        arfs.fit(X, y, w)
        grootcv_list = sorted(arfs.support_names_)

        assert bool(set(baseline_list) & set(grootcv_list)), "expect non-empty intersection"
