# Runs the pre-registered associations and summarizes them.

load_associations <- function(only = character()) {
  files <- sort(list.files("code/associations", pattern = "\\.R$", full.names = TRUE))
  associations <- lapply(files, function(f) {
    e <- new.env()
    sys.source(f, envir = e)
    e$association
  })
  if (length(only) > 0) associations <- Filter(function(a) a$id %in% only, associations)
  associations
}

summarize_fit <- function(fit) {
  list(estimate = fit$estimate, ci = fit$ci, b = fit$b, se = fit$se, df = fit$df, p = fit$p,
       n = fit$n, events = fit$events, weighted = fit$weighted, weights = fit$weights, dropped = fit$dropped,
       psu = fit$psu, strata = fit$strata)
}

# A 2021-2023 fit, or, if it can't be fitted, its error with p = 1: such an association counts as
# tested and not replicated, as the plan sets.
attempt <- function(fit) {
  tryCatch(summarize_fit(fit), error = function(e) {
    list(error = conditionMessage(e), estimate = NA, ci = c(NA, NA), b = NA, se = NA, df = NA, p = 1, n = 0)
  })
}

DEPARTURE_KINDS <- c("weighting", "sample", "coding", "model", "reporting")

# How a departure that bears on the headline estimate is established, added after registration in
# answer to review: "paper" where the paper alone shows it (two of its parts disagree, or arithmetic
# on its own printed numbers rules out what it describes); "data" where recomputing from the public
# files shows it and identifies what was computed (only another computation reproduces the paper's
# numbers); "unresolved" where recomputing rules out the described computation without identifying
# the one used. data/departure_coding.csv holds two independent codings of these.
DEPARTURE_EVIDENCE <- c("paper", "data", "unresolved")

# Where the coding check found a paper's analysis to have computed something otherwise than its
# text says (a$departures; an empty list if nowhere): each with its kind, what was found, whether
# it bears on the headline estimate, whether the paper version follows it (FALSE where it
# corrects it, NA where it bears on nothing the re-implementation computes), and, where it bears on
# the headline, its evidence.
check_departures <- function(a) {
  if (is.null(a$departures)) stop(a$id, ": no departures (an empty list if the coding check found none)")
  for (d in a$departures) {
    ok <- is.list(d) && isTRUE(d$kind %in% DEPARTURE_KINDS) && is.character(d$detail) && length(d$detail) == 1 && nzchar(d$detail) &&
      isTRUE(d$affects_headline %in% c(TRUE, FALSE)) && length(d$followed) == 1 && is.logical(d$followed) &&
      (!isTRUE(d$affects_headline) || isTRUE(d$evidence %in% DEPARTURE_EVIDENCE))
    if (!ok) stop(a$id, ": each departure needs a kind (", paste(DEPARTURE_KINDS, collapse = ", "), "), a detail, affects_headline, and followed, ",
                  "and one bearing on the headline its evidence (", paste(DEPARTURE_EVIDENCE, collapse = ", "), ")")
  }
  a$departures
}

# The kinds of departure the described version puts back as the text describes them: computations,
# not two of a paper's accounts disagreeing (reporting).
DESCRIBED_KINDS <- c("coding", "sample", "weighting", "model")

# The departures that qualify for the described version: bearing on the headline, followed by the
# re-implementation (one it corrects is already computed as described), and computations.
qualifying_departures <- function(a) {
  which(vapply(a$departures, function(d) isTRUE(d$affects_headline) && isTRUE(d$followed) && d$kind %in% DESCRIBED_KINDS, TRUE))
}

# Added in answer to the second version's reviews: the analysis as its paper describes it, the paper
# version with the qualifying departures in a$described$applies put back as the text describes them
# (code/described_rules.md), fitted on the paper's own cycles with
# its own analytic sample. Each qualifying departure is in applies or, where the text's version
# can't be built, in not_applied. With none applied, the described version is the paper version
# (fit NULL).
described_version <- function(a, constants) {
  d <- a$described
  applies <- if (is.null(d$applies)) integer() else as.integer(d$applies)
  not_applied <- if (is.null(d$not_applied)) integer() else vapply(d$not_applied, function(x) as.integer(x$departure), 0L)
  named <- c(applies, not_applied)
  qualifying <- qualifying_departures(a)
  if (anyDuplicated(named) || !setequal(named, qualifying)) {
    stop(a$id, ": described must name each qualifying departure (", paste(qualifying, collapse = ", "), ") once, in applies or not_applied")
  }
  if (length(applies) == 0) return(list(applies = applies, not_applied = d$not_applied, fit = NULL))
  changed <- modifyList(a, d$change)
  cycles <- if (!is.null(d$change$cycles)) d$change$cycles else a$cycles
  if (is.function(changed$constants)) constants <- changed$constants(pooled_data(changed, cycles))
  list(applies = applies, not_applied = d$not_applied, fit = fit_association(changed, cycles, constants, "paper"))
}

