import pandas as pd
import numpy as np

# 导入所需模块
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.visualize_route import visualize_route
from myffe.tools.interactive_utils import get_interactive_choice


def analyze_specify_needs(data):
    """分析数据，找出需要删除的列和行

    Args:
        data: pandas DataFrame，输入数据

    Returns:
        dict: 包含推荐删除列和行的字典
    """
    specify_params = {
        'drop_col': [],
        'drop_row': []
    }

    # 找出全为NaN的列
    all_nan_cols = data.columns[data.isna().all()].tolist()

    # 找出常量列（所有值都相同的列）
    constant_cols = []
    for col in data.columns:
        try:
            if data[col].nunique() == 1:
                constant_cols.append(col)
        except:
            pass

    # 合并推荐删除的列
    recommended_drop_col = list(set(all_nan_cols + constant_cols))

    # 找出全为NaN的行
    all_nan_rows = data.index[data.isna().all(axis=1)].tolist()

    specify_params['drop_col'] = recommended_drop_col
    specify_params['drop_row'] = all_nan_rows

    return specify_params


def display_missing_columns(data):
    """展示所有缺失列的信息

    Args:
        data: pandas DataFrame，输入数据

    Returns:
        list: 缺失列信息列表，格式为 [索引] 列名
    """
    missing_cols_info = []
    
    # 计算每列的缺失值数量和比例
    nan_counts = data.isna().sum()
    nan_ratios = data.isna().mean()
    
    # 筛选有缺失值的列
    missing_cols = nan_counts[nan_counts > 0]
    
    for idx, (col, nan_count) in enumerate(missing_cols.items()):
        nan_ratio = nan_ratios[col]
        info = f"[{idx}] {col} (缺失 {nan_count} 个值, 缺失率 {nan_ratio:.2%})"
        missing_cols_info.append(info)

    return missing_cols_info


def specify_py(data, auto_test=False):
    """Python实现的指定列删除分析
    
    Args:
        data: pandas DataFrame，输入数据
        auto_test: bool，是否自动测试模式，默认为 False
    
    Returns:
        dict: 包含指定删除参数的字典
    """
    # 展示缺失列信息
    ColorPrinter.print_info("=" * 60)
    ColorPrinter.print_info("数据缺失列信息：")
    ColorPrinter.print_info("=" * 60)

    missing_cols_info = display_missing_columns(data)

    if not missing_cols_info:
        ColorPrinter.print_info("数据中没有缺失列")
    else:
        # 展示缺失列信息
        for info in missing_cols_info:
            ColorPrinter.print_any(info, color='white')

    ColorPrinter.print_any("", color='white')

    # 对行进行智能推荐（自动删除全为NaN的行）
    recommend_params = analyze_specify_needs(data)
    drop_row = recommend_params['drop_row']

    if drop_row:
        ColorPrinter.print_success(f"智能推荐删除以下行（全为NaN的行）：{drop_row}")
    else:
        ColorPrinter.print_info("智能推荐：无需删除任何行")

    ColorPrinter.print_any("", color='white')

    # 询问是否使用智能推荐（针对列）
    inter_specify = get_interactive_choice(
        title="是否使用智能推荐列？",
        description="选择'ok'将自动推荐需要删除的列（全NaN列或常量列），选择'no'需要手动选择",
        auto_test=auto_test
    )

    drop_col = []

    if inter_specify:
        # 智能推荐模式（针对列）
        drop_col = recommend_params['drop_col']

        if drop_col:
            ColorPrinter.print_success(f"智能推荐删除以下列：{drop_col}")
        else:
            ColorPrinter.print_info("智能推荐：无需删除任何列")
    else:
        # 手动选择模式（只选择列）
        all_cols = list(data.columns)

        ColorPrinter.print_info("请选择要删除的列（按'm'键多选，按回车键确认）")
        selected_cols = interactive(all_cols, title="请按上下箭头选择，回车确认要删除的列,点击'm'键进行多选择", description="选择要删除的列，可以多选").run()

        drop_col = selected_cols

        ColorPrinter.print_success(f"你已选择删除以下列：{drop_col}")

    # 如果没有要删除的行和列，返回skip
    if not drop_col and not drop_row:
        ColorPrinter.print_info("没有需要删除的行和列")
        return 'skip'

    result = {
        'specify_preprocessing': {
            'specify_ways': [
                {
                    'params': {
                        'drop_col': drop_col,
                        'drop_row': drop_row
                    }
                }
            ]
        }
    }
    visualize_route(result)
    return result


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),
        'constant_col': 'same_value',
        'nan_col': [np.nan] * 20
    })

    # 随机添加缺失值
    for col in data.columns:
        if col not in ['constant_col', 'nan_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 = specify_py(data)
    ColorPrinter.print_success("处理结果:")
    ColorPrinter.print_any(result, color='white')
