# Example Sheet 1, Question 6: prepare the stepped-wedge data.
# Haines et al. (2017), PLOS Medicine 14(10): e1002412, supporting data S2.
# Article: https://doi.org/10.1371/journal.pmed.1002412
# Download this file to a working folder, set R's working directory there, then:
# install.packages("readxl")  # Once, if needed; the remaining code uses base R.
# source("stepped-wedge.R")

# Download the published raw workbook (reuse a local copy if present).
workbook <- "pmed.1002412.s002.xlsx"
url <- paste0("https://journals.plos.org/plosmedicine/article/file?",
              "type=supplementary&id=10.1371/journal.pmed.1002412.s002")
if (!file.exists(workbook)) download.file(url, workbook, mode = "wb")
raw <- as.data.frame(readxl::read_excel(workbook))
trial <- subset(raw, hospital == "Dandenong" & study1 == 1,
                select = c(index_ward, sw_step, acute_los))

# Keep all 7,986 trial-1 admissions in their original order.
stopifnot(nrow(trial) == 7986, !anyNA(trial),
          all(trial$index_ward %in% 1:6), all(trial$sw_step %in% 1:7),
          all(is.finite(trial$acute_los)), all(trial$acute_los > 0))
Z <- c(3, 5, 7, 6, 2, 4)  # Observed withdrawal months for wards 1,...,6
data <- data.frame(ward = as.integer(trial$index_ward),
                   month = as.integer(trial$sw_step),
                   log_los = log(trial$acute_los))
data$exposure <- as.integer(data$month >= Z[data$ward])
write.csv(data, "stepped-wedge-data.csv", row.names = FALSE)
# ward and month are W_i and T_i in the question; log_los is Y_i.
# exposure is 1 if weekend services had been withdrawn at admission, 0 otherwise.

# Generate all 720 schedules, so no separate schedule file is needed.
# Start with all six-tuples, then retain those using each month exactly once.
S <- as.matrix(expand.grid(rep(list(2:7), 6)))
S <- S[apply(S, 1, function(z) length(unique(z)) == 6), ]
colnames(S) <- paste0("ward", 1:6)
write.csv(S, "crossover-schedules.csv", row.names = FALSE)
# Row j of S is one schedule; column w gives ward w's withdrawal month.
cat("Prepared", nrow(data), "admissions and", nrow(S), "withdrawal schedules.\n")
print(head(data))

# Continue with the lm code in part (a), then extend it for (b) and (c).
# coef(fit)["exposure"] extracts the exposure coefficient from a fitted model.
# sample(data$month) shuffles the month entries without changing other columns.
# Choose the number of shuffles B; set a seed before shuffling for reproducibility.
# For each statistic, use the Monte Carlo p-value (1 + upper-tail count)/(B + 1).
# apply() or a for loop can visit the schedules.
# When counting ties between computed coefficients, allow rounding error of 1e-10.