import os

import pandas as pd
import pytest
from probatus.feature_elimination import EarlyStoppingShapRFECV, ShapRFECV
from probatus.utils import preprocess_labels
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import get_scorer
from sklearn.model_selection import RandomizedSearchCV
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVC
from sklearn.tree import DecisionTreeClassifier


@pytest.fixture(scope="function")
def X():
    """
    Fixture for X.
    """
    return pd.DataFrame(
        {"col_1": [1, 1, 1, 1, 1, 1, 1, 0], "col_2": [0, 0, 0, 0, 0, 0, 0, 1], "col_3": [1, 0, 1, 0, 1, 0, 1, 0]},
        index=[1, 2, 3, 4, 5, 6, 7, 8],
    )


@pytest.fixture(scope="session")
def catboost_classifier_class():
    """This fixture allows to reuse the import of the CatboostClassifier class across different tests.

    It is equivalent to importing the package at the beginning of the file.

    Importing catboost multiple times results in a ValueError: I/O operation on closed file.

    """
    from catboost import CatBoostClassifier

    return CatBoostClassifier


@pytest.fixture(scope="function")
def y():
    """
    Fixture for y.
    """
    return pd.Series([1, 0, 1, 0, 1, 0, 1, 0], index=[1, 2, 3, 4, 5, 6, 7, 8])


@pytest.fixture(scope="function")
def sample_weight():
    """
    Fixture for sample_weight.
    """
    return pd.Series([1, 1, 1, 1, 1, 1, 1, 1], index=[1, 2, 3, 4, 5, 6, 7, 8])


def test_shap_rfe_randomized_search(X, y, capsys):
    """
    Test with RandomizedSearchCV.
    """
    clf = DecisionTreeClassifier(max_depth=1)
    param_grid = {"criterion": ["gini"], "min_samples_split": [1, 2]}
    search = RandomizedSearchCV(clf, param_grid, cv=2, n_iter=2)
    with pytest.warns(None) as record:

        shap_elimination = ShapRFECV(search, step=0.8, cv=2, scoring="roc_auc", n_jobs=4, random_state=1)
        report = shap_elimination.fit_compute(X, y)

    assert shap_elimination.fitted
    shap_elimination._check_if_fitted()

    assert report.shape[0] == 2
    assert shap_elimination.get_reduced_features_set(1) == ["col_3"]

    _ = shap_elimination.plot(show=False)

    # Ensure that number of warnings was at least 2 for the verbose (2 generated by probatus + possibly more by SHAP)
    assert len(record) >= 2

    # Check if there is any prints
    out, _ = capsys.readouterr()
    assert len(out) == 0


def test_shap_rfe(X, y, sample_weight, capsys):
    """
    Test with ShapRFECV.
    """
    clf = DecisionTreeClassifier(max_depth=1, random_state=1)
    with pytest.warns(None) as record:
        shap_elimination = ShapRFECV(
            clf,
            random_state=1,
            step=1,
            cv=2,
            scoring="roc_auc",
            n_jobs=4,
        )
        shap_elimination = shap_elimination.fit(
            X, y, sample_weight=sample_weight, approximate=True, check_additivity=False
        )

    assert shap_elimination.fitted
    shap_elimination._check_if_fitted()

    report = shap_elimination.compute()

    assert report.shape[0] == 3
    assert shap_elimination.get_reduced_features_set(1) == ["col_3"]

    _ = shap_elimination.plot(show=False)

    # Ensure that number of warnings was 0
    assert len(record) == 0
    # Check if there is any prints
    out, _ = capsys.readouterr()
    assert len(out) == 0


def test_shap_pipeline_error(X, y, capsys):
    """
    Test with ShapRFECV for pipelines.
    """
    clf = Pipeline([("scaler", StandardScaler()), ("dt", DecisionTreeClassifier(max_depth=1, random_state=1))])
    with pytest.raises(TypeError):
        shap_elimination = ShapRFECV(
            clf,
            random_state=1,
            step=1,
            cv=2,
            scoring="roc_auc",
            n_jobs=4,
        )
        shap_elimination = shap_elimination.fit(X, y, approximate=True, check_additivity=False)


def test_shap_rfe_linear_model(X, y, capsys):
    """
    Test ShapRFECV with linear model.
    """
    clf = LogisticRegression(C=1, random_state=1)
    with pytest.warns(None) as record:
        shap_elimination = ShapRFECV(clf, random_state=1, step=1, cv=2, scoring="roc_auc", n_jobs=4)
        shap_elimination = shap_elimination.fit(X, y)

    assert shap_elimination.fitted
    shap_elimination._check_if_fitted()

    report = shap_elimination.compute()

    assert report.shape[0] == 3
    assert shap_elimination.get_reduced_features_set(1) == ["col_3"]

    _ = shap_elimination.plot(show=False)

    # Ensure that number of warnings was 0
    assert len(record) == 0
    # Check if there is any prints
    out, _ = capsys.readouterr()
    assert len(out) == 0


