# Week 08: Synthetic control, heterogeneity, and replication
# Synthetic regional panel plus a short subgroup example. No external data files required.
# ggplot2 is optional and used only for visualisation.

set.seed(20260908)

regions <- paste0("R", 1:8)
years <- 2014:2026
policy_year <- 2022

region_info <- data.frame(
  region = regions,
  base = c(55, 52, 57, 50, 60, 54, 48, 62),
  lambda1 = c(1.10, 1.00, 1.20, 0.80, 1.35, 1.05, 0.75, 1.45),
  lambda2 = c(3.0, 2.5, 3.5, 1.5, 4.0, 2.8, 1.0, 4.4)
)
time_info <- data.frame(
  year = years,
  f1 = years - min(years),
  f2 = sin(seq_along(years) / 2)
)

panel <- merge(
  expand.grid(region = regions, year = years, KEEP.OUT.ATTRS = FALSE),
  region_info,
  by = "region"
)
panel <- merge(panel, time_info, by = "year")
panel <- panel[order(panel$region, panel$year), ]
panel$treated_unit <- panel$region == "R1"
panel$post <- panel$year >= policy_year
panel$Y0 <- panel$base + panel$lambda1 * panel$f1 +
  panel$lambda2 * panel$f2 + rnorm(nrow(panel), sd = 0.7)
panel$effect <- ifelse(panel$treated_unit & panel$post,
                       -5 - 0.8 * (panel$year - policy_year), 0)
panel$Y <- panel$Y0 + panel$effect

wide <- reshape(panel[, c("region", "year", "Y")],
                idvar = "year", timevar = "region", direction = "wide")
wide <- wide[order(wide$year), ]
Y_mat <- as.matrix(wide[, paste0("Y.", regions)])
rownames(Y_mat) <- wide$year

theta_to_weights <- function(theta) {
  e <- exp(theta - max(theta))
  e / sum(e)
}

fit_synth <- function(Y_mat, treated_col, pre_rows,
                      donor_cols = setdiff(colnames(Y_mat), treated_col)) {
  y_treat_pre <- Y_mat[pre_rows, treated_col]
  y_donor_pre <- Y_mat[pre_rows, donor_cols, drop = FALSE]

  loss <- function(theta) {
    w <- theta_to_weights(theta)
    mean((as.vector(y_treat_pre) - as.vector(y_donor_pre %*% w))^2)
  }

  opt <- optim(rep(0, length(donor_cols)), loss, control = list(maxit = 2000))
  weights <- theta_to_weights(opt$par)
  names(weights) <- donor_cols
  synth <- as.vector(Y_mat[, donor_cols, drop = FALSE] %*% weights)

  list(weights = weights, synth = synth, donors = donor_cols)
}

pre_rows <- years < policy_year
main <- fit_synth(Y_mat, treated_col = "Y.R1", pre_rows = pre_rows)
gap <- as.vector(Y_mat[, "Y.R1"]) - main$synth
pre_mspe <- mean(gap[pre_rows]^2)
post_avg_gap <- mean(gap[!pre_rows])

weights <- data.frame(
  region = sub("Y\\.", "", names(main$weights)),
  weight = as.numeric(main$weights)
)
weights <- weights[order(-weights$weight), ]

placebo_rows <- lapply(colnames(Y_mat), function(treated_col) {
  fit <- fit_synth(Y_mat, treated_col = treated_col, pre_rows = pre_rows)
  placebo_gap <- as.vector(Y_mat[, treated_col]) - fit$synth
  pre <- mean(placebo_gap[pre_rows]^2)
  post <- mean(placebo_gap[!pre_rows]^2)
  data.frame(
    region = sub("Y\\.", "", treated_col),
    pre_mspe = pre,
    post_mspe = post,
    gap_ratio = post / pre,
    avg_post_gap = mean(placebo_gap[!pre_rows])
  )
})
placebo_summary <- do.call(rbind, placebo_rows)
placebo_summary <- placebo_summary[order(-placebo_summary$gap_ratio), ]
treated_rank <- which(placebo_summary$region == "R1")

