# Regularised Regression: A Machine Learning Toolkit for Econometrics
# Athanassios Stavrakoudis
# astavrak@uoi.gr
# with claude's assistance
#
# Complete Python code for the Lasso 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-lasso-optim]
import numpy as np
import wooldridge as woo
from sklearn.linear_model import Lasso

wp = woo.dataWoo("wagepan")
cols = [c for c in ["educ", "exper", "expersq", "married", "union", "hours"]
        if c in wp.columns]
X = wp[cols].to_numpy(dtype=float)
X = (X - X.mean(0)) / X.std(0)                 # standardise: comparable scales
y = wp["lwage"].to_numpy() - wp["lwage"].mean()  # centre: intercept drops out
n, p = X.shape

def soft_threshold(z, g):
    return np.sign(z) * max(abs(z) - g, 0.0)

def lasso_cd(X, y, lam, max_iter=1000, tol=1e-7):
    n, p = X.shape
    beta = np.zeros(p)                          # warm start at zero
    for it in range(max_iter):
        beta_old = beta.copy()
        for j in range(p):
            r_j = y - X @ beta + X[:, j] * beta[j]   # partial residual
            z_j = X[:, j] @ r_j / n
            beta[j] = soft_threshold(z_j, lam)
        if np.max(np.abs(beta - beta_old)) < tol:
            break
    return beta, it + 1

beta, iters = lasso_cd(X, y, lam=0.05)
print(f"From-scratch coordinate descent: converged in {iters} sweeps")
print({v: round(b, 4) for v, b in zip(cols, beta)})

# scikit-learn runs the same algorithm; these are its convergence controls:
m = Lasso(alpha=0.05, fit_intercept=False, max_iter=10000, tol=1e-7)
m.fit(X, y)
print("sklearn coefficients (same alpha):", np.round(m.coef_, 4))
# A "ConvergenceWarning: Objective did not converge" means max_iter was hit
# before tol — raise max_iter, loosen tol, or check standardisation.

# [py-lasso-wp-1]
import pandas as pd, numpy as np
import warnings; warnings.filterwarnings('ignore')
import wooldridge as woo
wp = woo.dataWoo("wagepan")

ctrl   = ["exper", "expersq", "married", "educ", "black", "hisp", "south"]
wp_use = wp[["nr", "year", "lwage", "union"] + ctrl].dropna()

# Within-demean: subtract individual (nr) means
num = [c for c in wp_use.columns if c not in ["nr", "year"]]
wp_dm = wp_use.copy()
wp_dm[num] = wp_dm[num] - wp_dm.groupby("nr")[num].transform("mean")

# TWFE baseline: OLS on demeaned data + year dummies
yr_dum = pd.get_dummies(wp_dm["year"], prefix="yr", drop_first=True).astype(float)
X_twfe = np.column_stack([wp_dm[["union", "exper", "expersq", "married", "educ"]].values,
                          yr_dum.values])
y_tw   = wp_dm["lwage"].values
b_twfe = np.linalg.lstsq(X_twfe, y_tw, rcond=None)[0][0]
print(f"TWFE union premium: {b_twfe:.4f}")

# [py-lasso-wp-2]
from sklearn.preprocessing import StandardScaler
num_cols = [c for c in wp_dm.columns if c not in ["nr", "year", "lwage", "union"]]
ctrl_all_py = np.column_stack([wp_dm[num_cols].values, yr_dum.values])
D_py = wp_dm["union"].values
y_py = wp_dm["lwage"].values
X_sc_py = StandardScaler().fit_transform(ctrl_all_py)
print(f"Controls: {X_sc_py.shape[1]} predictors, {X_sc_py.shape[0]} observations")

# [py-lasso-wp-3]
from sklearn.linear_model import LassoCV
lasso_y = LassoCV(cv=N_CV_FOLDS, max_iter=5000, n_jobs=N_CORES
                  ).fit(np.column_stack([D_py, X_sc_py]), y_py)
lasso_d = LassoCV(cv=N_CV_FOLDS, max_iter=5000, n_jobs=N_CORES).fit(X_sc_py, D_py)

sel_y = set(np.where(lasso_y.coef_[1:] != 0)[0])   # skip the D coefficient
sel_d = set(np.where(lasso_d.coef_ != 0)[0])
sel_u = sorted(sel_y | sel_d)                      # union selection rule
print(f"y-eq: {len(sel_y)}  D-eq: {len(sel_d)}  Union: {len(sel_u)}")

# [py-lasso-crime4]
import numpy as np
import pandas as pd
import wooldridge as woo
from sklearn.linear_model import LassoCV, LinearRegression

cr = woo.dataWoo("crime4").copy()
cr["lcrmrte"] = np.log(cr["crmrte"])
cr["lprbarr"] = np.log(cr["prbarr"])
cr["lpolpc"]  = np.log(cr["polpc"])

