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):
    """
    Evaluate dataset metrics and calculate comprehensive score
    
    Args:
        df: Processed dataset (pandas DataFrame)
        original_df: Original dataset (pandas DataFrame), used for comparing processing effects
        label_column: Label column name, used for evaluating class balance
    
    Returns:
        dict: Dictionary containing various metric scores and comprehensive score
    """
    scores = {}
    
    # 1. Data completeness metrics
    # Missing value ratio
    total_cells = df.shape[0] * df.shape[1]
    missing_cells = df.isnull().sum().sum()
    missing_ratio = missing_cells / total_cells
    # Lower missing ratio gets higher score
    scores['missing_ratio_score'] = max(0, 100 - missing_ratio * 100)
    
    # Number of rows
    if original_df is not None:
        row_keeping_ratio = len(df) / len(original_df)
        # Higher row keeping ratio gets higher score
        scores['row_keeping_score'] = min(100, row_keeping_ratio * 100)
    else:
        # If no original data, score based on absolute row count
        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
    
    # Number of columns
    scores['column_count_score'] = min(100, df.shape[1] * 5)  # 5 points per column, max 100 points
    
    # 2. Data quality metrics
    # Outlier handling effect (using IQR method to detect outliers)
    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
        # Lower outlier ratio gets higher score
        scores['outlier_score'] = max(0, 100 - outlier_ratio * 100)
    else:
        scores['outlier_score'] = 100
    
    # Data consistency (data type consistency)
    # Check if each column has consistent data types
    consistency_score = 100
    for col in df.columns:
        # Calculate number of data types for non-null values
        non_null_values = df[col].dropna()
        if len(non_null_values) > 0:
            # Check if all values are the same type
            types = non_null_values.apply(type).nunique()
            if types > 1:
                consistency_score -= 10  # Deduct 10 points for each inconsistent type
    scores['consistency_score'] = max(0, consistency_score)
    
    # 3. Feature quality metrics
    # Feature diversity (calculate unique value ratio for each column)
    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
    
    # Feature correlation (calculate correlation coefficients between numeric features)
    if len(numeric_cols) >= 2:
        corr_matrix = df[numeric_cols].corr()
        # Exclude diagonal (self-correlation)
        upper_tri = corr_matrix.where(np.triu(np.ones(corr_matrix.shape), k=1).astype(bool))
        # Calculate average absolute correlation coefficient
        mean_abs_corr = upper_tri.stack().abs().mean()
        # Lower correlation is better (avoid multicollinearity)
        scores['correlation_score'] = max(0, 100 - mean_abs_corr * 100)
    else:
        scores['correlation_score'] = 100
    
    # Feature information content (calculate variance of numeric features)
    variance_scores = []
    for col in numeric_cols:
        if len(df[col].dropna()) > 0:
            var = df[col].var()
            # Normalize variance score
            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. Data balance metrics
    # Class balance (if there is a label column)
    if label_column and label_column in df.columns:
        label_counts = df[label_column].value_counts()
        if len(label_counts) > 1:
            # Use Gini coefficient to evaluate balance
            total = len(df[label_column])
            proportions = label_counts / total
            gini = 1 - sum(proportions ** 2)
            # Gini coefficient closer to 1 (perfect balance) gets higher score
            scores['balance_score'] = gini * 100
        else:
            scores['balance_score'] = 50  # Only one class, give medium score
    else:
        scores['balance_score'] = 100  # No label column, give full score
    
    # 5. Feature range metrics
    # Check if numeric feature ranges are reasonable
    range_scores = []
    for col in numeric_cols:
        if len(df[col].dropna()) > 0:
            col_min = df[col].min()
            col_max = df[col].max()
            # Check for extreme values
            if col_max - col_min > 0:
                # Normalize range score
                range_score = min(100, 100 - abs(np.log10(col_max - col_min)) * 10)
                range_scores.append(range_score)
            else:
                range_scores.append(50)  # All values the same, give medium score
    
    if range_scores:
        scores['range_score'] = np.mean(range_scores)
    else:
        scores['range_score'] = 100
    
    # 6. Model compatibility metrics
    # Feature types
    # Check for feature types that are not suitable for models
    incompatible_types = 0
    for col in df.columns:
        col_dtype = df[col].dtype
        if col_dtype == 'object':
            # Check if it's a pure string column
            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  # Deduct 50 points for each incompatible type
        scores['compatibility_score'] = max(0, compatibility_score)
    else:
        scores['compatibility_score'] = 100
    
    # Calculate comprehensive score
    # Assign weights to different metrics
    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
    }
    
    # Calculate weighted average
    total_score = sum(scores[key] * weights[key] for key in scores)
    
    # Return result
    result = {
        'detailed_scores': scores,
        'total_score': round(total_score, 2),
        'data_shape': df.shape
    }
    
    return result


def print_evaluation_result(result):
    """
    Print evaluation result
    
    Args:
        result: Result dictionary returned by evaluate_dataset function
    """
    ColorPrinter.print_info("=" * 80)
    ColorPrinter.print_info("Dataset Evaluation Result")
    ColorPrinter.print_info("=" * 80)
    ColorPrinter.print_info(f"Total score: {result['total_score']}/100")
    ColorPrinter.print_info(f"Data shape: {result['data_shape']}")
    ColorPrinter.print_info("\nDetailed metric scores:")
    ColorPrinter.print_info("-" * 80)
    
    # Group by category
    categories = {
        'Data completeness': ['missing_ratio_score', 'row_keeping_score', 'column_count_score'],
        'Data quality': ['outlier_score', 'consistency_score'],
        'Feature quality': ['diversity_score', 'correlation_score', 'information_score'],
        'Data balance': ['balance_score', 'range_score'],
        'Model compatibility': ['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)