largest_donor_col <- paste0("Y.", weights$region[1])
loo_donors <- setdiff(main$donors, largest_donor_col)
loo <- fit_synth(Y_mat, treated_col = "Y.R1",
                 pre_rows = pre_rows, donor_cols = loo_donors)
gap_loo <- as.vector(Y_mat[, "Y.R1"]) - loo$synth
loo_post_avg_gap <- mean(gap_loo[!pre_rows])

# Short subgroup illustration: random assignment makes the interaction interpretable here.
n <- 600
X_subgroup <- rbinom(n, size = 1, prob = 0.45)
Z <- rbinom(n, size = 1, prob = 0.5)
D <- Z
tau <- ifelse(X_subgroup == 1, 6, 2)
Y0 <- 30 + 4 * X_subgroup + rnorm(n, sd = 5)
Y <- Y0 + tau * D
hte_fit <- lm(Y ~ D * X_subgroup)

cat("\nWeek 08: SCM, heterogeneity, and replication\n")
cat("--------------------------------------------\n")
cat("Largest donor weights for treated region R1:\n")
top_weights <- head(weights, 4)
top_weights$weight <- round(top_weights$weight, 3)
print(top_weights, row.names = FALSE)
cat("Pre-treatment MSPE: ", round(pre_mspe, 3), "\n", sep = "")
cat("Average post-treatment gap, R1 minus synthetic R1: ",
    round(post_avg_gap, 2), "\n", sep = "")
cat("Placebo gap-ratio rank for R1: ", treated_rank, " of ",
    nrow(placebo_summary), "\n", sep = "")

cat("\nPlacebo summary by gap ratio:\n")
placebo_print <- placebo_summary
placebo_print[, c("pre_mspe", "post_mspe", "gap_ratio", "avg_post_gap")] <-
  round(placebo_print[, c("pre_mspe", "post_mspe", "gap_ratio", "avg_post_gap")], 3)
print(placebo_print, row.names = FALSE)

cat("\nLeave-one-donor-out sensitivity:\n")
cat("Removed largest donor ", weights$region[1],
    "; average post-treatment gap becomes ",
    round(loo_post_avg_gap, 2), "\n", sep = "")

cat("\nSubgroup interaction from simulated randomised design:\n")
print(round(coef(summary(hte_fit)), 3))

cat("\nDataviz process:\n")
cat("Plot 1 compares the treated region with its synthetic counterfactual before and after the policy.\n")
cat("Plot 2 plots the treatment gap over time; the pre-period gap is a diagnostic for fit.\n")

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

  # Put the treated and synthetic series into one long data frame for ggplot.
  synth_plot_data <- rbind(
    data.frame(year = years, series = "Treated region R1", Y = as.vector(Y_mat[, "Y.R1"])),
    data.frame(year = years, series = "Synthetic R1", Y = main$synth)
  )

  # SCM fit plot: good pre-treatment fit is a diagnostic, not a proof of identification.
  plot_scm_fit <- ggplot(synth_plot_data, aes(x = year, y = Y, colour = series)) +
    geom_line(linewidth = 1) +
    geom_vline(xintercept = policy_year - 0.5, linetype = "dashed", colour = "grey35") +
    labs(
      title = "Treated region and synthetic control",
      x = "Year",
      y = "Outcome Y",
      colour = "Series"
    ) +
    theme_minimal()

  # Gap plot: after the policy, the treated-minus-synthetic gap is the estimated effect path.
  gap_plot_data <- data.frame(year = years, gap = gap)
  plot_scm_gap <- ggplot(gap_plot_data, aes(x = year, y = gap)) +
    geom_hline(yintercept = 0, colour = "grey55") +
    geom_line(linewidth = 1, colour = "#3266a8") +
    geom_point(size = 2, colour = "#3266a8") +
    geom_vline(xintercept = policy_year - 0.5, linetype = "dashed", colour = "grey35") +
    labs(
      title = "Treated minus synthetic gap",
      x = "Year",
      y = "Gap"
    ) +
    theme_minimal()

  if (interactive()) {
    print(plot_scm_fit)
    print(plot_scm_gap)
  } else {
    cat("Plots created as plot_scm_fit and plot_scm_gap. 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 SCM gap is credible only with a defensible donor pool, good pre-fit, no spillovers, and no differential post-treatment shocks.\n")