# Run before registration, from the bundle's root: checks that every association's replication
# runs on files shaped like NHANES August 2021-August 2023's, without reading any of them.
#
# Each 2021-2023 file is stood in for by the same component of one earlier release: the
# 2017-March 2020 files for an association that analyzed them, the 2017-2018 files for every
# other one. One release per association, because the 2017-March 2020 files renumbered every
# participant (SEQN), so their files can't be joined with 2017-2018's. A component that release
# lacks (2017-2018 released no sex steroid hormone file) is stood in for by its latest earlier
# release's records, drawn at random for the release's participants. Variables are renamed and
# recoded to their 2021-2023 names, and only the variables that 2021-2023's codebooks list
# (plan/variables_2021_2023.csv) are kept; a variable 2021-2023 has and the stand-in lacks is
# added as missing, except the weights 2021-2023 adds, which the stand-in's own weights stand in
# for. The estimates this gives mean nothing; what matters is that the code runs, reads only
# variables 2021-2023 has, and fits.
#
# Usage: Rscript plan/dry_run.R [--only=id,id]

suppressPackageStartupMessages(library(survey))
options(survey.lonely.psu = "adjust", warn = 1, stringsAsFactors = FALSE)
for (module in c("cycles", "io", "design", "model", "stats", "variables", "covariates", "indices", "liver", "pipeline", "summary")) {
  source(file.path("code", "lib", paste0(module, ".R")))
}

catalog <- read.csv("plan/variables_2021_2023.csv", colClasses = "character")
real_read <- read_nhanes
stand_in_cycle <- NULL
stand_in_files <- character()

recode <- function(component, data) {
  rename <- function(from, to) { if (from %in% names(data)) data[[to]] <<- data[[from]] }
  rename("WTINTPRP", "WTINT2YR"); rename("WTMECPRP", "WTMEC2YR"); rename("WTSAFPRP", "WTSAF2YR")
  rename("WTDRD1PP", "WTDRD1"); rename("WTDR2DPP", "WTDR2D")
  if (component == "TRIGLY") rename("LBXTR", "LBXTLG")
  # 2021-2023's three marital groups (as 2017-March 2020's): married or living with a partner,
  # widowed, divorced or separated, never married.
  if (component == "DEMO" && !"DMDMARTZ" %in% names(data) && "DMDMARTL" %in% names(data)) {
    data$DMDMARTZ <- ifelse(data$DMDMARTL %in% c(1, 6), 1, ifelse(data$DMDMARTL %in% 2:4, 2, ifelse(data$DMDMARTL %in% 5, 3, data$DMDMARTL)))
  }
  if (component == "BPQ") { rename("BPQ050A", "BPQ150"); rename("BPQ100D", "BPQ101D") }
  if (component == "KIQ_U") rename("KIQ480", "KIQ481")
  # 2021-2023's length of time in the US stops at 20 years or more.
  if (component == "DEMO" && !"DMDYRUSR" %in% names(data) && "DMDYRSUS" %in% names(data)) {
    data$DMDYRUSR <- ifelse(data$DMDYRSUS %in% 6:9, 6, data$DMDYRSUS)
  }
  # 2021-2023 codes the examination session as morning (0) or afternoon or evening (1).
  if (component == "FASTQX" && !"PHDSESNZ" %in% names(data) && "PHDSESN" %in% names(data)) {
    data$PHDSESNZ <- ifelse(data$PHDSESN %in% 0, 0, ifelse(data$PHDSESN %in% 1:2, 1, NA))
  }
  # Prescription medicine in the past month: RXQ033 in 2021-2023, RXDUSE before.
  if (component == "RXQ_RX") rename("RXDUSE", "RXQ033")
  # 2021-2023 asks about COPD, emphysema, and chronic bronchitis as one item.
  if (component == "MCQ" && !"MCQ160P" %in% names(data) && "MCQ160O" %in% names(data)) {
    any_yes <- data$MCQ160O %in% 1 | data$MCQ160G %in% 1 | data$MCQ160K %in% 1
    data$MCQ160P <- ifelse(any_yes, 1, data$MCQ160O)
  }
  # 2021-2023's kinds of usual place of care: 1 doctor's office or health center, 3 emergency
  # room, 5 some other place, 6 no one place (2017-2018's clinics, offices, and hospital
  # outpatient departments stand in for 1).
  if (component == "HUQ" && !"HUQ042" %in% names(data) && "HUQ041" %in% names(data)) {
    data$HUQ042 <- ifelse(data$HUQ041 %in% c(1, 2, 4), 1, data$HUQ041)
  }
  if (component == "PAQ" && "PAQ665" %in% names(data)) {
    data$PAD790Q <- ifelse(data$PAQ665 %in% 1, data$PAQ670, ifelse(data$PAQ665 %in% 2, 0, NA))
    data$PAD790U <- ifelse(data$PAQ665 %in% 1, "W", NA)
    data$PAD800 <- ifelse(data$PAQ665 %in% 1, data$PAD675, NA)
    data$PAD810Q <- ifelse(data$PAQ650 %in% 1, data$PAQ655, ifelse(data$PAQ650 %in% 2, 0, NA))
    data$PAD810U <- ifelse(data$PAQ650 %in% 1, "W", NA)
    data$PAD820 <- ifelse(data$PAQ650 %in% 1, data$PAD660, NA)
  }
  # Weights 2021-2023 adds: blood files' phlebotomy weight (the examination weight stands in) and
  # the supplement questionnaire's dietary day-one weight (the day-one recall's stands in).
  stand_in_weight <- function(variable, source, weight) {
    if (variable %in% names(data)) return()
    d <- recode(source, real_read(file_name(source, stand_in_cycle), stand_in_cycle))
    data[[variable]] <<- d[[weight]][match(data$SEQN, d$SEQN)]
  }
  if (component != "DEMO") stand_in_weight("WTPH2YR", "DEMO", "WTMEC2YR")
  if (component == "DSQTOT") stand_in_weight("WTDRD1", "DR1TOT", "WTDRD1")
  data
}

