test_conditional_feature_importance.py

from copy import deepcopy
from functools import partial

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pytest
from scipy.stats import ttest_1samp
from sklearn.base import clone
from sklearn.linear_model import (
    LinearRegression,
    LogisticRegression,
    LogisticRegressionCV,
    RidgeCV,
)
from sklearn.metrics import log_loss, mean_squared_error
from sklearn.model_selection import KFold, train_test_split
from sklearn.utils.estimator_checks import parametrize_with_checks

from hidimstat._utils.exception import InternalError
from hidimstat._utils.scenario import multivariate_simulation
from hidimstat.base_perturbation import BasePerturbation
from hidimstat.conditional_feature_importance import CFI, CFICV, cfi_importance
from hidimstat.statistical_tools.multiple_testing import fdp_power

from .conftest import SKLEARN_LT_1_6, check_estimator, fitted_linear_regression


def fitted_ridged_cv() -> RidgeCV:
    """Return a fitted RidgeCV model."""
    X, y, _, _ = multivariate_simulation(
        n_samples=500,
        n_features=50,
    )
    return RidgeCV(alphas=np.logspace(-3, 3, 13)).fit(X, y)


def list_fitted_ridge_cv() -> list[RidgeCV]:
    """Return a fitted RidgeCV model."""
    X, y, _, _ = multivariate_simulation(
        n_samples=500,
        n_features=50,
    )
    model = RidgeCV(alphas=np.logspace(-3, 3, 13))
    cv = KFold(n_splits=2, shuffle=True, random_state=0)
    return [
        clone(model.fit(X[train_index], y[train_index]))
        for train_index, _ in cv.split(X)
    ]


ESTIMATORS_TO_CHECK = [
    CFI(
        estimator=fitted_linear_regression(),
        imputation_model_continuous=LinearRegression(),
    ),
    CFI(
        estimator=LinearRegression(),
        imputation_model_continuous=LinearRegression(),
    ),
    CFICV(
        estimators=RidgeCV(),
        cv=KFold(n_splits=2),
    ),
    CFICV(
        estimators=fitted_ridged_cv(),
        cv=KFold(n_splits=2),
    ),
    CFICV(
        estimators=list_fitted_ridge_cv(),
        cv=KFold(n_splits=2, shuffle=True, random_state=0),
    ),
]


def expected_failed_checks(estimator):
    if isinstance(estimator, CFI):
        return {
            "check_estimator_sparse_array": "TODO",
            "check_estimator_sparse_matrix": "TODO",
            "check_estimator_sparse_tag": "TODO",
            "check_estimators_nan_inf": "TODO",
            "check_fit2d_1feature": "TODO",
            "check_fit2d_1sample": "TODO",
            "check_parameters_default_constructible": "TODO",
        }
    elif isinstance(estimator, CFICV):
        failed_checks = {
            "check_fit2d_1feature": "TODO",
            "check_parameters_default_constructible": "TODO",
            "check_dict_unchanged": "TODO",
            "check_dont_overwrite_parameters": "TODO",
            "check_estimator_sparse_tag": "TODO",
            "check_estimator_sparse_array": "TODO",
            "check_estimator_sparse_matrix": "TODO",
            "check_estimators_dtypes": "TODO",
            "check_estimators_nan_inf": "TODO",
            "check_estimators_fit_returns_self": "TODO",
            "check_estimators_overwrite_params": "TODO",
            "check_f_contiguous_array_estimator": "TODO",
            "check_fit_check_is_fitted": "TODO",
            "check_fit2d_predict1d": "TODO",
            "check_n_features_in": "TODO",
            "check_n_features_in_after_fitting": "TODO",
            "check_methods_sample_order_invariance": "TODO",
            "check_methods_subset_invariance": "TODO",
            "check_readonly_memmap_input": "TODO",
        }
        if isinstance(estimator.estimators, list):
            failed_checks |= {
                "check_complex_data": "TODO",
                "check_dtype_object": "TODO",
                "check_estimators_empty_data_messages": "TODO",
                "check_estimators_pickle": "TODO",
                "check_fit2d_1sample": "TODO",
                "check_fit_idempotent": "TODO",
                "check_fit_score_takes_y": "TODO",
                "check_pipeline_consistency": "TODO",
                "check_positive_only_tag_during_fit": "TODO",
            }
        return failed_checks


