from causalml.feature_selection.filters import FilterSelect

class FeatureImportance:

    def __init__(self, df, X_names, y_name, treatment_group):
        self.df = df
        self.X_names = X_names
        self.y_name = y_name
        self.treatment_group = treatment_group

    def compute_feature_importance(self):

        filter_method = FilterSelect()
        method = 'KL'
        kl_imp = filter_method.get_importance(self.df, self.X_names, self.y_name, method,
                            treatment_group = self.treatment_group,
                            n_bins=20)
        return kl_imp