# Module 8 Capstone
# Module 8 - Module 8 Capstone
# Dataset: MODULE8_INCOME_SAVINGS_SYNTHETIC

# Module 8 Capstone
#
# Module 8 notebook lab. This notebook uses original Ceteris Lab teaching data and does not report real empirical findings.

# Learning goal
# Complete a full diagnostic and correction workflow.
#
# Dataset: MODULE8_INCOME_SAVINGS_SYNTHETIC. Variables: household_id, income, savings, age, education, family_size.

# %% Cell 3
import pandas as pd
import statsmodels.api as sm
from statsmodels.stats.diagnostic import het_breuschpagan

df = pd.read_csv("/data/module-8/module8_income_savings_heteroskedastic.csv")
X = sm.add_constant(df[["income", "age", "education", "family_size"]])
ols = sm.OLS(df["savings"], X).fit()
robust = ols.get_robustcov_results(cov_type="HC1")
bp = het_breuschpagan(ols.resid, X)
wls = sm.WLS(df["savings"], X, weights=1 / (df["error_scale"] ** 2)).fit(cov_type="HC1")
print("Conventional OLS SE:", ols.bse.round(4).to_dict())
print("Robust OLS SE:", pd.Series(robust.bse, index=ols.params.index).round(4).to_dict())
print("BP p-value:", round(bp[1], 4))
print("Robust WLS coefficients:", wls.params.round(4).to_dict())
print("Recommendation: compare estimates, use robust inference, and explain the variance model.")

# Reflection
# Write two sentences: one sentence explaining what the diagnostic or robust result says, and one sentence explaining a limitation or next step.