# Week 01: Foundations of policy evaluation
# Synthetic tutoring data. No external data files required.
# ggplot2 is optional and used only for visualisation.

# Homework Exercise 2: a short treated-control comparison
dat <- data.frame(
  pupil = c("A", "B", "C", "D", "E", "F", "G", "H"),
  D = c(1, 1, 1, 1, 0, 0, 0, 0),
  X = c(62, 70, 58, 75, 65, 69, 60, 72),
  Y = c(73, 78, 69, 82, 68, 71, 64, 75)
)

head(dat)

treated_mean <- mean(dat$Y[dat$D == 1])
control_mean <- mean(dat$Y[dat$D == 0])
naive_difference <- treated_mean - control_mean

model <- lm(Y ~ D, data = dat)

cat("\nHomework Exercise 2: short R comparison\n")
cat("---------------------------------------\n")
cat("Mean Y, D = 1: ", round(treated_mean, 2), "\n", sep = "")
cat("Mean Y, D = 0: ", round(control_mean, 2), "\n", sep = "")
cat("Naive treated-control difference: ",
    round(naive_difference, 2), "\n", sep = "")

cat("\nRegression output for the homework comparison:\n")
print(summary(model))

cat("\nInterpretation prompt:\n")
cat("The coefficient on D is the treated-control comparison. It is causal only with an identification argument for the missing counterfactual.\n")

cat("\nDataviz process for the homework comparison:\n")
cat("This plot shows the raw treated-control outcome comparison before any causal claim is made.\n")

# ggplot2 is used only for visualisation. The comparison above works without it.
if (requireNamespace("ggplot2", quietly = TRUE)) {
  library(ggplot2)

  # Label D so the figure can be read without remembering that 1 means treated.
  dat$D_label <- ifelse(dat$D == 1, "D = 1", "D = 0")

  # The black point is the group mean. The small points show the individual pupils.
  plot_homework_comparison <- ggplot(dat, aes(x = D_label, y = Y, colour = D_label)) +
    geom_point(size = 2, alpha = 0.75) +
    stat_summary(fun = mean, geom = "point", size = 4, colour = "black") +
    labs(
      title = "Naive outcome comparison",
      x = "Treatment status",
      y = "End-of-year score Y"
    ) +
    theme_minimal() +
    theme(legend.position = "none")

  if (interactive()) {
    print(plot_homework_comparison)
  } else {
    cat("Plot created as plot_homework_comparison. Run the script in RStudio to display it.\n")
  }
} else {
  cat("ggplot2 is not installed. To draw the figure, run install.packages(\"ggplot2\") once, then rerun this script.\n")
}

set.seed(20260901)

n <- 500
id <- seq_len(n)
X <- rnorm(n, mean = 50, sd = 10)          # baseline score
motivation <- rnorm(n)                     # unobserved in a real study

Y0 <- 35 + 0.7 * X + 4 * motivation + rnorm(n, sd = 5)
tau <- 4 + 0.04 * (60 - X) + 1.5 * motivation
Y1 <- Y0 + tau

prob_D <- plogis(-6.4 + 0.08 * X + 0.9 * motivation)
D <- rbinom(n, size = 1, prob = prob_D)
Y <- ifelse(D == 1, Y1, Y0)

dat <- data.frame(id, X, motivation, D, Y, Y0, Y1, tau)

true_ate <- mean(dat$Y1 - dat$Y0)
true_att <- mean(dat$tau[dat$D == 1])
true_atu <- mean(dat$tau[dat$D == 0])
naive_difference <- with(dat, mean(Y[D == 1]) - mean(Y[D == 0]))
selection_in_Y0 <- with(dat, mean(Y0[D == 1]) - mean(Y0[D == 0]))

Z <- rbinom(n, size = 1, prob = 0.5)
Y_random_assignment <- ifelse(Z == 1, Y1, Y0)
randomised_difference <- mean(Y_random_assignment[Z == 1]) -
  mean(Y_random_assignment[Z == 0])

cat("\nWeek 01: Foundations\n")
cat("--------------------\n")
cat("True ATE: ", round(true_ate, 2), "\n", sep = "")
cat("True ATT: ", round(true_att, 2), "\n", sep = "")
cat("True ATU: ", round(true_atu, 2), "\n", sep = "")
cat("Naive observed difference: ", round(naive_difference, 2), "\n", sep = "")
cat("Selection in untreated potential outcomes: ",
    round(selection_in_Y0, 2), "\n", sep = "")
cat("Random-assignment benchmark difference: ",
    round(randomised_difference, 2), "\n", sep = "")

cat("\nSmall potential-outcomes table from the simulation:\n")
print(round(dat[1:6, c("id", "D", "Y", "Y0", "Y1", "tau")], 2))

cat("\nDataviz process for the simulation:\n")
cat("The simulation plot shows how selection can make treated and untreated pupils occupy different parts of the baseline-score distribution.\n")

if (requireNamespace("ggplot2", quietly = TRUE)) {
  dat$D_label <- ifelse(dat$D == 1, "D = 1", "D = 0")

  # This scatterplot connects three ideas: baseline X, observed Y, and voluntary treatment status D.
  plot_selection_simulation <- ggplot(dat, aes(x = X, y = Y, colour = D_label)) +
    geom_point(alpha = 0.55) +
    geom_smooth(method = "lm", se = FALSE) +
    labs(
      title = "Observed outcomes under voluntary participation",
      x = "Baseline score X",
      y = "Observed outcome Y",
      colour = "Treatment status"
    ) +
    theme_minimal()

  if (interactive()) {
    print(plot_selection_simulation)
  } else {
    cat("Plot created as plot_selection_simulation. Run the script in RStudio to display it.\n")
  }
}

cat("\nInterpretation prompt:\n")
cat("In real data only Y is observed. Y0 and Y1 are shown here only because this is a simulation.\n")