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

suppressPackageStartupMessages({
  library(glmnet)
  library(plm)
  library(sandwich)
  library(lmtest)
  library(wooldridge)
  library(xtable)
  library(tidyverse)
  library(patchwork)
})

# Reproducibility and parallel configuration
SEED       <- 14159L
N_CORES    <- 6L
N_THREADS  <- 12L
N_CV_FOLDS <- 10L
N_ALPHAS   <- 100L
N_MC_REPS  <- 200L
N_TREES    <- 1000L
set.seed(SEED)

# Lecture colour palette and ggplot theme
col_main   <- "#185FA5"
col_accent <- "#D85A30"
col_ok     <- "#1D9E75"
col_muted  <- "#6B7280"
col_warn   <- "#B22222"
theme_lecture <- theme_minimal(base_size = 13)


# [r-enet-wagepan]
library(wooldridge)
library(glmnet)
library(plm)
library(sandwich)

# - Load wagepan from the wooldridge package (hard fail if not installed)
wpe <- wagepan

# - Candidate controls for the union-premium specification
ctrl <- c("exper", "expersq", "married", "educ", "black", "hisp", "south")
wpe  <- wpe[, c("nr", "year", "lwage", "union", ctrl)]

# - Within-transform (remove individual fixed effect)
wpe_dm <- wpe %>%
  group_by(nr) %>%
  mutate(across(where(is.numeric),
                ~ . - mean(., na.rm = TRUE) + mean(wpe[[cur_column()]], na.rm = TRUE))) %>%
  ungroup()

# - Control matrix: demeaned covariates + year dummies
yr_dummies <- model.matrix(~ factor(year) - 1, data = wpe_dm)[, -1]
num_cols   <- setdiff(names(wpe_dm), c("nr", "year", "lwage", "union"))
ctrl_r2    <- cbind(as.matrix(wpe_dm[, num_cols]), yr_dummies)
y_r2       <- wpe_dm$lwage
D_r2       <- wpe_dm$union

# - TWFE baseline via plm (two-way within, cluster-robust SE)
pdate    <- pdata.frame(wpe, index = c("nr", "year"))
twfe_me  <- plm(lwage ~ union + exper + expersq + married + educ,
                data = pdate, model = "within", effect = "twoways")
b_twfe2  <- coef(twfe_me)["union"]
se_twfe2 <- sqrt(vcovHC(twfe_me, cluster = "group")["union", "union"])

cat(sprintf("TWFE: %.4f  |  Controls: %d obs × %d predictors\n",
            b_twfe2, nrow(ctrl_r2), ncol(ctrl_r2)))

# Lasso results (rebuild for comparison table)
set.seed(SEED)
cv_y_r  <- cv.glmnet(x = ctrl_r2, y = y_r2, alpha = 1, nfolds = N_CV_FOLDS)
cv_d_r  <- cv.glmnet(x = ctrl_r2, y = D_r2, alpha = 1, nfolds = N_CV_FOLDS)
sel_y_r <- which(coef(cv_y_r, s = "lambda.min")[-1] != 0)
sel_d_r <- which(coef(cv_d_r, s = "lambda.min")[-1] != 0)
sel_union <- union(sel_y_r, sel_d_r)
yres_l  <- y_r2 - predict(cv_y_r, ctrl_r2, s = "lambda.min")
Dres_l  <- D_r2 - predict(cv_d_r, ctrl_r2, s = "lambda.min")
lasso_ols <- lm(yres_l ~ Dres_l)
b_post   <- coef(lasso_ols)["Dres_l"]
se_post  <- sqrt(vcovHC(lasso_ols, "HC3")["Dres_l","Dres_l"])

# Ridge results (rebuild)
cv_ry <- cv.glmnet(x = ctrl_r2, y = y_r2, alpha = 0, nfolds = N_CV_FOLDS)
cv_rd <- cv.glmnet(x = ctrl_r2, y = D_r2, alpha = 0, nfolds = N_CV_FOLDS)
yres_r  <- y_r2 - predict(cv_ry, ctrl_r2, s = "lambda.min")
Dres_r  <- D_r2 - predict(cv_rd, ctrl_r2, s = "lambda.min")
ridge_ols2 <- lm(yres_r ~ Dres_r)
b_ridge2   <- coef(ridge_ols2)["Dres_r"]
se_ridge2  <- sqrt(vcovHC(ridge_ols2, "HC3")["Dres_r","Dres_r"])
d2_r2 <- svd(ctrl_r2, nu = 0, nv = 0)$d ^ 2
df_r2 <- sum(d2_r2 / (d2_r2 + cv_ry$lambda.min))

