import pandas as pd
import numpy as np
import time
# import the need modules
import sys
import os

# 添加MFM目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
# py_tools -> route -> FFEv3_2 -> FFE -> MFM -> 根目录
root_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(current_dir)))))
if root_dir not in sys.path:
    sys.path.append(root_dir)

from myffe.tools.interactive import interactive
from myffe.tools.color_printer import ColorPrinter
from myffe.tools.len_match import len_match
from myffe.tools.visualize_route import visualize_route
from myffe.tools.interactive_utils import get_interactive_choice
from myffe.tools.input_utils import ok_input

def analyze_nan_distribution(data, auto_test=False):
    """分析数据的缺失值分布，为每列推荐适合的缺失值处理方法
    
    Args:
        data: pandas DataFrame，输入数据
        auto_test: bool，是否自动测试模式，默认为 False
    
    Returns:
        dict: 包含缺失值处理参数的字典，使用嵌套字典列表结构
    """
    nan_params = {
        'nan_ways': []  # 存储每个列的处理参数
    }

    nan_missing_process_choice = ['drop_col_high_missing', 'fill_nan_knn', 'fill_nan_mean', 'fill_nan_mode', 'fill_nan_forward', 'fill_nan_backward', 'fill_nan_intelligent', 'fill_nan_ei', 'specify_fill', '随机填充','skip']
    # 计算每列的缺失值比例
    nan_ratio = data.isnull().mean() # 计算每列的缺失值比例
    nan_columns = nan_ratio[nan_ratio > 0].index.tolist() # 筛选出缺失值比例大于0的列
    
    if not nan_columns:
        ColorPrinter.print_info("数据中没有缺失值")
        return {'nan_preprocessing': 'skip'}
    
    
    inter_nan = get_interactive_choice(
        title="智能推荐？",
        description="选择'ok'将自动为所有缺失值列推荐最佳处理方法，选择'no'需要手动选择",
        auto_test=auto_test
    )
    
    if not inter_nan:
        selected_cols = interactive(nan_columns  , title="请按上下箭头选择，回车确认缺失值处理列,点击'm'键进行多选择", description="选择要处理的缺失值列，可以多选", min_select=1, max_select=len(nan_columns)).run()
        
        # 确保返回值是列表
        if selected_cols is None:
            time_cols = []
        elif isinstance(selected_cols, str):
            time_cols = [selected_cols]
        else:
            time_cols = selected_cols
        
        if time_cols:
            # 为每个选择的列单独设置处理方法和参数
            for col in time_cols:
                selected_method = interactive(nan_missing_process_choice,min_select = 1, max_select = 1, title=f"请按上下箭头选择，回车确认{col}的缺失值处理方法", description="drop_col_high_missing:删除缺失率高的列, fill_nan_knn:KNN填充, fill_nan_mean:均值填充, fill_nan_mode:众数填充, fill_nan_forward:前向填充, fill_nan_backward:后向填充, fill_nan_intelligent:智能填充, fill_nan_ei:集成填充, specify_fill:用户指定填充值, 随机填充:随机选择填充方式").run()
                # ensure the selected_method is a list
                method = selected_method[0]
                
                # 为当前列创建参数配置，包含所有可能的参数占位
                col_params = {
                    'col': col,
                    'method': method,
                    'drop_threshold': 0.5,  # 默认值
                    'knn_n_neighbors': 5,   # 默认值
                    'ei_model_type': 'linear',  # 默认值
                    'poly_degree': 2,  # 默认值
                    'fill_value': ''  # 默认值，使用空字符串代替 None
                }
                
                if method == '随机填充':
                    from random import choice
                    method = choice(nan_missing_process_choice)
                    col_params['method'] = method
                elif method == 'specify_fill':
                    # 让用户输入填充值
                    ColorPrinter.print_info(f'请为列 {col} 输入填充值：')
                    # 根据列的数据类型进行验证
                    col_dtype = data[col].dtype
                    fill_value = None
                    while fill_value is None:
                        if pd.api.types.is_numeric_dtype(col_dtype):
                            # 数值类型
                            if pd.api.types.is_integer_dtype(col_dtype):
                                fill_value = ok_input('int')
                            else:
                                fill_value = ok_input('float')
                        else:
                            # 非数值类型，使用字符串
                            fill_value = ok_input('str')
                    col_params['fill_value'] = fill_value
                elif method == 'fill_nan_ei':
                    # 为当前列选择EI模型类型
                    model_options = ['linear', 'ridge', 'lasso', 'random_forest', 'gradient_boosting', '自定义']
                    selected_model = interactive(model_options, min_select=1, max_select=1, title=f"请按上下箭头选择，回车确认列 {col} 的EI模型", description="选择EI方法中使用的机器学习模型").run()
                    # 确保返回值是单个值
                    model_type = selected_model[0]
                    if model_type == '自定义':
                        # 让用户输入自定义模型路径
                        ColorPrinter.print_info('请输入自定义模型完整路径（例如：sklearn.ensemble.RandomForestRegressor）：')
                        time.sleep(3)
                        custom_model = ok_input('str')
                        col_params['ei_model_type'] = custom_model
                        ColorPrinter.print_success(f'你已为列 {col} 选择自定义EI模型：' + custom_model)
                    else:
                        col_params['ei_model_type'] = model_type
                        ColorPrinter.print_success(f'你已为列 {col} 选择EI模型类型：' + model_type)
                    
                    # 为当前列选择多项式阶数
                    # 根据原始数据维度推荐多项式阶数
                    original_feature_count = len(data.columns)
                    data_rows = len(data)
                    
                    # 默认推荐阶数
                    default_degree = 2
                    if original_feature_count <= 5:
                        default_degree = 2
                    elif original_feature_count <= 10:
                        default_degree = 3
                    else:
                        default_degree = 4
                    
                    # 数据量特别大时可以适当增加阶数
                    if data_rows >= 5000:
                        default_degree = min(default_degree + 1, 5)
                    
                    # 构建选项，默认选项放在第一位
                    degree_options = [str(default_degree), '1', '2', '3', '4', '5', '自定义']
                    # 去重并保持顺序
                    degree_options = list(dict.fromkeys(degree_options))
                    
                    selected_degree = interactive(degree_options, min_select=1, max_select=1, title=f"请按上下箭头选择，回车确认列 {col} 的多项式特征阶数", description=f"设置EI方法中多项式回归的特征阶数，根据数据维度推荐：{default_degree}").run()
                    # 确保返回值是单个值
                    tem_degree = selected_degree[0]
                    if tem_degree == '自定义':
                        # 自动测试模式下，使用默认值
                        if auto_test:
                            degree = default_degree
                        else:
                            # 让用户输入自定义的多项式阶数
                            degree = ok_input('int', min_value=1, max_value=10)
                    else:
                        degree = int(tem_degree)
                    col_params['poly_degree'] = degree
                    ColorPrinter.print_success(f'你已为列 {col} 选择多项式特征阶数：' + str(degree))
                elif method == 'drop_col_high_missing':
                    threshold_options = ['0.3', '0.4', '0.5', '0.6', '0.7', '自定义']
                    selected_threshold = interactive(threshold_options, min_select=1, max_select=1, title=f"请按上下箭头选择，回车确认列 {col} 的高缺失值删除阈值", description="设置列的缺失值比例阈值，超过此比例的列将被删除").run()
                    # 确保返回值是单个值
                    tem_threshold = selected_threshold[0]
                    if tem_threshold == '自定义':
                        # 让用户输入自定义的阈值
                        drop_threshold = ok_input('float', min_value=0, max_value=1)
                    else:
                        drop_threshold = float(tem_threshold)
                    col_params['drop_threshold'] = drop_threshold
                    ColorPrinter.print_success(f'你已为列 {col} 选择高缺失值列删除阈值：' + str(drop_threshold))
                elif method == 'fill_nan_knn':
                    knn_options = ['3', '5', '7', '9', '10', '11', '13', '自定义']
                    selected_knn = interactive(knn_options, min_select=1, max_select=1, title=f"请按上下箭头选择，回车确认列 {col} 的KNN填充邻居数", description="设置KNN填充的邻居数量，选择'自定义'可输入任意值").run()
                    # 确保返回值是单个值
                    tem_knn = selected_knn[0]
                    if tem_knn == '自定义':
                        # 自动测试模式下，使用默认值 5
                        if auto_test:
                            knn_n_neighbors = 5
                        else:
                            # 让用户输入自定义的KNN邻居数
                            knn_n_neighbors = ok_input('int', min_value=1, max_value=50)
                    else:
                        knn_n_neighbors = int(tem_knn)
                    col_params['knn_n_neighbors'] = knn_n_neighbors
                    ColorPrinter.print_success(f'你已为列 {col} 选择KNN填充的邻居数：' + str(knn_n_neighbors))
                
                # 添加到nan_ways列表
                nan_params['nan_ways'].append(col_params)
                ColorPrinter.print_success(f'你已选择{col}的缺失值处理方法：{method}')
    else:
        # 只分析确实有缺失值的列
        for col in nan_columns:
            col_nan_ratio = nan_ratio[col]
            
            # 为当前列创建参数配置，包含所有可能的参数占位
            col_params = {
                'col': col,
                'drop_threshold': 0.5,  # 默认值
                'knn_n_neighbors': 5,   # 默认值
                'ei_model_type': 'linear',  # 默认值
                'poly_degree': 2,  # 默认值
                'fill_value': ''  # 默认值，使用空字符串代替 None
            }
            
            # 检查是否需要删除高缺失值列
            if col_nan_ratio > 0.5:
                col_params['method'] = 'drop_col_high_missing'
                col_params['drop_threshold'] = 0.5
            
            # 根据列的数据类型和缺失值比例推荐填充方法
            elif col_nan_ratio > 0:
                # 检查是否满足EI方法的推荐条件
                # 1. 计算整体缺失行比例
                total_missing_rows = data.isnull().any(axis=1).mean()
                # 2. 计算编码后的特征数量估计
                feature_count = 0
                for col_name in data.columns:
                    if col_name != col:  # 排除当前列
                        if pd.api.types.is_categorical_dtype(data[col_name]) or pd.api.types.is_object_dtype(data[col_name]):
                            # 分类列编码后特征数估计（假设平均每个分类列有10个类别）
                            feature_count += min(data[col_name].nunique(), 10)
                        elif pd.api.types.is_datetime64_any_dtype(data[col_name]):
                            # 时间列编码后特征数（年、月、日、小时等）
                            feature_count += 4
                        else:
                            # 数值列直接计数
                            feature_count += 1
                
                # 满足EI推荐条件
                # 1. 整体缺失行比例小于30%
                # 2. 编码后的特征数量不超过30个
                # 3. 目标列缺失比例在5%-40%之间
                # 4. 编码前的列数量达到3个以上
                # 5. 数据行数在1000行以上
                original_feature_count = len(data.columns)
                data_rows = len(data)
                if total_missing_rows < 0.3 and feature_count <= 30 and col_nan_ratio > 0.05 and col_nan_ratio < 0.4 and original_feature_count >= 3 and data_rows >= 1000:
                    col_params['method'] = 'fill_nan_ei'
                    # 根据数据特征推荐模型类型
                    # 1. 计算特征相关性，判断是否存在多重共线性
                    numeric_cols = [col_name for col_name in data.columns if col_name != col and pd.api.types.is_numeric_dtype(data[col_name])]
                    corr_matrix = None
                    if len(numeric_cols) >= 2:
                        try:
                            corr_matrix = data[numeric_cols].corr()
                        except:
                            pass
                    
                    # 2. 根据数据特征选择模型
                    model_type = 'linear'  # 默认模型
                    if corr_matrix is not None:
                        # 检查是否存在多重共线性（相关系数绝对值大于0.8）
                        high_corr = (abs(corr_matrix) > 0.8).sum().sum() - len(numeric_cols)  # 减去对角线元素
                        
                        # 检查是否存在非线性关系（通过特征与目标变量的关系）
                        has_nonlinearity = False
                        try:
                            # 检查每个数值特征与目标变量的关系
                            for num_col in numeric_cols:
                                if num_col != col:
                                    # 计算皮尔逊相关系数（线性关系）
                                    linear_corr = abs(data[num_col].corr(data[col]))
                                    # 计算斯皮尔曼相关系数（非线性关系）
                                    spearman_corr = abs(data[num_col].corr(data[col], method='spearman'))
                                    # 如果斯皮尔曼相关系数显著大于皮尔逊相关系数，说明存在非线性关系
                                    if spearman_corr > linear_corr + 0.2:
                                        has_nonlinearity = True
                                        break
                        except:
                            pass
                        
                        if high_corr > 0:
                            # 存在多重共线性，推荐使用Ridge回归
                            model_type = 'ridge'
                        elif len(numeric_cols) > 10:
                            # 特征数量较多，推荐使用Lasso回归进行特征选择
                            model_type = 'lasso'
                        elif has_nonlinearity and data_rows < 10000:
                            # 存在明显的非线性关系，且数据量适中，推荐使用梯度提升
                            model_type = 'gradient_boosting'
                        elif data_rows > 8000:
                            # 数据量很大，推荐使用随机森林
                            model_type = 'random_forest'
                        elif data_rows > 2000 and data_rows <= 8000:
                            # 数据量适中，推荐使用梯度提升
                            model_type = 'gradient_boosting'
                        else:
                            # 数据量较小，默认使用线性回归
                            model_type = 'linear'
                    else:
                        # 没有足够的数值特征，默认使用线性回归
                        model_type = 'linear'
                    
                    col_params['ei_model_type'] = model_type
                    
                    # 根据数据特征推荐多项式阶数
                    # 特征数量较少时使用较低阶数，特征数量较多时使用较高阶数
                    poly_degree = 2  # 默认阶数
                    if feature_count <= 10:
                        poly_degree = 2
                    elif feature_count <= 20:
                        poly_degree = 3
                    else:
                        poly_degree = 4
                    # 数据量特别大时可以适当增加阶数
                    if data_rows >= 5000:
                        poly_degree = min(poly_degree + 1, 5)
                    
                    col_params['poly_degree'] = poly_degree
                else:
                    # 根据列的数据类型推荐其他填充方法
                    if pd.api.types.is_numeric_dtype(data[col]):
                        if col_nan_ratio > 0.1: #当数值列缺失值比例大于0.1时，建议使用KNN填充
                            col_params['method'] = 'fill_nan_knn'
                            col_params['knn_n_neighbors'] = 5
                        else:
                            col_params['method'] = 'fill_nan_mean'
                    elif pd.api.types.is_categorical_dtype(data[col]) or pd.api.types.is_object_dtype(data[col]):
                        #当分类列缺失值比例大于0.1时，建议使用众数填充
                        col_params['method'] = 'fill_nan_mode'
                    elif pd.api.types.is_datetime64_any_dtype(data[col]):
                        # 检查日期时间列的单调性
                        is_increasing = data[col].dropna().is_monotonic_increasing
                        is_decreasing = data[col].dropna().is_monotonic_decreasing
                        if is_increasing:
                            # 递增数据建议使用向前填充
                            col_params['method'] = 'fill_nan_forward'
                        elif is_decreasing:
                            # 递减数据建议使用向后填充
                            col_params['method'] = 'fill_nan_backward'
                        else:
                            # 无明显单调性，默认使用向前填充
                            col_params['method'] = 'fill_nan_forward'
                    else:
                        col_params['method'] = 'fill_nan_intelligent'
            
            # 添加到nan_ways列表
            nan_params['nan_ways'].append(col_params)
    
    result = {'nan_preprocessing': nan_params}
    visualize_route(result)
    return result