if SKLEARN_LT_1_6:

    @pytest.mark.parametrize(
        "estimator, check, name",
        check_estimator(
            estimators=ESTIMATORS_TO_CHECK,
            return_expected_failed_checks=expected_failed_checks,
        ),
    )
    def test_check_estimator_sklearn_valid(estimator, check, name):  # noqa: ARG001
        """Check compliance with sklearn estimators."""
        check(estimator)

    @pytest.mark.xfail(reason="invalid checks should fail")
    @pytest.mark.parametrize(
        "estimator, check, name",
        check_estimator(
            estimators=ESTIMATORS_TO_CHECK,
            valid=False,
            return_expected_failed_checks=expected_failed_checks,
        ),
    )
    def test_check_estimator_sklearn_invalid(estimator, check, name):  # noqa: ARG001
        """Check compliance with sklearn estimators."""
        check(estimator)

else:

    @parametrize_with_checks(
        estimators=ESTIMATORS_TO_CHECK,
        expected_failed_checks=expected_failed_checks,
    )
    def test_check_estimator_sklearn(estimator, check):
        check(estimator)


def run_cfi(X, y, n_permutation, seed):
    """
    Configure Conditional Feature Importance (CFI) model with linear regression
    for feature importance analysis.

    Parameters
    ----------
    X : array-like of shape (n_samples, n_features)
        Input data matrix where each column represents a feature
        and each row a sample.
    y : array-like of shape (n_samples,)
        Target variable array.
    n_permutation : int
        Number of permutations to perform for the CFI analysis.
    seed : int
        Random seed for reproducibility.

    Returns
    -------
    importance : array-like
        Array containing importance scores for each feature.
        Higher values indicate greater feature importance in predicting
        the target variable.

    Notes
    -----
    The function performs the following steps:
    1. Splits data into training and test sets
    2. Fits a linear regression model on training data
    3. Configures CFI with linear regression as both estimator and imputer
    4. Calculates feature importance using the test set
    The CFI method uses permutation-based importance scoring with linear
    regression as the base model.
    """
    # split the data into training and test sets
    X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)

    # create and fit a linear regression model on the training set
    regression_model = LinearRegression()
    regression_model.fit(X_train, y_train)

    # instantiate CFI model with linear regression imputer
    cfi = CFI(
        estimator=regression_model,
        imputation_model_continuous=LinearRegression(),
        n_permutations=n_permutation,
        method="predict",
        features_groups=None,
        feature_types="auto",
        random_state=seed,
        n_jobs=1,
    )
    # fit the model using the training set
    cfi.fit(X_train)
    # calculate feature importance using the test set
    importance = cfi.importance(X_test, y_test)
    return importance


##############################################################################
## tests cfi on different type of data
parameter_exact = [
    ("HiDim", 150, 200, 10, 0.0, 42, 1.0, np.inf, 0.0),
    ("HiDim with noise", 150, 200, 10, 0.0, 42, 1.0, 10.0, 0.0),
    ("HiDim with correlated noise", 150, 200, 10, 0.0, 42, 1.0, 10.0, 0.2),
]


@pytest.mark.parametrize(
    "n_samples, n_features, support_size, rho, seed, value, signal_noise_ratio, rho_serial",
    zip(*(list(zip(*parameter_exact, strict=False))[1:]), strict=False),
    ids=next(zip(*parameter_exact, strict=False)),
)
@pytest.mark.parametrize(
    "n_permutation, cfi_seed", [(10, 5)], ids=["default_cfi"]
)
def test_linear_data_exact(data_generator, n_permutation, cfi_seed):
    """Tests the method on linear cases with noise and correlation"""
    X, y, important_features, _ = data_generator
    importance = run_cfi(X, y, n_permutation, cfi_seed)
    # check that importance scores are defined for each feature
    assert importance.shape == (X.shape[1],)
    # check that important features have the highest importance scores
    assert np.all(
        [int(i) in important_features for i in np.argsort(importance)[-10:]]
    )


parameter_partial = [
    (
        "HiDim with correlated features",
        150,
        200,
        10,
        0.2,
        42,
        1.0,
        np.inf,
        0.0,
    ),
    (
        "HiDim with correlated features and noise",
        150,
        200,
        10,
        0.2,
        42,
        1,
        10,
        0,
    ),
    (
        "HiDim with correlated features and correlated noise",
        150,
        200,
        10,
        0.2,
        42,
        1.0,
        10,
        0.2,
    ),
]


