# R-squared Comparison
# Module 3 - Goodness of Fit in Multiple Regression
# Dataset: WAGE1

# R-squared Comparison
#
# Compare R-squared across nested wage models and separate fit from causal interpretation.

# %% 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 model_stats(columns):
    y = df["lwage"].to_numpy()
    X = np.column_stack([np.ones(len(df)), df[columns].to_numpy()])
    beta = np.linalg.lstsq(X, y, rcond=None)[0]
    residuals = y - X @ beta
    r_squared = 1 - np.sum(residuals ** 2) / np.sum((y - y.mean()) ** 2)
    return r_squared, dict(zip(["intercept", *columns], beta))

for columns in [["educ"], ["educ", "exper"], ["educ", "exper", "tenure"]]:
    r2, beta = model_stats(columns)
    print("Model:", ", ".join(columns))
    print("  R-squared:", round(r2, 4))
    print("  education coefficient:", round(beta["educ"], 4))

# Interpretation
#
# R-squared usually rises when variables are added. That does not automatically mean the new model has a better causal interpretation.