def nan_py(data, auto_test=False):
    """Python实现的缺失值处理分析
    
    Args:
        data: pandas DataFrame，输入数据
        drop_threshold: 缺失值比例阈值，默认为0.5
        auto_test: bool，是否自动测试模式，默认为 False
    
    Returns:
        dict: 包含缺失值处理参数的字典
    """
    
    return analyze_nan_distribution(data, auto_test=auto_test)
if __name__ == "__main__":
    # 测试代码
    import pandas as pd
    import numpy as np
    
    # 创建测试数据
    np.random.seed(42)
    data = pd.DataFrame({
        'numeric_col1': np.random.normal(100, 10, 20),
        'numeric_col2': np.random.normal(50, 5, 20),
        'categorical_col': np.random.choice(['A', 'B', 'C'], 20),
        'datetime_col': pd.date_range('2020-01-01', periods=20)
    })
    
    # 随机添加缺失值
    for col in data.columns:
        if col != 'datetime_col':
            mask = np.random.choice([True, False], size=20, p=[0.2, 0.8])
            data.loc[mask, col] = np.nan
    
    ColorPrinter.print_info("测试数据:")
    ColorPrinter.print_any(data.head(), color='white')
    ColorPrinter.print_any("\n缺失值统计:", color='blue')
    ColorPrinter.print_any(data.isnull().sum(), color='white')
    
    # 测试缺失值处理
    ColorPrinter.print_info("\n测试缺失值处理:")
    result = nan_py(data)
    ColorPrinter.print_success("处理结果:")
    ColorPrinter.print_any(result, color='white')