def test_shap_rfe_svm(X, y, capsys):
    """
    Test with ShapRFECV with SVM.
    """
    clf = SVC(C=1, kernel="linear", probability=True)
    with pytest.warns(None) as record:
        shap_elimination = ShapRFECV(clf, random_state=1, step=1, cv=2, scoring="roc_auc", n_jobs=4)
        shap_elimination = shap_elimination.fit(X, y)

    assert shap_elimination.fitted
    shap_elimination._check_if_fitted()

    report = shap_elimination.compute()

    assert report.shape[0] == 3
    assert shap_elimination.get_reduced_features_set(1) == ["col_3"]

    _ = shap_elimination.plot(show=False)

    # Ensure that number of warnings was 0
    assert len(record) == 0
    # Check if there is any prints
    out, _ = capsys.readouterr()
    assert len(out) == 0


def test_shap_rfe_cols_to_keep(X, y, capsys):
    """
    Test for shap_rfe_cv with feautures to keep parameter.
    """
    clf = DecisionTreeClassifier(max_depth=1, random_state=1)
    with pytest.warns(None) as record:
        shap_elimination = ShapRFECV(
            clf, random_state=1, step=2, cv=2, scoring="roc_auc", n_jobs=4, min_features_to_select=1
        )
        shap_elimination = shap_elimination.fit(X, y, columns_to_keep=["col_2", "col_3"])

    assert shap_elimination.fitted
    shap_elimination._check_if_fitted()

    report = shap_elimination.compute()

    assert report.shape[0] == 2
    reduced_feature_set = set(shap_elimination.get_reduced_features_set(num_features=2))
    assert reduced_feature_set == set(["col_2", "col_3"])

    # Ensure that number of warnings was 0
    assert len(record) == 0
    # Check if there is any prints
    out, _ = capsys.readouterr()
    assert len(out) == 0


def test_shap_rfe_randomized_search_cols_to_keep(X, y, capsys):
    """
    Test with ShapRFECV with column to keep param.
    """
    clf = DecisionTreeClassifier(max_depth=1)
    param_grid = {"criterion": ["gini"], "min_samples_split": [1, 2]}
    search = RandomizedSearchCV(clf, param_grid, cv=2, n_iter=2)
    with pytest.warns(None) as record:

        shap_elimination = ShapRFECV(search, step=0.8, cv=2, scoring="roc_auc", n_jobs=4, random_state=1)
        report = shap_elimination.fit_compute(X, y, columns_to_keep=["col_2", "col_3"])

    assert shap_elimination.fitted
    shap_elimination._check_if_fitted()

    assert report.shape[0] == 2
    reduced_feature_set = set(shap_elimination.get_reduced_features_set(num_features=2))
    assert reduced_feature_set == set(["col_2", "col_3"])

    _ = shap_elimination.plot(show=False)

    # Ensure that number of warnings was at least 2 for the verbose (2 generated by probatus + possibly more by SHAP)
    assert len(record) >= 2

    # Check if there is any prints
    out, _ = capsys.readouterr()
    assert len(out) == 0


def test_calculate_number_of_features_to_remove():
    """
    Test with ShapRFECV with n features to remove.
    """
    assert 3 == ShapRFECV._calculate_number_of_features_to_remove(
        current_num_of_features=10, num_features_to_remove=3, min_num_features_to_keep=5
    )
    assert 3 == ShapRFECV._calculate_number_of_features_to_remove(
        current_num_of_features=8, num_features_to_remove=5, min_num_features_to_keep=5
    )
    assert 0 == ShapRFECV._calculate_number_of_features_to_remove(
        current_num_of_features=5, num_features_to_remove=1, min_num_features_to_keep=5
    )
    assert 4 == ShapRFECV._calculate_number_of_features_to_remove(
        current_num_of_features=5, num_features_to_remove=7, min_num_features_to_keep=1
    )


def test_get_feature_shap_values_per_fold(X, y):
    """
    Test with ShapRFECV with features per fold.
    """
    clf = DecisionTreeClassifier(max_depth=1)
    shap_elimination = ShapRFECV(clf)
    shap_values, train_score, test_score = shap_elimination._get_feature_shap_values_per_fold(
        X, y, clf, train_index=[2, 3, 4, 5, 6, 7], val_index=[0, 1], scorer=get_scorer("roc_auc")
    )
    assert test_score == 1
    assert train_score > 0.9
    assert shap_values.shape == (2, 3)


