# Week 03: Randomised experiments, design, and power
# Synthetic education trial. No external data files required.
# ggplot2 is optional and used only for visualisation.

set.seed(20260903)

make_trial <- function(n, stratified = FALSE) {
  id <- seq_len(n)
  baseline <- rnorm(n, mean = 50, sd = 10)
  high_baseline <- as.integer(baseline >= median(baseline))

  Z <- integer(n)
  if (stratified) {
    for (s in sort(unique(high_baseline))) {
      idx <- which(high_baseline == s)
      Z[idx] <- sample(rep(c(0, 1), length.out = length(idx)))
    }
  } else {
    Z <- sample(rep(c(0, 1), length.out = n))
  }

  D <- Z
  tau <- 5
  Y0 <- 45 + 0.55 * baseline + rnorm(n, sd = 8)
  Y <- Y0 + tau * D + rnorm(n, sd = 3)

  data.frame(id, baseline, high_baseline, Z, D, Y)
}

estimate_trial <- function(dat) {
  fit <- lm(Y ~ Z, data = dat)
  estimate <- unname(coef(fit)["Z"])
  se <- unname(coef(summary(fit))["Z", "Std. Error"])
  balance_diff <- mean(dat$baseline[dat$Z == 1]) -
    mean(dat$baseline[dat$Z == 0])
  c(
    estimate = estimate,
    se = se,
    balance_diff = balance_diff
  )
}

trial_simple <- make_trial(400, stratified = FALSE)
trial_stratified <- make_trial(400, stratified = TRUE)

simple_results <- estimate_trial(trial_simple)
stratified_results <- estimate_trial(trial_stratified)

simulate_se <- function(n, reps = 300) {
  estimates <- replicate(reps, estimate_trial(make_trial(n))["estimate"])
  sd(estimates)
}

sample_sizes <- c(100, 400, 1600)
simulated_se <- sapply(sample_sizes, simulate_se)
approx_mde <- 2.8 * simulated_se

cluster_size <- 20
icc <- 0.10
design_effect <- 1 + (cluster_size - 1) * icc

cat("\nWeek 03: RCT design and power\n")
cat("-----------------------------\n")
cat("Simple randomisation baseline difference: ",
    round(simple_results["balance_diff"], 2), "\n", sep = "")
cat("Stratified randomisation baseline difference: ",
    round(stratified_results["balance_diff"], 2), "\n", sep = "")
cat("Treatment estimate, simple trial: ",
    round(simple_results["estimate"], 2), "\n", sep = "")
cat("Standard error, simple trial: ",
    round(simple_results["se"], 2), "\n", sep = "")

cat("\nSimulated precision by sample size:\n")
precision_table <- data.frame(
  n = sample_sizes,
  simulated_se = round(simulated_se, 2),
  approx_mde_80_power = round(approx_mde, 2)
)
print(precision_table, row.names = FALSE)

cat("\nDataviz process:\n")
cat("Plot 1 checks baseline balance under simple and stratified randomisation.\n")
cat("Plot 2 shows how larger samples reduce the minimum detectable effect by improving precision.\n")

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

  # Combine the two trial designs into one plotting dataset.
  plot_trials <- rbind(
    transform(trial_simple, design = "Simple randomisation"),
    transform(trial_stratified, design = "Stratified randomisation")
  )
  plot_trials$Z_label <- ifelse(plot_trials$Z == 1, "Z = 1 assigned", "Z = 0 control")

  # Balance plot: randomisation should make groups comparable on average, not identical in every sample.
  plot_balance_design <- ggplot(plot_trials, aes(x = Z_label, y = baseline, fill = Z_label)) +
    geom_boxplot(outlier.alpha = 0.25) +
    facet_wrap(~ design) +
    labs(
      title = "Baseline balance by assignment group",
      x = "Assignment",
      y = "Baseline score"
    ) +
    theme_minimal() +
    theme(legend.position = "none")

  # Precision plot: a smaller MDE means the design can detect smaller effects.
  plot_mde <- ggplot(precision_table, aes(x = n, y = approx_mde_80_power)) +
    geom_line(colour = "#3266a8") +
    geom_point(size = 2.5, colour = "#3266a8") +
    labs(
      title = "Approximate MDE by sample size",
      x = "Sample size",
      y = "Approximate MDE"
    ) +
    theme_minimal()

  if (interactive()) {
    print(plot_balance_design)
    print(plot_mde)
  } else {
    cat("Plots created as plot_balance_design and plot_mde. 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("\nCluster design effect example:\n")
cat("Cluster size = ", cluster_size, ", ICC = ", icc,
    ", design effect = ", round(design_effect, 2), "\n", sep = "")

cat("\nInterpretation prompt:\n")
cat("Power changes precision around an estimator. It does not fix spillovers, attrition, or measurement problems.\n")