import sys
import os

# 添加项目根目录到Python路径
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..', '..', '..')))

import pandas as pd
import numpy as np
from myffe.tools.interactive import interactive
from myffe.tools.visualize_route import visualize_route
from myffe.tools.color_printer import ColorPrinter
from myffe.tools.param_utils import ensure_list, convert_to_list
from myffe.tools.interactive_utils import get_interactive_choice

def analyze_sample_distribution(data, label_columns):
    """分析数据的类别分布，推荐适合的采样处理方法"""
    # 初始化采样处理路线
    sample_route = {}
    
    # 对每个标签列进行分析
    for label_col in label_columns:
        # 检查该列是否存在于数据中
        if label_col not in data.columns:
            ColorPrinter.print_info(f"标签列 {label_col} 不存在于数据中，跳过采样处理")
            return {}
        
        # 检查该列是否适合作为标签列
        label_data = data[label_col]
        
        # 计算类别分布
        class_counts = label_data.value_counts()   # 计算每个类别的样本数量
        num_classes = len(class_counts)     # 类别数量
        total_samples = len(label_data)     # 总样本数量
        
        # 计算类别不平衡度
        if num_classes > 1:
            max_count = class_counts.max()
            min_count = class_counts.min()
            imbalance_ratio = max_count / min_count if min_count > 0 else float('inf')
        else:
            imbalance_ratio = 1.0
        
        # 分析数据特征，推荐适合的采样方法
        recommended_configs = []
        
        # 基于类别数量和不平衡度推荐方法
        if num_classes == 1:
            # 只有一个类别，不需要采样
            recommended_configs.append({
                'method': ['skip'],
                'balance_method': 'oversample',  # 占位
                'random_state': 42,  # 占位
                'high_sample_selection': 'RandomOverSampler',  # 占位
                'smote_k_neighbors': 5,  # 占位
                'adasyn_n_neighbors': 5,  # 占位
                'borderline_smote_k_neighbors': 5,  # 占位
                'svm_smote_k_neighbors': 5,  # 占位
                'enn_n_neighbors': 5,  # 占位
                'description': '跳过采样（只有一个类别）'
            })
        elif num_classes == 2:
            # 二分类问题
            if imbalance_ratio > 10:
                # 严重不平衡，推荐SMOTE或ADASYN
                recommended_configs.extend([
                    {
                        'method': ['MY_high_sample'],
                        'balance_method': 'oversample',  # 占位
                        'random_state': 42,  # 占位
                        'high_sample_selection': 'SMOTE',
                        'smote_k_neighbors': 5,
                        'adasyn_n_neighbors': 5,  # 占位
                        'borderline_smote_k_neighbors': 5,  # 占位
                        'svm_smote_k_neighbors': 5,  # 占位
                        'enn_n_neighbors': 5,  # 占位
                        'description': 'SMOTE采样（严重不平衡二分类）'
                    },
                    {
                        'method': ['MY_high_sample'],
                        'balance_method': 'oversample',  # 占位
                        'random_state': 42,  # 占位
                        'high_sample_selection': 'ADASYN',
                        'smote_k_neighbors': 5,  # 占位
                        'adasyn_n_neighbors': 5,
                        'borderline_smote_k_neighbors': 5,  # 占位
                        'svm_smote_k_neighbors': 5,  # 占位
                        'enn_n_neighbors': 5,  # 占位
                        'description': 'ADASYN采样（严重不平衡二分类）'
                    }
                ])
            elif imbalance_ratio > 2:
                # 中度不平衡，推荐RandomOverSampler或BorderlineSMOTE
                recommended_configs.extend([
                    {
                        'method': ['MY_high_sample'],
                        'balance_method': 'oversample',  # 占位
                        'random_state': 42,  # 占位
                        'high_sample_selection': 'RandomOverSampler',
                        'smote_k_neighbors': 5,  # 占位
                        'adasyn_n_neighbors': 5,  # 占位
                        'borderline_smote_k_neighbors': 5,  # 占位
                        'svm_smote_k_neighbors': 5,  # 占位
                        'enn_n_neighbors': 5,  # 占位
                        'description': 'RandomOverSampler采样（中度不平衡二分类）'
                    },
                    {
                        'method': ['MY_high_sample'],
                        'balance_method': 'oversample',  # 占位
                        'random_state': 42,  # 占位
                        'high_sample_selection': 'BorderlineSMOTE',
                        'smote_k_neighbors': 5,  # 占位
                        'adasyn_n_neighbors': 5,  # 占位
                        'borderline_smote_k_neighbors': 5,
                        'svm_smote_k_neighbors': 5,  # 占位
                        'enn_n_neighbors': 5,  # 占位
                        'description': 'BorderlineSMOTE采样（中度不平衡二分类）'
                    }
                ])
            else:
                # 平衡数据，推荐skip
                recommended_configs.append({
                    'method': ['skip'],
                    'balance_method': 'oversample',  # 占位
                    'random_state': 42,  # 占位
                    'high_sample_selection': 'RandomOverSampler',  # 占位
                    'smote_k_neighbors': 5,  # 占位
                    'adasyn_n_neighbors': 5,  # 占位
                    'borderline_smote_k_neighbors': 5,  # 占位
                    'svm_smote_k_neighbors': 5,  # 占位
                    'enn_n_neighbors': 5,  # 占位
                    'description': '跳过采样（平衡二分类）'
                })
        else:
            # 多分类问题
            if imbalance_ratio > 5:
                # 严重不平衡，推荐SMOTE或过采样
                recommended_configs.extend([     #extend方法将可迭代对象里面的东西添加到指定的列表的后面
                    {
                        'method': ['MY_high_sample'],
                        'balance_method': 'oversample',  # 占位
                        'random_state': 42,  # 占位
                        'high_sample_selection': 'SMOTE',
                        'smote_k_neighbors': 5,
                        'adasyn_n_neighbors': 5,  # 占位
                        'borderline_smote_k_neighbors': 5,  # 占位
                        'svm_smote_k_neighbors': 5,  # 占位
                        'enn_n_neighbors': 5,  # 占位
                        'description': 'SMOTE采样（严重不平衡多分类）'
                    },
                    {
                        'method': ['multi_class_balance'],
                        'balance_method': 'oversample',
                        'random_state': 42,
                        'high_sample_selection': 'RandomOverSampler',  # 占位
                        'smote_k_neighbors': 5,  # 占位
                        'adasyn_n_neighbors': 5,  # 占位
                        'borderline_smote_k_neighbors': 5,  # 占位
                        'svm_smote_k_neighbors': 5,  # 占位
                        'enn_n_neighbors': 5,  # 占位
                        'description': '过采样（严重不平衡多分类）'
                    }
                ])
            elif imbalance_ratio > 2:
                # 中度不平衡，推荐过采样
                recommended_configs.append({
                    'method': ['multi_class_balance'],
                    'balance_method': 'oversample',
                    'random_state': 42,
                    'high_sample_selection': 'RandomOverSampler',  # 占位
                    'smote_k_neighbors': 5,  # 占位
                    'adasyn_n_neighbors': 5,  # 占位
                    'borderline_smote_k_neighbors': 5,  # 占位
                    'svm_smote_k_neighbors': 5,  # 占位
                    'enn_n_neighbors': 5,  # 占位
                    'description': '过采样（中度不平衡多分类）'
                })
            else:
                # 平衡数据，推荐skip
                recommended_configs.append({
                    'method': ['skip'],
                    'balance_method': 'oversample',  # 占位
                    'random_state': 42,  # 占位
                    'high_sample_selection': 'RandomOverSampler',  # 占位
                    'smote_k_neighbors': 5,  # 占位
                    'adasyn_n_neighbors': 5,  # 占位
                    'borderline_smote_k_neighbors': 5,  # 占位
                    'svm_smote_k_neighbors': 5,  # 占位
                    'enn_n_neighbors': 5,  # 占位
                    'description': '跳过采样（平衡多分类）'
                })
        
        # 基于数据量大小调整推荐
        if total_samples < 1000:
            # 小数据集，优先推荐计算量较小的方法
            lightweight_configs = [config for config in recommended_configs if 
                                 ('multi_class_balance' in config.get('method', [])) or
                                 (config.get('high_sample_selection') in ['RandomOverSampler', 'RandomUnderSampler'])]
            if lightweight_configs:
                recommended_configs = lightweight_configs
        elif total_samples > 100000:
            # 大数据集，优先推荐高效的方法
            efficient_configs = [config for config in recommended_configs if 
                               ('multi_class_balance' in config.get('method', [])) or
                               (config.get('high_sample_selection') in ['RandomOverSampler', 'RandomUnderSampler'])]
            if efficient_configs:
                recommended_configs = efficient_configs
        
        # 为该标签列创建采样路线
        if recommended_configs:
            # 提取最佳配置
            best_config = recommended_configs[0]
            
            # 构建完整的参数配置，与参数配置文档和 sample_preprocessing.py 完全兼容
            # 新格式配置
            sample_params = {
                'sample_ways': [
                    {
                        'method': best_config.get('method', ['skip']),
                        'balance_method': best_config.get('balance_method', 'oversample'),
                        'random_state': best_config.get('random_state', 42),
                        'high_sample_selection': best_config.get('high_sample_selection', 'RandomOverSampler'),
                        'smote_k_neighbors': best_config.get('smote_k_neighbors', 5),
                        'adasyn_n_neighbors': best_config.get('adasyn_n_neighbors', 5),
                        'borderline_smote_k_neighbors': best_config.get('borderline_smote_k_neighbors', 5),
                        'svm_smote_k_neighbors': best_config.get('svm_smote_k_neighbors', 5),
                        'enn_n_neighbors': best_config.get('enn_n_neighbors', 5)
                    }
                ]
            }
            
            sample_route[label_col] = {
                'num_classes': num_classes,
                'total_samples': total_samples,
                'imbalance_ratio': round(imbalance_ratio, 2),
                'recommended_configs': recommended_configs,
                'best_config': best_config,
                'best_params': sample_params,
                'best_method': best_config['description']
            }
    
    return sample_route