@pytest.mark.skipif(os.environ.get("SKIP_LIGHTGBM") == "true", reason="LightGBM tests disabled")
def test_complex_dataset(complex_data, complex_lightgbm):
    """
    Test on complex dataset.
    """
    X, y = complex_data

    param_grid = {
        "n_estimators": [5, 7, 10],
        "num_leaves": [3, 5, 7, 10],
    }
    search = RandomizedSearchCV(complex_lightgbm, param_grid, n_iter=1)

    shap_elimination = ShapRFECV(clf=search, step=1, cv=10, scoring="roc_auc", n_jobs=3, verbose=50)
    with pytest.warns(None) as record:
        report = shap_elimination.fit_compute(X, y)

    assert report.shape[0] == X.shape[1]

    assert len(record) >= 2


@pytest.mark.skipif(os.environ.get("SKIP_LIGHTGBM") == "true", reason="LightGBM tests disabled")
def test_shap_rfe_early_stopping_lightGBM(complex_data, capsys):
    """
    Test EarlyStoppingShapRFECV with a LGBMClassifier.
    """
    from lightgbm import LGBMClassifier

    clf = LGBMClassifier(n_estimators=200, max_depth=3)
    X, y = complex_data

    with pytest.warns(None) as record:
        shap_elimination = EarlyStoppingShapRFECV(
            clf,
            random_state=1,
            step=1,
            cv=10,
            scoring="roc_auc",
            n_jobs=4,
            early_stopping_rounds=5,
            eval_metric="auc",
        )
        shap_elimination = shap_elimination.fit(X, y, approximate=False, check_additivity=False)

    assert shap_elimination.fitted
    shap_elimination._check_if_fitted()

    report = shap_elimination.compute()

    assert report.shape[0] == 5
    assert shap_elimination.get_reduced_features_set(1) == ["f5"]

    _ = shap_elimination.plot(show=False)

    # Ensure that number of warnings was 0
    assert len(record) == 0
    # Check if there is any prints
    out, _ = capsys.readouterr()
    assert len(out) == 0


@pytest.mark.skipif(os.environ.get("SKIP_LIGHTGBM") == "true", reason="LightGBM tests disabled")
def test_shap_rfe_early_stopping_XGBoost(complex_data, capsys):
    """
    Test EarlyStoppingShapRFECV with a LGBMClassifier.
    """
    from xgboost import XGBClassifier

    clf = XGBClassifier(n_estimators=200, max_depth=3, use_label_encoder=False, random_state=42)
    X, y = complex_data
    X["f1_categorical"] = X["f1_categorical"].astype(float)

    with pytest.warns(None) as record:
        shap_elimination = EarlyStoppingShapRFECV(
            clf,
            random_state=1,
            step=1,
            cv=10,
            scoring="roc_auc",
            n_jobs=4,
            early_stopping_rounds=5,
            eval_metric="auc",
        )
        shap_elimination = shap_elimination.fit(X, y, approximate=False, check_additivity=False)

    assert shap_elimination.fitted
    shap_elimination._check_if_fitted()

    report = shap_elimination.compute()

    assert report.shape[0] == 5
    assert shap_elimination.get_reduced_features_set(1) == ["f4"]

    _ = shap_elimination.plot(show=False)

    # Ensure that number of warnings was 0
    assert len(record) == 0
    # Check if there is any prints
    out, _ = capsys.readouterr()
    assert len(out) == 0


@pytest.mark.skipif(os.environ.get("SKIP_LIGHTGBM") == "true", reason="LightGBM tests disabled")
def test_shap_rfe_early_stopping_CatBoost(complex_data, capsys, catboost_classifier_class):
    """
    Test EarlyStoppingShapRFECV with a CatBoostClassifier.
    """

    clf = catboost_classifier_class(random_seed=42)
    X, y = complex_data
    X["f1_categorical"] = X["f1_categorical"].astype(str).astype("category")

    with pytest.warns(None) as record:
        shap_elimination = EarlyStoppingShapRFECV(
            clf,
            random_state=1,
            step=1,
            cv=10,
            scoring="roc_auc",
            n_jobs=4,
            early_stopping_rounds=5,
            eval_metric="auc",
        )
        shap_elimination = shap_elimination.fit(X, y, approximate=False, check_additivity=False)

    assert shap_elimination.fitted
    shap_elimination._check_if_fitted()

    report = shap_elimination.compute()

    assert report.shape[0] == 5
    assert shap_elimination.get_reduced_features_set(1)[0] in ["f4", "f5"]

    _ = shap_elimination.plot(show=False)

    # Ensure that number of warnings was 0
    assert len(record) == 0
    # Check if there is any prints
    out, _ = capsys.readouterr()
    assert len(out) == 0


