# Regularised Regression: A Machine Learning Toolkit for Econometrics
# Athanassios Stavrakoudis
# astavrak@uoi.gr
# with claude's assistance
#
# Complete R code for the Lasso 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-lasso-optim]
wp <- wagepan

# Standardised design from a few wagepan controls (coordinate descent needs
# comparable scales); centre the outcome so the intercept drops out.
vars <- c("educ", "exper", "expersq", "married", "union", "hours")
vars <- intersect(vars, names(wp))
X <- scale(as.matrix(wp[, vars]))
y <- wp$lwage - mean(wp$lwage)
n <- nrow(X)
p <- ncol(X)

# Soft-thresholding operator: S(z, g) = sign(z) * max(|z| - g, 0)
soft_threshold <- function(z, g) sign(z) * pmax(abs(z) - g, 0)

# Cyclic coordinate descent for the Lasso (glmnet objective (1/2n)RSS + lambda*||b||_1)
lasso_cd <- function(X, y, lambda, max_iter = 1000L, tol = 1e-7) {
  n <- nrow(X)
  p <- ncol(X)
  beta <- rep(0, p)                                  # warm start at zero
  for (iter in seq_len(max_iter)) {
    beta_old <- beta
    for (j in seq_len(p)) {
      r_j <- y - X[, -j, drop = FALSE] %*% beta[-j]  # partial residual
      z_j <- sum(X[, j] * r_j) / n                   # (1/n) x_j' r_j  (x_j standardised)
      beta[j] <- soft_threshold(z_j, lambda)
    }
    if (max(abs(beta - beta_old)) < tol) break       # convergence check
  }
  list(beta = beta, iters = iter)
}

fit <- lasso_cd(X, y, lambda = 0.05)
cat(sprintf("From-scratch coordinate descent: converged in %d sweeps\n", fit$iters))
print(round(setNames(fit$beta, vars), 4))

# glmnet runs the same algorithm; these are its convergence controls:
library(glmnet)
g <- glmnet(X, y, alpha = 1, lambda = 0.05,
            standardize = FALSE,   # X already standardised above
            thresh = 1e-7,         # convergence threshold (default 1e-7)
            maxit = 1e5)           # max passes over coordinates (default 1e5)
cat("glmnet coefficients (same lambda):\n")
print(round(as.numeric(coef(g))[-1], 4))

# [r-lasso-wp-1]
library(wooldridge)
library(plm)
library(lmtest)
library(sandwich)
wp <- wagepan

# Candidate controls used in the TWFE baseline
ctrl <- c("exper", "expersq", "married", "educ", "black", "hisp", "south")
wp   <- wp[, c("nr", "year", "lwage", "union", ctrl)]

# Within-transform: remove the individual (nr) fixed effect, add back the grand
# mean so variables stay on their original scale.
wp_dm <- wp %>%
  group_by(nr) %>%
  mutate(across(where(is.numeric),
                ~ . - mean(., na.rm = TRUE) + mean(wp[[cur_column()]], na.rm = TRUE))) %>%
  ungroup()

# TWFE baseline (two-way within, cluster-robust SE)
pdat    <- pdata.frame(wp, index = c("nr", "year"))
twfe_m  <- plm(lwage ~ union + exper + expersq + married + educ,
               data = pdat, model = "within", effect = "twoways")
b_twfe  <- coef(twfe_m)["union"]
se_twfe <- sqrt(vcovHC(twfe_m, cluster = "group")["union", "union"])
cat(sprintf("TWFE union premium: %.4f  (clustered SE: %.4f)\n", b_twfe, se_twfe))

# [r-lasso-wp-2]
yr_dummies <- model.matrix(~ factor(year) - 1, data = wp_dm)[, -1]
num_cols   <- setdiff(names(wp_dm), c("nr", "year", "lwage", "union"))
X_demeaned <- as.matrix(wp_dm[, num_cols])
ctrl_all   <- cbind(X_demeaned, yr_dummies)
y_wp       <- wp_dm$lwage
D_wp       <- wp_dm$union
cat(sprintf("Control matrix: %d obs × %d predictors\n",
            nrow(ctrl_all), ncol(ctrl_all)))

# [r-lasso-wp-3]
set.seed(SEED)
cv_y <- cv.glmnet(cbind(D_wp, ctrl_all), y_wp, alpha = 1,
                  nfolds = N_CV_FOLDS, standardize = TRUE)   # y on [D, X]
cv_d <- cv.glmnet(ctrl_all, D_wp, alpha = 1,
                  nfolds = N_CV_FOLDS, standardize = TRUE)   # D on X

