# -*- coding: utf-8 -*-
"""
survey.csv 多元线性回归：诊断 + 修正 + 稳健推断（可复现）
------------------------------------------------------------
流程：
  1) 缺失与描述统计
  2) 初始 OLS（未变换、含 ses_index）→ 四项诊断
  3) 逐项修正（去共线变量 / 年龄二次项 / 收入取对数）
  4) 稳健标准误 HC3 + 强影响点敏感性分析
  5) 效应量（标准化 beta、可解释的边际效应）、95% CI、图表
运行：python3 analysis.py
"""
import warnings
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import scipy.stats as st
import statsmodels.api as sm
import statsmodels.formula.api as smf
from statsmodels.stats.outliers_influence import variance_inflation_factor
from statsmodels.stats.diagnostic import het_breuschpagan, het_white, linear_reset
from statsmodels.stats.stattools import jarque_bera
from statsmodels.nonparametric.smoothers_lowess import lowess

warnings.filterwarnings("ignore")
OUT = "output"
plt.rcParams.update({"figure.dpi": 110, "font.size": 9, "axes.grid": True, "grid.alpha": 0.25})


def reset_p(model):
    """linear_reset 在 statsmodels>=0.13 返回 ContrastResults"""
    r = linear_reset(model, power=3, use_f=True)
    return getattr(r, "fvalue", r.statistic), getattr(r, "f_pvalue", r.pvalue)


def vif_table(xnames, X):
    return pd.DataFrame({
        "变量": xnames,
        "VIF": [variance_inflation_factor(X.values, i + 1) for i in range(len(xnames))]})


# ============================================================================
# 1. 读取、缺失、描述
# ============================================================================
df = pd.read_csv("survey.csv", encoding="utf-8-sig")
print("=" * 74)
print("原始数据维度：", df.shape)
miss = df.isna().sum()
print("缺失值：\n" + miss[miss > 0].to_string())
print("含缺失的行（任一变量）：", int(df.isna().any(axis=1).sum()))

# 变量编码
df["gender_female"] = df["gender"].map({"男": 0, "女": 1})
df["region"] = pd.Categorical(df["region"], categories=["一线城市", "二线城市", "三线及以下"], ordered=True)
rd = pd.get_dummies(df["region"], drop_first=True, dtype=int)
df = pd.concat([df, rd.rename(columns={"二线城市": "region_tier2", "三线及以下": "region_tier3"})], axis=1)

cont = ["age", "education_years", "monthly_income_k", "ses_index", "work_hours_week",
        "commute_min", "sleep_hours", "exercise_days", "social_support"]
cats = ["gender_female", "region_tier2", "region_tier3"]
Y = "life_satisfaction"

# listwise deletion
d = df[["respondent_id"] + cont + cats + [Y]].dropna().copy().reset_index(drop=True)
d["log_income"] = np.log(d["monthly_income_k"])
d["age_c"] = d["age"] - 45
d["age_c2"] = d["age_c"] ** 2
print("\n样本量：原始 %d，完整案例 n = %d（删除 %d 行）" % (len(df), len(d), len(df) - len(d)))

print("\n描述统计：")
desc = d[cont + [Y]].describe().T[["count", "mean", "std", "min", "50%", "max"]]
desc.columns = ["n", "均值", "标准差", "最小", "中位", "最大"]
print(desc.round(3).to_string())
print("\nlife_satisfaction：偏度 %.2f，超额峰度 %.2f；=100 的 %d 个；<10 的 %d 个" %
      (st.skew(d[Y]), st.kurtosis(d[Y]), int((d[Y] == 100).sum()), int((d[Y] < 10).sum())))

y = d[Y]

# ============================================================================
# 2. 初始模型 M0（未变换、含 ses_index）
# ============================================================================
print("\n" + "=" * 74)
print("模型 M0：初始线性模型（含 ses_index，收入原始尺度）")
x0 = ["age", "gender_female", "region_tier2", "region_tier3", "education_years",
      "monthly_income_k", "ses_index", "work_hours_week", "commute_min",
      "sleep_hours", "exercise_days", "social_support"]
