import sys
import os

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

from myffe.tools.Loggerv1_1 import log_class
from myffe.tools.color_printer import ColorPrinter
from myffe.tools.evaluation_decorator import with_evaluation

# 尝试导入可选依赖
try:
    from dask.dataframe import DataFrame as DaskDataFrame
    from MFM.FFE.FFEv3_2.preprocessing.dask_preprocessing import process_specify_with_dask
    dask_available = True
except ImportError:
    dask_available = False

@log_class(type='use')
class SpecifyPreprocessing:
    def __init__(self, shortcut_specify=None):
        self.shortcut_specify = shortcut_specify or {}

    @with_evaluation
    def __call__(self, data, data_label=None):
        self.data_label = data_label
        if not self.shortcut_specify:
            ColorPrinter.print_warning('没有指定任何处理参数')
            return data
        # 检查是否为Dask DataFrame
        is_dask = False
        if dask_available:
            is_dask = isinstance(data, DaskDataFrame)

        if is_dask:
            data = process_specify_with_dask(data, data_label, self.shortcut_specify)
        else:
            if self.shortcut_specify.get('drop_row'):
                data = self.specify_row(data)
            if self.shortcut_specify.get('drop_col'):
                data = self.specify_column(data)

        return data

    def specify_row(self, data):
        drop_row = self.shortcut_specify['drop_row']
        
        # 处理不同类型的输入
        if isinstance(drop_row, str):
            # 处理逗号分隔的字符串
            drop_row = [row.strip() for row in drop_row.split(',')]
        elif isinstance(drop_row, int):
            # 处理单个整数
            drop_row = [drop_row]
        elif not isinstance(drop_row, list):
            # 其他类型转换为列表
            drop_row = [drop_row]
        
        # 转换为整数
        try:
            drop_row = [int(row) for row in drop_row]
        except ValueError:
            ColorPrinter.print_warning("警告: 行索引必须是整数，跳过行删除")
            return data
        
        # 只删除存在的行
        existing_rows = [row for row in drop_row if row in data.index]
        if existing_rows:
            data.drop(existing_rows, axis=0, inplace=True)
            ColorPrinter.print_success(f'删除了行: {existing_rows}')
        else:
            ColorPrinter.print_warning('没有找到要删除的行')
        return data

    def specify_column(self, data):
        drop_col = self.shortcut_specify['drop_col']
        
        # 处理不同类型的输入
        if isinstance(drop_col, str):
            # 处理逗号分隔的字符串
            drop_col = [col.strip() for col in drop_col.split(',')]
        elif not isinstance(drop_col, list):
            # 其他类型转换为列表
            drop_col = [drop_col]
        
        # 只删除存在的列
        drop_col = [col for col in drop_col if col in data.columns]
        
        if drop_col:
            data.drop(drop_col, axis=1, inplace=True)
            ColorPrinter.print_success(f'删除了列: {drop_col}')
        else:
            ColorPrinter.print_warning('没有找到要删除的列')
        return data


"""
参数说明文档:

shortcut_specify 参数字典:
    - drop_row: 要删除的行索引
      类型: 索引列表、单个索引或逗号分隔的字符串
      可选值: 任意有效的行索引
      默认值: None（可选）
      说明: 删除指定的行
    
    - drop_col: 要删除的列名
      类型: 列名列表、单个列名或逗号分隔的字符串
      可选值: 任意有效的列名
      默认值: None（可选）
      说明: 删除指定的列

使用示例:
    # 删除列（列表形式）
    specify_params = {'drop_col': ['col1', 'col2']}
    specify_processor = SpecifyPreprocessing(shortcut_specify=specify_params)
    processed_data = specify_processor(data, 'label')
    
    # 删除列（字符串形式）
    specify_params = {'drop_col': 'col1, col2, col3'}
    specify_processor = SpecifyPreprocessing(shortcut_specify=specify_params)
    processed_data = specify_processor(data, 'label')
    
    # 删除行（列表形式）
    specify_params = {'drop_row': [0, 1, 2]}
    specify_processor = SpecifyPreprocessing(shortcut_specify=specify_params)
    processed_data = specify_processor(data, 'label')
    
    # 删除行（字符串形式）
    specify_params = {'drop_row': '0, 1, 2, 3'}
    specify_processor = SpecifyPreprocessing(shortcut_specify=specify_params)
    processed_data = specify_processor(data, 'label')
    
    # 同时删除行和列
    specify_params = {
        'drop_row': [0, 1],
        'drop_col': ['col1', 'col2']
    }
    specify_processor = SpecifyPreprocessing(shortcut_specify=specify_params)
    processed_data = specify_processor(data, 'label')
"""