b_y   <- coef(cv_y, s = "lambda.min")[-1]   # drop intercept; 1st = D, rest = controls
b_d   <- coef(cv_d, s = "lambda.min")[-1]
sel_y <- which(b_y[-1] != 0)                # controls in y-equation (drop D)
sel_d <- which(b_d != 0)                    # controls in D-equation
sel_union <- union(sel_y, sel_d)            # union selection rule
cat(sprintf("y-equation selects: %d  |  D-equation selects: %d  |  Union: %d\n",
            length(sel_y), length(sel_d), length(sel_union)))

# [r-lasso-wp-4]
X_sel_wp <- ctrl_all[, sel_union, drop = FALSE]
post_wp  <- lm(y_wp ~ D_wp + X_sel_wp)
b_post   <- coef(post_wp)["D_wp"]
se_post  <- sqrt(vcovHC(post_wp, "HC3")["D_wp", "D_wp"])
cat(sprintf("Post-Lasso union premium: %.4f  (HC3 SE: %.4f)\n", b_post, se_post))

# Which controls were selected, with the numbers behind the selection
ctrl_names <- colnames(ctrl_all)
post_coefs <- coef(post_wp)
post_named <- setNames(rep(NA_real_, length(sel_union)), ctrl_names[sel_union])
for (k in seq_along(sel_union))
  post_named[k] <- post_coefs[paste0("X_sel_wp", ctrl_names[sel_union][k])]

selected_tbl <- tibble(
  Control           = ctrl_names[sel_union],
  `In y-eqn`        = ifelse(sel_union %in% sel_y, "yes", ""),
  `In D-eqn`        = ifelse(sel_union %in% sel_d, "yes", ""),
  `Lasso b (y-eqn)` = round(b_y[-1][sel_union], 4),
  `Post-Lasso b`    = round(as.numeric(post_named), 4)
)
selected_tbl <- arrange(selected_tbl, desc(abs(`Post-Lasso b`)))
cat(sprintf("\nControls selected by the double-Lasso union rule: %d of %d\n",
            length(sel_union), ncol(ctrl_all)))
print(as.data.frame(selected_tbl), row.names = FALSE)

results_wp <- tibble(
  Estimator = c("TWFE (baseline)", "Post-Lasso OLS"),
  `Union premium` = c(b_twfe, b_post),
  `SE`           = c(se_twfe, se_post),
  `Controls`     = c("5 hand-picked", sprintf("%d Lasso-selected", length(sel_union))),
  `95% CI`       = sprintf("[%.4f, %.4f]",
                            c(b_twfe, b_post) - 1.96*c(se_twfe, se_post),
                            c(b_twfe, b_post) + 1.96*c(se_twfe, se_post))
)
cat("\nwagepan: TWFE vs Post-Lasso OLS — Union Wage Premium\n")
print(as.data.frame(results_wp), row.names = FALSE)

# [r-lasso-wp-latex]
library(xtable)
tab_wp <- data.frame(
  Estimator  = c("TWFE", "Post-Lasso OLS"),
  `$\\hat\\beta_{union}$`  = c(b_twfe, b_post),
  SE         = c(se_twfe, se_post),
  Controls   = c("5 (hand-picked)", sprintf("%d (Lasso-selected)", length(sel_union))),
  `$\\lambda^*$` = c("—", sprintf("%.4f", cv_y$lambda.min)),
  check.names = FALSE, stringsAsFactors = FALSE
)
print(xtable(tab_wp,
             caption = "wagepan: Union Wage Premium — TWFE vs Post-Lasso OLS",
             label   = "tab:lasso-wagepan",
             digits  = 4),
      include.rownames = FALSE,
      booktabs         = TRUE,
      sanitize.text.function = identity,
      comment          = FALSE)

# [r-lasso-wp-tests]
# Post-Lasso OLS residuals
resid_pl <- residuals(post_wp)

# 1. Heteroskedasticity: Breusch-Pagan test
bp_test <- lmtest::bptest(post_wp)
cat(sprintf("Breusch-Pagan (heteroskedasticity): χ²(df=%d) = %.4f  p = %.4f\n",
            bp_test$parameter, bp_test$statistic, bp_test$p.value))
cat(sprintf("  → %s\n",
            if(bp_test$p.value < 0.05) "REJECT H0: heteroskedastic residuals → use HC3 SEs (already applied)"
            else "FAIL TO REJECT H0: residuals approximately homoskedastic"))

# 2. Serial correlation: Durbin-Watson
dw_test <- lmtest::dwtest(post_wp)
cat(sprintf("Durbin-Watson (serial corr.): DW = %.4f  p = %.4f\n",
            dw_test$statistic, dw_test$p.value))

# 3. RESET test: functional form
reset_test <- lmtest::resettest(post_wp, power = 2:3)
cat(sprintf("RESET (functional form): F(%d,%d) = %.4f  p = %.4f\n",
            reset_test$parameter[1], reset_test$parameter[2],
            reset_test$statistic, reset_test$p.value))
