# Week 05: Difference-in-differences and simple TWFE
# Synthetic two-period regional panel. No external data files required.
# ggplot2 is optional and used only for visualisation.

set.seed(20260905)

n_units <- 80
years <- c(2024, 2025)
unit_info <- data.frame(
  id = seq_len(n_units),
  treated = as.integer(seq_len(n_units) <= n_units / 2),
  alpha = rnorm(n_units, sd = 6)
)

panel <- merge(
  expand.grid(id = seq_len(n_units), year = years, KEEP.OUT.ATTRS = FALSE),
  unit_info,
  by = "id"
)
panel$post <- as.integer(panel$year == 2025)

Y0 <- 50 + panel$alpha + 4 * panel$post + rnorm(nrow(panel), sd = 4)
tau_att <- 6
panel$Y <- Y0 + tau_att * panel$treated * panel$post

means <- aggregate(Y ~ treated + post, data = panel, mean)
get_mean <- function(treated_value, post_value) {
  means$Y[means$treated == treated_value & means$post == post_value]
}

treated_pre <- get_mean(1, 0)
treated_post <- get_mean(1, 1)
control_pre <- get_mean(0, 0)
control_post <- get_mean(0, 1)
did_hand <- (treated_post - treated_pre) - (control_post - control_pre)

fit <- lm(Y ~ treated * post, data = panel)
did_regression <- coef(fit)["treated:post"]

cat("\nWeek 05: Difference-in-differences\n")
cat("----------------------------------\n")
cat("Group-time means:\n")
print(round(means[order(means$treated, means$post), ], 2), row.names = FALSE)
cat("Hand-computed DiD: ", round(did_hand, 2), "\n", sep = "")
cat("Regression interaction coefficient: ",
    round(did_regression, 2), "\n", sep = "")

cat("\nRegression table for the simple 2x2 model:\n")
print(round(coef(summary(fit)), 3))

cat("\nDataviz process:\n")
cat("The line plot shows the two ingredients of DiD: the treated trend and the control trend.\n")
cat("Parallel trends is an assumption about the missing untreated trend for the treated group.\n")

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

  # Convert post back into a calendar year so the x-axis reads naturally.
  means$year <- ifelse(means$post == 1, 2025, 2024)
  means$group <- ifelse(means$treated == 1, "Treated", "Control")

  # In a 2x2 design, the vertical gap between the two trend changes is the DiD estimate.
  plot_did_means <- ggplot(means, aes(x = year, y = Y, colour = group, group = group)) +
    geom_line(linewidth = 1) +
    geom_point(size = 3) +
    scale_x_continuous(breaks = years) +
    labs(
      title = "Observed group means over time",
      x = "Year",
      y = "Mean outcome Y",
      colour = "Group"
    ) +
    theme_minimal()

  if (interactive()) {
    print(plot_did_means)
  } else {
    cat("Plot created as plot_did_means. 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")
}

cat("\nInterpretation prompt:\n")
cat("The interaction equals the hand DiD in this simple two-period design.\n")
cat("Causal interpretation still rests on parallel trends.\n")