# Survey weights and designs.
#
# An analysis names its weight by its 2-year name in 2001-2018 files ("WTMEC2YR", "WTINT2YR",
# "WTSAF2YR", "WTDRD1", "WTDR2D", or a subsample weight). The 2017-March 2020 files and the
# 1999-2002 four-year weights name theirs differently. In 2021-2023 two components moved to smaller
# subsamples, and NCHS directs their weights: blood analytes to the phlebotomy subsample (WTPH2YR,
# whenever a blood analyte is in the analysis), and the 30-day dietary supplement questionnaire to
# the telephone interview after the first dietary recall (WTDRD1, for any analysis using it, "either
# alone or when Dietary Supplement data are used in conjunction with MEC data").

PREPANDEMIC_WEIGHTS <- c(WTMEC2YR = "WTMECPRP", WTINT2YR = "WTINTPRP", WTSAF2YR = "WTSAFPRP",
                         WTDRD1 = "WTDRD1PP", WTDR2D = "WTDR2DPP")
FOUR_YEAR_WEIGHTS <- c(WTMEC2YR = "WTMEC4YR", WTINT2YR = "WTINT4YR", WTSAF2YR = "WTSAF4YR", WTDRD1 = "WTDR4YR")

# The components of the 30-day dietary supplement questionnaire.
SUPPLEMENT_COMPONENTS <- c("DSQTOT", "DSQIDS")

# The weight variable an analysis uses in a cycle. `blood` says whether a blood analyte is in it,
# `supplements` whether the 30-day supplement questionnaire is. The supplement weight replaces the
# examination and interview weights; the fasting and dietary weights already cover subsamples at
# least as small.
weight_variable <- function(weight, cycle, blood = FALSE, supplements = FALSE) {
  if (cycle == "2017-2020") {
    if (!weight %in% names(PREPANDEMIC_WEIGHTS)) stop("No 2017-March 2020 weight for ", weight)
    return(unname(PREPANDEMIC_WEIGHTS[weight]))
  }
  if (cycle == REPLICATION_CYCLE && supplements && weight %in% c("WTMEC2YR", "WTINT2YR")) return("WTDRD1")
  if (cycle == REPLICATION_CYCLE && blood && weight == "WTMEC2YR") return("WTPH2YR")
  weight
}

# Pooled weights: each cycle's weight times the share of the pooled years the cycle covers,
# as NCHS's analytic guidelines say (a 2-year cycle among k 2-year cycles gets 1/k). When both
# 1999-2000 and 2001-2002 are pooled, NCHS's four-year weights replace their 2-year weights.
# rule = "equal" divides every cycle's weight by the number of cycles instead, as some papers do.
pooled_weight <- function(data, weight, cycles, rule = "nchs") {
  four_year <- all(c("1999-2000", "2001-2002") %in% cycles) && weight %in% names(FOUR_YEAR_WEIGHTS)
  total <- sum(sapply(cycles, function(cycle) cycle_info(cycle)$years))
  w <- rep(NA_real_, nrow(data))
  for (cycle in cycles) {
    rows <- data$cycle == cycle
    if (rule == "equal") {
      w[rows] <- data$weight_2yr[rows] / length(cycles)
    } else if (four_year && cycle %in% c("1999-2000", "2001-2002")) {
      w[rows] <- data$weight_4yr[rows] * 4 / total
    } else {
      w[rows] <- data$weight_2yr[rows] * cycle_info(cycle)$years / total
    }
  }
  w
}

# The survey design over everyone with a positive weight, with Taylor-linearized variance,
# strata and PSUs nested within cycles' pseudo-strata.
survey_design <- function(data) {
  data <- data[!is.na(data$w) & data$w > 0, ]
  survey::svydesign(ids = ~SDMVPSU, strata = ~SDMVSTRA, weights = ~w, nest = TRUE, data = data)
}

# Where each weight lives, by component, and its four-year name in 1999-2002.
WEIGHT_FILES <- c(WTMEC2YR = "DEMO", WTINT2YR = "DEMO", WTSAF2YR = "GLU", WTDRD1 = "DR1TOT", WTDR2D = "DR1TOT")

# SEQN with the cycle's weight for an analysis (weight_2yr, and weight_4yr in 1999-2002).
# In 2021-2023 a blood analysis takes the phlebotomy weight from the blood file it reads, and an
# analysis using the supplement questionnaire takes WTDRD1 from DSQTOT_L.
weights_for <- function(weight, cycle, blood_file = NULL, file = NULL, supplements = FALSE) {
  variable <- weight_variable(weight, cycle, blood = !is.null(blood_file), supplements = supplements)
  source <- if (variable == "WTPH2YR") blood_file else if (variable != weight && variable == "WTDRD1") "DSQTOT" else if (!is.null(file)) file else WEIGHT_FILES[[weight]]
  data <- component(source, cycle)
  out <- data.frame(SEQN = data$SEQN, weight_2yr = data[[variable]])
  four <- FOUR_YEAR_WEIGHTS[weight]
  out$weight_4yr <- if (cycle %in% c("1999-2000", "2001-2002") && !is.na(four) && four %in% names(data)) data[[four]] else NA_real_
  out
}