import pandas as pd
import numpy as np
from scipy import stats
from sklearn.preprocessing import StandardScaler
from myffe.tools.color_printer import ColorPrinter


def evaluate_dataset(df, original_df=None, label_column=None):
    """
    评价数据集的各项指标并计算综合得分
    
    Args:
        df: 处理后的数据集 (pandas DataFrame)
        original_df: 原始数据集 (pandas DataFrame)，用于对比处理效果
        label_column: 标签列名称，用于评估类别平衡性
    
    Returns:
        dict: 包含各项指标得分和综合得分的字典
    """
    scores = {}
    
    # 1. 数据完整性指标
    # 缺失值比例
    total_cells = df.shape[0] * df.shape[1]
    missing_cells = df.isnull().sum().sum()
    missing_ratio = missing_cells / total_cells
    # 缺失值比例越低得分越高
    scores['missing_ratio_score'] = max(0, 100 - missing_ratio * 100)
    
    # 数据行数
    if original_df is not None:
        row_keeping_ratio = len(df) / len(original_df)
        # 保留的数据行数比例越高得分越高
        scores['row_keeping_score'] = min(100, row_keeping_ratio * 100)
    else:
        # 如果没有原始数据，基于绝对行数评分
        if len(df) >= 10000:
            scores['row_keeping_score'] = 100
        elif len(df) >= 1000:
            scores['row_keeping_score'] = 80
        elif len(df) >= 100:
            scores['row_keeping_score'] = 60
        else:
            scores['row_keeping_score'] = 40
    
    # 数据列数
    scores['column_count_score'] = min(100, df.shape[1] * 5)  # 每列5分，最高100分
    
    # 2. 数据质量指标
    # 异常值处理效果（使用IQR方法检测异常值）
    numeric_cols = df.select_dtypes(include=[np.number]).columns
    outlier_count = 0
    total_numeric_cells = 0
    
    for col in numeric_cols:
        if len(df[col].dropna()) > 0:
            Q1 = df[col].quantile(0.25)
            Q3 = df[col].quantile(0.75)
            IQR = Q3 - Q1
            lower_bound = Q1 - 1.5 * IQR
            upper_bound = Q3 + 1.5 * IQR
            outlier_count += ((df[col] < lower_bound) | (df[col] > upper_bound)).sum()
            total_numeric_cells += len(df[col].dropna())
    
    if total_numeric_cells > 0:
        outlier_ratio = outlier_count / total_numeric_cells
        # 异常值比例越低得分越高
        scores['outlier_score'] = max(0, 100 - outlier_ratio * 100)
    else:
        scores['outlier_score'] = 100
    
    # 数据一致性（数据类型一致性）
    # 检查每列的数据类型是否一致
    consistency_score = 100
    for col in df.columns:
        # 计算非空值的数据类型数量
        non_null_values = df[col].dropna()
        if len(non_null_values) > 0:
            # 检查是否所有值都是相同类型
            types = non_null_values.apply(type).nunique()
            if types > 1:
                consistency_score -= 10  # 每种不一致类型扣10分
    scores['consistency_score'] = max(0, consistency_score)
    
    # 3. 特征质量指标
    # 特征多样性（计算每列的唯一值比例）
    diversity_scores = []
    for col in df.columns:
        if len(df[col].dropna()) > 0:
            unique_ratio = df[col].nunique() / len(df[col].dropna())
            diversity_scores.append(min(1, unique_ratio) * 100)
    
    if diversity_scores:
        scores['diversity_score'] = np.mean(diversity_scores)
    else:
        scores['diversity_score'] = 0
    
    # 特征相关性（计算数值特征之间的相关系数）
    if len(numeric_cols) >= 2:
        corr_matrix = df[numeric_cols].corr()
        # 排除对角线（自相关）
        upper_tri = corr_matrix.where(np.triu(np.ones(corr_matrix.shape), k=1).astype(bool))
        # 计算平均绝对相关系数
        mean_abs_corr = upper_tri.stack().abs().mean()
        # 相关性越低越好（避免多重共线性）
        scores['correlation_score'] = max(0, 100 - mean_abs_corr * 100)
    else:
        scores['correlation_score'] = 100
    
    # 特征信息量（计算数值特征的方差）
    variance_scores = []
    for col in numeric_cols:
        if len(df[col].dropna()) > 0:
            var = df[col].var()
            # 标准化方差得分
            if var > 0:
                variance_scores.append(min(100, var * 10))
            else:
                variance_scores.append(0)
    
    if variance_scores:
        scores['information_score'] = np.mean(variance_scores)
    else:
        scores['information_score'] = 0
    
    # 4. 数据平衡指标
    # 类别平衡性（如果有标签列）
    if label_column and label_column in df.columns:
        label_counts = df[label_column].value_counts()
        if len(label_counts) > 1:
            # 使用基尼系数评估平衡性
            total = len(df[label_column])
            proportions = label_counts / total
            gini = 1 - sum(proportions ** 2)
            # 基尼系数越接近1（完全平衡）得分越高
            scores['balance_score'] = gini * 100
        else:
            scores['balance_score'] = 50  # 只有一个类别，给予中等分数
    else:
        scores['balance_score'] = 100  # 无标签列，给予满分
    
    # 5. 特征范围指标
    # 检查数值特征的范围是否合理
    range_scores = []
    for col in numeric_cols:
        if len(df[col].dropna()) > 0:
            col_min = df[col].min()
            col_max = df[col].max()
            # 检查是否有极端值
            if col_max - col_min > 0:
                # 标准化范围得分
                range_score = min(100, 100 - abs(np.log10(col_max - col_min)) * 10)
                range_scores.append(range_score)
            else:
                range_scores.append(50)  # 所有值相同，给予中等分数
    
    if range_scores:
        scores['range_score'] = np.mean(range_scores)
    else:
        scores['range_score'] = 100
    
    # 6. 模型兼容性指标
    # 特征类型
    # 检查是否有不适合模型的特征类型
    incompatible_types = 0
    for col in df.columns:
        col_dtype = df[col].dtype
        if col_dtype == 'object':
            # 检查是否为纯字符串列
            if all(isinstance(x, str) for x in df[col].dropna()):
                incompatible_types += 1
    
    total_cols = len(df.columns)
    if total_cols > 0:
        compatibility_score = 100 - (incompatible_types / total_cols) * 50  # 每个不兼容类型扣50分
        scores['compatibility_score'] = max(0, compatibility_score)
    else:
        scores['compatibility_score'] = 100
    
    # 计算综合得分
    # 为不同指标分配权重
    weights = {
        'missing_ratio_score': 0.15,
        'row_keeping_score': 0.1,
        'column_count_score': 0.1,
        'outlier_score': 0.1,
        'consistency_score': 0.1,
        'diversity_score': 0.1,
        'correlation_score': 0.05,
        'information_score': 0.1,
        'balance_score': 0.1,
        'range_score': 0.05,
        'compatibility_score': 0.05
    }
    
    # 计算加权平均分
    total_score = sum(scores[key] * weights[key] for key in scores)
    
    # 返回结果
    result = {
        'detailed_scores': scores,
        'total_score': round(total_score, 2),
        'data_shape': df.shape
    }
    
    return result