def sample_py(data, label_columns, auto_test=False):
    """Python实现的采样处理分析"""
    # 确保label_columns是列表
    label_columns = ensure_list(label_columns)
    if not label_columns:
        ColorPrinter.print_info("没有提供标签列，跳过采样处理")
        return {'sample_preprocessing': 'skip'}
    
    #region 分析样本分布,分析错误直接跳过
    sample_route = analyze_sample_distribution(data, label_columns)
    
    if not sample_route:
        ColorPrinter.print_info("无法分析样本分布，跳过采样处理")
        return {'sample_preprocessing': 'skip'}
    #endregion
    
    # 获取第一个标签列的分析结果
    label_col = list(sample_route.keys())[0]
    analysis = sample_route[label_col]
    
#region 显示分析信息
    ColorPrinter.print_info(f"标签列: {label_col}")
    ColorPrinter.print_info(f"类别数量: {analysis['num_classes']}")
    ColorPrinter.print_info(f"总样本数量: {analysis['total_samples']}")
    ColorPrinter.print_info(f"类别不平衡度: {analysis['imbalance_ratio']}")
    ColorPrinter.print_info(f"推荐方法: {analysis['best_method']}")
#endregion
    #region 智能推荐菜单
    inter_choice = get_interactive_choice(
        title="采样处理",
        description="选择'ok'使用推荐方案，选择'no'自定义处理方案",
        auto_test=auto_test
    )
    #endregion
    
    if inter_choice:
        ColorPrinter.print_success("使用推荐方案")
        # 使用推荐参数
        sample_params = analysis['best_params']
        # 构建结果
        result = {'sample_preprocessing': sample_params}
        
        # 可视化处理路线
        visualize_route(result)
        return result
    else:
        # 自定义处理方案
        ColorPrinter.print_info("自定义处理方案")
        