X0 = sm.add_constant(d[x0])
m0 = sm.OLS(y, X0).fit()
print("R²=%.3f  adjR²=%.3f  AIC=%.1f  BIC=%.1f  n=%d" %
      (m0.rsquared, m0.rsquared_adj, m0.aic, m0.bic, int(m0.nobs)))
print("F(%d,%d)=%.2f, p=%.3g" % (int(m0.df_model), int(m0.df_resid), m0.fvalue, m0.f_pvalue))

print("\n[共线性] VIF：")
print(vif_table(x0, X0).round(2).to_string(index=False))

f0, p0 = reset_p(m0)
bp = het_breuschpagan(m0.resid, m0.model.exog)
wh = het_white(m0.resid, m0.model.exog)
jb = jarque_bera(m0.resid)
sw = st.shapiro(m0.resid)
print("\n[线性] RESET(power=3)：F=%.2f, p=%.4f" % (f0, p0))
print("[异方差] Breusch-Pagan：LM=%.2f, p=%.4g；White：LM=%.2f, p=%.4g" % (bp[0], bp[1], wh[0], wh[1]))
print("[正态] Jarque-Bera p=%.4g；Shapiro-Wilk p=%.4g；残差偏度 %.2f，峰度 %.2f" %
      (jb[1], sw[1], jb[2], jb[3]))

inf0 = m0.get_influence().summary_frame()
print("\n[影响点] Cook's D 前 5：")
print(inf0.nlargest(5, "cooks_d")[["cooks_d", "student_resid", "hat_diag"]].round(3).to_string())

# ============================================================================
# 3. 逐项修正：模型对比
# ============================================================================
x_no_ses = ["age", "gender_female", "region_tier2", "region_tier3", "education_years",
            "monthly_income_k", "work_hours_week", "commute_min", "sleep_hours",
            "exercise_days", "social_support"]
x_age2 = ["age_c", "age_c2", "gender_female", "region_tier2", "region_tier3", "education_years",
          "monthly_income_k", "work_hours_week", "commute_min", "sleep_hours",
          "exercise_days", "social_support"]
x_final = ["age_c", "age_c2", "gender_female", "region_tier2", "region_tier3", "education_years",
           "log_income", "work_hours_week", "commute_min", "sleep_hours",
           "exercise_days", "social_support"]

specs = [
    ("M0 原始（含ses）", x0),
    ("M1 去掉ses", x_no_ses),
    ("M2 +年龄二次项", x_age2),
    ("M3 +收入取对数", x_final),
]
fits, comp = {}, []
for name, xs in specs:
    fit = sm.OLS(y, sm.add_constant(d[xs])).fit()
    fits[name] = fit
    comp.append({"模型": name, "k": len(xs) + 1, "n": int(fit.nobs),
                 "R²": fit.rsquared, "adj R²": fit.rsquared_adj,
                 "AIC": fit.aic, "BIC": fit.bic})
print("\n" + "=" * 74)
print("逐项修正的模型比较：")
comp_df = pd.DataFrame(comp)
print(comp_df.round(3).to_string(index=False))

# 年龄二次项嵌套 F 检验
from statsmodels.stats.anova import anova_lm
aov = anova_lm(fits["M1 去掉ses"], fits["M2 +年龄二次项"])
print("\n年龄二次项嵌套 F 检验：F=%.2f, p=%.3g" %
      (aov.loc[1, "F"], aov.loc[1, "Pr(>F)"]))

m1 = fits["M3 +收入取对数"]          # 修正后的 OLS（经典 SE）
m1_hc3 = sm.OLS(y, sm.add_constant(d[x_final])).fit(cov_type="HC3")
X1 = sm.add_constant(d[x_final])

print("\n[修正模型 M3] R²=%.3f  adjR²=%.3f  AIC=%.1f  BIC=%.1f" %
      (m1.rsquared, m1.rsquared_adj, m1.aic, m1.bic))
print("\n[共线性] 修正后 VIF：")
print(vif_table(x_final, X1).round(2).to_string(index=False))