focal = "lprbarr"
ctrl  = [c for c in ["prbconv", "prbpris", "avgsen", "lpolpc", "density", "taxpc",
                     "pctmin80", "pctymle", "west", "central", "urban",
                     "wcon", "wtuc", "wtrd", "wfir", "wser", "wmfg",
                     "wfed", "wsta", "wloc"] if c in cr.columns]
cr = cr.dropna(subset=["lcrmrte", focal] + ctrl)

# Two-way FE: within-county demean + year dummies
num = ["lcrmrte", focal] + ctrl
cr[num] = cr.groupby("county")[num].transform(lambda v: v - v.mean()) + cr[num].mean()
yr = pd.get_dummies(cr["year"], drop_first=True).astype(float).to_numpy()
Xc = np.column_stack([cr[ctrl].to_numpy(float), yr])

cv_y = LassoCV(cv=10).fit(Xc, cr["lcrmrte"].to_numpy())
cv_d = LassoCV(cv=10).fit(Xc, cr[focal].to_numpy())
sel  = np.where((cv_y.coef_ != 0) | (cv_d.coef_ != 0))[0]

Xpds = np.column_stack([cr[focal].to_numpy(), Xc[:, sel]])
pds  = LinearRegression().fit(Xpds, cr["lcrmrte"].to_numpy())
print(f"Controls selected (union): {len(sel)} of {Xc.shape[1]}")
print(f"Deterrence elasticity (PDS Lasso): {pds.coef_[0]:.3f}")

# [py-cv-lasso]
import numpy as np
import pandas as pd
from sklearn.linear_model import LassoCV, RidgeCV
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import make_pipeline

rng  = np.random.default_rng(14159)
n, p_main, p_signal = 500, 50, 5
beta_signal = np.array([1.5, -1.2, 0.8, -0.5, 1.0])

X = rng.standard_normal((n, p_main))
D = 0.5*X[:,0] - 0.4*X[:,1] + 0.3*X[:,2] + rng.standard_normal(n)
y = 2.0*D + X[:,:p_signal] @ beta_signal + rng.standard_normal(n)

# Pipeline: StandardScaler → LassoCV
#   StandardScaler   : centres and scales each feature to μ=0, σ=1 before Lasso
#                      (same role as standardize=TRUE in glmnet — required for fair penalisation)
#   LassoCV          : fits Lasso over a grid of alpha values (sklearn calls λ "alpha")
#     cv=10          : 10-fold CV; same as nfolds=10 in R's cv.glmnet
#     max_iter=5000  : solver (coordinate descent) iterations; increase if convergence warnings appear
#     n_alphas=100   : number of λ values on the regularisation path (denser = smoother path)
lasso_cv = make_pipeline(StandardScaler(),
                         LassoCV(cv=N_CV_FOLDS, max_iter=5000, n_alphas=N_ALPHAS))
lasso_cv.fit(X, y)

# RidgeCV: alpha_ is the λ chosen by CV
#   alphas = logspace(-3, 4, 100): tests 100 λ values from 0.001 to 10,000
#     lower bound: near-OLS; upper bound: heavy shrinkage toward zero
#   cv=10: 10-fold CV (same as Lasso for fair comparison)
alphas   = np.logspace(-3, 4, 100)
ridge_cv = make_pipeline(StandardScaler(),
                         RidgeCV(alphas=alphas, cv=N_CV_FOLDS))
ridge_cv.fit(X, y)

print(f"Lasso  λ_min = {lasso_cv['lassocv'].alpha_:.4f}")
print(f"Ridge  λ_min = {ridge_cv['ridgecv'].alpha_:.4f}")

# sklearn's LassoCV exposes alpha_ (= λ.min in R notation)
# There is no built-in λ.1se; use cross_val_score manually if needed
coef_l = lasso_cv['lassocv'].coef_
nnz    = np.sum(coef_l != 0)
sel_idx = np.where(coef_l != 0)[0]
print(f"\nLasso selected {nnz}/{p_main} regressors (true non-zero: {p_signal})")
print(f"  Selected indices: {sel_idx}")
print(f"  True signal indices: 0 to {p_signal-1}")
print(f"  Correctly identified: {sum(i < p_signal for i in sel_idx)}/{p_signal}")

# - CV error path plot — mirrors R's ggplot
import matplotlib.pyplot as plt
from sklearn.linear_model import Ridge
from sklearn.model_selection import cross_val_score

fig, axes = plt.subplots(1, 2, figsize=(12, 4))

# Lasso panel: LassoCV stores the full path (alphas_, mse_path_)
lcv = lasso_cv['lassocv']
log_a    = np.log(lcv.alphas_)
mse_mean = lcv.mse_path_.mean(axis=1)
mse_std  = lcv.mse_path_.std(axis=1)
_ = axes[0].fill_between(log_a, mse_mean - mse_std, mse_mean + mse_std,
                         alpha=0.15, color="#1a6ea8")