#region 初始化参数
        sample_params = {
            'sample_ways': [
                {
                    'method': ['skip'],
                    'balance_method': 'oversample',
                    'random_state': 42,
                    'high_sample_selection': 'RandomOverSampler',
                    'smote_k_neighbors': 5,
                    'adasyn_n_neighbors': 5,
                    'borderline_smote_k_neighbors': 5,
                    'svm_smote_k_neighbors': 5,
                    'enn_n_neighbors': 5
                }
            ]
        }
#endregion
        
        # 选择处理方法
        method_options = ["multi_class_balance", "MY_high_sample", "skip"]
        method_selected = interactive(method_options, title="选择处理方法", description="请选择采样方法:", min_select=1, max_select=1).run()
        selected_method = method_selected[0] if method_selected else "skip"
        
        # 更新新格式
        sample_params['sample_ways'][0]['method'] = [selected_method]
        
        if selected_method == "multi_class_balance":
            # 选择平衡方法
            balance_options = ["oversample", "undersample"]
            balance_selected = interactive(balance_options, title="选择平衡方法", description="请选择平衡方法:", min_select=1, max_select=1).run()
            balance_method = balance_selected[0] if balance_selected else "oversample"
            
            # 更新新格式
            sample_params['sample_ways'][0]['balance_method'] = balance_method
            
            # 输入随机种子
            random_state = input("请输入随机种子 (默认42): ")
            random_state_val = int(random_state) if random_state else 42
            
            # 更新新格式
            sample_params['sample_ways'][0]['random_state'] = random_state_val
        elif selected_method == "MY_high_sample":
            # 选择高级采样方法
            high_sample_options = [
                "SMOTE",
                "ADASYN",
                "BorderlineSMOTE",
                "SVMSMOTE",
                "RandomOverSampler",
                "RandomUnderSampler",
                "TomekLinks",
                "EditedNearestNeighbours",
                "ClusterCentroids"
            ]
            high_sample_selected = interactive(high_sample_options, title="选择高级采样方法", description="请选择高级采样方法:", min_select=1, max_select=1).run()
            high_sample_selection = high_sample_selected[0] if high_sample_selected else "RandomOverSampler"
            
            # 更新新格式
            sample_params['sample_ways'][0]['high_sample_selection'] = high_sample_selection
            
            # 根据选择的方法添加相应参数
            if high_sample_selection == "SMOTE":
                smote_k = input("请输入SMOTE邻居数 (默认5): ")
                smote_k_val = int(smote_k) if smote_k else 5
                # 更新新格式
                sample_params['sample_ways'][0]['smote_k_neighbors'] = smote_k_val
            elif high_sample_selection == "ADASYN":
                adasyn_k = input("请输入ADASYN邻居数 (默认5): ")
                adasyn_k_val = int(adasyn_k) if adasyn_k else 5
                # 更新新格式
                sample_params['sample_ways'][0]['adasyn_n_neighbors'] = adasyn_k_val
            elif high_sample_selection == "BorderlineSMOTE":
                borderline_k = input("请输入BorderlineSMOTE邻居数 (默认5): ")
                borderline_k_val = int(borderline_k) if borderline_k else 5
                # 更新新格式
                sample_params['sample_ways'][0]['borderline_smote_k_neighbors'] = borderline_k_val
            elif high_sample_selection == "SVMSMOTE":
                svm_k = input("请输入SVMSMOTE邻居数 (默认5): ")
                svm_k_val = int(svm_k) if svm_k else 5
                # 更新新格式
                sample_params['sample_ways'][0]['svm_smote_k_neighbors'] = svm_k_val
            elif high_sample_selection == "EditedNearestNeighbours":
                enn_k = input("请输入EditedNearestNeighbours邻居数 (默认5): ")
                enn_k_val = int(enn_k) if enn_k else 5
                # 更新新格式
                sample_params['sample_ways'][0]['enn_n_neighbors'] = enn_k_val
    
    # 构建最终结果
    result = {'sample_preprocessing': sample_params}
        
    # 可视化处理路线
    visualize_route(result)
        
    return result


