# Regularised Regression: A Machine Learning Toolkit for Econometrics
# Athanassios Stavrakoudis
# astavrak@uoi.gr
# with claude's assistance
#
# Complete Python code for the Ridge part
# Standalone: libraries imported and configuration hard-coded below.

import numpy as np
import pandas as pd
import wooldridge as woo
import matplotlib.pyplot as plt
from sklearn.linear_model import (Lasso, LassoCV, Ridge, RidgeCV,
                                  ElasticNet, ElasticNetCV, LinearRegression)
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.metrics import mean_squared_error, r2_score
import warnings
warnings.filterwarnings("ignore")

SEED       = 14159
N_CORES    = 6
N_CV_FOLDS = 10
N_ALPHAS   = 100
rng = np.random.default_rng(SEED)


# [py-ridge-s1]
import numpy as np
import pandas as pd
import wooldridge as woo
from sklearn.preprocessing import StandardScaler

# - Real data: wagepan. Ridge as prediction of log wage from collinear controls.
wp_py = woo.dataWoo("wagepan")
cand = [c for c in ["exper","expersq","educ","married","union","black","hisp",
                    "south","nrthcen","nrtheast","rur",
                    "agric","bus","construc","ent","fin","manuf","min",
                    "pro","pub","tra","trad",
                    "occ1","occ2","occ3","occ4","occ5","occ6","occ7","occ8","occ9",
                    "poorhlth","hours"]
        if c in wp_py.columns]
dat_py = wp_py[["lwage"] + cand].dropna()
X_py = dat_py[cand].to_numpy(dtype=float)
y_py = dat_py["lwage"].to_numpy()
n, p = X_py.shape

sc   = StandardScaler()
X_sc = sc.fit_transform(X_py)   # centre + scale to σ=1 per predictor
kappa_py = np.linalg.cond(X_sc.T @ X_sc)
print(f"n = {n}  |  p = {p} controls")
print(f"κ(X'X) = {kappa_py:.0f}  [>30 problematic, >1000 severe]")

# 70/30 split — same seed as R
rng2   = np.random.default_rng(SEED)
tr_idx = rng2.choice(n, size=int(0.7 * n), replace=False)
te_idx = np.setdiff1d(np.arange(n), tr_idx)
X_tr_py, X_te_py = X_sc[tr_idx], X_sc[te_idx]
y_tr_py, y_te_py = y_py[tr_idx], y_py[te_idx]
print(f"Train: {len(tr_idx)}  |  Test: {len(te_idx)}")

# [py-ridge-s2]
from sklearn.linear_model import ElasticNetCV, LinearRegression
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
import numpy as np

# OLS on training sample
ols_py = LinearRegression().fit(X_tr_py, y_tr_py)

# Ridge via Pipeline — the ISLP textbook approach (James et al. 2023, §6.5.2)
# Pipeline separates standardisation from model fitting:
#   'scaler': StandardScaler — uses ONLY training-set mean/std (no leakage)
#   'ridge' : ElasticNetCV with l1_ratio=0 → pure Ridge (ℓ₂ penalty only)
#             l1_ratio=1 would give Lasso; 0 < l1_ratio < 1 → Elastic Net
# sklearn calls λ 'alpha'; large alpha = heavy regularisation
alphas   = np.logspace(4, -2, 100)   # search grid: 10^4 down to 10^-2
scaler_r = StandardScaler(with_mean=True, with_std=True)
ridgeCV  = ElasticNetCV(alphas=alphas, l1_ratio=0.0,
                         cv=N_CV_FOLDS, max_iter=10000)
pipe_r   = Pipeline([('scaler', scaler_r), ('ridge', ridgeCV)])
pipe_r.fit(X_tr_py, y_tr_py)

lam_py = ridgeCV.alpha_
print(f"Ridge CV λ (alpha): {lam_py:.4f}")
print(f"Pipeline standardises internally — no manual scaling needed")

# Effective df: df(λ) = Σ d_j² / (d_j² + λ)
# Degrees of freedom decreases from p (OLS) toward 0 as λ → ∞
X_sc_py = scaler_r.transform(X_tr_py)
d2_py   = np.linalg.svd(X_sc_py, compute_uv=False) ** 2
df_py   = np.sum(d2_py / (d2_py + lam_py))
print(f"OLS df: {p}  |  Ridge df(λ): {df_py:.2f}")

# [py-ridge-s3]
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.linear_model import LinearRegression, RidgeCV
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import numpy as np, pandas as pd
import wooldridge as woo

# Self-contained on real wagepan: predict log wage from collinear controls
wp_r3 = woo.dataWoo("wagepan")
cand = [c for c in ["exper","expersq","educ","married","union","black","hisp",
                    "south","nrthcen","nrtheast","rur",
                    "agric","bus","construc","ent","fin","manuf","min",
                    "pro","pub","tra","trad",
                    "occ1","occ2","occ3","occ4","occ5","occ6","occ7","occ8","occ9",
                    "poorhlth","hours"]
        if c in wp_r3.columns]
