# Covariates as papers in this literature define them. Where a paper leaves a definition
# unstated, its association file says which of these it uses and why.

# Drinking status. 1999-2016 asked whether someone had at least 12 drinks in any one year
# (ALQ101) or in their life (ALQ110) and how often they drank in the past 12 months (ALQ120Q,
# ALQ120U); 2017 on asks whether they ever had a drink (ALQ111) and how often in the past 12
# months (ALQ121, 0 = never in the last year).
# three = 1 never, 2 former, 3 current.
alcohol <- function(cycle) {
  d <- component("ALQ", cycle)
  if (has(d, "ALQ121")) {
    ever <- yes(d$ALQ111)
    past_year <- ifelse(d$ALQ121 %in% 1:10, 1, ifelse(d$ALQ121 %in% 0, 0, NA))
    drinks <- ifelse(d$ALQ130 %in% 1:15, d$ALQ130, NA)
    # days a year from ALQ121's categories (midpoints)
    days <- c(`0` = 0, `1` = 365, `2` = 300, `3` = 182, `4` = 104, `5` = 52, `6` = 30, `7` = 12, `8` = 9, `9` = 4.5, `10` = 1.5)
    days_per_year <- unname(days[as.character(d$ALQ121)])
    days_per_year[ever %in% 0] <- 0
    binge_days <- if (has(d, "ALQ142")) unname(days[as.character(d$ALQ142)]) else NA
  } else {
    # At least 12 drinks in any one year: ALQ101 (ALQ100 in 1999-2000, ALD100 in 2001-2002); in
    # life: ALQ110.
    year12 <- yes(d[[intersect(c("ALQ101", "ALD100", "ALQ100"), names(d))[1]]])
    lifetime12 <- ifelse(year12 %in% 1 | yes(d$ALQ110) %in% 1, 1,
                         ifelse(yes(d$ALQ110) %in% 0 | (year12 %in% 0 & is.na(d$ALQ110)), 0, NA))
    ever <- lifetime12
    q <- ifelse(d$ALQ120Q %in% 0:365, d$ALQ120Q, NA)
    per_year <- ifelse(d$ALQ120U %in% 1, 52, ifelse(d$ALQ120U %in% 2, 12, ifelse(d$ALQ120U %in% 3, 1, NA)))
    days_per_year <- ifelse(q %in% 0, 0, q * per_year)
    past_year <- ifelse(days_per_year > 0, 1, ifelse(days_per_year %in% 0, 0, NA))
    # those with fewer than 12 drinks ever skip the frequency question
    past_year[ever %in% 0] <- 0
    days_per_year[ever %in% 0] <- 0
    # ALQ130 counts drinks; 77/99 and 777/999 code refused or don't know.
    drinks <- ifelse(d$ALQ130 >= 1 & d$ALQ130 < 77, d$ALQ130, NA)
    # Days of heavy drinking in the past 12 months: 5 or more drinks (ALQ140Q/U, 1999-2010) or
    # 4 or more for women and 5 or more for men (ALQ141Q/U, 2011-2016).
    q_var <- if (has(d, "ALQ141Q")) "ALQ141Q" else if (has(d, "ALQ140Q")) "ALQ140Q" else NA
    u_var <- if (has(d, "ALQ141U")) "ALQ141U" else if (has(d, "ALQ140U")) "ALQ140U" else NA
    bq <- if (!is.na(q_var)) ifelse(d[[q_var]] %in% 0:365, d[[q_var]], NA) else NA
    bu <- if (!is.na(u_var)) ifelse(d[[u_var]] %in% 1, 52, ifelse(d[[u_var]] %in% 2, 12, ifelse(d[[u_var]] %in% 3, 1, NA))) else NA
    binge_days <- ifelse(bq %in% 0, 0, bq * bu)
  }
  three <- ifelse(ever %in% 0, 1, ifelse(ever %in% 1 & past_year %in% 0, 2, ifelse(past_year %in% 1, 3, NA)))
  data.frame(SEQN = d$SEQN, alcohol3 = three, drinks_per_day = drinks, drinking_days = days_per_year,
             binge_days = binge_days,
             # Ever drank 4 or 5 or more drinks almost every day: ALQ151 (2011 on), ALQ150 (1999-2010).
             daily_binge_ever = if (has(d, "ALQ151")) yes(d$ALQ151) else if (has(d, "ALQ150")) yes(d$ALQ150) else NA)
}

# Drinking in five groups (never, former, mild, moderate, heavy) as Rattan and colleagues'
# NHANES definitions that much of this literature cites: heavy is 3 or more drinks a day for
# women, 4 or more for men, or binge drinking (4 or 5 drinks on one occasion) on 5 or more days a
# month; moderate is 2 or more a day for women, 3 or more for men, or binge drinking on 2 or more
# days a month; mild is any other current drinking.
alcohol5 <- function(alq, sex) {
  binge_month <- alq$binge_days / 12
  heavy <- (sex == 2 & alq$drinks_per_day >= 3) | (sex == 1 & alq$drinks_per_day >= 4) | binge_month >= 5
  moderate <- (sex == 2 & alq$drinks_per_day >= 2) | (sex == 1 & alq$drinks_per_day >= 3) | binge_month >= 2
  ifelse(alq$alcohol3 %in% 1, 1, ifelse(alq$alcohol3 %in% 2, 2,
    ifelse(alq$alcohol3 %in% 3 & heavy %in% TRUE, 5, ifelse(alq$alcohol3 %in% 3 & moderate %in% TRUE, 4,
      ifelse(alq$alcohol3 %in% 3, 3, NA)))))
}