@pytest.mark.parametrize(
    "n_samples, n_features, support_size, rho, seed, value, signal_noise_ratio, rho_serial",
    zip(*(list(zip(*parameter_partial, strict=False))[1:]), strict=False),
    ids=next(zip(*parameter_partial, strict=False)),
)
@pytest.mark.parametrize(
    "n_permutation, cfi_seed", [(10, 5)], ids=["default_cfi"]
)
def test_linear_data_partial(data_generator, n_permutation, cfi_seed):
    """Tests the method on linear cases with noise and correlation"""
    X, y, important_features, _ = data_generator
    importance = run_cfi(X, y, n_permutation, cfi_seed)
    # check that importance scores are defined for each feature
    assert importance.shape == (X.shape[1],)
    # check that important features have the highest importance scores
    min_rank = 0
    importance_sort = np.flip(np.argsort(importance))
    for index in important_features:
        rank = np.where(importance_sort == index)[0]
        min_rank = max(min_rank, rank)
    # accept missing ranking of 15 elements
    assert min_rank < 25


@pytest.mark.parametrize(
    "n_samples, n_features, support_size, rho, seed, value, signal_noise_ratio, rho_serial",
    [(150, 200, 10, 0.2, 42, 1.0, 1.0, 0.0)],
    ids=["high level noise"],
)
@pytest.mark.parametrize(
    "n_permutation, cfi_seed", [(20, 5)], ids=["default_cfi"]
)
def test_linear_data_fail(data_generator, n_permutation, cfi_seed):
    """Tests when the method doesn't identify all important features"""
    X, y, important_features, not_important_features = data_generator
    importance = run_cfi(X, y, n_permutation, cfi_seed)
    # check that importance is defined for each feature
    assert importance.shape == (X.shape[1],)
    # check that mean importance of important features is
    # higher than mean importance of other features
    assert (
        importance[important_features].mean()
        > importance[not_important_features].mean()
    )
    # Verify that not all important features are detected
    assert np.sum(
        [int(i) in important_features for i in np.argsort(importance)[-10:]]
    ) != len(important_features)


##############################################################################
## Test specific options of cfi
@pytest.mark.parametrize(
    "n_samples, n_features, support_size, rho, seed, value, signal_noise_ratio, rho_serial",
    [(150, 200, 10, 0.0, 42, 1.0, np.inf, 0.0)],
    ids=["high dimension"],
)
def test_group(data_generator):
    """Test CFI with groups using pandas objects"""
    X, y, important_features, not_important_features = data_generator

    # Create groups and convert to pandas DataFrame
    groups = {
        "group_0": [f"col_{i}" for i in important_features],
        "the_group_1": [f"col_{i}" for i in not_important_features],
    }
    X_df = pd.DataFrame(X, columns=[f"col_{i}" for i in range(X.shape[1])])
    # Split data into training and test sets
    X_train_df, X_test_df, y_train, y_test = train_test_split(
        X_df, y, random_state=0
    )

    # Create and fit linear regression model on training set
    regression_model = LinearRegression()
    regression_model.fit(X_train_df, y_train)

    cfi = CFI(
        estimator=regression_model,
        imputation_model_continuous=LinearRegression(),
        n_permutations=20,
        method="predict",
        features_groups=groups,
        feature_types="auto",
        random_state=0,
        n_jobs=1,
    )
    cfi.fit(X_train_df)
    # Warning expected since column names in pandas are not considered
    with pytest.warns(
        UserWarning, match="X does not have valid feature names, but"
    ):
        importance = cfi.importance(X_test_df, y_test)

    # Check if importance scores are computed for each feature
    assert importance.shape == (2,)
    # Verify that important feature group has higher score
    # than non-important feature group
    assert importance[0] > importance[1]


@pytest.mark.parametrize(
    "n_samples, n_features, support_size, rho, seed, value, signal_noise_ratio, rho_serial",
    zip(*(list(zip(*parameter_exact, strict=False))[1:]), strict=False),
    ids=next(zip(*parameter_exact, strict=False)),
)
def test_classication(data_generator):
    """Test CFI for a classification problem"""
    X, y, important_features, not_important_features = data_generator
    # Create categories
    y_clf = deepcopy(y)
    y_clf[np.where(y > 4)] = 0
    y_clf[np.where(np.logical_and(y <= 4, y > 0))] = 1
    y_clf[np.where(np.logical_and(y <= 0, y > -4))] = 2
    y_clf[np.where(y <= -4)] = 3
    y_clf = np.array(y_clf, dtype=int)

    # Split the data into training and test sets
    X_train, X_test, y_train_clf, y_test_clf = train_test_split(
        X, y_clf, random_state=0
    )

    # Create and fit a logistic regression model on the training set
    logistic_model = LogisticRegression()
    logistic_model.fit(X_train, y_train_clf)

    cfi = CFI(
        estimator=logistic_model,
        imputation_model_continuous=LinearRegression(),
        n_permutations=20,
        method="predict_proba",
        loss=log_loss,
        features_groups=None,
        feature_types=["continuous"] * X.shape[1],
        random_state=0,
        n_jobs=1,
    )
    cfi.fit(X_train)
    importance = cfi.importance(X_test, y_test_clf)
    # Check that importance scores are defined for each feature
    assert importance.shape == (X.shape[1],)
    # Check that important features have higher mean importance scores
    assert (
        importance[important_features].mean()
        > importance[not_important_features].mean()
    )