@pytest.mark.skipif(os.environ.get("SKIP_LIGHTGBM") == "true", reason="LightGBM tests disabled")
def test_shap_rfe_randomized_search_early_stopping_lightGBM(complex_data):
    """
    Test EarlyStoppingShapRFECV with RandomizedSearchCV and a LGBMClassifier on complex dataset.
    """
    from lightgbm import LGBMClassifier

    clf = LGBMClassifier(n_estimators=200)
    X, y = complex_data

    param_grid = {
        "max_depth": [3, 4, 5],
    }
    search = RandomizedSearchCV(clf, param_grid, cv=2, n_iter=2)
    with pytest.warns(None) as record:
        shap_elimination = EarlyStoppingShapRFECV(
            search,
            step=1,
            cv=10,
            scoring="roc_auc",
            early_stopping_rounds=5,
            eval_metric="auc",
            n_jobs=4,
            verbose=50,
            random_state=1,
        )
        report = shap_elimination.fit_compute(X, y)

    assert shap_elimination.fitted
    shap_elimination._check_if_fitted()

    assert report.shape[0] == X.shape[1]
    assert shap_elimination.get_reduced_features_set(1) == ["f5"]

    _ = shap_elimination.plot(show=False)

    # Ensure that number of warnings was at least 3 for the verbose (2 generated by probatus + possibly more by SHAP)
    assert len(record) >= 3


@pytest.mark.skipif(os.environ.get("SKIP_LIGHTGBM") == "true", reason="LightGBM tests disabled")
def test_get_feature_shap_values_per_fold_early_stopping_lightGBM(complex_data):
    """
    Test with ShapRFECV with features per fold.
    """
    from lightgbm import LGBMClassifier

    clf = LGBMClassifier(n_estimators=200, max_depth=3)
    X, y = complex_data
    y = preprocess_labels(y, y_name="y", index=X.index)

    shap_elimination = EarlyStoppingShapRFECV(clf, early_stopping_rounds=5)
    shap_values, train_score, test_score = shap_elimination._get_feature_shap_values_per_fold(
        X,
        y,
        clf,
        train_index=list(range(5, 50)),
        val_index=[0, 1, 2, 3, 4],
        scorer=get_scorer("roc_auc"),
    )
    assert test_score > 0.6
    assert train_score > 0.6
    assert shap_values.shape == (5, 5)


@pytest.mark.skipif(os.environ.get("SKIP_LIGHTGBM") == "true", reason="LightGBM tests disabled")
def test_get_feature_shap_values_per_fold_early_stopping_CatBoost(complex_data, catboost_classifier_class):
    """
    Test with ShapRFECV with features per fold.
    """
    clf = catboost_classifier_class(random_seed=42)
    X, y = complex_data
    X["f1_categorical"] = X["f1_categorical"].astype(str).astype("category")
    y = preprocess_labels(y, y_name="y", index=X.index)

    shap_elimination = EarlyStoppingShapRFECV(clf, early_stopping_rounds=5)
    shap_values, train_score, test_score = shap_elimination._get_feature_shap_values_per_fold(
        X,
        y,
        clf,
        train_index=list(range(5, 50)),
        val_index=[0, 1, 2, 3, 4],
        scorer=get_scorer("roc_auc"),
    )
    assert test_score > 0
    assert train_score > 0.6
    assert shap_values.shape == (5, 5)


@pytest.mark.skipif(os.environ.get("SKIP_LIGHTGBM") == "true", reason="LightGBM tests disabled")
def test_get_feature_shap_values_per_fold_early_stopping_XGBoost(complex_data):
    """
    Test with ShapRFECV with features per fold.
    """
    from xgboost import XGBClassifier

    clf = XGBClassifier(n_estimators=200, max_depth=3, use_label_encoder=False, random_state=42)
    X, y = complex_data
    X["f1_categorical"] = X["f1_categorical"].astype(float)
    y = preprocess_labels(y, y_name="y", index=X.index)

    shap_elimination = EarlyStoppingShapRFECV(clf, early_stopping_rounds=5)
    shap_values, train_score, test_score = shap_elimination._get_feature_shap_values_per_fold(
        X,
        y,
        clf,
        train_index=list(range(5, 50)),
        val_index=[0, 1, 2, 3, 4],
        scorer=get_scorer("roc_auc"),
    )
    assert test_score > 0
    assert train_score > 0.6
    assert shap_values.shape == (5, 5)