# - EN: alpha grid with lapply (5 calls — lapply is fine, no parallel needed)
# future_map with plan(multisession) cannot see variables from the parent session.
# For 5 alpha values lapply is fast enough and avoids all scoping issues.
alpha_vals_wp <- c(0.1, 0.25, 0.5, 0.75, 0.9)
set.seed(SEED)

# lapply sees ctrl_r2/y_r2/D_r2 directly (same R session — no scoping issue)
cv_en_wp   <- lapply(alpha_vals_wp, \(a)
  cv.glmnet(x = ctrl_r2, y = y_r2, alpha = a, nfolds = N_CV_FOLDS, standardize = TRUE))
cv_en_d_wp <- lapply(alpha_vals_wp, \(a)
  cv.glmnet(x = ctrl_r2, y = D_r2, alpha = a, nfolds = N_CV_FOLDS, standardize = TRUE))

best_idx_wp  <- which.min(sapply(cv_en_wp,   \(cv) min(cv$cvm)))
best_d_wp    <- which.min(sapply(cv_en_d_wp, \(cv) min(cv$cvm)))
cv_best_wp   <- cv_en_wp[[best_idx_wp]]
cv_d_best    <- cv_en_d_wp[[best_d_wp]]
alpha_best_wp<- alpha_vals_wp[best_idx_wp]

# Partial-out using best EN
yres_en <- y_r2 - predict(cv_best_wp, ctrl_r2, s="lambda.min")
Dres_en <- D_r2 - predict(cv_d_best,  ctrl_r2, s="lambda.min")

en_ols   <- lm(yres_en ~ Dres_en)
b_en_wp  <- coef(en_ols)["Dres_en"]
se_en_wp <- sqrt(vcovHC(en_ols,"HC3")["Dres_en","Dres_en"])

# How many controls selected by EN?
b_en_coef  <- coef(cv_best_wp, s="lambda.min")[-1]
nsel_en_wp <- sum(b_en_coef != 0)

cat(sprintf("EN best α* = %.2f  λ* = %.4f\n", alpha_best_wp, cv_best_wp$lambda.min))
cat(sprintf("Controls selected: %d/%d\n", nsel_en_wp, ncol(ctrl_r2)))
cat(sprintf("EN union premium: %.4f  (HC3 SE: %.4f)\n", b_en_wp, se_en_wp))

# - Which controls Elastic Net selected, with their coefficients
# EN both zeroes and shrinks; show the survivors and tag occupation/industry
# membership so the GROUP structure (the reason for EN) is visible.
occ_ind <- c("agric","bus","construc","ndurman","trcommpu","trade",
             "services","profserv","profocc","clerocc","servocc")
sel_en_idx <- which(b_en_coef != 0)
en_sel_tbl <- tibble(
  Control     = colnames(ctrl_r2)[sel_en_idx],
  `EN b`      = round(as.numeric(b_en_coef[sel_en_idx]), 4),
  Group       = ifelse(colnames(ctrl_r2)[sel_en_idx] %in% occ_ind,
                       "occ/industry", "base/year")
)
en_sel_tbl <- arrange(en_sel_tbl, desc(abs(`EN b`)))
cat(sprintf("\nElastic Net selected %d controls (alpha* = %.2f):\n",
            nsel_en_wp, alpha_best_wp))
print(as.data.frame(en_sel_tbl), row.names = FALSE)

# Final comparison: all three methods
tibble(
  Estimator        = c("TWFE","Post-Lasso OLS","Ridge (partial)","Elastic Net (partial)"),
  `Union premium`  = c(b_twfe2, b_post, b_ridge2, b_en_wp),
  `SE`             = c(se_twfe2, se_post, se_ridge2, se_en_wp),
  `α*`             = c("—","1.00","0.00", sprintf("%.2f", alpha_best_wp)),
  `Controls/df`    = c("5 fixed",
                       sprintf("%d selected",length(sel_union)),
                       sprintf("%.1f eff. df", df_r2),
                       sprintf("%d selected",nsel_en_wp))
) %>%
  kbl(caption="wagepan: Union Premium — TWFE, Lasso, Ridge, Elastic Net", digits=4) %>%
  kable_styling(font_size=19, full_width=TRUE) %>%
  row_spec(4, bold=TRUE, color="white", background=col_main)