def print_evaluation_result(result):
    """
    打印评估结果
    
    Args:
        result: evaluate_dataset函数返回的结果字典
    """
    ColorPrinter.print_info("=" * 80)
    ColorPrinter.print_info("数据集评估结果")
    ColorPrinter.print_info("=" * 80)
    ColorPrinter.print_info(f"综合得分: {result['total_score']}/100")
    ColorPrinter.print_info(f"数据形状: {result['data_shape']}")
    ColorPrinter.print_info("\n详细指标得分:")
    ColorPrinter.print_info("-" * 80)
    
    # 按类别分组显示
    categories = {
        '数据完整性': ['missing_ratio_score', 'row_keeping_score', 'column_count_score'],
        '数据质量': ['outlier_score', 'consistency_score'],
        '特征质量': ['diversity_score', 'correlation_score', 'information_score'],
        '数据平衡': ['balance_score', 'range_score'],
        '模型兼容性': ['compatibility_score']
    }
    
    for category, metrics in categories.items():
        ColorPrinter.print_info(f"\n{category}:")
        for metric in metrics:
            if metric in result['detailed_scores']:
                score = result['detailed_scores'][metric]
                ColorPrinter.print_info(f"  • {metric.replace('_score', '')}: {score:.2f}/100")
    
    ColorPrinter.print_info("=" * 80)


if __name__ == "__main__":
    # 测试代码
    # 创建示例数据
    data = {
        'age': [25, 30, 35, 40, 45, 50, 55, 60, 65, 70],
        'income': [50000, 60000, 70000, 80000, 90000, 100000, 110000, 120000, 130000, 140000],
        'score': [85, 90, 75, 80, 95, 70, 85, 90, 80, 75],
        'label': ['A', 'B', 'A', 'B', 'A', 'B', 'A', 'B', 'A', 'B']
    }
    
    df = pd.DataFrame(data)
    
    # 评估数据集
    result = evaluate_dataset(df, label_column='label')
    
    # 打印结果
    print_evaluation_result(result)