##############################################################################
@pytest.mark.parametrize(
    "n_samples, n_features, support_size, rho, seed, value, signal_noise_ratio, rho_serial",
    [(150, 200, 10, 0.0, 42, 1.0, 0.0, 0.0)],
    ids=["default data"],
)
class TestCFIClass:
    """Test the element of the class"""

    def test_init(self, data_generator):
        """Test CFI initialization"""
        X, y, _, _ = data_generator
        fitted_model = LinearRegression().fit(X, y)
        cfi = CFI(
            estimator=fitted_model,
            method="predict",
        )
        assert cfi.n_jobs == 1
        assert cfi.n_permutations == 50
        assert cfi.loss == mean_squared_error
        assert cfi.method == "predict"
        assert cfi.categorical_max_cardinality == 10
        assert isinstance(
            cfi.imputation_model_categorical, LogisticRegressionCV
        )
        assert isinstance(cfi.imputation_model_continuous, RidgeCV)

    def test_fit(self, data_generator):
        """Test fitting CFI"""
        X, y, _, _ = data_generator
        fitted_model = LinearRegression().fit(X, y)
        cfi = CFI(
            estimator=fitted_model,
            imputation_model_continuous=LinearRegression(),
            random_state=42,
        )

        # Test fit with auto var_type
        cfi.fit(X)
        assert len(cfi._list_imputation_models) == X.shape[1]
        assert cfi.n_features_groups_ == X.shape[1]

    def test_fit_group(self, data_generator):
        """Test fitting CFI with group"""
        X, y, _, _ = data_generator
        fitted_model = LinearRegression().fit(X, y)
        # Test with specified groups
        groups = {"g1": [0, 1], "g2": [2, 3, 4]}
        cfi = CFI(
            estimator=fitted_model,
            imputation_model_continuous=LinearRegression(),
            features_groups=groups,
            random_state=42,
        )
        cfi.fit(X)

        assert len(cfi._list_imputation_models) == 2
        assert cfi.n_features_groups_ == 2

    def test_categorical(
        self,
        n_samples,
        n_features,  # noqa: ARG002
        support_size,  # noqa: ARG002
        rho,  # noqa: ARG002
        seed,
        value,  # noqa: ARG002
        signal_noise_ratio,  # noqa: ARG002
        rho_serial,  # noqa: ARG002
    ):
        """Test CFI with categorical variables"""
        rng = np.random.default_rng(seed)
        X_cont = rng.random((n_samples, 2))
        X_cat = rng.integers(low=0, high=3, size=(n_samples, 1))
        X = np.hstack([X_cont, X_cat])
        y = rng.random((n_samples, 1))
        fitted_model = LinearRegression().fit(X, y)

        feature_types = ["continuous", "continuous", "categorical"]
        cfi = CFI(
            estimator=fitted_model,
            imputation_model_continuous=LinearRegression(),
            imputation_model_categorical=LogisticRegression(),
            feature_types=feature_types,
            random_state=0,
        )
        cfi.fit(X, y)

        importances = cfi.importance(X, y)
        assert len(importances) == 3