if __name__ == "__main__":
    import pandas as pd
    import numpy as np
    
    print("=" * 60)
    print("指定行列删除测试程序")
    print("=" * 60)
    
    # 创建测试数据
    np.random.seed(42)
    test_data = pd.DataFrame({
        'feature1': np.random.randint(1, 100, 10),
        'feature2': np.random.randint(1, 100, 10),
        'feature3': np.random.randint(1, 100, 10),
        'feature4': np.random.randint(1, 100, 10),
        'label': ['A', 'B', 'A', 'B', 'A', 'B', 'A', 'B', 'A', 'B']
    })
    
    print("\n原始数据:")
    print(test_data)
    print(f"\n数据形状: {test_data.shape}")
    
    # 测试不同的删除方式
    test_cases = [
        {
            'params': {'drop_col': ['feature1', 'feature2']},
            'description': '删除指定列'
        },
        {
            'params': {'drop_col': 'feature3'},
            'description': '删除单个列'
        },
        {
            'params': {'drop_row': [0, 1, 2]},
            'description': '删除指定行'
        },
        {
            'params': {'drop_row': 5},
            'description': '删除单个行'
        },
        {
            'params': {'drop_col': ['feature4']},
            'description': '删除标签列以外的列'
        }
    ]
    
    for i, test_case in enumerate(test_cases, 1):
        print(f"\n{'=' * 60}")
        print(f"测试 {i}: {test_case['description']}")
        print('=' * 60)
        print(f"参数: {test_case['params']}")
        
        specify_processor = SpecifyPreprocessing(shortcut_specify=test_case['params'])
        try:
            processed_data = specify_processor(test_data.copy(), 'label')
            
            print("\n处理后的数据:")
            print(processed_data)
            print(f"\n处理后的数据形状: {processed_data.shape}")
        except Exception as e:
            print(f"\n处理出错: {e}")
    
    print("\n" + "=" * 60)
    print("测试完成！")
    print("=" * 60)
    
    # 深度渗透测试套件
    print(f"\n{'=' * 80}")
    print("深度渗透测试套件")
    print('=' * 80)
    
    # 创建更大的测试数据
    np.random.seed(42)
    large_test_data = pd.DataFrame({
        'col1': np.random.randint(1, 1000, 100),
        'col2': np.random.randint(1, 1000, 100),
        'col3': np.random.randint(1, 1000, 100),
        'col4': np.random.randint(1, 1000, 100),
        'col5': np.random.randint(1, 1000, 100),
        'label': np.random.choice(['A', 'B', 'C'], 100)
    })
    
    # 测试1: 边界条件测试 - 删除不存在的列
    print(f"\n{'-' * 60}")
    print("测试1: 边界条件测试 - 删除不存在的列")
    print('-' * 60)
    
    params_nonexistent_col = {'drop_col': ['nonexistent_col', 'col1']}
    specify_processor_1 = SpecifyPreprocessing(shortcut_specify=params_nonexistent_col)
    result_1 = specify_processor_1(large_test_data.copy(), 'label')
    
    print(f"\n数据形状变化: {large_test_data.shape} -> {result_1.shape}")
    print("期望：只删除存在的 col1，忽略不存在的列")
    
    # 测试2: 边界条件测试 - 删除不存在的行
    print(f"\n{'-' * 60}")
    print("测试2: 边界条件测试 - 删除不存在的行")
    print('-' * 60)
    
    params_nonexistent_row = {'drop_row': [999, -1, 0, 1, 2]}
    specify_processor_2 = SpecifyPreprocessing(shortcut_specify=params_nonexistent_row)
    result_2 = specify_processor_2(large_test_data.copy(), 'label')
    
    print(f"\n数据形状变化: {large_test_data.shape} -> {result_2.shape}")
    print("期望：只删除存在的行 0,1,2，忽略不存在的行")
    
    # 测试3: 边界条件测试 - 空参数
    print(f"\n{'-' * 60}")
    print("测试3: 边界条件测试 - 空参数")
    print('-' * 60)
    
    params_empty = {}
    specify_processor_3 = SpecifyPreprocessing(shortcut_specify=params_empty)
    result_3 = specify_processor_3(large_test_data.copy(), 'label')
    
    print(f"\n数据形状变化: {large_test_data.shape} -> {result_3.shape}")
    print("期望：数据保持不变，没有任何删除")
    
    # 测试4: 边界条件测试 - 所有行删除
    print(f"\n{'-' * 60}")
    print("测试4: 边界条件测试 - 所有行删除")
    print('-' * 60)
    
    params_all_rows = {'drop_row': list(range(100))}
    specify_processor_4 = SpecifyPreprocessing(shortcut_specify=params_all_rows)
    result_4 = specify_processor_4(large_test_data.copy(), 'label')
    
    print(f"\n数据形状变化: {large_test_data.shape} -> {result_4.shape}")
    print("期望：所有行被删除，返回空DataFrame")
    
    # 测试5: 边界条件测试 - 所有列删除（除了标签列）
    print(f"\n{'-' * 60}")
    print("测试5: 边界条件测试 - 所有列删除（除了标签列）")
    print('-' * 60)
    
    params_all_cols = {'drop_col': ['col1', 'col2', 'col3', 'col4', 'col5']}
    specify_processor_5 = SpecifyPreprocessing(shortcut_specify=params_all_cols)
    result_5 = specify_processor_5(large_test_data.copy(), 'label')
    
    print(f"\n数据形状变化: {large_test_data.shape} -> {result_5.shape}")
    print("期望：只保留标签列，其他列被删除")
    
    # 测试6: 组合测试 - 同时删除行和列
    print(f"\n{'-' * 60}")
    print("测试6: 组合测试 - 同时删除行和列")
    print('-' * 60)
    
    params_combined = {
        'drop_row': [0, 1, 2, 3, 4],
        'drop_col': ['col1', 'col2']
    }
    specify_processor_6 = SpecifyPreprocessing(shortcut_specify=params_combined)
    result_6 = specify_processor_6(large_test_data.copy(), 'label')
    
    print(f"\n数据形状变化: {large_test_data.shape} -> {result_6.shape}")
    print("期望：删除前5行和前2列")
    
    # 测试7: 字符串格式参数测试
    print(f"\n{'-' * 60}")
    print("测试7: 字符串格式参数测试")
    print('-' * 60)
    
    params_string = {
        'drop_row': '0, 1, 2, 3, 4',
        'drop_col': 'col1, col2, col3'
    }
    specify_processor_7 = SpecifyPreprocessing(shortcut_specify=params_string)
    result_7 = specify_processor_7(large_test_data.copy(), 'label')
    
    print(f"\n数据形状变化: {large_test_data.shape} -> {result_7.shape}")
    print("期望：删除前5行和前3列（使用逗号分隔字符串格式）")
    
    # 测试8: 单一参数测试
    print(f"\n{'-' * 60}")
    print("测试8: 单一参数测试")
    print('-' * 60)
    
    params_single = {
        'drop_row': 5,
        'drop_col': 'col1'
    }
    specify_processor_8 = SpecifyPreprocessing(shortcut_specify=params_single)
    result_8 = specify_processor_8(large_test_data.copy(), 'label')
    
    print(f"\n数据形状变化: {large_test_data.shape} -> {result_8.shape}")
    print("期望：删除第5行和col1列")
    
    # 测试9: 空列表参数测试
    print(f"\n{'-' * 60}")
    print("测试9: 空列表参数测试")
    print('-' * 60)
    
    params_empty_list = {
        'drop_row': [],
        'drop_col': []
    }
    specify_processor_9 = SpecifyPreprocessing(shortcut_specify=params_empty_list)
    result_9 = specify_processor_9(large_test_data.copy(), 'label')
    
    print(f"\n数据形状变化: {large_test_data.shape} -> {result_9.shape}")
    print("期望：数据保持不变")
    
    # 测试10: 混合数据类型测试
    print(f"\n{'-' * 60}")
    print("测试10: 混合数据类型测试")
    print('-' * 60)
    
    mixed_test_data = pd.DataFrame({
        'int_col': [1, 2, 3, 4, 5],
        'float_col': [1.1, 2.2, 3.3, 4.4, 5.5],
        'str_col': ['a', 'b', 'c', 'd', 'e'],
        'bool_col': [True, False, True, False, True],
        'label': ['X', 'Y', 'X', 'Y', 'X']
    })
    
    params_mixed = {
        'drop_row': [0, 2],
        'drop_col': ['str_col', 'bool_col']
    }
    specify_processor_10 = SpecifyPreprocessing(shortcut_specify=params_mixed)
    result_10 = specify_processor_10(mixed_test_data.copy(), 'label')
    
    print(f"\n原始数据类型: {list(mixed_test_data.dtypes)}")
    print(f"处理后数据类型: {list(result_10.dtypes)}")
    print(f"\n数据形状变化: {mixed_test_data.shape} -> {result_10.shape}")
    print("期望：删除指定的行和列，其他列的数据类型保持不变")
    
    print(f"\n{'=' * 80}")
    print("深度渗透测试完成！")
    print('=' * 80)
