# Simple vs Multiple Regression
# Module 3 - Simple Regression versus Multiple Regression
# Dataset: WAGE1

# Simple vs Multiple Regression
#
# Compare the education coefficient in a short wage regression and a controlled multiple-regression model.

# %% Cell 2
from pathlib import Path

DATASET = "WAGE1.DTA"
VARIABLES = ["lwage", "educ", "exper", "tenure"]
DATA_FOLDER = "https://drive.google.com/drive/folders/1_STdcydIcst-opcbwOKRFzUXsgxQgBoS?usp=sharing"

if not Path(DATASET).exists():
    raise FileNotFoundError(
        "Dataset file not installed yet\n"
        f"Dataset: WAGE1\n"
        f"Variables needed: {', '.join(VARIABLES)}\n"
        f"Course data folder: {DATA_FOLDER}\n"
        "Admin upload instruction: upload the dataset in Admin -> Datasets and make it available to this notebook."
    )

# %% Cell 3
import pandas as pd
import numpy as np

df = pd.read_stata(DATASET)[VARIABLES].dropna()

def ols(y, columns):
    X = np.column_stack([np.ones(len(df)), df[columns].to_numpy()])
    beta = np.linalg.lstsq(X, y.to_numpy(), rcond=None)[0]
    fitted = X @ beta
    residuals = y.to_numpy() - fitted
    r_squared = 1 - np.sum(residuals ** 2) / np.sum((y.to_numpy() - y.mean()) ** 2)
    return dict(zip(["intercept", *columns], beta)), r_squared

simple_beta, simple_r2 = ols(df["lwage"], ["educ"])
multiple_beta, multiple_r2 = ols(df["lwage"], ["educ", "exper", "tenure"])

print("Simple educ coefficient:", round(simple_beta["educ"], 4))
print("Multiple educ coefficient:", round(multiple_beta["educ"], 4))
print("Simple R-squared:", round(simple_r2, 4))
print("Multiple R-squared:", round(multiple_r2, 4))

# Interpretation
#
# If the education coefficient changes after adding controls, the simple regression was mixing education with differences in the added variables.