# VIF Calculation
# Module 3 - Multicollinearity and VIF
# Dataset: WAGE1

# VIF Calculation
#
# Calculate variance inflation factors for education, experience, and tenure using auxiliary regressions.

# %% Cell 2
from pathlib import Path

DATASET = "WAGE1.DTA"
VARIABLES = ["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 r_squared(target, controls):
    y = df[target].to_numpy()
    X = np.column_stack([np.ones(len(df)), df[controls].to_numpy()])
    beta = np.linalg.lstsq(X, y, rcond=None)[0]
    residuals = y - X @ beta
    return 1 - np.sum(residuals ** 2) / np.sum((y - y.mean()) ** 2)

for target in VARIABLES:
    controls = [name for name in VARIABLES if name != target]
    r2 = r_squared(target, controls)
    vif = 1 / (1 - r2)
    print(f"{target}: auxiliary R-squared={r2:.4f}, VIF={vif:.3f}")

# Interpretation
#
# VIF summarizes how strongly one regressor can be predicted by the other regressors. It is a precision warning, not automatic proof that a model is invalid.