##############################################################################
@pytest.mark.parametrize(
    "n_samples, n_features, support_size, rho, seed, value, signal_noise_ratio, rho_serial",
    [(150, 200, 10, 0.0, 42, 1.0, 0.0, 0.0)],
    ids=["default data"],
)
class TestCFIExceptions:
    """Test class for CFI exceptions"""

    def test_unknown_predict_method(self, data_generator):
        """Test when an unknown prediction method is provided"""
        X, y, _, _ = data_generator
        fitted_model = LinearRegression().fit(X, y)

        cfi = CFI(
            estimator=fitted_model,
            method="unknown method",
        )

        with pytest.raises(ValueError):
            cfi.fit(X, y)

    def test_unfitted_importance(self, data_generator):
        """Test importance method with unfitted model"""
        X, y, _, _ = data_generator
        fitted_model = LinearRegression().fit(X, y)
        cfi = CFI(
            estimator=fitted_model,
            method="predict",
        )

        with pytest.raises(
            ValueError, match="This CFI instance is not fitted yet"
        ):
            cfi.importance(X, y)

    def test_unfitted_base_perturbation(self, data_generator):
        """Test base perturbation with unfitted estimators"""
        X, y, _, _ = data_generator
        fitted_model = LinearRegression().fit(X, y)
        cfi = CFI(
            estimator=fitted_model,
            method="predict",
        )
        BasePerturbation.fit(cfi, X, y)

        with pytest.raises(
            ValueError,
            match="The imputation models require to be fitted before being used",
        ):
            cfi.importance(X, y)

    def test_invalid_type(self, data_generator):
        """Test invalid type of data"""
        X, y, _, _ = data_generator
        fitted_model = LinearRegression().fit(X, y)
        cfi = CFI(estimator=fitted_model, feature_types="invalid")

        # Test error when passing invalid var_type
        with pytest.raises(
            ValueError, match="feature_types support only the string 'auto'"
        ):
            cfi.fit(X)

    def test_invalid_n_permutations(self, data_generator):
        """Test when invalid number of permutations is provided"""
        X, y, _, _ = data_generator
        fitted_model = LinearRegression().fit(X, y)

        cfi = CFI(estimator=fitted_model, n_permutations=-1, method="predict")
        with pytest.raises(
            AssertionError, match="n_permutations must be positive"
        ):
            cfi.fit(X, y)

    def test_not_good_type_X(self, data_generator):
        """Test when X is wrong type"""
        X, y, _, _ = data_generator
        fitted_model = LinearRegression().fit(X, y)
        cfi = CFI(
            estimator=fitted_model,
            imputation_model_continuous=LinearRegression(),
            features_groups=None,
            feature_types="auto",
            method="predict",
        )
        cfi.fit(X)

        with pytest.raises(
            ValueError,
            match="X should be a pandas dataframe or a numpy array",
        ):
            cfi.importance(X.tolist(), y)

    def test_mismatched_features(self, data_generator):
        """Test when number of features doesn't match between fit and predict"""
        X, y, _, _ = data_generator
        fitted_model = LinearRegression().fit(X, y)
        cfi = CFI(
            estimator=fitted_model,
            imputation_model_continuous=LinearRegression(),
            method="predict",
            features_groups=None,
            feature_types="auto",
        )
        cfi.fit(X)

        with pytest.raises(
            AssertionError, match="X does not correspond to the fitting data"
        ):
            cfi.importance(X[:, :-1], y)

    def test_mismatched_features_string(self, data_generator):
        """Test when name of features doesn't match between fit and predict"""
        X, y, _, _ = data_generator
        X = pd.DataFrame({"col_" + str(i): X[:, i] for i in range(X.shape[1])})
        subgroups = {
            "group1": ["col_" + str(i) for i in range(int(X.shape[1] / 2))],
            "group2": [
                "col_" + str(i)
                for i in range(int(X.shape[1] / 2), X.shape[1] - 3)
            ],
        }
        fitted_model = LinearRegression().fit(X, y)
        cfi = CFI(
            estimator=fitted_model,
            imputation_model_continuous=LinearRegression(),
            method="predict",
            features_groups=subgroups,
            feature_types="auto",
        )
        cfi.fit(X)

        with pytest.raises(
            AssertionError,
            match=r"The array is missing at least one of the following columns \['col_100', 'col_101', 'col_102',",
        ):
            cfi.importance(
                X[
                    np.concatenate(
                        [subgroups["group1"], subgroups["group2"][:-2]]
                    )
                ],
                y,
            )

    def test_internal_error(self, data_generator):
        """Test when name of features doesn't match between fit and predict"""
        X, y, _, _ = data_generator
        X = pd.DataFrame({"col_" + str(i): X[:, i] for i in range(X.shape[1])})
        subgroups = {
            "group1": ["col_" + str(i) for i in range(int(X.shape[1] / 2))],
            "group2": [
                "col_" + str(i)
                for i in range(int(X.shape[1] / 2), X.shape[1] - 3)
            ],
        }
        fitted_model = LinearRegression().fit(X, y)
        cfi = CFI(
            estimator=fitted_model,
            imputation_model_continuous=LinearRegression(),
            method="predict",
            features_groups=subgroups,
            feature_types="auto",
        )
        cfi.fit(X)
        cfi.features_groups["group1"] = [None for i in range(100)]

        X = X.to_records(index=False)
        X = np.array(X, dtype=X.dtype.descr)
        with pytest.raises(
            InternalError,
            match="A problem with indexing has happened during the fit",
        ):
            cfi.importance(X, y)

    def test_invalid_var_type(self, data_generator):
        """Test when invalid variable type is provided"""
        X, y, _, _ = data_generator
        fitted_model = LinearRegression().fit(X, y)
        cfi = CFI(
            estimator=fitted_model,
            method="predict",
            features_groups=None,
            feature_types=["invalid_type"] * X.shape[1],
        )

        with pytest.raises(
            ValueError, match="type of data 'invalid_type' unknown"
        ):
            cfi.fit(X)

    def test_incompatible_imputer(self, data_generator):
        """Test when incompatible imputer is provided"""
        X, y, _, _ = data_generator
        fitted_model = LinearRegression().fit(X, y)

        cfi = CFI(
            estimator=fitted_model,
            imputation_model_continuous="invalid_imputer",
            method="predict",
        )
        with pytest.raises(
            AssertionError, match="Continuous imputation model invalid"
        ):
            cfi.fit(X)

        cfi = CFI(
            estimator=fitted_model,
            imputation_model_categorical="invalid_imputer",
            method="predict",
        )
        with pytest.raises(
            AssertionError, match="Categorial imputation model invalid"
        ):
            cfi.fit(X)

    def test_invalid_groups_format(self, data_generator):
        """Test when groups are provided in invalid format"""
        X, y, _, _ = data_generator
        invalid_groups = ["group1", "group2"]  # Should be dictionary
        fitted_model = LinearRegression().fit(X, y)
        cfi = CFI(
            estimator=fitted_model,
            method="predict",
            features_groups=invalid_groups,
            feature_types="auto",
        )

        with pytest.raises(
            ValueError, match="features_groups needs to be a dictionary"
        ):
            cfi.fit(X)

    def test_groups_warning(self, data_generator):
        """Test if a subgroup raise a warning"""
        X, y, _, _ = data_generator
        subgroups = {"group1": [0, 1], "group2": [2, 3]}
        fitted_model = LinearRegression().fit(X, y)
        cfi = CFI(
            estimator=fitted_model,
            imputation_model_continuous=LinearRegression(),
            method="predict",
            features_groups=subgroups,
            feature_types="auto",
        )
        cfi.fit(X, y)

        with pytest.warns(
            UserWarning,
            match="The number of features in X: 200 differs from the"
            " number of features for which importance is computed: 4",
        ):
            cfi.importance(X, y)

    def test_assert_dimension_pvalue(self, data_generator):
        """Test that assert is raise if function stat is not good"""
        X, y, _, _ = data_generator
        fitted_model = LinearRegression().fit(X, y)
        cfi = CFI(
            estimator=fitted_model,
            imputation_model_continuous=LinearRegression(),
            statistical_test=partial(ttest_1samp, popmean=0, axis=0),
        )
        cfi.fit(X, y)
        with pytest.raises(
            AssertionError,
            match="The statistical test doesn't provide the correct dimension",
        ):
            cfi.importance(X, y)