dat_r3 = wp_r3[["lwage"] + cand].dropna()
X_r3 = StandardScaler().fit_transform(dat_r3[cand].to_numpy(dtype=float))
y_r3 = dat_r3["lwage"].to_numpy()
X_tr_py, X_te_py, y_tr_py, y_te_py = train_test_split(
    X_r3, y_r3, test_size=0.3, random_state=SEED)

ols_py = LinearRegression().fit(X_tr_py, y_tr_py)
rcv    = RidgeCV(alphas=np.logspace(-2, 4, 100)).fit(X_tr_py, y_tr_py)

def gof(m, Xtr, ytr, Xte, yte):
    tr = mean_squared_error(ytr, m.predict(Xtr)) ** 0.5
    te = mean_squared_error(yte, m.predict(Xte)) ** 0.5
    r2 = r2_score(yte, m.predict(Xte))
    return {"Train RMSE": tr, "Test RMSE": te, "Test R²": r2,
            "Overfit ratio": round(tr / te, 3)}

rows = {
    "OLS"          : gof(ols_py, X_tr_py, y_tr_py, X_te_py, y_te_py),
    "Ridge (λ.min)": gof(rcv,    X_tr_py, y_tr_py, X_te_py, y_te_py),
}
gof_df = pd.DataFrame(rows).T.round(4)
print(gof_df.to_string())

# [py-ridge-wagepan]
import pandas as pd, numpy as np, warnings
from sklearn.linear_model import RidgeCV
from sklearn.preprocessing import StandardScaler
warnings.filterwarnings('ignore')

# Self-contained: reload wagepan, demean, build controls
try:
    import wooldridge as woo
    wp_r = woo.data("wagepan")
except Exception:
    wp_r = pd.read_csv("../data/wagepan.csv")

ctrl_c = [c for c in ["exper","expersq","married","educ","black","hisp","south","smsa",
                       "agric","bus","construc","ndurman","trcommpu","trade",
                       "services","profserv","profocc","clerocc","servocc"]
          if c in wp_r.columns]
use_c = ["nr","year","lwage","union"] + ctrl_c
wp_r  = wp_r[[c for c in use_c if c in wp_r.columns]].dropna()
num_c = [c for c in wp_r.columns if c not in ["nr","year"]]
wp_r[num_c] = wp_r[num_c] - wp_r.groupby("nr")[num_c].transform("mean")
yr_d  = pd.get_dummies(wp_r["year"], prefix="yr", drop_first=True).astype(float)
ctrl_all_r = pd.concat([wp_r[ctrl_c], yr_d], axis=1).values
y_r_py = wp_r["lwage"].values
D_r_py = wp_r["union"].values

# TWFE baseline
X_twfe_r = np.column_stack([D_r_py, ctrl_all_r])
b_twfe_r  = np.linalg.lstsq(X_twfe_r, y_r_py, rcond=None)[0][0]
print(f"TWFE union premium: {b_twfe_r:.4f}")

# Ridge partial-out
sc2 = StandardScaler().fit(ctrl_all_r)
X_sc2 = sc2.transform(ctrl_all_r)
alphas_r = np.logspace(-2, 4, 100)
ridge_y_py = RidgeCV(alphas=alphas_r, cv=N_CV_FOLDS).fit(X_sc2, y_r_py)
ridge_d_py = RidgeCV(alphas=alphas_r, cv=N_CV_FOLDS).fit(X_sc2, D_r_py)

yres_py = y_r_py - ridge_y_py.predict(X_sc2)
Dres_py = D_r_py - ridge_d_py.predict(X_sc2)
b_ridge_py  = np.dot(Dres_py, yres_py) / np.dot(Dres_py, Dres_py)
resid_rr    = yres_py - b_ridge_py * Dres_py
psi_rr      = Dres_py * resid_rr
se_ridge_py = ((Dres_py**2).mean()**(-2) * (psi_rr**2).mean() / len(y_r_py))**0.5

d2_py2 = np.linalg.svd(X_sc2, compute_uv=False)**2
df_py2 = np.sum(d2_py2 / (d2_py2 + ridge_y_py.alpha_))
print(f"Ridge union premium: {b_ridge_py:.4f}  SE: {se_ridge_py:.4f}")
print(f"Ridge λ*: {ridge_y_py.alpha_:.4f}  effective df: {df_py2:.1f}")

from tabulate import tabulate
print(tabulate([
    ["TWFE",                f"{b_twfe_r:.4f}", "demeaned OLS", "—"],
    ["Ridge (partial-out)", f"{b_ridge_py:.4f}", f"df {df_py2:.1f}", f"{ridge_y_py.alpha_:.4f}"]],
    headers=["Estimator","Union premium","Controls/df","λ*"],
    tablefmt="rounded_outline"))