import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import numpy as np

# 设置绘图风格
sns.set(style="whitegrid")
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False 

# 读取数据
df = pd.read_csv('healthcare-dataset-stroke-data.csv')

# 1. 生成 EDA 截图 (特征相关性热力图 + 箱线图)
plt.figure(figsize=(18, 8))

# 子图1: 相关性热力图
plt.subplot(1, 2, 1)
numeric_cols = df.select_dtypes(include=[np.number]).columns
# 简单填充缺失值以计算相关性
temp_df = df.copy()
temp_df['bmi'] = temp_df['bmi'].fillna(temp_df['bmi'].mean()) 
sns.heatmap(temp_df[numeric_cols].corr(), annot=True, cmap='coolwarm', fmt=".2f")
plt.title('Feature Correlation Heatmap')

# 子图2: 年龄分布
plt.subplot(1, 2, 2)
sns.kdeplot(data=df[df['stroke'] == 0]['age'], label='No Stroke', shade=True, color='blue')
sns.kdeplot(data=df[df['stroke'] == 1]['age'], label='Stroke', shade=True, color='red')
plt.title('Age Distribution by Stroke Risk')
plt.legend()

plt.tight_layout()
plt.savefig('images/notebook_eda.png', dpi=100)
print("Generated images/notebook_eda.png")