# Total physical activity in MET-minutes per week from the Global Physical Activity
# Questionnaire (2007-2018 and 2017-March 2020): 8 METs for vigorous and 4 for moderate work and
# leisure, 4 for walking or cycling to get places. NA in cycles without it (1999-2006, 2021-2023).
met_minutes <- function(cycle) {
  d <- component("PAQ", cycle)
  if (!has(d, "PAQ605")) return(data.frame(SEQN = d$SEQN, met = NA_real_))
  part <- function(answer, days, minutes, met) {
    m <- ifelse(minutes %in% c(7777, 9999), NA, minutes)
    n <- ifelse(days %in% 1:7, days, NA)
    ifelse(answer %in% 2, 0, ifelse(answer %in% 1, met * n * m, NA))
  }
  parts <- cbind(
    part(d$PAQ605, d$PAQ610, d$PAD615, 8), part(d$PAQ620, d$PAQ625, d$PAD630, 4),
    part(d$PAQ635, d$PAQ640, d$PAD645, 4), part(d$PAQ650, d$PAQ655, d$PAD660, 8),
    part(d$PAQ665, d$PAQ670, d$PAD675, 4))
  data.frame(SEQN = d$SEQN, met = rowSums(parts))
}

# Nutrient totals from the 24-hour recalls: the first day's (days = 1), the mean of both days for
# those with two reliable recalls (days = 2, missing otherwise), or that mean where there are two
# and the first day's otherwise (days = "mean_or_one"). A recall counts if its status is 1
# (reliable and met the minimum criteria). `variables` are names after the prefix, such as
# "KCAL"; 1999-2002 name the day-one file DRXTOT and its variables DRXT...
dietary_totals <- function(cycle, variables, days = 1) {
  d1 <- component("DR1TOT", cycle)
  prefix1 <- if (has(d1, "DR1TKCAL")) "DR1T" else "DRXT"
  # Recall status: DR1DRSTZ, DRDDRSTZ in 2001-2002, DRDDRSTS in 1999-2000.
  status1 <- d1[[intersect(c("DR1DRSTZ", "DRDDRSTZ", "DRDDRSTS"), names(d1))[1]]]
  absent <- setdiff(paste0(prefix1, variables), names(d1))
  if (length(absent) > 0) stop(cycle, "'s day-one dietary file has no ", paste(absent, collapse = ", "))
  out <- data.frame(SEQN = d1$SEQN, recall_day1 = as.integer(status1 %in% 1))
  if (!identical(days, 1)) {
    d2 <- component("DR2TOT", cycle)
    rows <- match(d1$SEQN, d2$SEQN)
    out$recall_day2 <- as.integer(d2$DR2DRSTZ[rows] %in% 1)
  }
  for (v in variables) {
    one <- ifelse(status1 %in% 1, d1[[paste0(prefix1, v)]], NA)
    if (identical(days, 1)) {
      out[[v]] <- one
      next
    }
    two <- ifelse(out$recall_day2 %in% 1, d2[[paste0("DR2T", v)]][rows], NA)
    both <- !is.na(one) & !is.na(two)
    out[[v]] <- ifelse(both, (one + two) / 2, if (identical(days, "mean_or_one")) one else NA)
  }
  out
}

# Diabetes from any of: told by a doctor (DIQ010 = 1), HbA1c 6.5% or more, fasting glucose
# 126 mg/dL or more, or taking insulin or diabetes pills. `parts` picks which; a part a person
# lacks counts as not met. The glucose part reads every value in the fasting-subsample file, as
# papers in this literature do, including those of people NCHS doesn't count as fasted
# (fasting_glucose() flags them). An OGTT criterion, where a paper uses one, is its own file's.
diabetes_status <- function(cycle, parts = c("told", "hba1c", "glucose", "medication")) {
  q <- diabetes_questions(cycle)
  out <- q[, "SEQN", drop = FALSE]
  met <- matrix(FALSE, nrow(q), 0)
  known <- matrix(FALSE, nrow(q), 0)
  add <- function(flag) { met <<- cbind(met, flag %in% 1); known <<- cbind(known, !is.na(flag)) }
  if ("told" %in% parts) add(q$told_diabetes)
  if ("medication" %in% parts) add(ifelse(q$insulin %in% 1 | q$pills %in% 1, 1, 0))
  if ("hba1c" %in% parts) {
    h <- hba1c(cycle)
    add(as.integer(h$hba1c[match(q$SEQN, h$SEQN)] >= 6.5))
  }
  if ("glucose" %in% parts) {
    g <- fasting_glucose(cycle)
    add(as.integer(g$glucose[match(q$SEQN, g$SEQN)] >= 126))
  }
  out$diabetes <- ifelse(rowSums(met) > 0, 1, ifelse(rowSums(known) > 0, 0, NA))
  out
}