carried <- function(name, cycle) any(SOURCES$cycle == cycle & SOURCES$file == name)

# A component's records from its latest release before the stand-in release, drawn at random
# (with replacement) for each of the stand-in release's participants.
drawn <- function(component) {
  earlier <- rev(CYCLES$cycle[seq_len(match(stand_in_cycle, CYCLES$cycle) - 1)])
  cycle <- Find(function(cycle) carried(file_name(component, cycle), cycle), earlier)
  if (is.null(cycle)) stop("no release of ", component, " to stand in for its 2021-2023 file")
  records <- real_read(file_name(component, cycle), cycle)
  people <- real_read(file_name("DEMO", stand_in_cycle), stand_in_cycle)$SEQN
  set.seed(20261006)
  data <- records[sample(nrow(records), length(people), replace = TRUE), , drop = FALSE]
  data$SEQN <- people
  structure(data, from = paste(file_name(component, cycle), "drawn for", stand_in_cycle))
}

read_nhanes <- function(name, cycle) {
  if (cycle != REPLICATION_CYCLE) return(real_read(name, cycle))
  component <- sub("_L$", "", name)
  from <- file_name(component, stand_in_cycle)
  if (carried(from, stand_in_cycle)) {
    data <- recode(component, real_read(from, stand_in_cycle))
  } else {
    data <- drawn(component)
    from <- attr(data, "from")
    data <- recode(component, data)
  }
  listed <- catalog$variable[catalog$file == name]
  if (length(listed) == 0) stop(name, " is not a 2021-2023 file")
  for (v in setdiff(listed, names(data))) data[[v]] <- NA
  stand_in_files <<- union(stand_in_files, paste(name, "from", from))
  data[, listed, drop = FALSE]
}

only <- unlist(strsplit(sub("^--only=", "", grep("^--only=", commandArgs(trailingOnly = TRUE), value = TRUE)), ","))
failed <- character()
results <- lapply(load_associations(only), function(a) {
  stand_in_cycle <<- if ("2017-2020" %in% a$cycles) "2017-2020" else "2017-2018"
  # The variants run only on the paper's own cycles, which this doesn't check.
  a$variants <- list()
  tryCatch(run_association(a, replication = TRUE), error = function(e) {
    failed <<- c(failed, paste0(a$id, ": ", conditionMessage(e)))
    NULL
  })
})
results <- correct_all(Filter(Negate(is.null), results), replication = TRUE)
summary <- summarize_results(results, replication = TRUE)
cat("The replications ran on stand-ins for:\n", paste(" ", sort(stand_in_files), collapse = "\n"), "\n")
for (r in results) cat(sprintf("%-8s n = %5d  estimate %.3f  weights %s\n", r$id, r$replication$n, r$replication$estimate, paste(r$replication$weights, collapse = ",")))
if (length(failed) > 0) {
  cat("Failed:\n", paste(" ", failed, collapse = "\n"), "\n")
  quit(status = 1)
}