"""
Stratified Analysis - 30 Features with Hadlock 4 Formula
=========================================================
对比Hadlock 4公式与30特征ML模型在不同超声医师经验和出生体重分层中的预测误差。
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy import stats
import joblib
import os
import sys

# 导入主脚本中的特征工程函数
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from fetal_weight_prediction_30features import create_30_features, BASE_FEATURES

RANDOM_STATE = 42
OUTPUT_DIR = '/Users/wangkai/Desktop/kimi/2/30特征'

# ============================================================================
# Hadlock 4 公式（最新版，使用BPD+HC+AC+FL）
# ============================================================================
def hadlock_4_efw(bpd_mm, hc_mm, ac_mm, fl_mm):
    """
    Hadlock 4 公式 (EFW in grams)
    log10(weight) = 1.3596 - 0.00386*AC*FL + 0.0064*HC + 0.00061*BPD*AC + 0.0424*AC + 0.174*FL
    输入单位为 mm，内部转换为 cm。
    """
    bpd_cm = bpd_mm / 10.0
    hc_cm = hc_mm / 10.0
    ac_cm = ac_mm / 10.0
    fl_cm = fl_mm / 10.0

    log_efw = (1.3596
               - 0.00386 * ac_cm * fl_cm
               + 0.0064 * hc_cm
               + 0.00061 * bpd_cm * ac_cm
               + 0.0424 * ac_cm
               + 0.174 * fl_cm)
    return 10 ** log_efw


# ============================================================================
# 数据加载与处理
# ============================================================================
def load_data():
    """加载数据，过滤超声-分娩间隔≤7天的样本，计算Hadlock 4和ML预测值。"""
    df = pd.read_csv('processed_data.csv')
    df_filtered = df[df['days_last_US_to_delivery'] <= 7].copy()
    print(f"[INFO] Filtered to {len(df_filtered)} samples (US-to-delivery ≤7 days)")

    # Hadlock 4 EFW
    df_filtered['EFW_g'] = hadlock_4_efw(
        df_filtered['BPD'], df_filtered['HC'], df_filtered['AC'], df_filtered['FL']
    )

    # 加载30特征模型
    model = joblib.load(f'{OUTPUT_DIR}/best_model_30features.pkl')
    scaler = joblib.load(f'{OUTPUT_DIR}/scaler_30features.pkl')

    X_30 = create_30_features(df_filtered)
    # 确保列顺序与训练时一致
    X_30_scaled = scaler.transform(X_30)
    df_filtered['ML_g'] = model.predict(X_30_scaled)

    # 经验分组（0=Junior<3y, 1=Senior≥3y）
    df_filtered['experience'] = df_filtered['sonographer_experience'].map(
        {0: 'Junior (<3y)', 1: 'Senior (≥3y)'}
    )

    # 出生体重分层
    def weight_strata(bw):
        if bw < 2500:
            return 'LBW\n(<2500g)'
        elif bw < 3000:
            return 'Normal-Low\n(2500-3000g)'
        elif bw < 3500:
            return 'Normal-Mid\n(3000-3500g)'
        elif bw < 4000:
            return 'Normal-High\n(3500-4000g)'
        else:
            return 'Macrosomia\n(>=4000g)'

    df_filtered['weight_strata'] = df_filtered['birth_weight_g'].apply(weight_strata)
    df_filtered['EFW_AE'] = np.abs(df_filtered['EFW_g'] - df_filtered['birth_weight_g'])
    df_filtered['ML_AE'] = np.abs(df_filtered['ML_g'] - df_filtered['birth_weight_g'])

    return df_filtered


def calculate_metrics(df_filtered):
    """按经验和体重分层计算MAE、RMSE及改善百分比。"""
    results = []
    for exp in ['Junior (<3y)', 'Senior (≥3y)']:
        exp_data = df_filtered[df_filtered['experience'] == exp]
        # Overall
        efw_mae = exp_data['EFW_AE'].mean()
        ml_mae = exp_data['ML_AE'].mean()
        efw_rmse = np.sqrt(np.mean(exp_data['EFW_AE'] ** 2))
        ml_rmse = np.sqrt(np.mean(exp_data['ML_AE'] ** 2))
        results.append({
            'experience': exp, 'weight_strata': 'Overall', 'n_samples': len(exp_data),
            'EFW_MAE': efw_mae, 'ML_MAE': ml_mae,
            'EFW_RMSE': efw_rmse, 'ML_RMSE': ml_rmse,
            'MAE_improvement_%': (efw_mae - ml_mae) / efw_mae * 100,
            'RMSE_improvement_%': (efw_rmse - ml_rmse) / efw_rmse * 100,
            'absolute_reduction_g': efw_mae - ml_mae,
        })

        for strata in ['LBW\n(<2500g)', 'Normal-Low\n(2500-3000g)',
                       'Normal-Mid\n(3000-3500g)', 'Normal-High\n(3500-4000g)',
                       'Macrosomia\n(>=4000g)']:
            sub = exp_data[exp_data['weight_strata'] == strata]
            if len(sub) == 0:
                continue
            efw_mae_s = sub['EFW_AE'].mean()
            ml_mae_s = sub['ML_AE'].mean()
            efw_rmse_s = np.sqrt(np.mean(sub['EFW_AE'] ** 2))
            ml_rmse_s = np.sqrt(np.mean(sub['ML_AE'] ** 2))
            results.append({
                'experience': exp, 'weight_strata': strata, 'n_samples': len(sub),
                'EFW_MAE': efw_mae_s, 'ML_MAE': ml_mae_s,
                'EFW_RMSE': efw_rmse_s, 'ML_RMSE': ml_rmse_s,
                'MAE_improvement_%': (efw_mae_s - ml_mae_s) / efw_mae_s * 100,
                'RMSE_improvement_%': (efw_rmse_s - ml_rmse_s) / efw_rmse_s * 100,
                'absolute_reduction_g': efw_mae_s - ml_mae_s,
            })

    return pd.DataFrame(results)


def create_stratified_figure(df_metrics, output_path):
    """生成5-panel综合分层分析图。"""
    plt.rcParams['font.family'] = 'Arial'
    plt.rcParams['axes.unicode_minus'] = False

    strata_order = ['LBW\n(<2500g)', 'Normal-Low\n(2500-3000g)',
                    'Normal-Mid\n(3000-3500g)', 'Normal-High\n(3500-4000g)',
                    'Macrosomia\n(>=4000g)']
    colors = {'Junior (<3y)': '#E74C3C', 'Senior (≥3y)': '#3498DB'}
    methods = ['EFW_MAE', 'ML_MAE']
    method_labels = ['Hadlock 4 EFW', '30-Feature ML']
    method_colors = ['#95A5A6', '#2ECC71']

    fig, axes = plt.subplots(2, 3, figsize=(18, 10))
    axes = axes.flatten()

    # Panel 1: MAE by strata (Junior)
    ax = axes[0]
    df_junior = df_metrics[df_metrics['experience'] == 'Junior (<3y)']
    df_junior = df_junior[df_junior['weight_strata'] != 'Overall']
    df_junior = df_junior.set_index('weight_strata').reindex(strata_order).reset_index()
    x = np.arange(len(strata_order))
    width = 0.35
    ax.bar(x - width/2, df_junior['EFW_MAE'], width, label='Hadlock 4 EFW', color=method_colors[0])
    ax.bar(x + width/2, df_junior['ML_MAE'], width, label='30-Feature ML', color=method_colors[1])
    ax.set_ylabel('MAE (g)', fontsize=11)
    ax.set_title('(a) Junior Sonographers (<3y)', fontsize=12, fontweight='bold')
    ax.set_xticks(x)
    ax.set_xticklabels(strata_order, fontsize=9)
    ax.legend()
    ax.set_ylim(0, max(df_junior['EFW_MAE'].max(), df_junior['ML_MAE'].max()) * 1.2)

    # Panel 2: MAE by strata (Senior)
    ax = axes[1]
    df_senior = df_metrics[df_metrics['experience'] == 'Senior (≥3y)']
    df_senior = df_senior[df_senior['weight_strata'] != 'Overall']
    df_senior = df_senior.set_index('weight_strata').reindex(strata_order).reset_index()
    ax.bar(x - width/2, df_senior['EFW_MAE'], width, label='Hadlock 4 EFW', color=method_colors[0])
    ax.bar(x + width/2, df_senior['ML_MAE'], width, label='30-Feature ML', color=method_colors[1])
    ax.set_ylabel('MAE (g)', fontsize=11)
    ax.set_title('(b) Senior Sonographers (≥3y)', fontsize=12, fontweight='bold')
    ax.set_xticks(x)
    ax.set_xticklabels(strata_order, fontsize=9)
    ax.legend()
    ax.set_ylim(0, max(df_senior['EFW_MAE'].max(), df_senior['ML_MAE'].max()) * 1.2)

    # Panel 3: Overall MAE comparison
    ax = axes[2]
    overall = df_metrics[df_metrics['weight_strata'] == 'Overall']
    experiences = overall['experience'].values
    x_overall = np.arange(len(experiences))
    ax.bar(x_overall - width/2, overall['EFW_MAE'], width, label='Hadlock 4 EFW', color=method_colors[0])
    ax.bar(x_overall + width/2, overall['ML_MAE'], width, label='30-Feature ML', color=method_colors[1])
    ax.set_ylabel('MAE (g)', fontsize=11)
    ax.set_title('(c) Overall MAE by Experience', fontsize=12, fontweight='bold')
    ax.set_xticks(x_overall)
    ax.set_xticklabels(experiences, fontsize=10)
    ax.legend()
    for i, row in overall.iterrows():
        idx = list(overall.index).index(i)
        ax.annotate(f"{row['MAE_improvement_%']:.1f}%",
                    xy=(idx, max(row['EFW_MAE'], row['ML_MAE']) + 5),
                    ha='center', fontsize=9, color='darkgreen', fontweight='bold')

    # Panel 4: MAE improvement % by strata
    ax = axes[3]
    df_plot = df_metrics[df_metrics['weight_strata'] != 'Overall']
    for exp in ['Junior (<3y)', 'Senior (≥3y)']:
        sub = df_plot[df_plot['experience'] == exp]
        sub = sub.set_index('weight_strata').reindex(strata_order).reset_index()
        ax.plot(range(len(strata_order)), sub['MAE_improvement_%'],
                marker='o', label=exp, color=colors[exp], linewidth=2, markersize=8)
    ax.axhline(0, color='black', linestyle='--', linewidth=1)
    ax.set_ylabel('MAE Improvement (%)', fontsize=11)
    ax.set_title('(d) ML Improvement over Hadlock 4 by Strata', fontsize=12, fontweight='bold')
    ax.set_xticks(range(len(strata_order)))
    ax.set_xticklabels(strata_order, fontsize=9)
    ax.legend()

    # Panel 5: Absolute reduction in inter-operator gap
    ax = axes[4]
    junior_overall = overall[overall['experience'] == 'Junior (<3y)'].iloc[0]
    senior_overall = overall[overall['experience'] == 'Senior (≥3y)'].iloc[0]
    methods_gap = ['Hadlock 4', '30-Feature ML']
    efw_gap = junior_overall['EFW_MAE'] - senior_overall['EFW_MAE']
    ml_gap = junior_overall['ML_MAE'] - senior_overall['ML_MAE']
    gap_values = [efw_gap, ml_gap]
    gap_colors = ['#95A5A6', '#2ECC71']
    bars = ax.bar(methods_gap, gap_values, color=gap_colors, width=0.5)
    ax.set_ylabel('MAE Gap (g)', fontsize=11)
    ax.set_title(f'(e) Inter-Operator Gap Reduction\n({(1 - ml_gap/efw_gap)*100:.0f}% reduction)',
                 fontsize=12, fontweight='bold')
    for bar, val in zip(bars, gap_values):
        ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 2,
                f'{val:.1f} g', ha='center', fontsize=10, fontweight='bold')

    # Hide the 6th panel
    axes[5].set_visible(False)

    plt.tight_layout()
    fig.savefig(f"{output_path}.png", dpi=300, bbox_inches='tight')
    fig.savefig(f"{output_path}.pdf", bbox_inches='tight')
    fig.savefig(f"{output_path}.eps", bbox_inches='tight')
    plt.close()
    print(f"[INFO] Saved stratified figure to {output_path}.{{png,pdf,eps}}")


def main():
    print("=" * 70)
    print("Stratified Analysis - 30 Features with Hadlock 4")
    print("=" * 70)

    df_filtered = load_data()
    df_metrics = calculate_metrics(df_filtered)

    output_csv = f'{OUTPUT_DIR}/stratified_analysis_metrics.csv'
    df_metrics.to_csv(output_csv, index=False)
    print(f"[INFO] Saved: {output_csv}")

    # 打印关键发现
    overall = df_metrics[df_metrics['weight_strata'] == 'Overall']
    junior = overall[overall['experience'] == 'Junior (<3y)'].iloc[0]
    senior = overall[overall['experience'] == 'Senior (≥3y)'].iloc[0]
    efw_gap = junior['EFW_MAE'] - senior['EFW_MAE']
    ml_gap = junior['ML_MAE'] - senior['ML_MAE']
    print(f"\n[KEY FINDINGS]")
    print(f"  Junior: Hadlock 4 MAE = {junior['EFW_MAE']:.1f} g -> ML MAE = {junior['ML_MAE']:.1f} g")
    print(f"  Senior: Hadlock 4 MAE = {senior['EFW_MAE']:.1f} g -> ML MAE = {senior['ML_MAE']:.1f} g")
    print(f"  Gap: Hadlock 4 = {efw_gap:.1f} g -> ML = {ml_gap:.1f} g (reduced by {(1-ml_gap/efw_gap)*100:.0f}%)")

    create_stratified_figure(df_metrics, f'{OUTPUT_DIR}/Figure_Experience_Stratified_Analysis')

    print("\n" + "=" * 70)
    print("Stratified analysis completed!")
    print("=" * 70)


if __name__ == '__main__':
    main()