"""
Generate Figures - 30-Feature Results
=====================================
生成基于30特征实验结果的SCI级别图片（300 DPI，Arial字体）。
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patheffects as pe
import seaborn as sns
from sklearn.metrics import roc_curve, auc, precision_recall_curve, average_precision_score
import os
import sys
import joblib

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from fetal_weight_prediction_30features import create_30_features

plt.rcParams['font.family'] = 'Arial'
plt.rcParams['axes.unicode_minus'] = False
sns.set_palette("husl")

OUTPUT_DIR = '/Users/wangkai/Desktop/kimi/2/30特征'
SCI_DIR = f'{OUTPUT_DIR}/SCI_Figures'
os.makedirs(SCI_DIR, exist_ok=True)


def prepare_data():
    """加载数据、模型和预测结果。"""
    df = pd.read_csv('processed_data.csv')
    X = create_30_features(df)
    y = df['birth_weight_g']
    y_class = pd.cut(y, bins=[0, 2500, 4000, 10000], labels=[0, 1, 2]).astype(int)

    # 使用与主脚本相同的split（需固定random_state和stratify）
    from sklearn.model_selection import train_test_split
    _, X_test, _, y_test, _, y_class_test = train_test_split(
        X, y, y_class, test_size=0.2, random_state=42, stratify=y_class)

    model = joblib.load(f'{OUTPUT_DIR}/best_model_30features.pkl')
    scaler = joblib.load(f'{OUTPUT_DIR}/scaler_30features.pkl')
    X_test_s = scaler.transform(X_test)
    y_pred = model.predict(X_test_s)

    return X_test, y_test, y_class_test, y_pred


def fig2_regression_performance(y_test, y_pred):
    """Figure 2: Scatter + Bland-Altman"""
    fig, axes = plt.subplots(1, 2, figsize=(14, 5.5))

    # Scatter
    ax = axes[0]
    ax.scatter(y_test, y_pred, alpha=0.6, edgecolors='k', linewidth=0.5, s=60)
    ax.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--', lw=2, label='Identity')
    z = np.polyfit(y_test, y_pred, 1)
    p = np.poly1d(z)
    ax.plot(y_test.sort_values(), p(y_test.sort_values()), 'b-', lw=1.5, label='Fitted line')
    ax.set_xlabel('Actual Birth Weight (g)', fontsize=12)
    ax.set_ylabel('Predicted Birth Weight (g)', fontsize=12)
    ax.set_title('(a) Predicted vs Actual Birth Weight', fontsize=13, fontweight='bold')
    ax.legend()

    # Bland-Altman
    ax = axes[1]
    diff = y_pred - y_test
    mean = (y_test + y_pred) / 2
    md = np.mean(diff)
    sd = np.std(diff, ddof=1)
    ax.scatter(mean, diff, alpha=0.6, edgecolors='k', linewidth=0.5, s=60)
    ax.axhline(md, color='red', linestyle='--', lw=2, label=f'Mean diff: {md:.1f} g')
    ax.axhline(md + 1.96*sd, color='gray', linestyle=':', lw=1.5)
    ax.axhline(md - 1.96*sd, color='gray', linestyle=':', lw=1.5)
    ax.fill_between([mean.min(), mean.max()], md - 1.96*sd, md + 1.96*sd, alpha=0.1, color='gray')
    ax.set_xlabel('Mean of Actual and Predicted (g)', fontsize=12)
    ax.set_ylabel('Difference (Predicted - Actual) (g)', fontsize=12)
    ax.set_title('(b) Bland-Altman Plot', fontsize=13, fontweight='bold')
    ax.legend()

    plt.tight_layout()
    fig.savefig(f'{SCI_DIR}/Fig2_Regression_Performance.png', dpi=300, bbox_inches='tight')
    fig.savefig(f'{SCI_DIR}/Fig2_Regression_Performance.pdf', bbox_inches='tight')
    plt.close()
    print("[INFO] Saved Fig2_Regression_Performance")


def fig3_confusion_matrix(y_class_test, y_pred):
    """Figure 3: Confusion Matrix (推荐策略 P10/P90+FH)"""
    from sklearn.metrics import confusion_matrix
    # 重新计算推荐策略的分类结果
    y_train_pred = None
    # 需要训练集预测值来计算P10/P90阈值
    df_full = pd.read_csv('processed_data.csv')
    X_full = create_30_features(df_full)
    y_full = df_full['birth_weight_g']
    y_class_full = pd.cut(y_full, bins=[0, 2500, 4000, 10000], labels=[0, 1, 2]).astype(int)
    from sklearn.model_selection import train_test_split
    X_train, X_test, y_train, _, _, _ = train_test_split(
        X_full, y_full, y_class_full, test_size=0.2, random_state=42, stratify=y_class_full)

    model = joblib.load(f'{OUTPUT_DIR}/best_model_30features.pkl')
    scaler = joblib.load(f'{OUTPUT_DIR}/scaler_30features.pkl')
    X_train_s = scaler.transform(X_train)
    y_train_pred = model.predict(X_train_s)

    p10 = np.percentile(y_train_pred, 10)
    p90 = np.percentile(y_train_pred, 90)
    p80 = np.percentile(y_train_pred, 80)
    fh_test = X_test['fundal_height_cm'].values

    y_class_pred = np.ones(len(y_pred), dtype=int)
    y_class_pred[y_pred < p10] = 0
    macro_condition = (y_pred > p90) | ((y_pred > p80) & (fh_test > 36))
    y_class_pred[macro_condition] = 2

    cm = confusion_matrix(y_class_test, y_class_pred, labels=[0, 1, 2])

    fig, ax = plt.subplots(figsize=(8, 7))
    sns.heatmap(cm, annot=False, fmt='d', cmap='Blues', ax=ax,
                xticklabels=['LBW', 'Normal', 'Macro'],
                yticklabels=['LBW', 'Normal', 'Macro'],
                linewidths=1.5, linecolor='gray')
    # 手动添加标注，确保所有数值在任何背景下都清晰可辨
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            ax.text(j + 0.5, i + 0.5, str(cm[i, j]),
                    ha='center', va='center', fontsize=16, weight='bold', color='black',
                    path_effects=[pe.withStroke(linewidth=3, foreground='white')])
    ax.set_xlabel('Predicted', fontsize=13)
    ax.set_ylabel('Actual', fontsize=13)
    ax.set_title('Confusion Matrix (P10/P90 + FH>36cm)', fontsize=14, fontweight='bold')

    plt.tight_layout()
    fig.savefig(f'{SCI_DIR}/Fig3_Confusion_Matrix.png', dpi=300, bbox_inches='tight')
    fig.savefig(f'{SCI_DIR}/Fig3_Confusion_Matrix.pdf', bbox_inches='tight')
    plt.close()
    print("[INFO] Saved Fig3_Confusion_Matrix")


def fig4_roc_pr_curves(y_class_test, y_pred):
    """Figure 4: ROC and PR curves"""
    fig, axes = plt.subplots(1, 2, figsize=(14, 5.5))

    # ROC
    ax = axes[0]
    for label, name, pos_label in [(0, 'LBW vs Others', 0), (2, 'Macrosomia vs Others', 2)]:
        y_binary = (y_class_test == label).astype(int)
        fpr, tpr, _ = roc_curve(y_binary, -y_pred if label == 0 else y_pred)
        roc_auc = auc(fpr, tpr)
        ax.plot(fpr, tpr, lw=2, label=f'{name} (AUC = {roc_auc:.3f})')
    ax.plot([0, 1], [0, 1], 'k--', lw=1)
    ax.set_xlabel('False Positive Rate', fontsize=12)
    ax.set_ylabel('True Positive Rate', fontsize=12)
    ax.set_title('(a) ROC Curves', fontsize=13, fontweight='bold')
    ax.legend(loc='lower right')

    # PR
    ax = axes[1]
    for label, name in [(0, 'LBW'), (2, 'Macrosomia')]:
        y_binary = (y_class_test == label).astype(int)
        precision, recall, _ = precision_recall_curve(y_binary, -y_pred if label == 0 else y_pred)
        ap = average_precision_score(y_binary, -y_pred if label == 0 else y_pred)
        ax.plot(recall, precision, lw=2, label=f'{name} (AP = {ap:.3f})')
    ax.set_xlabel('Recall', fontsize=12)
    ax.set_ylabel('Precision', fontsize=12)
    ax.set_title('(b) Precision-Recall Curves', fontsize=13, fontweight='bold')
    ax.legend()

    plt.tight_layout()
    fig.savefig(f'{SCI_DIR}/Fig4_ROC_PR_Curves.png', dpi=300, bbox_inches='tight')
    fig.savefig(f'{SCI_DIR}/Fig4_ROC_PR_Curves.pdf', bbox_inches='tight')
    plt.close()
    print("[INFO] Saved Fig4_ROC_PR_Curves")


def figs1_feature_importance():
    """Supplementary Figure S1: Ridge coefficient bar chart"""
    importance = pd.read_csv(f'{OUTPUT_DIR}/feature_importance.csv')
    fig, ax = plt.subplots(figsize=(10, 10))
    colors = ['#E74C3C' if c > 0 else '#3498DB' for c in importance['Coefficient']]
    ax.barh(range(len(importance)), importance['Coefficient'], color=colors)
    ax.set_yticks(range(len(importance)))
    ax.set_yticklabels(importance['Feature'], fontsize=10)
    ax.set_xlabel('Ridge Coefficient', fontsize=12)
    ax.set_title('Feature Importance (Ridge Coefficients)', fontsize=13, fontweight='bold')
    ax.invert_yaxis()
    plt.tight_layout()
    fig.savefig(f'{SCI_DIR}/FigS1_Feature_Importance.png', dpi=300, bbox_inches='tight')
    fig.savefig(f'{SCI_DIR}/FigS1_Feature_Importance.pdf', bbox_inches='tight')
    plt.close()
    print("[INFO] Saved FigS1_Feature_Importance")


def figs3_dca_curves():
    """Supplementary Figure S3: DCA Curves"""
    # DCA需要重新计算
    # 我们用测试集数据重新计算
    from sklearn.model_selection import train_test_split
    df = pd.read_csv('processed_data.csv')
    X = create_30_features(df)
    y = df['birth_weight_g']
    y_class = pd.cut(y, bins=[0, 2500, 4000, 10000], labels=[0, 1, 2]).astype(int)
    _, X_test, _, y_test, _, y_class_test = train_test_split(
        X, y, y_class, test_size=0.2, random_state=42, stratify=y_class)

    model = joblib.load(f'{OUTPUT_DIR}/best_model_30features.pkl')
    scaler = joblib.load(f'{OUTPUT_DIR}/scaler_30features.pkl')
    y_pred = model.predict(scaler.transform(X_test))

    y_binary = (y_class_test == 0).astype(int)
    n = len(y_binary)
    thresholds = np.linspace(0.01, 0.99, 100)
    nb_model, nb_all = [], []

    for thresh in thresholds:
        pred_pos = y_pred < np.percentile(y_pred, thresh * 100)
        tp = np.sum((y_binary == 1) & pred_pos)
        fp = np.sum((y_binary == 0) & pred_pos)
        nb = (tp / n) - (fp / n) * (thresh / (1 - thresh))
        nb_model.append(nb)
        nb_all.append(np.mean(y_binary) - (1 - np.mean(y_binary)) * (thresh / (1 - thresh)))

    fig, ax = plt.subplots(figsize=(8, 6))
    ax.plot(thresholds, nb_model, lw=2, label='30-Feature ML Model')
    ax.plot(thresholds, nb_all, lw=2, linestyle='--', label='Treat All')
    ax.axhline(0, color='black', linestyle=':', label='Treat None')
    ax.set_xlabel('Threshold Probability', fontsize=12)
    ax.set_ylabel('Net Benefit', fontsize=12)
    ax.set_title('Decision Curve Analysis (LBW)', fontsize=13, fontweight='bold')
    ax.legend()
    plt.tight_layout()
    fig.savefig(f'{SCI_DIR}/FigS3_DCA_Curves.png', dpi=300, bbox_inches='tight')
    fig.savefig(f'{SCI_DIR}/FigS3_DCA_Curves.pdf', bbox_inches='tight')
    plt.close()
    print("[INFO] Saved FigS3_DCA_Curves")


def main():
    print("=" * 70)
    print("Generating Figures - 30-Feature Results")
    print("=" * 70)

    X_test, y_test, y_class_test, y_pred = prepare_data()

    fig2_regression_performance(y_test, y_pred)
    fig3_confusion_matrix(y_class_test, y_pred)
    fig4_roc_pr_curves(y_class_test, y_pred)
    figs1_feature_importance()
    figs3_dca_curves()

    print("\n" + "=" * 70)
    print("All figures generated successfully!")
    print(f"Saved to: {SCI_DIR}")
    print("=" * 70)


if __name__ == '__main__':
    import json
    main()