#!/usr/bin/env Rscript
# -----------------------------------------------------------------------------
# Script 04: DiD with Staggered Adoption (Callaway & Sant'Anna)
# Atlas: Causality Atlas — Ch 09 (Econometric / Quasi-Experimental)
# -----------------------------------------------------------------------------

library(did); library(ggplot2)

cat("=== Difference-in-Differences: Staggered Adoption ===\n")
set.seed(42)
N <- 500; T_periods <- 5

# Generate staggered treatment timing
# Units first treated in periods 2, 3, 4, or never
dat <- matrix(NA, nrow = N * T_periods, ncol = 5)
colnames(dat) <- c("id", "year", "Y", "first_treat", "A")
idx <- 1

for (i in 1:N) {
  first_treat <- sample(c(2, 3, 4, 0), 1, prob = c(0.3, 0.3, 0.2, 0.2))
  unit_effect <- rnorm(1, 0, 0.5)

  for (t in 1:T_periods) {
    year <- 2019 + t
    treat <- ifelse(first_treat > 0 & t >= first_treat, 1, 0)
    Y <- unit_effect + 0.5 * year + 2 * treat + rnorm(1, 0, 1)
    dat[idx, ] <- c(i, year, Y, first_treat, treat)
    idx <- idx + 1
  }
}

dat <- as.data.frame(dat)

# Callaway & Sant'Anna staggered DiD
out <- att_gt(yname = "Y", tname = "year", idname = "id",
              gname = "first_treat", data = dat,
              control_group = "notyettreated",
              bstrap = TRUE, cband = TRUE)

# Aggregate results
cat("\nGroup-time ATTs:\n")
summary(out)

# Simple aggregated ATT
agg_simple <- aggte(out, type = "simple")
cat(sprintf("\nSimple weighted average ATT = %.3f\n", agg_simple$overall.att))

# Event study aggregation
agg_event <- aggte(out, type = "dynamic")
cat(sprintf("Event study ATT (overall) = %.3f\n", agg_event$overall.att))

# Plot event study
pdf("../outputs/figures/did_event_study.pdf", width = 8, height = 5)
ggdid(agg_event) + ggtitle("DiD Event Study (Callaway & Sant'Anna)")
dev.off()
cat("Saved: did_event_study.pdf\n")
cat("\n=== Script complete. True ATT = 2.0 ===\n")