# Week 02: Selection, regression, matching, and inference
# Synthetic employment programme data. No external data files required.
# ggplot2 is optional and used only for visualisation.

set.seed(20260902)

n <- 800
id <- seq_len(n)
X <- rnorm(n, mean = 50, sd = 10)          # observed baseline skill
female <- rbinom(n, size = 1, prob = 0.55)
U <- rnorm(n)                              # unobserved motivation/network

Y0 <- 1600 + 22 * X + 120 * female + 350 * U + rnorm(n, sd = 350)
tau <- 220 + 60 * (X < 50)
prob_D <- plogis(-5.2 + 0.075 * X + 0.85 * U - 0.25 * female)
D <- rbinom(n, size = 1, prob = prob_D)
Y <- Y0 + tau * D

dat <- data.frame(id, X, female, U, D, Y, tau)

naive <- coef(lm(Y ~ D, data = dat))["D"]
adjusted <- coef(lm(Y ~ D + X + female, data = dat))["D"]

ps_model <- glm(D ~ X + female, data = dat, family = binomial())
dat$ps <- fitted(ps_model)

treated_range <- range(dat$ps[dat$D == 1])
control_range <- range(dat$ps[dat$D == 0])
support_lower <- max(treated_range[1], control_range[1])
support_upper <- min(treated_range[2], control_range[2])
dat_support <- subset(dat, ps >= support_lower & ps <= support_upper)

breaks <- unique(quantile(dat_support$ps, probs = seq(0, 1, 0.2), na.rm = TRUE))
dat_support$ps_bin <- cut(dat_support$ps, breaks = breaks, include.lowest = TRUE)

bins <- split(dat_support, dat_support$ps_bin)
bin_att <- sapply(bins, function(b) {
  if (length(unique(b$D)) < 2) return(NA_real_)
  mean(b$Y[b$D == 1]) - mean(b$Y[b$D == 0])
})
bin_weights <- sapply(bins, function(b) sum(b$D == 1))
matched_att <- weighted.mean(bin_att, bin_weights, na.rm = TRUE)

true_att <- mean(dat$tau[dat$D == 1])
true_att_support <- mean(dat_support$tau[dat_support$D == 1])

balance_table <- rbind(
  untreated = c(
    mean_X = mean(dat$X[dat$D == 0]),
    mean_female = mean(dat$female[dat$D == 0]),
    mean_ps = mean(dat$ps[dat$D == 0])
  ),
  treated = c(
    mean_X = mean(dat$X[dat$D == 1]),
    mean_female = mean(dat$female[dat$D == 1]),
    mean_ps = mean(dat$ps[dat$D == 1])
  )
)

cat("\nWeek 02: Selection, regression, and matching\n")
cat("--------------------------------------------\n")
cat("Naive difference/regression coefficient: ", round(naive, 2), "\n", sep = "")
cat("Adjusted regression coefficient: ", round(adjusted, 2), "\n", sep = "")
cat("Propensity-score subclassification estimate: ",
    round(matched_att, 2), "\n", sep = "")
cat("True ATT in full simulated data: ", round(true_att, 2), "\n", sep = "")
cat("True ATT among treated in common support: ",
    round(true_att_support, 2), "\n", sep = "")

cat("\nPropensity-score ranges:\n")
print(round(rbind(treated = treated_range, control = control_range), 3))
cat("Common support interval: [", round(support_lower, 3), ", ",
    round(support_upper, 3), "]\n", sep = "")
cat("Treated observations outside common support: ",
    sum(dat$D == 1 & (dat$ps < support_lower | dat$ps > support_upper)),
    "\n", sep = "")

cat("\nObserved balance table:\n")
print(round(balance_table, 2))

cat("\nDataviz process:\n")
cat("The propensity-score plot checks overlap: do treated and untreated observations exist in the same regions of X?\n")
cat("The dashed lines mark the common-support interval used before subclassification.\n")

# ggplot2 is used only for visualisation. The matching calculations above use base R.
if (requireNamespace("ggplot2", quietly = TRUE)) {
  library(ggplot2)

  # Label D so the legend reads like the empirical comparison.
  dat$D_label <- ifelse(dat$D == 1, "D = 1 treated", "D = 0 control")

  # Overlap plot: poor overlap warns us that the estimand may become more local.
  plot_propensity_overlap <- ggplot(dat, aes(x = ps, fill = D_label)) +
    geom_histogram(position = "identity", bins = 30, alpha = 0.55) +
    geom_vline(xintercept = c(support_lower, support_upper),
               linetype = "dashed", colour = "grey30") +
    labs(
      title = "Propensity-score overlap",
      x = "Estimated propensity score",
      y = "Number of observations",
      fill = "Treatment status"
    ) +
    theme_minimal()

  # Balance plot: this shows observed X before any claim about unobserved U.
  plot_x_by_treatment <- ggplot(dat, aes(x = D_label, y = X, colour = D_label)) +
    geom_boxplot(outlier.alpha = 0.35) +
    labs(
      title = "Observed baseline skill by treatment status",
      x = "Treatment status",
      y = "Observed covariate X"
    ) +
    theme_minimal() +
    theme(legend.position = "none")

  if (interactive()) {
    print(plot_propensity_overlap)
    print(plot_x_by_treatment)
  } else {
    cat("Plots created as plot_propensity_overlap and plot_x_by_treatment. Run the script in RStudio to display them.\n")
  }
} else {
  cat("ggplot2 is not installed. To draw the figures, run install.packages(\"ggplot2\") once, then rerun this script.\n")
}

cat("\nInterpretation prompt:\n")
cat("The script adjusts for observed X and female, but U remains unobserved in a real study.\n")