# Controlled interrupted time-series analysis: AM vs. control, annual EP data
# Put arsenic_am_data_wqtc_demo.csv in the working directory (downloaded above).

# pacman package and p_load command makes it easy to install and load commands
if (!require(pacman)) install.packages("pacman")
p_load(dplyr, ggplot2, sandwich, lmtest)

dat <- read.csv("arsenic_am_data_wqtc_demo.csv", stringsAsFactors = FALSE)
names(dat)[1] <- "year"  # guards against a UTF-8 BOM corrupting the first column name

# --- Validate before analysis ---
expected_years <- 1998:2019
stopifnot(setequal(unique(dat$technology), c("AM", "control")))
stopifnot(!anyDuplicated(dat[c("year", "technology")]))
stopifnot(all(vapply(split(dat$year, dat$technology),
  function(x) identical(sort(x), expected_years), logical(1))))

# --- Build the intervention variables (2006 = first post-intervention year) ---
dat <- dat |>
  mutate(
    technology = factor(technology, levels = c("control", "AM")),
    treated    = as.integer(technology == "AM"),
    time_index = year - 1998,
    post       = as.integer(year >= 2006),
    post_time  = pmax(year - 2006, 0),
    pre_period = year < 2006
  )

# --- Descriptive pre/post means (not causal, just context) ---
dat |>
  group_by(technology, post) |>
  summarise(annual_mean = mean(annual_mean), .groups = "drop") |>
  mutate(period = ifelse(post == 0, "1998-2005 (pre)", "2006-2019 (post)")) |>
  select(technology, period, annual_mean) |>
  print()

# --- Plot the two series with the intervention boundary marked ---
ggplot(dat, aes(year, annual_mean, colour = technology)) +
  geom_line(linewidth = 0.9) +
  geom_point(size = 2) +
  geom_vline(xintercept = 2006, linetype = "dashed", colour = "#b45309") +
  labs(title = "Annual EP mean: AM vs. control",
       subtitle = "Dashed line = 2006, first post-intervention year",
       x = "Year", y = "annual_mean", colour = "Technology") +
  theme_minimal(base_size = 12)

# --- Primary controlled segmented regression ---
model <- lm(
  annual_mean ~ treated + time_index + post + post_time +
    treated:time_index + treated:post + treated:post_time,
  data = dat
)
summary(model)
cbind(estimate = coef(model), confint(model))

# --- Newey-West robust inference (lag = 2 as an annual-data starting point;
#     sensitivity-test other lags, there are few years per series) ---
nw_vcov <- NeweyWest(model, lag = 2, prewhite = FALSE)
coeftest(model, vcov. = nw_vcov)

nw_se <- sqrt(diag(nw_vcov))
nw_ci <- data.frame(
  term     = names(coef(model)),
  estimate = coef(model),
  lower    = coef(model) - 1.96 * nw_se,
  upper    = coef(model) + 1.96 * nw_se,
  row.names = NULL
)

key <- filter(nw_ci, term %in% c("treated:post", "treated:post_time"))
print(key)

# --- AM-specific effect over time relative to the no-intervention counterfactual ---
data.frame(year = 2006:2019) |>
  mutate(am_specific_effect = coef(model)["treated:post"] +
           coef(model)["treated:post_time"] * (year - 2006)) |>
  print()

# --- Pre-trend check (parallel trends) ---
pre_model <- lm(annual_mean ~ treated + time_index + treated:time_index,
                 data = filter(dat, pre_period))
coef(summary(pre_model))["treated:time_index", , drop = FALSE]

# --- Placebo test: move the boundary to 2004, using only real pre-period years ---
placebo <- dat |>
  filter(pre_period) |>
  mutate(
    placebo_post      = as.integer(year >= 2004),
    placebo_post_time = pmax(year - 2004, 0)
  )
placebo_model <- lm(
  annual_mean ~ treated + time_index + placebo_post + placebo_post_time +
    treated:time_index + treated:placebo_post + treated:placebo_post_time,
  data = placebo
)
coef(summary(placebo_model))[c("treated:placebo_post", "treated:placebo_post_time"), , drop = FALSE]

# --- Residuals over time (watch for patterns a 2006-only story wouldn't predict) ---
plot(dat$year, residuals(model), pch = 16, xlab = "Year", ylab = "Residual",
     main = "Residuals from controlled segmented model")
abline(h = 0, lty = 2)

# --- Plain-language summary, filled in from the fitted model at runtime ---
fmt   <- function(x) formatC(x, digits = 2, format = "f")
level <- filter(key, term == "treated:post")
slope <- filter(key, term == "treated:post_time")
cat(sprintf(
  "At the 2006 boundary, the estimated AM-specific differential level change was %s units (95%% CI %s, %s). The estimated additional annual slope change after 2006 was %s units/year (95%% CI %s, %s).\n",
  fmt(level$estimate), fmt(level$lower), fmt(level$upper),
  fmt(slope$estimate), fmt(slope$lower), fmt(slope$upper)
))

# --- Optional sensitivity analysis: inverse-variance weighting ---
# annual_mean is already an aggregated value. The primary model above does NOT
# weight by annual_sd / n_systems / se_mean. As a sensitivity check only:
#
# weighted_model <- lm(
#   annual_mean ~ treated + time_index + post + post_time +
#     treated:time_index + treated:post + treated:post_time,
#   data = dat, weights = 1 / dat$se_mean^2
# )
# summary(weighted_model)
#
# Compare coefficients and uncertainty against the unweighted primary model
# rather than reporting the weighted version alone.