# 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
}

# One association: the paper's analysis and the harmonized one on the paper's own cycles, the
# variants tried there, 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)
  )
  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$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"))
    }
    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)
    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)
}

# Benjamini-Hochberg over the replication p-values (the primary outcome), over the heterogeneity
# tests, and over the replication with the paper's own weights.
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")
    p_paper_weight <- vapply(results, function(r) (if (is.null(r$replication_paper_weight)) r$replication else r$replication_paper_weight)$p, 0)
    q_paper_weight <- stats::p.adjust(p_paper_weight, method = "BH")
    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]]
      b_paper_weight <- (if (is.null(r$replication_paper_weight)) r$replication else r$replication_paper_weight)$b
      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"
      results[[i]]$replicated_paper_weight <- isTRUE(sign(b_paper_weight) == sign(r$published_b[1])) && q_paper_weight[i] < 0.05
      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)
}