# One association: the paper's analysis and the harmonized one on the paper's own cycles, the
# variants tried there, the analysis as the paper describes it, and, when `replication` is TRUE,
# the harmonized analysis on 2021-2023.
run_association <- function(a, replication) {
  message("== ", a$id, " (Table A row ", a$row, ")")
  constants <- if (is.function(a$constants)) a$constants(pooled_data(a, a$cycles)) else NULL
  original <- fit_association(a, a$cycles, constants, "paper")
  harmonized <- fit_association(a, a$cycles, constants, "harmonized")
  # Other ways of running the paper's analysis, tried on its own cycles to settle choices it left
  # unstated. A variant that changes the model keeps the paper model's analytic sample.
  variants <- lapply(a$variants, function(v) {
    changed <- modifyList(a, v[names(v) != "label"])
    if (is.null(a$sample_formula) && is.null(v$sample_formula) && !is.null(v$formula)) changed$sample_formula <- a$formula
    cycles <- if (!is.null(v$cycles)) v$cycles else a$cycles
    c(list(label = v$label), summarize_fit(fit_association(changed, cycles, constants, "paper")))
  })
  published <- published_b(a$published, FAMILIES[[a$family]]$ratio)
  years <- sum(vapply(a$cycles, function(cycle) cycle_info(cycle)$years, 0))
  out <- list(
    id = a$id, row = a$row, doi = a$doi, measure = a$published$measure, contrast = a$published$contrast,
    published = a$published, published_b = published, constants = constants,
    original = summarize_fit(original),
    harmonized = summarize_fit(harmonized),
    reproduced = a$published$low <= original$estimate && original$estimate <= a$published$high,
    same_sign_original = sign(original$b) == sign(published[1]),
    ratio_original = original$b / published[1],
    # The published estimate's z, and the harmonized estimate against the paper version's on the
    # same cycles (what leaving out what 2021-2023 lacks changes).
    published_z = published[1] / published_se(published),
    # The published interval's width against ours on the same cycles (analysis scale): below 1
    # where the published interval is narrower than our analysis of the same data gives.
    interval_ratio_original = (published[3] - published[2]) / (original$ci_b[2] - original$ci_b[1]),
    ratio_harmonized = harmonized$b / original$b,
    variants = variants, left_out = a$left_out, choices = a$choices, departures = check_departures(a),
    projected_power = projected_power(published, years)
  )
  described <- described_version(a, constants)
  described_fit <- if (is.null(described$fit)) original else described$fit
  # A described version whose fit doesn't converge, or gives no finite estimate, has no estimate, and
  # like a 2021-2023 model that can't be fitted, it counts as meeting neither criterion.
  fitted <- isTRUE(described_fit$converged) && is.finite(described_fit$b) && is.finite(described_fit$se)
  if (!fitted) described_fit[c("estimate", "ci", "b", "se", "p")] <- list(NA, c(NA, NA), NA, NA, NA)
  out$described <- c(list(applies = as.list(described$applies), not_applied = described$not_applied, same_as_paper_version = is.null(described$fit),
                          converged = fitted), summarize_fit(described_fit))
  # Whether the analysis as described would meet the registered reproduction criterion, and whether
  # it gives a significant estimate with the published sign, beside the paper version's.
  out$described_in_published_ci <- fitted && a$published$low <= described_fit$estimate && described_fit$estimate <= a$published$high
  out$described_significant <- fitted && isTRUE(sign(described_fit$b) == sign(published[1]) && described_fit$p < 0.05)
  out$paper_version_significant <- isTRUE(out$same_sign_original && original$p < 0.05)
  out$described_ratio <- if (fitted) described_fit$b / published[1] else NA
  if (replication) {
    fit <- tryCatch(fit_association(a, REPLICATION_CYCLE, constants, "harmonized"), error = function(e) e)
    if (inherits(fit, "error")) {
      out$replication <- attempt(stop(fit))
      out$power <- 0
      out$power_half <- 0
      out$same_sign <- FALSE
      out$in_published_ci <- FALSE
      out$ratio <- NA
      out$ratio_own <- NA
      out$heterogeneity_p <- 1
      out$published_difference_p <- 1
      out$heterogeneity_p_t <- 1
      out$published_difference_p_t <- 1
      return(out)
    }
    out$replication <- summarize_fit(fit)
    # Sensitivity: the paper's own weight in place of the weights NCHS directs for 2021-2023's
    # subsamples (phlebotomy, supplement questionnaire).
    if (fit$weighted && !identical(fit$weights, weight_variable(a$weight, REPLICATION_CYCLE))) {
      out$replication_paper_weight <- attempt(fit_association(modifyList(a, list(cycle_weights = FALSE)), REPLICATION_CYCLE, constants, "harmonized"))
    }
    # A paper re-run unweighted, because only that reproduces its estimates, is also run weighted.
    if (!original$weighted) {
      out$replication_weighted <- attempt(fit_association(modifyList(a, list(weighted = TRUE)), REPLICATION_CYCLE, constants, "harmonized"))
      # Added after registration, in answer to review: its unweighted estimate with design-based
      # variance (equal weights within the strata and PSUs, t on the design degrees of freedom).
      out$replication_design <- attempt(fit_association(modifyList(a, list(equal_weights = TRUE)), REPLICATION_CYCLE, constants, "harmonized"))
    }
    # Added after registration, in answer to review: a blood analyte analyzed with the interview
    # weight, which the registered rule keeps, analyzed with the phlebotomy weight instead.
    if (fit$weighted && !is.null(a$blood_file) && "WTINT2YR" %in% fit$weights) {
      out$replication_phlebotomy <- attempt(fit_association(modifyList(a, list(blood_replaces = c("WTMEC2YR", "WTINT2YR"))), REPLICATION_CYCLE, constants, "harmonized"))
    }
    z <- (fit$b - harmonized$b) / sqrt(fit$se^2 + harmonized$se^2)
    # The 2021-2023 estimate against the published one: independent samples, z on the analysis
    # scale with the published standard error from its interval.
    z_published <- (fit$b - published[1]) / sqrt(fit$se^2 + published_se(published)^2)
    out$published_difference_p <- 2 * stats::pnorm(-abs(z_published))
    out$power <- power_to_detect(published[1], fit$se, fit$df)
    # Added after registration, in answer to review: power to detect half the published effect,
    # since effects published for crossing a threshold are expected to be inflated.
    out$power_half <- power_to_detect(published[1] / 2, fit$se, fit$df)
    out$same_sign <- sign(fit$b) == sign(published[1])
    out$in_published_ci <- a$published$low <= fit$estimate && fit$estimate <= a$published$high
    out$ratio <- fit$b / published[1]
    out$ratio_own <- fit$b / harmonized$b
    out$heterogeneity_p <- 2 * stats::pnorm(-abs(z))
    # Sensitivity added after registration, in answer to review: the same two tests referred to a t
    # distribution on the design degrees of freedom of the fits they compare (15 for a weighted
    # 2021-2023 fit) rather than to the normal.
    out$published_difference_p_t <- 2 * stats::pt(-abs(z_published), fit$df)
    out$heterogeneity_p_t <- 2 * stats::pt(-abs(z), min(fit$df, harmonized$df))
  }
  out
}