_ = axes[0].plot(log_a, mse_mean, color="#1a6ea8", lw=1.4, label="Lasso")
_ = axes[0].axvline(np.log(lcv.alpha_), color="#1a6ea8", ls="--", lw=1.2,
                    label=f"λ_min = {lcv.alpha_:.4f}")
_ = axes[0].set_xlabel(r"log(λ)"); axes[0].set_ylabel("CV Mean Squared Error")
axes[0].set_title("Lasso: 10-fold CV Error Path\n(band = ±1 SD across folds)",
                  fontweight="bold")
_ = axes[0].legend(fontsize=9); axes[0].grid(True, color="#e8e8e8")

# Ridge panel: RidgeCV does NOT expose a per-alpha MSE path, so compute it
# explicitly with cross_val_score over the same alpha grid.
ridge_mse = [-cross_val_score(Ridge(alpha=a), X, y, cv=N_CV_FOLDS,
                              scoring="neg_mean_squared_error").mean()
             for a in alphas]
log_ra    = np.log(alphas)
best_ra   = ridge_cv['ridgecv'].alpha_
_ = axes[1].plot(log_ra, ridge_mse, color="#e8521a", lw=1.4, label="Ridge")
_ = axes[1].axvline(np.log(best_ra), color="#e8521a", ls="--", lw=1.2,
                    label=f"λ_min = {best_ra:.4f}")
_ = axes[1].set_xlabel(r"log(λ)"); axes[1].set_ylabel("CV Mean Squared Error")
axes[1].set_title("Ridge: 10-fold CV Error Path",
                  fontweight="bold")
_ = axes[1].legend(fontsize=9); axes[1].grid(True, color="#e8e8e8")

_ = fig.suptitle("Cross-Validation Path: Lasso vs Ridge (n=500, p=50, s=5)",
             fontsize=11, fontweight="bold")
fig.tight_layout(); plt.show()

# [py-pcr-pls]
from sklearn.decomposition import PCA
from sklearn.cross_decomposition import PLSRegression
from sklearn.linear_model import LinearRegression
from sklearn.pipeline import Pipeline
from sklearn.model_selection import cross_val_score
from sklearn.preprocessing import StandardScaler
import numpy as np
import wooldridge as woo

# Self-contained design from wagepan (the deck's running dataset): predict lwage
# from a standardised block of controls — the high-dimensional setting where
# PCR and PLS are natural alternatives to Ridge.
wp_pcr = woo.dataWoo("wagepan")
cols = [c for c in ["exper","expersq","married","educ","black","hisp","south",
                    "agric","bus","construc","manuf","fin","tra","trad","pub",
                    "occ1","occ2","occ3","occ4","occ5","occ6","occ7","occ8","occ9"]
        if c in wp_pcr.columns]
X_pcr_py = StandardScaler().fit_transform(wp_pcr[cols].to_numpy(dtype=float))
y_pcr_py = (wp_pcr["lwage"] - wp_pcr["lwage"].mean()).to_numpy()

# PCR: PCA + LinearRegression in a Pipeline
# n_components M is the tuning parameter — choose by CV
M_grid = range(1, min(21, X_pcr_py.shape[1] + 1))
cv_mse_pcr, cv_mse_pls = [], []

for M in M_grid:
    # PCR pipeline
    pipe_pcr = Pipeline([('pca', PCA(n_components=M)),
                          ('ols', LinearRegression())])
    score_pcr = -cross_val_score(pipe_pcr, X_pcr_py, y_pcr_py,
                                  cv=N_CV_FOLDS, scoring='neg_mean_squared_error').mean()
    # PLS
    pls = PLSRegression(n_components=M)
    score_pls = -cross_val_score(pls, X_pcr_py, y_pcr_py,
                                  cv=N_CV_FOLDS, scoring='neg_mean_squared_error').mean()
    cv_mse_pcr.append(score_pcr)
    cv_mse_pls.append(score_pls)

M_pcr_py = np.argmin(cv_mse_pcr) + 1
M_pls_py = np.argmin(cv_mse_pls) + 1
print(f"PCR optimal M = {M_pcr_py} | PLS optimal M = {M_pls_py}")

import matplotlib.pyplot as plt
fig, ax = plt.subplots(figsize=(9, 3.6))
_ = ax.plot(M_grid, cv_mse_pcr, color="#1a6ea8", lw=1.8, label="PCR")
_ = ax.plot(M_grid, cv_mse_pls, color="#e8521a", lw=1.8, label="PLS")
_ = ax.axvline(M_pcr_py, color="#1a6ea8", ls="--", lw=1)
_ = ax.axvline(M_pls_py, color="#e8521a", ls="--", lw=1)
_ = ax.set_xlabel("M (number of components)"); ax.set_ylabel("10-fold CV MSE")
ax.set_title("PCR vs PLS: CV error by number of components",
             fontsize=10, fontweight="bold")
_ = ax.legend(fontsize=9); ax.grid(True, color="#e8e8e8")
plt.tight_layout(); plt.show()