"""
这个文件输入一个data
输出一个字典
格式是：
{
    'nan_preprocessing': {
        'nan_ways': [
            {
                'col': 'col1',
                'method': 'fill_nan_mean'
            },
            {
                'col': 'col2',
                'method': 'fill_nan_knn',
                'knn_n_neighbors': 5
            },
            {
                'col': 'col3',
                'method': 'fill_nan_ei',
                'ei_model_type': 'linear',
                'poly_degree': 2
            },
            {
                'col': 'col4',
                'method': 'drop_col_high_missing',
                'drop_threshold': 0.5
            },
            {
                'col': 'col5',
                'method': 'specify_fill',
                'fill_value': 0
            }
        ]
    }
}

method的取值有

    'drop_col_high_missing'   # 删除缺失值比例超过阈值的列
    'fill_nan_knn'   # 使用KNN填充缺失值
    'fill_nan_mean'   # 使用均值填充缺失值
    'fill_nan_mode'   # 使用众数填充缺失值
    'fill_nan_forward'   # 前向填充缺失值
    'fill_nan_backward'   # 后向填充缺失值
    'fill_nan_intelligent'   # 智能填充缺失值
    'fill_nan_ei'   # 集成填充缺失值
    'specify_fill'   # 用户指定填充值
    '随机填充'   # 随机选择填充方式
    'skip'   # 跳过缺失值处理

参数处理逻辑：
- 每个列的处理配置是独立的，包含在nan_ways列表中的一个字典里
- 无论方法是否需要特定参数，所有可能的参数都会在配置中进行占位
- 未被选中的参数会使用默认值进行占位，确保配置结构的一致性
- 这种设计使得配置结构更加统一，便于后续处理和维护
"""
