# -*- coding:utf-8 -*-
"""
@author: yq
@time: 2023/12/28
@desc:  特征工具类
"""
import pandas as pd
from sklearn.preprocessing import KBinsDiscretizer
from entitys import DataSplitEntity
from enums import BinsStrategyEnum
import toad as td


def f_get_bins(data: DataSplitEntity, feat: str, strategy: str='quantile', nbins: int=10) -> pd.DataFrame:
    # 等频分箱
    if strategy == BinsStrategyEnum.QUANTILE.value:
        kbin_encoder = KBinsDiscretizer(n_bins=nbins, encode='ordinal', strategy='quantile')
        feature_binned = kbin_encoder.fit_transform(data[feat])
        return feature_binned.astype(int).astype(str)
    # 等宽分箱
    if strategy == BinsStrategyEnum.WIDTH.value:
        bin_width = (data[feat].max() - data[feat].min()) / nbins
        return pd.cut(data[feat], bins=nbins, labels=[f'Bin_{i}' for i in range(1, nbins + 1)])
    # 使用toad分箱
    '''
    c = td.transfrom.Combiner()
    # method参数需要根据toad指定的几种方法名称选择
    c.fit(data, y = 'target', method = strategy, min_samples=None, n_bins = nbins, empty_separate = False)
    # 返回toad分箱combiner，用于训练集和测试集的分箱
    # 可使用c.export()[feature]查看某一特征的分箱临界值
    return c
    '''
def f_get_woe(data: DataSplitEntity, c: td.transform.Combiner, to_drop:list) -> pd.DataFrame:
    transer = td.transform.WOETransformer()
    # 根据训练数据来训练woe转换器，并选择目标变量和排除变量
    train_woe = transer.fit_transform(c.transform(data.train_data()), data.train_data()['target'],exclude=to_drop+['target'])
    test_woe = transer.transform(c.transfrom(data.test_data()))
    oot_woe = transer.transform(c.transform(data.val_data()))
    return train_woe, test_woe, oot_woe

def f_get_iv(data: DataSplitEntity) -> pd.DataFrame:
    # 计算前，先排除掉不需要计算IV的cols
    return td.quality(data, 'target',iv_only=True)

def f_get_psi(train_data: DataSplitEntity, oot_data: DataSplitEntity) -> pd.DataFrame:
    # 计算前，先排除掉不需要的cols
    return td.metrics.PSI(train_data, oot_data)


def f_get_corr(data: DataSplitEntity) -> pd.DataFrame:
    pass


def f_get_ivf(data: DataSplitEntity) -> pd.DataFrame:
    pass