f1, p1 = reset_p(m1)
bp1 = het_breuschpagan(m1.resid, m1.model.exog)
jb1 = jarque_bera(m1.resid)
sw1 = st.shapiro(m1.resid)
print("\n[线性] RESET：F=%.2f, p=%.4f" % (f1, p1))
print("[异方差] Breusch-Pagan：LM=%.2f, p=%.4g" % (bp1[0], bp1[1]))
print("[正态] Jarque-Bera p=%.4g；Shapiro p=%.4g；偏度 %.2f，峰度 %.2f" %
      (jb1[1], sw1[1], jb1[2], jb1[3]))

b_a, b_a2 = m1.params["age_c"], m1.params["age_c2"]
vertex = 45 - b_a / (2 * b_a2)
print("\n年龄效应：开口向上抛物线，最低点 age = %.1f 岁" % vertex)

# ============================================================================
# 4. 强影响点明细 + 敏感性分析
# ============================================================================
inf1 = m1.get_influence().summary_frame()
inf1["respondent_id"] = d["respondent_id"]
inf1["life_satisfaction"] = d[Y]
inf1["fitted"] = m1.fittedvalues
top = inf1.nlargest(10, "cooks_d")
print("\n[影响点明细] Cook's D 前 10：")
print(top[["respondent_id", "life_satisfaction", "fitted", "student_resid", "hat_diag", "cooks_d"]]
      .round(3).to_string(index=False))

n = len(d)
keep = inf1["cooks_d"] <= 4 / n
print("\nCook's D > 4/n (%.4f) 的样本：%d 个" % (4 / n, int((~keep).sum())))
m_trim = sm.OLS(y[keep], X1[keep]).fit(cov_type="HC3")

# ============================================================================
# 5. 最终模型结果（HC3）：系数、CI、标准化效应量
# ============================================================================
print("\n" + "=" * 74)
print("最终模型（M3，HC3 稳健标准误）")
ci = m1_hc3.conf_int()
sd_y = d[Y].std(ddof=1)
res = pd.DataFrame({
    "b": m1_hc3.params, "SE_HC3": m1_hc3.bse, "t": m1_hc3.tvalues,
    "p": m1_hc3.pvalues, "CI_low": ci.iloc[:, 0], "CI_up": ci.iloc[:, 1]})
res["std_beta"] = np.nan
for v in ["age_c", "age_c2", "education_years", "log_income", "work_hours_week",
          "commute_min", "sleep_hours", "exercise_days", "social_support"]:
    res.loc[v, "std_beta"] = res.loc[v, "b"] * d[v].std(ddof=1) / sd_y
res = res.drop("const", errors="ignore")
print(res.round(3).to_string())
res.round(4).to_csv(f"{OUT}/final_model_coef.csv", encoding="utf-8-sig")

# 可解释的边际效应（含 95% CI）
print("\n可解释效应（点估计 [95% CI]）：")
def eff(var, delta, label):
    b, lo, hi = res.loc[var, "b"], res.loc[var, "CI_low"], res.loc[var, "CI_up"]
    print("  %-28s %.2f [%.2f, %.2f]" % (label, b * delta, lo * delta, hi * delta))
eff("log_income", np.log(2), "收入翻倍 (log 差 ln2)")
eff("log_income", np.log(1.1), "收入 +10%")
eff("work_hours_week", 10, "每周工作 +10 小时")
eff("commute_min", 30, "通勤 +30 分钟")
eff("social_support", 1, "社会支持 +1 分")
eff("exercise_days", 1, "每周锻炼 +1 天")
eff("sleep_hours", 1, "睡眠 +1 小时")
print("  年龄 25→45 岁（二次项）       %.2f" % (res.loc["age_c2", "b"] * (20 ** 2)))
print("  年龄 45→65 岁（二次项）       %.2f" % (res.loc["age_c2", "b"] * (20 ** 2)))

# 剔除强影响点后的系数对比
print("\n敏感性分析（剔除 %d 个高 Cook's D 样本，n=%d）：" % (int((~keep).sum()), int(keep.sum())))
cmp = pd.DataFrame({"全部 b": m1_hc3.params, "剔除后 b": m_trim.params,
                    "全部 SE": m1_hc3.bse, "剔除后 SE": m_trim.bse})
print(cmp.drop("const", errors="ignore").round(3).to_string())

