from sklearn.preprocessing import StandardScaler, MinMaxScaler
import warnings
from pandas.core.common import SettingWithCopyWarning
warnings.simplefilter(action='ignore', category=SettingWithCopyWarning)


class DenseTransform(object):
    def __init__(self, name='dense_preprocess', dense_encode='raw', dense_list=[], dense_specific_encode={}, inplace=False, **params):
        '''连续变量处理
        dense_encode:  编码方式
            raw: 不作处理
            std_scale: StandardScaler
            min_max_scale: MinMaxScaler()
        dense_list: 需要处理的特征列表
        dense_specific_encode: 单独为每个特征指定编码类型
        return_encoder：返回训练后的encoder
        '''
        self.name = name
        self.default_encode=dense_encode
        self.feature_list = dense_list
        self.specific_encode = dense_specific_encode
        self.encoder = dict()
        self.inplace = inplace
    
    def fit(self, X, y=None, **tracker):
        df = X if self.inplace else X.copy()
        self.feature_list += [i for i in tracker.get('dense_list', []) if i not in self.feature_list]
        self.feature_list += [i for i in self.specific_encode.keys() if i not in self.feature_list]

        for col in self.feature_list:
            choice = self.specific_encode.get(col, self.default_encode)
            if choice == 'raw':
                continue
            elif choice == 'std_scale':
                stdsc = StandardScaler()
                df[col] = stdsc.fit_transform(df[[col]])
                self.encoder[col] = stdsc
            elif choice == 'min_max_scale':
                stdsc = MinMaxScaler()
                df[col] = stdsc.fit_transform(df[[col]])
                self.encoder[col] = stdsc
            elif choice == 'qcut':
                pass
            elif choice == 'cut':
                pass
            else:
                raise ValueError(f'illegal choice={choice}')

        tracker['dense_list'] = self.feature_list
        return df, y, tracker
    
    def predict(self, X):
        for col, trans in self.encoder.items():
            X[col] = trans.transform(X[[col]])
        return X

