# Simple Regression in Python
# Module 2 - Regression in Python
# Dataset: /data/wage_sample.csv

# Simple Regression in Python
#
# Estimate a simple regression of wage on education. This notebook runs in the browser and can also be downloaded as a regular `.ipynb` file. If `statsmodels` is available, the notebook uses it. If not, it uses the same OLS formulas directly.

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

df = pd.read_csv("wage_sample.csv")
print(df.head())
print("Rows:", len(df))

# %% Cell 3
plt.scatter(df["education"], df["wage"])
plt.xlabel("Years of education")
plt.ylabel("Hourly wage")
plt.title("Wage and education")
plt.show()

# %% Cell 4
x = df["education"]
y = df["wage"]

try:
    import statsmodels.api as sm
    X = sm.add_constant(x)
    model = sm.OLS(y, X).fit()
    intercept = model.params["const"]
    slope = model.params["education"]
    r_squared = model.rsquared
    print(model.summary())
except Exception as error:
    print("statsmodels is not available in this browser runtime.")
    print("Using the same OLS formulas directly instead.")
    x_bar = x.mean()
    y_bar = y.mean()
    slope = ((x - x_bar) * (y - y_bar)).sum() / ((x - x_bar) ** 2).sum()
    intercept = y_bar - slope * x_bar
    fitted = intercept + slope * x
    residuals = y - fitted
    ssr = (residuals ** 2).sum()
    sst = ((y - y_bar) ** 2).sum()
    r_squared = 1 - ssr / sst

print("Intercept:", round(intercept, 2))
print("Education slope:", round(slope, 2))
print("R-squared:", round(r_squared, 3))

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

plt.scatter(df["education"], df["wage"], label="Actual workers")
plt.plot(education_grid, predicted_wage, label="Fitted regression line")
plt.xlabel("Years of education")
plt.ylabel("Hourly wage")
plt.title("Simple regression fitted line")
plt.legend()
plt.show()

# Interpretation practice
#
# In this teaching sample, one more year of education is associated with about the estimated slope dollars higher predicted hourly wage. This is an association from a simple model, not automatic proof of causation.