# Week 04: Compliance, IV, LATE, and external validity
# Synthetic encouragement design. No external data files required.
# ggplot2 is optional and used only for visualisation.

set.seed(20260904)

n <- 1000
id <- seq_len(n)
X <- rnorm(n)
type <- sample(
  c("never-taker", "complier", "always-taker"),
  size = n,
  replace = TRUE,
  prob = c(0.25, 0.50, 0.25)
)

D0 <- ifelse(type == "always-taker", 1, 0)
D1 <- ifelse(type == "never-taker", 0, 1)

Z <- rbinom(n, size = 1, prob = 0.5)
D <- ifelse(Z == 1, D1, D0)

tau <- ifelse(type == "complier", 8,
              ifelse(type == "always-taker", 5, 3)) + 0.5 * X
Y0 <- 40 + 3 * X + 2 * (type == "always-taker") -
  2 * (type == "never-taker") + rnorm(n, sd = 5)
Y1 <- Y0 + tau
Y <- ifelse(D == 1, Y1, Y0)

dat <- data.frame(id, X, type, Z, D0, D1, D, Y, tau)

itt_y <- with(dat, mean(Y[Z == 1]) - mean(Y[Z == 0]))
first_stage <- with(dat, mean(D[Z == 1]) - mean(D[Z == 0]))
wald <- itt_y / first_stage
naive_ols <- coef(lm(Y ~ D, data = dat))["D"]
true_late <- mean(dat$tau[dat$type == "complier"])

# Covariate-adjusted IV point estimate by residualising Y, D, and Z on X.
# This gives the point estimate intuition. It is not a recipe for valid IV standard errors.
resid_y <- resid(lm(Y ~ X, data = dat))
resid_d <- resid(lm(D ~ X, data = dat))
resid_z <- resid(lm(Z ~ X, data = dat))
iv_adjusted_point <- sum(resid_z * resid_y) / sum(resid_z * resid_d)

cat("\nWeek 04: IV and LATE\n")
cat("--------------------\n")
cat("ITT on Y: ", round(itt_y, 2), "\n", sep = "")
cat("First stage: ", round(first_stage, 3), "\n", sep = "")
cat("Wald/LATE estimate: ", round(wald, 2), "\n", sep = "")
cat("Naive OLS coefficient on D: ", round(naive_ols, 2), "\n", sep = "")
cat("True complier average effect in the simulation: ",
    round(true_late, 2), "\n", sep = "")
cat("Covariate-adjusted IV point estimate: ",
    round(iv_adjusted_point, 2), "\n", sep = "")

cat("\nCompliance-type shares shown only because this is a simulation:\n")
print(round(prop.table(table(dat$type)), 3))

cat("\nDataviz process:\n")
cat("Plot 1 shows the first stage: assignment Z changes treatment receipt D.\n")
cat("Plot 2 shows the ITT: assignment Z changes the outcome Y before dividing by the first stage.\n")

# ggplot2 is used only for visualisation. The ITT, first stage, and Wald ratio above use base R.
if (requireNamespace("ggplot2", quietly = TRUE)) {
  library(ggplot2)

  # Aggregate means by assignment status. In real data, Z and D are observed; compliance type is not.
  first_stage_data <- aggregate(D ~ Z, data = dat, mean)
  first_stage_data$Z_label <- ifelse(first_stage_data$Z == 1, "Z = 1 encouraged", "Z = 0 not encouraged")

  itt_data <- aggregate(Y ~ Z, data = dat, mean)
  itt_data$Z_label <- ifelse(itt_data$Z == 1, "Z = 1 encouraged", "Z = 0 not encouraged")

  # First-stage plot: the height difference is Delta_D(c) in the IV logic.
  plot_first_stage <- ggplot(first_stage_data, aes(x = Z_label, y = D, fill = Z_label)) +
    geom_col(width = 0.65) +
    labs(
      title = "First stage: assignment and treatment receipt",
      x = "Assignment Z",
      y = "Mean treatment receipt D"
    ) +
    theme_minimal() +
    theme(legend.position = "none")

  # ITT plot: the height difference is the effect of assignment on the outcome.
  plot_itt <- ggplot(itt_data, aes(x = Z_label, y = Y, fill = Z_label)) +
    geom_col(width = 0.65) +
    labs(
      title = "ITT: assignment and outcomes",
      x = "Assignment Z",
      y = "Mean outcome Y"
    ) +
    theme_minimal() +
    theme(legend.position = "none")

  if (interactive()) {
    print(plot_first_stage)
    print(plot_itt)
  } else {
    cat("Plots created as plot_first_stage and plot_itt. 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("In real data D0 and D1 are not observed. LATE refers to compliers moved by Z.\n")
cat("Do not report ordinary second-stage OLS standard errors from a manual two-step fit.\n")