# Every association, then the corrections over all of them.
run_all <- function(associations, replication) {
  correct_all(lapply(associations, run_association, replication = replication), replication)
}

# The sensitivity analyses that refit some associations in 2021-2023, by the name of their fit: the
# paper's own weight (registered), the survey design and weights for an analysis re-run unweighted
# (registered), and, added after registration in answer to review, the design-based variance of its
# unweighted estimate, and the phlebotomy weight for a blood analyte kept on the interview weight.
SENSITIVITIES <- c(paper_weight = "replication_paper_weight", weighted = "replication_weighted",
                   design = "replication_design", phlebotomy = "replication_phlebotomy")

# An association's 2021-2023 fit in a sensitivity analysis: its refit there, or the registered fit
# where the sensitivity doesn't refit it.
sensitivity_fit <- function(r, sensitivity) {
  refit <- r[[SENSITIVITIES[[sensitivity]]]]
  if (is.null(refit)) r$replication else refit
}

# Benjamini-Hochberg over the replication p-values (the primary outcome), over the heterogeneity
# tests, and over the replication in each sensitivity analysis.
correct_all <- function(results, replication) {
  names(results) <- vapply(results, function(r) r$id, "")
  if (replication) {
    q <- stats::p.adjust(vapply(results, function(r) r$replication$p, 0), method = "BH")
    q_het <- stats::p.adjust(vapply(results, function(r) r$heterogeneity_p, 0), method = "BH")
    q_published <- stats::p.adjust(vapply(results, function(r) r$published_difference_p, 0), method = "BH")
    q_sensitivity <- lapply(names(SENSITIVITIES), function(s) {
      stats::p.adjust(vapply(results, function(r) sensitivity_fit(r, s)$p, 0), method = "BH")
    })
    names(q_sensitivity) <- names(SENSITIVITIES)
    # The 2021-2023 estimate against the published one in each sensitivity analysis, as registered
    # for the primary fit (z on the analysis scale with the published standard error).
    q_published_sensitivity <- lapply(names(SENSITIVITIES), function(s) {
      stats::p.adjust(vapply(results, function(r) {
        fit <- sensitivity_fit(r, s)
        if (is.null(fit$error)) 2 * stats::pnorm(-abs((fit$b - r$published_b[1]) / sqrt(fit$se^2 + published_se(r$published_b)^2))) else 1
      }, 0), method = "BH")
    })
    names(q_published_sensitivity) <- names(SENSITIVITIES)
    q_published_sensitivity_t <- lapply(names(SENSITIVITIES), function(s) {
      stats::p.adjust(vapply(results, function(r) {
        fit <- sensitivity_fit(r, s)
        if (is.null(fit$error)) 2 * stats::pt(-abs((fit$b - r$published_b[1]) / sqrt(fit$se^2 + published_se(r$published_b)^2)), fit$df) else 1
      }, 0), method = "BH")
    })
    names(q_published_sensitivity_t) <- names(SENSITIVITIES)
    q_published_t <- stats::p.adjust(vapply(results, function(r) r$published_difference_p_t, 0), method = "BH")
    q_het_t <- stats::p.adjust(vapply(results, function(r) r$heterogeneity_p_t, 0), method = "BH")
    # Power at the Bonferroni level, 0.05 over the number of tests, added after registration in
    # answer to review: Benjamini-Hochberg rejects every p-value below that level and none above
    # 0.05, so its power for a test lies between this and the registered power at 0.05.
    alpha_bonferroni <- 0.05 / length(results)
    for (i in seq_along(results)) {
      r <- results[[i]]
      results[[i]]$q <- q[i]
      results[[i]]$replicated <- isTRUE(r$same_sign) && q[i] < 0.05
      results[[i]]$informative <- r$power >= 0.8
      results[[i]]$heterogeneity_q <- q_het[i]
      # Where the 2021-2023 estimate differs significantly from the published one, and how.
      results[[i]]$published_difference_q <- q_published[i]
      results[[i]]$differs_from_published <- q_published[i] < 0.05
      b <- r$replication$b
      results[[i]]$difference <- if (!isTRUE(q_published[i] < 0.05)) "none" else if (sign(b) != sign(r$published_b[1])) "opposite sign" else
        if (abs(b) < abs(r$published_b[1])) "smaller" else "larger"
      for (s in names(SENSITIVITIES)) {
        refit <- sensitivity_fit(r, s)
        results[[i]][[paste0("replicated_", s)]] <- isTRUE(sign(refit$b) == sign(r$published_b[1])) && q_sensitivity[[s]][i] < 0.05
        results[[i]][[paste0("published_difference_q_", s)]] <- q_published_sensitivity[[s]][i]
        results[[i]][[paste0("published_difference_q_t_", s)]] <- q_published_sensitivity_t[[s]][i]
        results[[i]][[paste0("power_", s)]] <- if (is.null(refit$error)) power_to_detect(r$published_b[1], refit$se, refit$df) else 0
      }
      results[[i]]$published_difference_q_t <- q_published_t[i]
      results[[i]]$differs_from_published_t <- q_published_t[i] < 0.05
      results[[i]]$heterogeneity_q_t <- q_het_t[i]
      fit <- r$replication
      results[[i]]$power_bonferroni <- if (is.null(fit$error)) power_to_detect(r$published_b[1], fit$se, fit$df, alpha_bonferroni) else 0
      results[[i]]$informative_bonferroni <- results[[i]]$power_bonferroni >= 0.8
    }
  }
  results
}

write_json <- function(x, path) {
  writeLines(jsonlite::toJSON(x, auto_unbox = TRUE, digits = NA, pretty = TRUE, null = "null", na = "null"), path)
}