import pymrmr
import pandas as pd


class MRMRFeatureSelection:
    def __init__(self, X, y):
        """
        初始化 mRMR 特征选择模块。

        :param X: 输入特征矩阵 (DataFrame)，每列为一个特征。
        :param y: 目标变量 (Series)，与特征矩阵 X 对应。
        """
        self.X = X
        self.y = y
        self.selected_features = None

    def select_features(self, k=18, method='MIQ'):
        """
        执行 mRMR 特征选择。

        :param k: 选择的特征个数。
        :param method: 选择的 mRMR 方法 ('MIQ' 或 'MRMR')。
        :return: 选择的特征列表
        """
        # 拼接特征和目标变量
        df = pd.concat([self.y, self.X], axis=1)

        # 使用 pymrmr 进行 mRMR 特征选择
        self.selected_features = pymrmr.mRMR(df, method, k)

        return self.selected_features

    def get_selected_features(self):
        """
        获取已选择的特征。

        :return: 选择的特征
        """
        return self.selected_features

    def get_selected_feature_names(self):
        """
        获取已选择特征的列名

        :return: 选择的特征列名列表
        """
        if self.selected_features is None:
            return None
        return self.selected_features