"""
Fetal Birth Weight Prediction - 30-Feature Pipeline (No Data Leakage)
=====================================================================
基于30个特征（17基础+13工程）的胎儿出生体重预测完整流程。
所有特征均从原始临床/超声测量值衍生，不含任何与出生体重直接相关的泄露特征。

Version: 1.0
Date: 2026-04-14

依赖:
    pip install numpy pandas scikit-learn matplotlib seaborn joblib shap xgboost lightgbm
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import (
    train_test_split, cross_val_score, StratifiedKFold, GridSearchCV
)
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import Ridge, Lasso, ElasticNet
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.metrics import (
    confusion_matrix, mean_absolute_error, mean_squared_error, r2_score
)
import json
import warnings
import os
import joblib
from typing import Dict, Tuple, List, Any

warnings.filterwarnings('ignore')

# ============================================================================
# 配置
# ============================================================================
RANDOM_STATE = 42
N_SPLITS = 5
TEST_SIZE = 0.2
OUTPUT_DIR = '/Users/wangkai/Desktop/kimi/2/30特征'
DATA_PATH = 'processed_data.csv'

np.random.seed(RANDOM_STATE)


# ============================================================================
# 数据加载与特征工程（30特征，无泄露）
# ============================================================================
LEAKED_FEATURES = [
    'birth_weight_category_encoded',
    'gestational_category_encoded',
    'comorbidity_category_encoded',
    'ultrasound_composite',
    'maternal_size_index',
    'gestational_delivery_gap',
    'FL_AC_ratio',
    'BPD_HC_ratio',
    'complication_score',
]

BASE_FEATURES = [
    'gestational_weeks', 'AC', 'BPD', 'HC', 'FL', 'AFI',
    'parity', 'prepregnancy_BMI', 'fundal_height_cm', 'hypertension',
    'diabetes', 'prepregnancy_weight_kg', 'height_cm', 'maternal_AC_cm',
    'placental_grade', 'delivery_mode', 'age'
]


def create_30_features(df: pd.DataFrame) -> pd.DataFrame:
    """
    创建30个无泄露特征（17基础 + 13工程）。
    所有特征仅基于原始临床/超声测量值计算，不涉及出生体重。
    """
    df_features = df[BASE_FEATURES].copy()

    # 1-2. 多项式特征
    df_features['AC_squared'] = df['AC'] ** 2
    df_features['HC_squared'] = df['HC'] ** 2

    # 3-5. 比率特征
    df_features['AC_FL_ratio'] = df['AC'] / df['FL']
    df_features['HC_AC_ratio'] = df['HC'] / df['AC']
    df_features['BPD_FL_ratio'] = df['BPD'] / df['FL']

    # 6-9. 交互特征
    df_features['GW_x_AC'] = df['gestational_weeks'] * df['AC']
    df_features['GW_x_FL'] = df['gestational_weeks'] * df['FL']
    df_features['AC_x_BMI'] = df['AC'] * df['prepregnancy_BMI']
    df_features['FH_x_AC'] = df['fundal_height_cm'] * df['AC']

    # 10-11. 复合特征
    df_features['AC_HC_FL_composite'] = (df['AC'] + df['HC']) / df['FL']
    df_features['Estimated_Volume'] = (df['AC'] ** 2) * df['FL'] / 1000

    # 12-13. 母体-胎儿特征
    df_features['Maternal_Fetal_AC_ratio'] = df['maternal_AC_cm'] / df['AC']
    df_features['Weight_Height_Ratio'] = df['prepregnancy_weight_kg'] / df['height_cm']

    return df_features


def load_and_preprocess_data(filepath: str = DATA_PATH) -> Tuple[pd.DataFrame, pd.DataFrame, pd.Series]:
    """加载数据，确保无泄露特征混入，返回30特征矩阵和目标变量。"""
    df = pd.read_csv(filepath)

    # 检查泄露特征是否意外存在于BASE_FEATURES中
    for leaked in LEAKED_FEATURES:
        if leaked in BASE_FEATURES:
            raise ValueError(f"CRITICAL: Leaked feature '{leaked}' found in BASE_FEATURES!")

    X = create_30_features(df)
    y = df['birth_weight_g']

    # 再次确认X中没有泄露列
    for leaked in LEAKED_FEATURES:
        if leaked in X.columns:
            raise ValueError(f"CRITICAL: Leaked feature '{leaked}' found in engineered features!")

    print(f"[INFO] Loaded {len(df)} samples with {len(X.columns)} features (no leakage)")
    print(f"[INFO] Features: {list(X.columns)}")
    return df, X, y


# ============================================================================
# 模型训练与评估
# ============================================================================
def train_model_with_tuning(X_train: pd.DataFrame, y_train: pd.Series) -> Tuple[Ridge, StandardScaler, Dict]:
    """使用GridSearchCV训练Ridge模型并进行超参数调优。"""
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)

    param_grid = {'alpha': [0.0001, 0.001, 0.01, 0.1, 0.5, 1.0, 2.0, 5.0, 10.0, 50.0, 100.0]}
    grid_search = GridSearchCV(
        Ridge(random_state=RANDOM_STATE),
        param_grid,
        cv=StratifiedKFold(n_splits=5, shuffle=True, random_state=RANDOM_STATE),
        scoring='neg_mean_absolute_error',
        n_jobs=-1
    )

    # 用于StratifiedKFold，需要按出生体重分层
    y_class_train = pd.cut(y_train, bins=[0, 2500, 4000, 10000], labels=[0, 1, 2]).astype(int)
    grid_search.fit(X_train_scaled, y_train)

    best_alpha = grid_search.best_params_['alpha']
    best_score = -grid_search.best_score_

    print(f"[INFO] Optimal alpha: {best_alpha}")
    print(f"[INFO] Best CV MAE: {best_score:.2f} g")

    model = Ridge(alpha=best_alpha, random_state=RANDOM_STATE)
    model.fit(X_train_scaled, y_train)

    tuning_info = {
        'optimal_alpha': float(best_alpha),
        'cv_mae': float(best_score),
        'all_cv_results': {str(k): float(-v) for k, v in zip(param_grid['alpha'], grid_search.cv_results_['mean_test_score'])}
    }
    return model, scaler, tuning_info


def evaluate_regression_cv_stratified(X: pd.DataFrame, y: pd.Series, alpha: float) -> Dict:
    """使用StratifiedKFold评估回归性能。"""
    y_class = pd.cut(y, bins=[0, 2500, 4000, 10000], labels=[0, 1, 2]).astype(int)
    skf = StratifiedKFold(n_splits=N_SPLITS, shuffle=True, random_state=RANDOM_STATE)
    maes, rmses, r2s, mapes = [], [], [], []

    for fold, (train_idx, val_idx) in enumerate(skf.split(X, y_class), 1):
        X_train, X_val = X.iloc[train_idx], X.iloc[val_idx]
        y_train, y_val = y.iloc[train_idx], y.iloc[val_idx]

        scaler = StandardScaler()
        X_train_s = scaler.fit_transform(X_train)
        X_val_s = scaler.transform(X_val)

        model = Ridge(alpha=alpha, random_state=RANDOM_STATE)
        model.fit(X_train_s, y_train)
        y_pred = model.predict(X_val_s)

        maes.append(mean_absolute_error(y_val, y_pred))
        rmses.append(np.sqrt(mean_squared_error(y_val, y_pred)))
        r2s.append(r2_score(y_val, y_pred))
        mapes.append(np.mean(np.abs((y_val - y_pred) / y_val)) * 100)

    return {
        'cv_mae_mean': float(np.mean(maes)),
        'cv_mae_std': float(np.std(maes)),
        'cv_rmse_mean': float(np.mean(rmses)),
        'cv_rmse_std': float(np.std(rmses)),
        'cv_r2_mean': float(np.mean(r2s)),
        'cv_r2_std': float(np.std(r2s)),
        'cv_mape_mean': float(np.mean(mapes)),
        'cv_mape_std': float(np.std(mapes)),
        'fold_results': [{'fold': i+1, 'mae': m, 'rmse': r, 'r2': r2, 'mape': mp}
                         for i, (m, r, r2, mp) in enumerate(zip(maes, rmses, r2s, mapes))]
    }


def evaluate_regression_test(y_true: np.ndarray, y_pred: np.ndarray) -> Dict[str, float]:
    """测试集回归评估。"""
    return {
        'mae': float(mean_absolute_error(y_true, y_pred)),
        'rmse': float(np.sqrt(mean_squared_error(y_true, y_pred))),
        'r2': float(r2_score(y_true, y_pred)),
        'mape': float(np.mean(np.abs((y_true - y_pred) / y_true)) * 100),
    }


def compute_bootstrap_ci(X_train, y_train, X_test, y_test, models_dict, n_bootstrap=1000):
    """对所有模型计算Bootstrap 95% CI。"""
    from sklearn.base import clone

    results = []
    rng = np.random.RandomState(RANDOM_STATE)
    n_test = len(y_test)

    for name, model_template in models_dict.items():
        print(f"  Bootstrapping {name}...")
        # 先拟合点估计
        scaler = StandardScaler()
        X_train_s = scaler.fit_transform(X_train)
        X_test_s = scaler.transform(X_test)
        model = clone(model_template)
        model.fit(X_train_s, y_train)
        y_pred_point = model.predict(X_test_s)
        point = evaluate_regression_test(y_test.values, y_pred_point)

        # Bootstrap
        boot_maes, boot_rmses, boot_r2s, boot_mapes = [], [], [], []
        for _ in range(n_bootstrap):
            idx = rng.choice(n_test, size=n_test, replace=True)
            y_true_b = y_test.values[idx]
            y_pred_b = y_pred_point[idx]
            boot_maes.append(mean_absolute_error(y_true_b, y_pred_b))
            boot_rmses.append(np.sqrt(mean_squared_error(y_true_b, y_pred_b)))
            boot_r2s.append(r2_score(y_true_b, y_pred_b))
            boot_mapes.append(np.mean(np.abs((y_true_b - y_pred_b) / y_true_b)) * 100)

        def percentile_ci(arr):
            return float(np.percentile(arr, 2.5)), float(np.percentile(arr, 97.5))

        mae_ci = percentile_ci(boot_maes)
        rmse_ci = percentile_ci(boot_rmses)
        r2_ci = percentile_ci(boot_r2s)
        mape_ci = percentile_ci(boot_mapes)

        results.append({
            'model': name,
            'test_mae_point': point['mae'],
            'test_mae_ci_lower': mae_ci[0], 'test_mae_ci_upper': mae_ci[1],
            'test_rmse_point': point['rmse'],
            'test_rmse_ci_lower': rmse_ci[0], 'test_rmse_ci_upper': rmse_ci[1],
            'test_r2_point': point['r2'],
            'test_r2_ci_lower': r2_ci[0], 'test_r2_ci_upper': r2_ci[1],
            'test_mape_point': point['mape'],
            'test_mape_ci_lower': mape_ci[0], 'test_mape_ci_upper': mape_ci[1],
        })

    return results


# ============================================================================
# 分类策略
# ============================================================================
def classify_by_percentiles(y_pred: np.ndarray, y_train_pred: np.ndarray, fh_values: np.ndarray) -> np.ndarray:
    """
    基于训练集百分位数的分类策略 (无数据泄露):
    - < P10: LBW (0)
    - > P90 或 (P80-P90 且 宫高>36): Macrosomia (2)
    - 其余: Normal (1)
    """
    p10 = np.percentile(y_train_pred, 10)
    p90 = np.percentile(y_train_pred, 90)
    p80 = np.percentile(y_train_pred, 80)

    y_class = np.ones(len(y_pred), dtype=int)
    y_class[y_pred < p10] = 0

    high_fh = fh_values > 36
    macro_condition = (y_pred > p90) | ((y_pred > p80) & high_fh)
    y_class[macro_condition] = 2

    return y_class


def evaluate_classification(y_true: np.ndarray, y_pred: np.ndarray) -> Dict[str, Any]:
    """分类性能评估。"""
    cm = confusion_matrix(y_true, y_pred, labels=[0, 1, 2])
    results = {
        'confusion_matrix': cm.tolist(),
        'accuracy': float(np.sum(np.diag(cm)) / np.sum(cm)),
    }

    tp_lbw, fn_lbw = cm[0, 0], cm[0, 1] + cm[0, 2]
    fp_lbw = cm[1, 0] + cm[2, 0]
    results['lbw_recall'] = float(tp_lbw / (tp_lbw + fn_lbw)) if (tp_lbw + fn_lbw) > 0 else 0.0
    results['lbw_precision'] = float(tp_lbw / (tp_lbw + fp_lbw)) if (tp_lbw + fp_lbw) > 0 else 0.0

    tp_macro, fn_macro = cm[2, 2], cm[2, 0] + cm[2, 1]
    fp_macro = cm[0, 2] + cm[1, 2]
    results['macro_recall'] = float(tp_macro / (tp_macro + fn_macro)) if (tp_macro + fn_macro) > 0 else 0.0
    results['macro_precision'] = float(tp_macro / (tp_macro + fp_macro)) if (tp_macro + fp_macro) > 0 else 0.0

    tn_normal = cm[1, 1]
    fp_normal = cm[1, 0] + cm[1, 2]
    results['fp_rate'] = float(fp_normal / (fp_normal + tn_normal)) if (fp_normal + tn_normal) > 0 else 0.0
    results['specificity'] = 1.0 - results['fp_rate']

    return results


def compare_classification_strategies(X_train, X_test, y_train, y_test,
                                      y_class_train, y_class_test, model, scaler):
    """对比多种分类策略。"""
    X_train_s = scaler.transform(X_train)
    X_test_s = scaler.transform(X_test)

    y_train_pred = model.predict(X_train_s)
    y_test_pred = model.predict(X_test_s)
    fh_test = X_test['fundal_height_cm'].values

    p10 = np.percentile(y_train_pred, 10)
    p90 = np.percentile(y_train_pred, 90)
    p80 = np.percentile(y_train_pred, 80)

    strategies = {}

    # 硬阈值
    y_hard = np.where(y_test_pred < 2500, 0, np.where(y_test_pred > 4000, 2, 1))
    strategies['hard_threshold'] = evaluate_classification(y_class_test, y_hard)
    strategies['hard_threshold']['name'] = "Hard Threshold"

    # P10/P90 only
    y_p10p90 = np.where(y_test_pred < p10, 0, np.where(y_test_pred > p90, 2, 1))
    strategies['p10p90_only'] = evaluate_classification(y_class_test, y_p10p90)
    strategies['p10p90_only']['name'] = "P10/P90 Only"

    # P10/P90 + FH (推荐)
    y_fh = classify_by_percentiles(y_test_pred, y_train_pred, fh_test)
    strategies['p10p90_fh'] = evaluate_classification(y_class_test, y_fh)
    strategies['p10p90_fh']['name'] = "P10/P90 + Fundal Height (Recommended)"

    # 仅宫高
    y_fh_only = np.where(fh_test < 30, 0, np.where(fh_test > 36, 2, 1))
    strategies['fh_only'] = evaluate_classification(y_class_test, y_fh_only)
    strategies['fh_only']['name'] = "Fundal Height Only"

    return strategies, {'p10': p10, 'p80': p80, 'p90': p90}


# ============================================================================
# 特征重要性
# ============================================================================
def calculate_feature_importance(model: Ridge, feature_names: List[str]) -> pd.DataFrame:
    """基于Ridge系数计算归一化特征重要性。"""
    coefficients = model.coef_
    abs_coef = np.abs(coefficients)
    importance_scores = (abs_coef / abs_coef.max()) * 100

    return pd.DataFrame({
        'Feature': feature_names,
        'Coefficient': coefficients,
        'Abs_Coefficient': abs_coef,
        'Importance_Score': importance_scores
    }).sort_values('Importance_Score', ascending=False)


# ============================================================================
# 扩展模型对比（30特征，相同数据划分）
# ============================================================================
def compare_models_30_features(X_train, y_train, X_test, y_test) -> pd.DataFrame:
    """对比7个模型在30特征上的性能。"""
    models = {
        'Ridge': Ridge(alpha=5.0, random_state=RANDOM_STATE),
        'Lasso': Lasso(alpha=1.0, random_state=RANDOM_STATE, max_iter=10000),
        'ElasticNet': ElasticNet(alpha=1.0, l1_ratio=0.5, random_state=RANDOM_STATE, max_iter=10000),
        'RandomForest': RandomForestRegressor(n_estimators=300, max_depth=12, min_samples_split=5, random_state=RANDOM_STATE, n_jobs=-1),
        'GradientBoosting': GradientBoostingRegressor(n_estimators=300, max_depth=5, learning_rate=0.1, random_state=RANDOM_STATE),
    }
    try:
        import xgboost as xgb
        models['XGBoost'] = xgb.XGBRegressor(n_estimators=300, max_depth=6, learning_rate=0.1, subsample=0.8, random_state=RANDOM_STATE, n_jobs=-1)
    except ImportError:
        pass
    try:
        import lightgbm as lgb
        models['LightGBM'] = lgb.LGBMRegressor(n_estimators=300, max_depth=6, learning_rate=0.1, subsample=0.8, random_state=RANDOM_STATE, n_jobs=-1, verbose=-1)
    except ImportError:
        pass

    results = []
    y_class = pd.cut(y_train, bins=[0, 2500, 4000, 10000], labels=[0, 1, 2]).astype(int)

    for name, model in models.items():
        cv_scores = cross_val_score(model, X_train, y_train,
                                     cv=StratifiedKFold(n_splits=5, shuffle=True, random_state=RANDOM_STATE),
                                     scoring='neg_mean_absolute_error', n_jobs=-1)
        cv_mae = -cv_scores.mean()
        cv_std = cv_scores.std()
        model.fit(X_train, y_train)
        y_pred = model.predict(X_test)
        test_mae = mean_absolute_error(y_test, y_pred)
        test_rmse = np.sqrt(mean_squared_error(y_test, y_pred))
        test_r2 = r2_score(y_test, y_pred)
        test_mape = np.mean(np.abs((y_test - y_pred) / y_test)) * 100
        results.append({
            'Model': name, 'CV_MAE_Mean': cv_mae, 'CV_MAE_Std': cv_std,
            'Test_MAE': test_mae, 'Test_RMSE': test_rmse, 'Test_R2': test_r2, 'Test_MAPE': test_mape
        })
        print(f"  {name:<18} CV MAE: {cv_mae:.2f}+-{cv_std:.2f}g | Test MAE: {test_mae:.2f}g | R2: {test_r2:.4f}")

    return pd.DataFrame(results).sort_values('Test_MAE')


# ============================================================================
# 主流程
# ============================================================================
def main():
    """执行完整分析流程。"""
    print("=" * 70)
    print("Fetal Birth Weight Prediction - 30-Feature Pipeline v1.0")
    print("=" * 70)

    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # 1. 加载数据（30特征，无泄露）
    print("\n[1/8] Loading data with 30 features (no leakage)...")
    df, X, y = load_and_preprocess_data()
    y_class = pd.cut(y, bins=[0, 2500, 4000, 10000], labels=[0, 1, 2]).astype(int)
    feature_names = list(X.columns)

    # 2. 划分数据集
    print("\n[2/8] Splitting data...")
    X_train, X_test, y_train, y_test, y_class_train, y_class_test = train_test_split(
        X, y, y_class, test_size=TEST_SIZE, random_state=RANDOM_STATE, stratify=y_class)
    print(f"  Train: {len(X_train)}, Test: {len(X_test)}")

    # 3. 训练模型 (带调参)
    print("\n[3/8] Training Ridge model with hyperparameter tuning...")
    model, scaler, tuning_info = train_model_with_tuning(X_train, y_train)
    best_alpha = tuning_info['optimal_alpha']

    # 4. 交叉验证
    print("\n[4/8] Running stratified 5-fold cross-validation...")
    cv_results = evaluate_regression_cv_stratified(X, y, alpha=best_alpha)
    print(f"  CV MAE: {cv_results['cv_mae_mean']:.2f} +- {cv_results['cv_mae_std']:.2f} g")
    print(f"  CV R2:  {cv_results['cv_r2_mean']:.4f} +- {cv_results['cv_r2_std']:.4f}")

    X_test_s = scaler.transform(X_test)
    y_test_pred = model.predict(X_test_s)
    test_results = evaluate_regression_test(y_test.values, y_test_pred)
    print(f"  Test MAE: {test_results['mae']:.2f} g")
    print(f"  Test R2:  {test_results['r2']:.4f}")

    # 5. 分类评估
    print("\n[5/8] Evaluating classification performance...")
    strategies, thresholds = compare_classification_strategies(
        X_train, X_test, y_train, y_test, y_class_train, y_class_test, model, scaler)
    print(f"\n  Thresholds: P10={thresholds['p10']:.0f}g, P90={thresholds['p90']:.0f}g")
    print(f"\n  {'Strategy':<45} {'Accuracy':>10} {'LBW Rec':>10} {'Macro Rec':>10}")
    print(f"  {'-'*75}")
    for key, strat in strategies.items():
        print(f"  {strat['name']:<45} {strat['accuracy']:>9.1%} {strat['lbw_recall']:>9.1%} {strat['macro_recall']:>9.1%}")
    recommended = strategies['p10p90_fh']

    # 6. 特征重要性
    print("\n[6/8] Feature importance...")
    importance_df = calculate_feature_importance(model, feature_names)
    print("  Top 5 features:")
    for _, row in importance_df.head(5).iterrows():
        print(f"    {row['Feature']}: {row['Importance_Score']:.1f}%")

    # 7. 7模型横向对比（含Bootstrap CI）
    print("\n[7/8] Comparing 7 models on identical 30 features with Bootstrap 95% CI...")
    scaler_all = StandardScaler()
    X_train_scaled = scaler_all.fit_transform(X_train)
    X_test_scaled = scaler_all.transform(X_test)

    ext_results_df = compare_models_30_features(X_train_scaled, y_train, X_test_scaled, y_test)
    ext_results_df.to_csv(f'{OUTPUT_DIR}/model_comparison_30features.csv', index=False)
    print(f"  Saved: {OUTPUT_DIR}/model_comparison_30features.csv")

    # Bootstrap CI
    models_for_bootstrap = {
        'Ridge': Ridge(alpha=5.0, random_state=RANDOM_STATE),
        'Lasso': Lasso(alpha=1.0, random_state=RANDOM_STATE, max_iter=10000),
        'ElasticNet': ElasticNet(alpha=1.0, l1_ratio=0.5, random_state=RANDOM_STATE, max_iter=10000),
        'RandomForest': RandomForestRegressor(n_estimators=300, max_depth=12, min_samples_split=5, random_state=RANDOM_STATE, n_jobs=-1),
        'GradientBoosting': GradientBoostingRegressor(n_estimators=300, max_depth=5, learning_rate=0.1, random_state=RANDOM_STATE),
    }
    try:
        import xgboost as xgb
        models_for_bootstrap['XGBoost'] = xgb.XGBRegressor(n_estimators=300, max_depth=6, learning_rate=0.1, subsample=0.8, random_state=RANDOM_STATE, n_jobs=-1)
    except ImportError:
        pass
    try:
        import lightgbm as lgb
        models_for_bootstrap['LightGBM'] = lgb.LGBMRegressor(n_estimators=300, max_depth=6, learning_rate=0.1, subsample=0.8, random_state=RANDOM_STATE, n_jobs=-1, verbose=-1)
    except ImportError:
        pass

    bootstrap_results = compute_bootstrap_ci(
        pd.DataFrame(X_train_scaled, columns=feature_names), y_train,
        pd.DataFrame(X_test_scaled, columns=feature_names), y_test,
        models_for_bootstrap, n_bootstrap=1000)

    with open(f'{OUTPUT_DIR}/bootstrap_ci_all_models.json', 'w', encoding='utf-8') as f:
        json.dump({'n_bootstrap': 1000, 'results': bootstrap_results}, f, indent=2, ensure_ascii=False)
    print(f"  Saved: {OUTPUT_DIR}/bootstrap_ci_all_models.json")

    # 8. 保存结果
    print("\n[8/8] Saving all results...")

    all_results = {
        'experiment_info': {
            'version': '1.0 (30-Feature No Leakage)',
            'date': '2026-04-14',
            'random_state': RANDOM_STATE,
            'n_samples': len(df),
            'n_features': len(feature_names),
            'features': feature_names,
            'leaked_features_excluded': LEAKED_FEATURES,
        },
        'tuning_info': tuning_info,
        'regression_cv': cv_results,
        'regression_test': test_results,
        'classification_strategies': strategies,
        'recommended_strategy': recommended,
        'thresholds': {k: float(v) for k, v in thresholds.items()},
        'feature_importance': importance_df.to_dict('records'),
    }

    with open(f'{OUTPUT_DIR}/final_metrics.json', 'w', encoding='utf-8') as f:
        json.dump(all_results, f, indent=2, ensure_ascii=False)
    print(f"  Saved: {OUTPUT_DIR}/final_metrics.json")

    importance_df.to_csv(f'{OUTPUT_DIR}/feature_importance.csv', index=False)
    print(f"  Saved: {OUTPUT_DIR}/feature_importance.csv")

    joblib.dump(model, f'{OUTPUT_DIR}/best_model_30features.pkl')
    joblib.dump(scaler_all, f'{OUTPUT_DIR}/scaler_30features.pkl')
    print(f"  Saved: {OUTPUT_DIR}/best_model_30features.pkl")
    print(f"  Saved: {OUTPUT_DIR}/scaler_30features.pkl")

    # 生成用于后续分析的测试集预测文件
    test_pred_df = pd.DataFrame({
        'actual': y_test.values,
        'predicted': y_test_pred,
        'fundal_height_cm': X_test['fundal_height_cm'].values,
        'birth_weight_category': y_class_test.values,
    })
    test_pred_df.to_csv(f'{OUTPUT_DIR}/test_predictions.csv', index=False)
    print(f"  Saved: {OUTPUT_DIR}/test_predictions.csv")

    print("\n" + "=" * 70)
    print("30-Feature Pipeline completed successfully!")
    print("=" * 70)
    return all_results


if __name__ == '__main__':
    results = main()