@pytest.mark.parametrize(
    "n_samples, n_features, support_size, rho, seed, value, signal_noise_ratio, rho_serial",
    [(150, 200, 10, 0.2, 42, 1.0, 1.0, 0.0)],
    ids=["high level noise"],
)
@pytest.mark.parametrize(
    "n_permutation, cfi_seed", [(20, 0)], ids=["default_cfi"]
)
def test_function_cfi(data_generator, n_permutation, cfi_seed):
    """Test CFI function"""
    X, y, _, _ = data_generator
    cfi_importance(
        LinearRegression().fit(X, y),
        X,
        y,
        imputation_model_continuous=LinearRegression(),
        n_permutations=n_permutation,
        random_state=cfi_seed,
    )


@pytest.mark.parametrize(
    "n_samples, n_features, support_size, rho, seed, value, signal_noise_ratio, rho_serial",
    [(10, 10, 1, 0.2, 0, 1.0, 1.0, 0.0)],
    ids=["10 features"],
)
@pytest.mark.mpl_image_compare
def test_cfi_plot(data_generator):
    """Test CFI plot function"""
    X, y, _, _ = data_generator
    X_train, _, y_train, _ = train_test_split(
        X, y, test_size=0.5, random_state=0
    )
    fitted_model = LinearRegression().fit(X_train, y_train)
    cfi = CFI(
        estimator=fitted_model,
        imputation_model_continuous=LinearRegression(),
        feature_types="continuous",
        random_state=0,
    )
    cfi.fit(X_train, y_train)
    cfi.loss_reference_ = []
    cfi.loss_ = []
    # Make the plot independent of data / randomness to test only the plotting function
    cfi.importances_ = np.arange(X.shape[1])
    fig, ax = plt.subplots(figsize=(6, 3))
    ax = cfi.plot_importance(ax=ax)
    return fig