# ============================================================================
# 6. 图
# ============================================================================
# 图1：初始模型残差诊断
fig, ax = plt.subplots(1, 3, figsize=(13, 3.8))
ax[0].scatter(m0.fittedvalues, m0.resid, s=12, alpha=0.5)
lo = lowess(m0.resid, m0.fittedvalues, return_sorted=True)
ax[0].plot(lo[:, 0], lo[:, 1], color="tab:red", lw=2); ax[0].axhline(0, color="k", lw=0.7)
ax[0].set_xlabel("Fitted"); ax[0].set_ylabel("Residual"); ax[0].set_title("M0: residuals vs fitted")
st.probplot(m0.resid / m0.resid.std(ddof=1), dist="norm", plot=ax[1])
ax[1].set_title("M0: normal Q-Q (standardized)")
ax[2].scatter(d["monthly_income_k"], m0.resid, s=12, alpha=0.5)
lo = lowess(m0.resid, d["monthly_income_k"], return_sorted=True)
ax[2].plot(lo[:, 0], lo[:, 1], color="tab:red", lw=2); ax[2].axhline(0, color="k", lw=0.7)
ax[2].set_xlabel("monthly_income_k"); ax[2].set_ylabel("Residual"); ax[2].set_title("M0: residuals vs income")
fig.tight_layout(); fig.savefig(f"{OUT}/fig1_M0_residuals.png"); plt.close(fig)

# 图2：成分残差图（线性检查）
fig, axes = plt.subplots(2, 4, figsize=(14, 6)); axes = axes.ravel()
plot_cont = ["age", "education_years", "monthly_income_k", "work_hours_week",
             "commute_min", "sleep_hours", "exercise_days", "social_support"]
for i, v in enumerate(plot_cont):
    pr = m0.resid + m0.params[v] * d[v]
    axes[i].scatter(d[v], pr, s=10, alpha=0.5)
    lo = lowess(pr, d[v], return_sorted=True)
    axes[i].plot(lo[:, 0], lo[:, 1], color="tab:red", lw=2)
    axes[i].set_title(v); axes[i].set_ylabel("Component + residual")
fig.suptitle("M0 component+residual plots (lowess): curvature = nonlinearity")
fig.tight_layout(); fig.savefig(f"{OUT}/fig2_M0_component_residual.png"); plt.close(fig)

# 图3：修正模型残差诊断
fig, ax = plt.subplots(1, 3, figsize=(13, 3.8))
ax[0].scatter(m1.fittedvalues, m1.resid, s=12, alpha=0.5)
lo = lowess(m1.resid, m1.fittedvalues, return_sorted=True)
ax[0].plot(lo[:, 0], lo[:, 1], color="tab:red", lw=2); ax[0].axhline(0, color="k", lw=0.7)
ax[0].set_xlabel("Fitted"); ax[0].set_ylabel("Residual"); ax[0].set_title("M3: residuals vs fitted")
st.probplot(m1.resid / m1.resid.std(ddof=1), dist="norm", plot=ax[1])
ax[1].set_title("M3: normal Q-Q (standardized)")
sr = m1.get_influence().resid_studentized_external
ax[2].scatter(np.sqrt(np.abs(sr)), m1.fittedvalues, s=12, alpha=0.5)
lo = lowess(np.sqrt(np.abs(sr)), m1.fittedvalues, return_sorted=True)
ax[2].plot(lo[:, 0], lo[:, 1], color="tab:red", lw=2)
ax[2].set_xlabel("sqrt(|studentized residual|)"); ax[2].set_ylabel("Fitted")
ax[2].set_title("M3: scale-location")
fig.tight_layout(); fig.savefig(f"{OUT}/fig3_M3_residuals.png"); plt.close(fig)

# 图4：影响诊断
fig, ax = plt.subplots(1, 2, figsize=(10, 4))
ax[0].scatter(inf1["hat_diag"], inf1["student_resid"], s=14, alpha=0.6)
ax[0].axhline(3, color="r", lw=0.8, ls="--"); ax[0].axhline(-3, color="r", lw=0.8, ls="--")
ax[0].set_xlabel("Leverage (hat)"); ax[0].set_ylabel("Studentized residual")
ax[0].set_title("M3: studentized residual vs leverage")
ax[1].stem(np.arange(n), inf1["cooks_d"], markerfmt=" ", basefmt=" ")
ax[1].axhline(4 / n, color="r", lw=0.8, ls="--", label="4/n")
ax[1].set_xlabel("Observation index"); ax[1].set_ylabel("Cook's D")
ax[1].set_title("M3: Cook's distance"); ax[1].legend()
fig.tight_layout(); fig.savefig(f"{OUT}/fig4_M3_influence.png"); plt.close(fig)