cat(sprintf("  → %s\n",
            if(reset_test$p.value < 0.05)
              "REJECT H0: non-linear terms significant → model may be misspecified"
            else "FAIL TO REJECT H0: no evidence of functional form misspecification"))

# 4. Residual diagnostic plots — ggplot
diag_df <- tibble(
  fitted   = fitted(post_wp),
  residual = resid_pl
)
p_rvf <- ggplot(diag_df, aes(fitted, residual)) +
  geom_point(colour = col_main, alpha = 0.4, size = 0.6) +
  geom_hline(yintercept = 0, linetype = "dashed", colour = "grey50") +
  labs(x = "Fitted values", y = "Residuals",
       title = "Residuals vs Fitted (Post-Lasso OLS)") +
  theme_lecture

# Q-Q via ggplot stat_qq
p_qq <- ggplot(diag_df, aes(sample = residual)) +
  stat_qq(colour = col_main, alpha = 0.4, size = 0.6) +
  stat_qq_line(colour = col_accent, linewidth = 0.8) +
  labs(x = "Theoretical quantiles", y = "Sample quantiles",
       title = "Normal Q-Q (Post-Lasso residuals)") +
  theme_lecture

p_rvf + p_qq    # patchwork side-by-side

# [r-lasso-crime4]
cr <- crime4

cr$lcrmrte <- log(cr$crmrte)
cr$lprbarr <- log(cr$prbarr)
cr$lpolpc  <- log(cr$polpc)

focal     <- "lprbarr"
ctrl_vars <- c("prbconv", "prbpris", "avgsen", "lpolpc", "density", "taxpc",
               "pctmin80", "pctymle", "west", "central", "urban",
               "wcon", "wtuc", "wtrd", "wfir", "wser", "wmfg", "wfed", "wsta", "wloc")
ctrl_vars <- intersect(ctrl_vars, names(cr))

keep <- c("county", "year", "lcrmrte", focal, ctrl_vars)
cr   <- cr[complete.cases(cr[, keep]), keep]

# Two-way FE = within-county demean + year dummies
cr_dm <- group_by(cr, county)
cr_dm <- mutate(cr_dm, across(where(is.numeric),
                ~ . - mean(.) + mean(cr[[cur_column()]])))
cr_dm <- ungroup(cr_dm)
yr_d  <- model.matrix(~ factor(year) - 1, data = cr_dm)[, -1]
Xc    <- cbind(as.matrix(cr_dm[, ctrl_vars]), yr_d)

# Post-double-selection: Lasso the outcome on controls, and the focal regressor
# on controls; keep the union of selected controls; then OLS.
library(glmnet)
cv_y <- cv.glmnet(Xc, cr_dm$lcrmrte, alpha = 1, nfolds = 10)
cv_d <- cv.glmnet(Xc, cr_dm[[focal]], alpha = 1, nfolds = 10)
sel  <- union(which(coef(cv_y, s = "lambda.min")[-1] != 0),
              which(coef(cv_d, s = "lambda.min")[-1] != 0))

Xsel <- Xc[, sel, drop = FALSE]
pds  <- lm(cr_dm$lcrmrte ~ cr_dm[[focal]] + Xsel)
b    <- coef(pds)[2]
se   <- sqrt(sandwich::vcovHC(pds)[2, 2])
cat(sprintf("Controls selected (union): %d of %d\n", length(sel), ncol(Xc)))
cat(sprintf("Deterrence elasticity (PDS Lasso): %.3f  (HC SE %.3f)\n", b, se))

# [r-cv-selection-table]
# How many variables does each rule select for Lasso?
coef_min <- coef(cv_lasso, s = "lambda.min")   # coefficients at λ.min
coef_1se <- coef(cv_lasso, s = "lambda.1se")   # coefficients at λ.1se

nnz_min  <- sum(coef_min[-1] != 0)  # [-1] drops the intercept
nnz_1se  <- sum(coef_1se[-1] != 0)

tibble(
  Rule        = c("λ.min (lowest CV-MSE)", "λ.1se (1 SE rule)"),
  `λ value`   = c(round(cv_lasso$lambda.min, 4), round(cv_lasso$lambda.1se, 4)),
  `Variables selected` = c(nnz_min, nnz_1se),
  `True non-zero`      = 5L,    # P_SIGNAL = 5 in the DGP
  `CV-MSE`    = round(c(cv_lasso$cvm[cv_lasso$lambda == cv_lasso$lambda.min],
                        cv_lasso$cvm[cv_lasso$lambda == cv_lasso$lambda.1se]), 4)
) %>%
  kbl(caption = "Lasso selection: λ.min aggressively includes more variables; λ.1se is sparser") %>%
  kable_styling(font_size = 20, full_width = TRUE)