# Hyperlipidemia as the NCEP ATP III thresholds define it: total cholesterol 200 mg/dL or more,
# triglycerides 150 or more, LDL 130 or more, HDL under 40 (men) or 50 (women), or taking
# cholesterol-lowering medication. Triglycerides and LDL are every value in the fasting-subsample
# file, fasted or not (triglycerides() flags who fasted).
hyperlipidemia_status <- function(cycle, demo) {
  tc <- total_cholesterol(cycle)
  hdl <- hdl_cholesterol(cycle)
  tg <- triglycerides(cycle)
  bpq <- bp_questions(cycle)
  s <- demo[, c("SEQN", "sex")]
  s$tc <- tc$tc[match(s$SEQN, tc$SEQN)]
  s$hdl <- hdl$hdl[match(s$SEQN, hdl$SEQN)]
  s$tg <- tg$tg[match(s$SEQN, tg$SEQN)]
  s$ldl <- tg$ldl[match(s$SEQN, tg$SEQN)]
  s$med <- bpq$cholesterol_medication[match(s$SEQN, bpq$SEQN)]
  flags <- cbind(s$tc >= 200, s$tg >= 150, s$ldl >= 130, (s$sex == 1 & s$hdl < 40) | (s$sex == 2 & s$hdl < 50), s$med == 1)
  data.frame(SEQN = s$SEQN, hyperlipidemia = ifelse(rowSums(flags, na.rm = TRUE) > 0, 1, ifelse(rowSums(!is.na(flags)) > 0, 0, NA)))
}

# Hypertension: ever told (BPQ020), taking medication, or a mean measured pressure at or above the
# thresholds (blood_pressure() averages every reading; an association that averages others builds
# its own).
hypertension_status <- function(cycle, sbp_at = 140, dbp_at = 90, bp_method = "auto") {
  q <- bp_questions(cycle)
  bp <- blood_pressure(cycle, method = bp_method)
  out <- q[, "SEQN", drop = FALSE]
  sbp <- bp$sbp[match(q$SEQN, bp$SEQN)]
  dbp <- bp$dbp[match(q$SEQN, bp$SEQN)]
  flags <- cbind(q$told_hypertension == 1, q$bp_medication == 1, sbp >= sbp_at, dbp >= dbp_at)
  out$hypertension <- ifelse(rowSums(flags, na.rm = TRUE) > 0, 1, ifelse(rowSums(!is.na(flags)) > 0, 0, NA))
  out
}

# Chronic kidney disease: eGFR under 60 mL/min/1.73 m2 or urine albumin-to-creatinine ratio of
# 30 mg/g or more.
ckd_status <- function(cycle, demo, equation = "2009") {
  b <- biochemistry(cycle)
  a <- urine_albumin_creatinine(cycle)
  s <- demo[, c("SEQN", "age", "sex", "race")]
  s$creatinine <- b$creatinine[match(s$SEQN, b$SEQN)]
  s$acr <- a$acr[match(s$SEQN, a$SEQN)]
  s$egfr <- egfr(standard_creatinine(s$creatinine, cycle), s$age, s$sex, s$race, equation)
  flags <- cbind(s$egfr < 60, s$acr >= 30)
  data.frame(SEQN = s$SEQN, egfr = s$egfr, ckd = ifelse(rowSums(flags, na.rm = TRUE) > 0, 1, ifelse(rowSums(!is.na(flags)) > 0, 0, NA)))
}

# Prescription medicines taken in the past 30 days, by Multum Lexicon therapeutic category
# (RXQ_RX's drug codes joined to RXQ_DRUG's categories). Returns, per SEQN, whether any medicine
# in each of `categories` (a named list of category ID vectors, any of levels 1-3) was taken.
# 2021-2023 releases no drug names, so this is unavailable there.
medication_use <- function(cycle, categories) {
  rx <- component("RXQ_RX", cycle)
  drugs <- read_nhanes("RXQ_DRUG", "lexicon")
  ids <- drugs[, c("RXDDRGID", grep("^RXDDCI", names(drugs), value = TRUE))]
  taken <- rx[!is.na(rx$RXDDRGID) & rx$RXDDRGID != "", c("SEQN", "RXDDRGID")]
  taken <- merge(taken, ids, by = "RXDDRGID")
  code_columns <- grep("^RXDDCI", names(taken), value = TRUE)
  people <- data.frame(SEQN = unique(rx$SEQN))
  answered <- rx$SEQN[rx$RXDUSE %in% 1:2]
  for (name in names(categories)) {
    hit <- apply(taken[, code_columns], 1, function(codes) any(codes %in% categories[[name]]))
    users <- unique(taken$SEQN[hit])
    people[[name]] <- ifelse(people$SEQN %in% users, 1, ifelse(people$SEQN %in% answered, 0, NA))
  }
  people
}