# 图5：系数森林图（HC3, 95% CI）
show = ["log_income", "social_support", "exercise_days", "sleep_hours", "education_years",
        "age_c2", "age_c", "gender_female", "region_tier2", "region_tier3",
        "commute_min", "work_hours_week"]
bv = res.loc[show, "b"].values; lo_ = res.loc[show, "CI_low"].values; hi_ = res.loc[show, "CI_up"].values
yp = np.arange(len(show))[::-1]
fig, ax = plt.subplots(figsize=(7, 5.5))
ax.errorbar(bv, yp, xerr=[bv - lo_, hi_ - bv], fmt="o", color="tab:blue",
            ecolor="tab:gray", elinewidth=1.2, capsize=3)
ax.axvline(0, color="k", lw=0.8)
ax.set_yticks(yp); ax.set_yticklabels(show)
ax.set_xlabel("Coefficient (95% CI, HC3)"); ax.set_title("Final model coefficients")
fig.tight_layout(); fig.savefig(f"{OUT}/fig5_final_coef.png"); plt.close(fig)

# 图6：非线性形状（调整后预测）
base = {c: d[c].mean() for c in ["education_years", "log_income", "work_hours_week",
                                  "commute_min", "sleep_hours", "exercise_days", "social_support"]}
base.update({"const": 1.0, "gender_female": 0, "region_tier2": 0, "region_tier3": 0,
             "age_c": d["age_c"].mean(), "age_c2": d["age_c2"].mean()})
cov = m1_hc3.cov_params()
names = list(X1.columns)

def adj_pred(rows):
    Xg = pd.DataFrame(rows)[names]
    p = Xg.dot(m1_hc3.params).to_numpy()
    se = np.sqrt(np.einsum("ij,jk,ik->i", Xg.to_numpy(), cov.to_numpy(), Xg.to_numpy()))
    return p, se

# 年龄
ages = np.arange(19, 69)
rows = []
for a in ages:
    r = dict(base); r["age_c"] = a - 45; r["age_c2"] = (a - 45) ** 2; rows.append(r)
p_age, se_age = adj_pred(rows)
# 收入
incs = np.linspace(d["monthly_income_k"].min(), d["monthly_income_k"].max(), 200)
rows = []
for inc in incs:
    r = dict(base); r["log_income"] = np.log(inc); rows.append(r)
p_inc, se_inc = adj_pred(rows)

fig, ax = plt.subplots(1, 2, figsize=(11, 4))
ax[0].plot(ages, p_age, color="tab:blue", lw=2)
ax[0].fill_between(ages, p_age - 1.96 * se_age, p_age + 1.96 * se_age, alpha=0.2, color="tab:blue")
ax[0].axvline(vertex, color="r", ls="--", lw=1)
ax[0].annotate("min at %.1f" % vertex, (vertex, p_age.min()), xytext=(vertex + 2, p_age.min() - 3),
               color="r", fontsize=9)
ax[0].set_xlabel("Age (years)"); ax[0].set_ylabel("Predicted life satisfaction")
ax[0].set_title("Adjusted age effect (U-shape)")
ax[1].plot(incs, p_inc, color="tab:green", lw=2)
ax[1].fill_between(incs, p_inc - 1.96 * se_inc, p_inc + 1.96 * se_inc, alpha=0.2, color="tab:green")
ax[1].set_xscale("log")
ax[1].set_xlabel("Monthly income (k, log scale)"); ax[1].set_ylabel("Predicted life satisfaction")
ax[1].set_title("Adjusted income effect (diminishing returns)")
fig.tight_layout(); fig.savefig(f"{OUT}/fig6_nonlinear_effects.png"); plt.close(fig)

print("\n完成：图表写入 %s/，系数表写入 %s/final_model_coef.csv" % (OUT, OUT))