# Prediction Intervals with Changing Variance
# Module 8 - Prediction under Heteroskedasticity
# Dataset: MODULE8_HOUSING_VARIANCE_SYNTHETIC

# Prediction Intervals with Changing Variance
#
# Module 8 notebook lab. This notebook uses original Ceteris Lab teaching data and does not report real empirical findings.

# Learning goal
# Plot prediction intervals under changing variance.
#
# Dataset: MODULE8_HOUSING_VARIANCE_SYNTHETIC. Variables: house_id, price, log_price, lotsize, sqrft, bedrooms.

# %% Cell 3
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import statsmodels.api as sm

df = pd.read_csv("/data/module-8/module8_housing_variance_synthetic.csv")
X = sm.add_constant(df[["sqrft", "lotsize"]])
model = sm.OLS(df["price"], X).fit(cov_type="HC1")
grid = pd.DataFrame({"sqrft": np.linspace(df["sqrft"].min(), df["sqrft"].max(), 30), "lotsize": df["lotsize"].median()})
pred = model.get_prediction(sm.add_constant(grid, has_constant="add")).summary_frame()
plt.plot(grid["sqrft"], pred["mean"], label="mean prediction")
plt.fill_between(grid["sqrft"], pred["obs_ci_lower"], pred["obs_ci_upper"], alpha=0.2, label="prediction interval")
plt.xlabel("Square feet")
plt.ylabel("Price")
plt.legend()
plt.title("Prediction intervals widen with uncertainty")
plt.show()
print(pred[["mean", "obs_ci_lower", "obs_ci_upper"]].head().round(2))

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