import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib
from ipywidgets import interact, FloatSlider, IntSlider
import ipywidgets as widgets
matplotlib.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'SimHei', 'STHeiti']
matplotlib.rcParams['axes.unicode_minus'] = False
# 导入交互式组件（可选）
try:

    INTERACTIVE_MODE = True
    print("✅ 交互式模式已启用")
except ImportError:
    INTERACTIVE_MODE = False
    print("⚠️  请安装 ipywidgets: pip install ipywidgets")
    print("   或使用手动修改参数的方式")

if not INTERACTIVE_MODE:
    print("⚠️  请先安装 ipywidgets:")
    print("   pip install ipywidgets")
    print("   jupyter nbextension enable --py widgetsnbextension")
else:
    def simulate_and_plot(n_cities=5000, true_drug_effect=-2.0, confound_strength=3.0, random_seed=42):
        """交互式模拟函数"""

        np.random.seed(random_seed)
        severity = np.random.normal(0, 1, n_cities)
        treatment = (severity > 0).astype(int)
        recovery = 10 + severity * confound_strength + treatment * true_drug_effect + np.random.normal(0, 1, n_cities)

        data = pd.DataFrame({
            'severity': severity,
            'treatment': treatment,
            'recovery': recovery
        })

        control_mean = data[data['treatment'] == 0]['recovery'].mean()
        treated_mean = data[data['treatment'] == 1]['recovery'].mean()
        observed_effect = treated_mean - control_mean
        selection_bias = observed_effect - true_drug_effect

        # 简化版可视化（4图）
        fig, axes = plt.subplots(2, 2, figsize=(11, 9))

        # 图1
        axes[0, 0].hist(data[data['treatment'] == 0]['severity'], bins=25, alpha=0.6, color='blue', label='未获药')
        axes[0, 0].hist(data[data['treatment'] == 1]['severity'], bins=25, alpha=0.6, color='red', label='获得药')
        axes[0, 0].set_title('疫情严重程度分布', fontsize=11, fontweight='bold')
        axes[0, 0].legend()

        # 图2
        axes[0, 1].bar(['未获药', '获得药'], [control_mean, treated_mean], color=['blue', 'red'], alpha=0.7)
        axes[0, 1].axhline(10, color='gray', linestyle='--', alpha=0.5)
        axes[0, 1].set_ylabel('平均天数')
        axes[0, 1].set_title('康复时间对比', fontsize=11, fontweight='bold')

        # 图3
        axes[1, 0].barh(['观察', '真实', '偏误'],
                       [observed_effect, true_drug_effect, selection_bias],
                       color=['orange', 'green', 'red'], alpha=0.7)
        axes[1, 0].axvline(0, color='black', linewidth=1)
        axes[1, 0].set_xlabel('天数')
        axes[1, 0].set_title('效应分解', fontsize=11, fontweight='bold')

        # 图4
        axes[1, 1].scatter(data[data['treatment'] == 0]['severity'],
                          data[data['treatment'] == 0]['recovery'],
                          alpha=0.3, s=15, color='blue', label='未获药')
        axes[1, 1].scatter(data[data['treatment'] == 1]['severity'],
                          data[data['treatment'] == 1]['recovery'],
                          alpha=0.3, s=15, color='red', label='获得药')
        axes[1, 1].set_xlabel('严重程度')
        axes[1, 1].set_ylabel('康复天数')
        axes[1, 1].set_title('混淆效应', fontsize=11, fontweight='bold')
        axes[1, 1].legend()

        plt.tight_layout()
        plt.show()

        print(f"\n📊 观察效应: {observed_effect:.2f} 天")
        print(f"✅ 真实效应: {true_drug_effect:.2f} 天")
        print(f"⚠️  选择偏误: {selection_bias:.2f} 天\n")

    # 创建交互式界面
    interact(
        simulate_and_plot,
        n_cities=IntSlider(value=5000, min=1000, max=10000, step=500, description='城市数量:'),
        true_drug_effect=FloatSlider(value=-2.0, min=-5.0, max=2.0, step=0.5, description='真实药效:'),
        confound_strength=FloatSlider(value=3.0, min=0.0, max=6.0, step=0.5, description='混淆强度:'),
        random_seed=IntSlider(value=42, min=1, max=100, step=1, description='随机种子:')
    )