@pytest.mark.parametrize(
    "n_samples, n_features, support_size, rho, seed, value, signal_noise_ratio, rho_serial",
    [(10, 5, 1, 0.2, 0, 1.0, 1.0, 0.0)],
    ids=["5_features"],
)
@pytest.mark.mpl_image_compare
def test_cfi_plot_2d_imp(data_generator):
    """Test CFI plot function"""
    X, y, _, _ = data_generator
    X_train, _, y_train, _ = train_test_split(
        X, y, test_size=0.5, random_state=0
    )
    fitted_model = LinearRegression().fit(X_train, y_train)
    cfi = CFI(
        estimator=fitted_model,
        imputation_model_continuous=LinearRegression(),
        feature_types="continuous",
        random_state=0,
    )
    cfi.fit(X_train, y_train)
    cfi.loss_reference_ = []
    cfi.loss_ = []
    # Make the plot independent of data / randomness to test only the plotting function
    cfi.importances_ = np.stack(
        [
            np.arange(X_train.shape[1]),
            np.arange(X_train.shape[1]) - 1,
            np.arange(X_train.shape[1]) + 1,
        ],
        axis=0,
    )
    fig, ax = plt.subplots(figsize=(6, 3))
    ax = cfi.plot_importance(ax=ax)
    return fig


@pytest.mark.parametrize(
    "n_samples, n_features, support_size, rho, seed, value, signal_noise_ratio, rho_serial",
    [(10, 3, 1, 0.2, 0, 1.0, 1.0, 0.0)],
)
def test_cfi_plot_coverage(data_generator, rng):
    """Add arguments combinations to test coverage of the plot function"""
    X, y, _, _ = data_generator
    X_train, _, y_train, _ = train_test_split(
        X, y, test_size=0.5, random_state=0
    )
    fitted_model = LinearRegression().fit(X_train, y_train)
    cfi = CFI(
        estimator=fitted_model,
        imputation_model_continuous=LinearRegression(),
        feature_types="continuous",
        random_state=0,
    )
    cfi.fit(X_train, y_train)
    cfi.loss_reference_ = []
    cfi.loss_ = []
    # Make the plot independent of data / randomness to test only the plotting function
    cfi.importances_ = np.arange(X.shape[1])
    _, ax = plt.subplots(figsize=(6, 3))

    ax = cfi.plot_importance(ax=None)
    assert isinstance(ax, plt.Axes)

    _, ax = plt.subplots()
    cfi.importances_ = rng.standard_normal((3, X.shape[1]))
    ax = cfi.plot_importance(ax=ax)
    assert isinstance(ax, plt.Axes)


@pytest.fixture(scope="module")
def cfi_test_data():
    """
    Fixture to generate test data and a fitted LinearRegression model for CFI
    reproducibility tests.
    """
    X, y, _, _ = multivariate_simulation(
        n_samples=100,
        n_features=5,
        support_size=2,
        rho=0,
        value=1,
        signal_noise_ratio=4,
        rho_serial=0,
        shuffle=False,
        seed=0,
    )
    X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)
    model = LinearRegression()
    model.fit(X_train, y_train)
    cfi_default_parameters = {
        "estimator": model,
        "imputation_model_continuous": LinearRegression(),
        "n_permutations": 20,
        "method": "predict",
        "n_jobs": 1,
    }
    return X_train, X_test, y_test, cfi_default_parameters


def test_cfi_repeatibility(cfi_test_data):
    """
    Test that multiple calls of .importance() when CFI is seeded provide deterministic
    results.
    """
    X_train, X_test, y_test, cfi_default_parameters = cfi_test_data
    cfi = CFI(**cfi_default_parameters)
    cfi.fit(X_train)
    vim = cfi.importance(X_test, y_test)
    # repeat
    vim_repeat = cfi.importance(X_test, y_test)
    assert not np.array_equal(vim, vim_repeat)


