# OLS and R-squared by Formula
# Module 2 - Ordinary Least Squares intuition
# Dataset: /data/wage_sample.csv

# OLS and R-squared by Formula
#
# Compute the simple regression slope, intercept, residual sum of squares, total sum of squares, and R-squared step by step.

# %% 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"]

x_bar = x.mean()
y_bar = y.mean()
numerator = ((x - x_bar) * (y - y_bar)).sum()
denominator = ((x - x_bar) ** 2).sum()
slope = numerator / denominator
intercept = y_bar - slope * x_bar

print("x mean:", round(x_bar, 2))
print("y mean:", round(y_bar, 2))
print("slope numerator:", round(numerator, 2))
print("slope denominator:", round(denominator, 2))
print("intercept:", round(intercept, 2))
print("slope:", round(slope, 2))

# %% Cell 3
fitted = intercept + slope * x
residuals = y - fitted
ssr = (residuals ** 2).sum()
sst = ((y - y_bar) ** 2).sum()
sse = sst - ssr
r_squared = 1 - ssr / sst

print("SSR, residual variation:", round(ssr, 2))
print("SST, total variation:", round(sst, 2))
print("SSE, explained variation:", round(sse, 2))
print("R-squared:", round(r_squared, 3))

# %% Cell 4
trial_slopes = np.linspace(1.0, 3.6, 14)
records = []

for trial_slope in trial_slopes:
    trial_intercept = y_bar - trial_slope * x_bar
    trial_fitted = trial_intercept + trial_slope * x
    trial_ssr = ((y - trial_fitted) ** 2).sum()
    records.append({"trial_slope": trial_slope, "SSR": trial_ssr})

comparison = pd.DataFrame(records)
print(comparison.round(2))

plt.plot(comparison["trial_slope"], comparison["SSR"], marker="o")
plt.axvline(slope, linestyle="--", label="OLS slope")
plt.xlabel("Trial slope")
plt.ylabel("Sum of squared residuals")
plt.title("OLS chooses the slope with the lowest SSR")
plt.legend()
plt.show()

# Interpretation practice
#
# OLS picks the line with the smallest sum of squared residuals. R-squared compares residual variation with total variation in the outcome.