if __name__ == "__main__":
    """
    测试采样处理功能
    """
    import pandas as pd
    import numpy as np
    
    print("=" * 80)
    print("采样处理模块测试")
    print("=" * 80)
    
    # 创建测试数据
    np.random.seed(42)
    data = pd.DataFrame({
        'feature1': np.random.normal(50, 10, 100),
        'feature2': np.random.normal(100, 20, 100),
        'feature3': np.random.normal(200, 30, 100),
        'label': ['A'] * 70 + ['B'] * 20 + ['C'] * 10
    })
    
    print("测试数据:")
    print(data)
    print(f"\n数据形状: {data.shape}")
    print(f"列名: {list(data.columns)}")
    print(f"类别分布:")
    print(data['label'].value_counts())
    
    print("\n" + "=" * 80)
    print("测试标签列检测")
    print("=" * 80)
    
    # 测试标签列检测
    label_cols = ['label']  # 直接使用已知的标签列
    print(f"使用的标签列: {label_cols}")
    
    print("\n" + "=" * 80)
    print("测试样本分布分析")
    print("=" * 80)
    
    # 测试样本分布分析
    sample_route = analyze_sample_distribution(data, label_cols)
    if sample_route:
        label_col = list(sample_route.keys())[0]
        analysis = sample_route[label_col]
        print(f"标签列: {label_col}")
        print(f"类别数量: {analysis['num_classes']}")
        print(f"总样本数量: {analysis['total_samples']}")
        print(f"类别不平衡度: {analysis['imbalance_ratio']}")
        print(f"推荐方法: {analysis['best_method']}")
        print(f"推荐参数: {analysis['best_params']}")
    
    print("\n" + "=" * 80)
    print("测试主处理函数")
    print("=" * 80)
    
    # 测试主处理函数
    try:
        result = sample_py(data, label_columns=['label'])
        print("\n处理结果:")
        if result == 'skip':
            print("跳过处理")
        else:
            for key, value in result.items():
                print(f"\n{key}:")
                for param_key, param_value in value.items():
                    print(f"  {param_key}: {param_value}")
    except Exception as e:
        print(f"测试过程中出错: {e}")
    
    print("\n" + "=" * 80)
    print("测试完成!")
    print("=" * 80)