def test_cfi_randomness_with_none(cfi_test_data):
    """
    Test that multiple calls of .importance() when CFI has random_state=None
    """
    X_train, X_test, y_test, cfi_default_parameters = cfi_test_data
    cfi = CFI(random_state=None, **cfi_default_parameters)
    cfi.fit(X_train)
    vim = cfi.importance(X_test, y_test)
    # repeat importance
    vim_repeat = cfi.importance(X_test, y_test)
    assert not np.array_equal(vim, vim_repeat)

    # refit
    cfi.fit(X_train)
    vim_refit = cfi.importance(X_test, y_test)
    assert not np.array_equal(vim, vim_refit)

    # Reproducibility
    cfi_2 = CFI(random_state=None, **cfi_default_parameters)
    cfi_2.fit(X_train)
    vim_reproducibility = cfi_2.importance(X_test, y_test)
    assert not np.array_equal(vim, vim_reproducibility)


def test_cfi_reproducibility_with_integer(cfi_test_data):
    """
    Test that multiple calls of .importance() when CFI has random_state=42
    """
    X_train, X_test, y_test, cfi_default_parameters = cfi_test_data
    cfi = CFI(random_state=42, **cfi_default_parameters)
    cfi.fit(X_train)
    vim = cfi.importance(X_test, y_test)
    # repeat importance
    vim_repeat = cfi.importance(X_test, y_test)
    assert np.array_equal(vim, vim_repeat)

    # refit
    cfi.fit(X_train)
    vim_refit = cfi.importance(X_test, y_test)
    assert np.array_equal(vim, vim_refit)

    # Reproducibility
    cfi_2 = CFI(random_state=42, **cfi_default_parameters)
    cfi_2.fit(X_train)
    vim_reproducibility = cfi_2.importance(X_test, y_test)
    assert np.array_equal(vim, vim_reproducibility)


def test_cfi_reproducibility_with_rng(cfi_test_data):
    """
    Test that:
     1. Mmultiple calls of .importance() when CFI has random_state=rng are random
     2. refit with same rng provides same result
    """
    X_train, X_test, y_test, cfi_default_parameters = cfi_test_data
    rng = np.random.default_rng(0)
    cfi = CFI(random_state=rng, **cfi_default_parameters)
    cfi.fit(X_train)
    vim = cfi.importance(X_test, y_test)
    # repeat importance
    vim_repeat = cfi.importance(X_test, y_test)
    assert not np.array_equal(vim, vim_repeat)

    # refit
    cfi.fit(X_train)
    vim_refit = cfi.importance(X_test, y_test)
    assert not np.array_equal(vim, vim_refit)

    # refit repeatability
    rng = np.random.default_rng(0)
    cfi.random_state = rng
    cfi.fit(X_train)
    vim_refit_2 = cfi.importance(X_test, y_test)
    assert np.array_equal(vim, vim_refit_2)

    # Reproducibility
    cfi_2 = CFI(
        random_state=np.random.default_rng(0), **cfi_default_parameters
    )
    cfi_2.fit(X_train)
    vim_reproducibility = cfi_2.importance(X_test, y_test)
    assert np.array_equal(vim, vim_reproducibility)


@pytest.mark.parametrize(
    "n_samples, n_features, support_size, rho, seed, value, signal_noise_ratio, rho_serial",
    [(500, 50, 5, 0.1, 0, 8.0, 4, 0.0)],
    ids=["default data"],
)
def test_cfi_cv(data_generator):
    """
    Test that CFI with cross-validated estimator works as expected. In particular,
     - Empirical FDP is below the target FDR level
     - Power is above 0.8, which is an arbitrary threshold

    Note: Although only the expected FDP should be controlled, in practice
    the simulation setting is simple enough to satisfy this stronger condition.
    """
    X, y, important_features, _ = data_generator

    model = RidgeCV(alphas=np.logspace(-3, 3, 13))
    cv = KFold(n_splits=5, shuffle=True, random_state=0)
    estimators = [
        clone(model).fit(X[train_index], y[train_index])
        for train_index, _ in cv.split(X)
    ]
    cfi_cv = CFICV(
        estimators=estimators,
        cv=cv,
        random_state=0,
        n_jobs=5,
    )
    cfi_cv.fit_importance(X, y)
    selected = cfi_cv.fdr_selection(fdr=0.05)

    alpha = 0.05
    gt_mask = np.zeros(X.shape[1], dtype=int)
    gt_mask[important_features] = 1
    fdp, power = fdp_power(selected=selected, ground_truth=gt_mask)
    assert fdp < alpha
    assert power > 0.8