# Fitted Values and Residuals
# Module 2 - Fitted values and residuals
# Dataset: /data/wage_sample.csv

# Fitted Values and Residuals
#
# Use the wage regression to calculate predicted wages, residuals, and squared residuals. Then visualize the vertical gaps between actual and fitted values.

# %% Cell 2
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

df = pd.read_csv("wage_sample.csv")
x = df["education"]
y = df["wage"]

slope = ((x - x.mean()) * (y - y.mean())).sum() / ((x - x.mean()) ** 2).sum()
intercept = y.mean() - slope * x.mean()

df["fitted_wage"] = intercept + slope * df["education"]
df["residual"] = df["wage"] - df["fitted_wage"]
df["squared_residual"] = df["residual"] ** 2

print("Fitted equation: wage_hat =", round(intercept, 2), "+", round(slope, 2), "* education")
print(df[["education", "wage", "fitted_wage", "residual", "squared_residual"]].round(2))

# %% Cell 3
education_grid = np.linspace(df["education"].min(), df["education"].max(), 50)
line = intercept + slope * education_grid

plt.scatter(df["education"], df["wage"], label="Actual wage")
plt.plot(education_grid, line, label="Fitted line")

for _, row in df.iterrows():
    plt.plot([row["education"], row["education"]], [row["wage"], row["fitted_wage"]], linestyle="--")

plt.xlabel("Years of education")
plt.ylabel("Hourly wage")
plt.title("Residuals are vertical gaps")
plt.legend()
plt.show()

# %% Cell 4
plt.axhline(0, linestyle="--")
plt.scatter(df["education"], df["residual"])
plt.xlabel("Years of education")
plt.ylabel("Residual")
plt.title("Residual plot")
plt.show()

print("Largest absolute residual:")
print(df.loc[df["residual"].abs().idxmax(), ["education", "wage", "fitted_wage", "residual"]].round(2))

# Interpretation practice
#
# A positive residual means the worker earned more than the fitted line predicted. A negative residual means the worker earned less than the fitted line predicted.