# The review's registered analysis: does unconditional cash in high-income countries improve recipients' mental
# health, does it change whether they work, and does a mental-health benefit fade while the payments go on?
#
#   Rscript code/analyze.R [--blind] [data dir] [results dir]
#
# Reads data/trials.json (each trial's decided design: its comparisons, its waves with their timing, its
# registration), data/extraction.json (every result its reports print, as decided), data/parameters.json, and
# data/registry.json. Writes results/R*.json (declared results), results/estimates.csv (every usable entry with
# its effect), and results/analyses.csv (every model fitted). --blind flips the sign of each comparison's effects
# at random (seed in parameters.json), so the code can be run and checked before the numbers mean anything.

suppressPackageStartupMessages({
  library(metafor)
  library(jsonlite)
})

args <- commandArgs(trailingOnly = TRUE)
blind <- "--blind" %in% args
args <- setdiff(args, "--blind")
data_dir <- if (length(args) >= 1) args[1] else "data"
out_dir <- if (length(args) >= 2) args[2] else "results"
dir.create(out_dir, showWarnings = FALSE, recursive = TRUE)

P <- fromJSON(file.path(data_dir, "parameters.json"), simplifyVector = FALSE)
trials <- fromJSON(file.path(data_dir, "trials.json"), simplifyVector = FALSE)
extraction <- fromJSON(file.path(data_dir, "extraction.json"), simplifyVector = FALSE)
registry <- fromJSON(file.path(data_dir, "registry.json"), simplifyVector = FALSE)

`%||%` <- function(a, b) if (is.null(a) || length(a) == 0) b else a
num <- function(x) if (is.null(x) || length(x) == 0) NA_real_ else suppressWarnings(as.numeric(x))
`%|na|%` <- function(a, b) if (length(a) == 0 || is.na(a[1])) b else a
set.seed(P$seed)

MH <- unlist(P$mental_health_classes)          # in order of preference
SCALES <- c(MH, "stress", "life_satisfaction")
WORK <- c("employed_any", "employed_full_time")

# ---------------------------------------------------------------------------------------------------
# Comparisons and waves.

eligible <- Filter(function(t) isTRUE(t$eligible), trials)
comparisons <- do.call(rbind, lapply(eligible, function(t) do.call(rbind, lapply(t$comparisons, function(c) data.frame(
  trial = t$trial, name = t$name, comparison = c$id, kind = c$kind, country = t$country,
  currency = c$currency %||% NA_character_, monthly_amount = num(c$monthly_amount), lump_sum = num(c$lump_sum),
  months_paid = num(c$months_paid), n_treated = num(c$n_treated), n_control = num(c$n_control),
  control = c$control %||% NA_character_,
  registered = any(vapply(t$registrations %||% list(), function(r) isTRUE(r$registered_before_payments), TRUE)),
  integrity_ok = !isFALSE(t$integrity$ok), stringsAsFactors = FALSE)))))
sign_of <- setNames(ifelse(blind & runif(nrow(comparisons)) < 0.5, -1, 1), comparisons$comparison)

waves <- do.call(rbind, lapply(eligible, function(t) do.call(rbind, lapply(t$waves, function(w) data.frame(
  trial = t$trial, wave = w$label, months = num(w$months_since_first_payment), during = isTRUE(w$during_payments),
  stringsAsFactors = FALSE)))))
result_rank <- function(trial, source) {
  t <- Filter(function(x) identical(x$trial, trial), eligible)[[1]]
  m <- match(source, unlist(t$result_reports %||% list()))
  if (is.na(m)) 99 else m
}

# ---------------------------------------------------------------------------------------------------
# Effects. Scales become standardized mean differences (Hedges' g), positive when the cash arm is better off;
# shares employed become differences in percentage points, positive when more of the cash arm work; hours become
# differences in hours per week.

z_of <- function(level) qnorm(1 - (1 - (num(level) %|na|% 95) / 100) / 2)
se_from <- function(e, log_scale = FALSE) {
  if (is.null(e)) return(c(NA, NA))
  v <- num(e$value); se <- num(e$se); lo <- num(e$lower); hi <- num(e$upper); p <- num(e$p)
  if (log_scale) { v <- log(v); lo <- log(lo); hi <- log(hi) }
  if (!is.na(se) && se > 0) return(c(se, 1))
  if (!is.na(lo) && !is.na(hi) && hi > lo) return(c((hi - lo) / (2 * z_of(e$level)), 2))
  if (!is.na(p) && p > 0 && p < 1 && !is.na(v) && v != 0) return(c(abs(v) / qnorm(1 - p / 2), 3))
  c(NA, NA)
}
se_how <- c("se", "interval", "p value")
frac <- function(x, unit) if (identical(unit, "percent")) num(x) / 100 else num(x)
J <- function(n) ifelse(is.na(n) | n < 4, 1, 1 - 3 / (4 * n - 9))

# Wave sizes: the entry's own, else the attrition table's, else those randomized.
sizes <- function(trial, e) {
  nt <- num(e$treated$n); nc <- num(e$control$n)
  a <- Filter(function(x) identical(x$comparison, e$comparison) && identical(x$wave, e$wave),
              (Filter(function(x) identical(x$trial, trial), extraction)[[1]])$attrition %||% list())
  if (is.na(nt) && length(a)) nt <- num(a[[1]]$responded_treated)
  if (is.na(nc) && length(a)) nc <- num(a[[1]]$responded_control)
  cm <- comparisons[comparisons$comparison == e$comparison, ]
  if (is.na(nt) && nrow(cm)) nt <- cm$n_treated[1]
  if (is.na(nc) && nrow(cm)) nc <- cm$n_control[1]
  c(nt, nc)
}

# Every standard deviation printed for an instrument, for the instrument-level fallback.
printed_sds <- do.call(rbind, lapply(extraction, function(tr) do.call(rbind, lapply(tr$entries, function(e) {
  s <- c(num(e$treated$sd), num(e$control$sd), num(e$baseline_sd))
  s <- s[!is.na(s) & s > 0]
  if (!length(s) || isTRUE(e$binary)) return(NULL)
  data.frame(trial = tr$trial, instrument = e$instrument %||% "other", measure = e$measure %||% "", sd = s, stringsAsFactors = FALSE)
}))))
if (is.null(printed_sds)) printed_sds <- data.frame(trial = character(), instrument = character(), measure = character(), sd = numeric())

sd_for <- function(trial, e, nt, nc, se_diff) {
  st <- num(e$treated$sd); sc <- num(e$control$sd)
  if (!is.na(st) && !is.na(sc)) {
    w <- if (!is.na(nt) && !is.na(nc)) c(nt - 1, nc - 1) else c(1, 1)
    return(list(sd = sqrt((w[1] * st^2 + w[2] * sc^2) / sum(w)), rule = "wave"))
  }
  if (!is.na(sc)) return(list(sd = sc, rule = "wave"))
  if (!is.na(num(e$baseline_sd))) return(list(sd = num(e$baseline_sd), rule = "baseline"))
  same <- printed_sds$sd[printed_sds$trial == trial & printed_sds$measure == (e$measure %||% "")]
  if (length(same)) return(list(sd = median(same), rule = "same_trial"))
  inst <- e$instrument %||% "other"
  if (inst != "other") {
    other <- printed_sds$sd[printed_sds$instrument == inst]
    if (length(other)) return(list(sd = median(other), rule = "instrument"))
  }
  if (!is.na(se_diff) && !is.na(nt) && !is.na(nc)) return(list(sd = se_diff / sqrt(1 / nt + 1 / nc), rule = "from_se"))
  list(sd = NA_real_, rule = "none")
}

scale_effect <- function(trial, e) {
  dir <- if (identical(e$higher_is, "worse")) -1 else 1
  n <- sizes(trial, e); nt <- n[1]; nc <- n[2]
  est <- e$estimate; kind <- est$kind %||% "none"
  out <- function(g, v, how, rule = NA_character_) list(yi = g, vi = v, how = how, sd_rule = rule, nt = nt, nc = nc)
  if (isTRUE(e$binary)) {
    # A share above a cut-off: log odds ratio, then Chinn's conversion to a standardized difference.
    pu <- e$proportion_unit %||% "fraction"
    pt <- frac(e$treated$proportion, pu); pc <- frac(e$control$proportion, pu)
    if (identical(kind, "odds_ratio") && !is.na(num(est$value)) && num(est$value) > 0) {
      s <- se_from(est, log_scale = TRUE)
      if (!is.na(s[1])) return(out(dir * log(num(est$value)) * sqrt(3) / pi, s[1]^2 * 3 / pi^2, paste("odds ratio,", se_how[s[2]])))
    }
    if (!is.na(pc) && kind %in% c("difference", "adjusted_difference") && !is.na(num(est$value))) {
      d <- num(est$value); if (identical(est$scale, "percentage_points")) d <- d / 100
      pt2 <- pc + d; s <- se_from(est)
      sd_d <- if (!is.na(s[1])) (if (identical(est$scale, "percentage_points")) s[1] / 100 else s[1]) else NA
      if (pt2 > 0 && pt2 < 1 && pc > 0 && pc < 1) {
        lor <- qlogis(pt2) - qlogis(pc)
        v <- if (!is.na(sd_d)) (sd_d / (pt2 * (1 - pt2)))^2 else 1 / (nt * pt2 * (1 - pt2)) + 1 / (nc * pc * (1 - pc))
        if (!is.na(v)) return(out(dir * lor * sqrt(3) / pi, v * 3 / pi^2, "share difference"))
      }
    }
    if (!is.na(pt) && !is.na(pc) && !is.na(nt) && !is.na(nc) && pt > 0 && pt < 1 && pc > 0 && pc < 1) {
      lor <- qlogis(pt) - qlogis(pc)
      v <- 1 / (nt * pt * (1 - pt)) + 1 / (nc * pc * (1 - pc))
      return(out(dir * lor * sqrt(3) / pi, v * 3 / pi^2, "shares"))
    }
    return(NULL)
  }
  if (identical(kind, "standardized_difference") && !is.na(num(est$value))) {
    g <- dir * num(est$value); s <- se_from(est)
    v <- if (!is.na(s[1])) s[1]^2 else (1 / nt + 1 / nc + g^2 / (2 * (nt + nc)))
    if (is.na(v)) return(NULL)
    return(out(g, v, paste("standardized,", if (!is.na(s[1])) se_how[s[2]] else "sizes")))
  }
  if (kind %in% c("difference", "adjusted_difference") && !is.na(num(est$value))) {
    d <- num(est$value); s <- se_from(est)
    sdv <- sd_for(trial, e, nt, nc, s[1])
    if (is.na(sdv$sd) || sdv$sd <= 0) return(NULL)
    g <- dir * d / sdv$sd
    v <- if (!is.na(s[1])) (s[1] / sdv$sd)^2 else (1 / nt + 1 / nc + g^2 / (2 * (nt + nc)))
    if (is.na(v)) return(NULL)
    return(out(g, v, paste("difference,", if (!is.na(s[1])) se_how[s[2]] else "sizes"), sdv$rule))
  }
  mt <- num(e$treated$mean); mc <- num(e$control$mean)
  if (!is.na(mt) && !is.na(mc) && !is.na(nt) && !is.na(nc)) {
    sdv <- sd_for(trial, e, nt, nc, NA)
    if (is.na(sdv$sd) || sdv$sd <= 0) return(NULL)
    g <- J(nt + nc) * dir * (mt - mc) / sdv$sd
    return(out(g, (nt + nc) / (nt * nc) + g^2 / (2 * (nt + nc)), "means", sdv$rule))
  }
  NULL
}

work_effect <- function(trial, e) {
  n <- sizes(trial, e); nt <- n[1]; nc <- n[2]
  est <- e$estimate; kind <- est$kind %||% "none"
  pu <- e$proportion_unit %||% "fraction"
  pt <- frac(e$treated$proportion, pu); pc <- frac(e$control$proportion, pu)
  out <- function(d, se, how) list(yi = d, vi = se^2, how = how, sd_rule = NA_character_, nt = nt, nc = nc)
  if (kind %in% c("difference", "adjusted_difference") && !is.na(num(est$value)) && (est$scale %||% "") %in% c("percentage_points", "proportion")) {
    k <- if (identical(est$scale, "proportion")) 100 else 1
    s <- se_from(est)
    if (!is.na(s[1])) return(out(k * num(est$value), k * s[1], paste("difference,", se_how[s[2]])))
    p0 <- if (!is.na(pc)) pc else pt
    if (!is.na(p0) && !is.na(nt) && !is.na(nc)) return(out(k * num(est$value), 100 * sqrt(p0 * (1 - p0) * (1 / nt + 1 / nc)), "difference, sizes"))
  }
  if (identical(kind, "odds_ratio") && !is.na(pc) && !is.na(num(est$value)) && num(est$value) > 0) {
    s <- se_from(est, log_scale = TRUE)
    if (!is.na(s[1])) {
      o <- num(est$value) * pc / (1 - pc); p1 <- o / (1 + o)
      return(out(100 * (p1 - pc), 100 * p1 * (1 - p1) * s[1], "odds ratio"))
    }
  }
  if (!is.na(pt) && !is.na(pc) && !is.na(nt) && !is.na(nc)) {
    return(out(100 * (pt - pc), 100 * sqrt(pt * (1 - pt) / nt + pc * (1 - pc) / nc), "shares"))
  }
  NULL
}

hours_effect <- function(trial, e) {
  k <- switch(e$per %||% "week", week = 1, month = 12 / 52, year = 1 / 52, NA)
  if (is.na(k)) return(NULL)
  n <- sizes(trial, e); nt <- n[1]; nc <- n[2]
  est <- e$estimate
  if ((est$kind %||% "none") %in% c("difference", "adjusted_difference") && !is.na(num(est$value))) {
    s <- se_from(est)
    if (!is.na(s[1])) return(list(yi = k * num(est$value), vi = (k * s[1])^2, how = "difference", sd_rule = NA_character_, nt = nt, nc = nc))
  }
  mt <- num(e$treated$mean); mc <- num(e$control$mean); st <- num(e$treated$sd); sc <- num(e$control$sd)
  if (!any(is.na(c(mt, mc, st, sc, nt, nc)))) return(list(yi = k * (mt - mc), vi = k^2 * (st^2 / nt + sc^2 / nc), how = "means", sd_rule = NA_character_, nt = nt, nc = nc))
  NULL
}

kind_rank <- c(adjusted_difference = 1, difference = 2, standardized_difference = 3, odds_ratio = 4)
rows <- list()
for (tr in extraction) {
  for (i in seq_along(tr$entries)) {
    e <- tr$entries[[i]]
    cl <- e$class %||% NA_character_
    if (is.na(cl) || !(e$comparison %in% comparisons$comparison)) next
    w <- waves[waves$trial == tr$trial & waves$wave == e$wave, ]
    if (nrow(w) == 0) next
    r <- if (cl %in% SCALES) scale_effect(tr$trial, e) else if (cl %in% WORK) work_effect(tr$trial, e) else if (cl == "hours") hours_effect(tr$trial, e) else NULL
    if (is.null(r) || is.na(r$yi) || is.na(r$vi) || r$vi <= 0) next
    rows[[length(rows) + 1]] <- data.frame(
      trial = tr$trial, entry = paste0(tr$trial, "#", i), entry_n = i, comparison = e$comparison, wave = e$wave,
      months = w$months[1], during = w$during[1], class = cl, instrument = e$instrument %||% NA_character_,
      measure = e$measure %||% NA_character_, binary = isTRUE(e$binary), subgroup = e$subgroup %||% NA_character_,
      kind = e$estimate$kind %||% "arms", adjusted = isTRUE(e$adjusted), source = e$source %||% NA_character_,
      source_rank = result_rank(tr$trial, e$source %||% ""), yi = r$yi, vi = r$vi, how = r$how, sd_rule = r$sd_rule,
      nt = r$nt, nc = r$nc, stringsAsFactors = FALSE)
  }
}
est <- do.call(rbind, rows)
est$yi <- est$yi * unname(sign_of[est$comparison])
write.csv(est, file.path(out_dir, "estimates.csv"), row.names = FALSE)

# ---------------------------------------------------------------------------------------------------
# One estimate per comparison, outcome, and wave: the whole sample before a subgroup; a score before a share
# above a cut-off; an adjusted difference, then an unadjusted one, then a standardized one, then an odds
# ratio, then arms' means or shares, and within each an estimate adjusted for covariates first; the trial's
# first-listed results report first; then the entry listed first.

pick <- function(d) {
  if (is.null(d) || nrow(d) == 0) return(NULL)
  kr <- unname(kind_rank[d$kind]); kr[is.na(kr)] <- 5
  d[order(!is.na(d$subgroup), d$binary, kr, !d$adjusted, d$source_rank, d$entry_n), ][1, ]
}

# The mental-health class of each comparison: the first in order of preference with an estimate at a wave
# during payments (or at any wave, for a lump sum); at a wave without it, the next class there.
mh_class_of <- sapply(comparisons$comparison, function(cmp) {
  d <- est[est$comparison == cmp & est$class %in% MH, ]
  dd <- if (any(d$during)) d[d$during, ] else d
  for (cl in MH) if (cl %in% dd$class) return(cl)
  NA_character_
})
mh_at <- function(cmp, wave) {
  d <- est[est$comparison == cmp & est$wave == wave & est$class %in% MH, ]
  if (nrow(d) == 0) return(NULL)
  first <- mh_class_of[[cmp]]
  order_cl <- if (is.na(first)) MH else c(first, setdiff(MH, first))
  for (cl in order_cl) if (cl %in% d$class) return(pick(d[d$class == cl, ]))
  NULL
}
class_at <- function(cmp, wave, cls) pick(est[est$comparison == cmp & est$wave == wave & est$class %in% cls, ])

cmp_waves <- function(cmp) {
  t <- comparisons$trial[comparisons$comparison == cmp]
  w <- waves[waves$trial == t, ]
  w[order(w$months), ]
}
latest_during <- function(cmp, getter) {
  w <- cmp_waves(cmp); w <- w[w$during & !is.na(w$months), ]
  for (i in rev(seq_len(nrow(w)))) { r <- getter(cmp, w$wave[i]); if (!is.null(r)) return(r) }
  NULL
}
# The wave during payments closest to `primary_months` after the first payment (the later of two as close).
closest_during <- function(cmp, getter) {
  w <- cmp_waves(cmp); w <- w[w$during & !is.na(w$months), ]
  w <- w[order(abs(w$months - P$primary_months), -w$months), ]
  for (i in seq_len(nrow(w))) { r <- getter(cmp, w$wave[i]); if (!is.null(r)) return(r) }
  NULL
}
earliest_during <- function(cmp, getter) {
  w <- cmp_waves(cmp); w <- w[w$during & !is.na(w$months), ]
  for (i in seq_len(nrow(w))) { r <- getter(cmp, w$wave[i]); if (!is.null(r)) return(r) }
  NULL
}
first_after <- function(cmp, getter) {
  w <- cmp_waves(cmp); mp <- comparisons$months_paid[comparisons$comparison == cmp]
  w <- w[!w$during & !is.na(w$months) & (is.na(mp) | w$months - mp <= P$after_max_months), ]
  for (i in seq_len(nrow(w))) { r <- getter(cmp, w$wave[i]); if (!is.null(r)) return(r) }
  NULL
}
# A lump sum has no "during": its latest wave at most `lump_max_months` after the payment.
latest_lump <- function(cmp, getter) {
  w <- cmp_waves(cmp); w <- w[!is.na(w$months) & w$months <= P$lump_max_months, ]
  for (i in rev(seq_len(nrow(w)))) { r <- getter(cmp, w$wave[i]); if (!is.null(r)) return(r) }
  NULL
}

collect <- function(cmps, chooser, getter) {
  out <- lapply(cmps, function(c) chooser(c, getter))
  out <- Filter(Negate(is.null), out)
  if (!length(out)) return(NULL)
  d <- do.call(rbind, out)
  split_shared(d)
}

# Comparisons of one trial that enter one analysis share its control arm: each comparison's variance is
# recomputed as if the control's participants were split equally among them (Cochrane Handbook 23.3.4).
split_shared <- function(d) {
  d$shared <- 1
  for (t in unique(d$trial)) {
    i <- which(d$trial == t)
    if (length(i) < 2) next
    ctl <- comparisons$control[match(d$comparison[i], comparisons$comparison)]
    for (g in unique(ctl)) {
      j <- i[ctl == g]
      if (length(j) < 2) next
      m <- length(j)
      f <- (1 / d$nt[j] + m / d$nc[j]) / (1 / d$nt[j] + 1 / d$nc[j])
      f[is.na(f)] <- m
      d$vi[j] <- d$vi[j] * f
      d$shared[j] <- m
    }
  }
  d
}

# ---------------------------------------------------------------------------------------------------
# Models.

analyses <- list()
fmt_p <- function(p) if (is.na(p)) NA_character_ else if (p < 0.001) "p < 0.001" else sprintf("p = %.3f", p)
summ <- function(f, label, d, digits = 3) {
  pi <- tryCatch(predict(f), error = function(e) NULL)
  r <- list(k = f$k, trials = length(unique(d$trial)), estimate = round(as.numeric(f$beta), digits),
            lower = round(f$ci.lb, digits), upper = round(f$ci.ub, digits), p = signif(f$pval, 3), p_text = fmt_p(f$pval),
            tau2 = signif(f$tau2 %||% 0, 3), i2_pct = round(f$I2 %||% 0, 1),
            pi_lower = if (!is.null(pi) && !is.null(pi$pi.lb)) round(pi$pi.lb, digits) else NULL,
            pi_upper = if (!is.null(pi) && !is.null(pi$pi.ub)) round(pi$pi.ub, digits) else NULL,
            participants = sum(d$nt + d$nc, na.rm = TRUE))
  analyses[[length(analyses) + 1]] <<- data.frame(analysis = label, k = f$k, estimate = r$estimate, lower = r$lower, upper = r$upper,
                                                 p = r$p, tau2 = r$tau2, i2 = r$i2_pct, stringsAsFactors = FALSE)
  r
}
re <- function(d, label, digits = 3) {
  if (is.null(d) || nrow(d) == 0) return(NULL)
  if (nrow(d) == 1) {
    se <- sqrt(d$vi); p <- 2 * pnorm(-abs(d$yi / se))
    analyses[[length(analyses) + 1]] <<- data.frame(analysis = label, k = 1, estimate = round(d$yi, digits), lower = round(d$yi - 1.96 * se, digits),
                                                   upper = round(d$yi + 1.96 * se, digits), p = signif(p, 3), tau2 = NA, i2 = NA, stringsAsFactors = FALSE)
    return(list(k = 1, trials = 1, estimate = round(d$yi, digits), lower = round(d$yi - 1.96 * se, digits), upper = round(d$yi + 1.96 * se, digits),
                p = signif(p, 3), p_text = fmt_p(p), participants = sum(d$nt + d$nc, na.rm = TRUE), single = TRUE))
  }
  f <- rma(d$yi, d$vi, method = "REML", test = "knha", control = list(stepadj = 0.5, maxiter = 1000))
  summ(f, label, d, digits)
}
fe <- function(d, label, digits = 3) {
  if (is.null(d) || nrow(d) < 2) return(NULL)
  f <- rma(d$yi, d$vi, method = "FE")
  r <- summ(f, label, d, digits); r$tau2 <- NULL; r$i2_pct <- NULL; r$pi_lower <- NULL; r$pi_upper <- NULL
  r
}
# Bayesian random effects: half-normal prior on tau (scale from parameters.json), flat prior on the mean;
# the posterior of the mean by integration over a grid of tau.
bayes <- function(d, scale, digits = 3) {
  if (is.null(d) || nrow(d) < 2) return(NULL)
  taus <- seq(0, 6 * scale, length.out = 3001)[-1]
  lw <- sapply(taus, function(tau) {
    w <- 1 / (d$vi + tau^2); m <- sum(w * d$yi) / sum(w)
    dnorm(tau, 0, scale, log = TRUE) + 0.5 * sum(log(w)) - 0.5 * sum(w * (d$yi - m)^2) - 0.5 * log(sum(w))
  })
  pw <- exp(lw - max(lw)); pw <- pw / sum(pw)
  grid <- seq(min(d$yi) - 4 * max(sqrt(d$vi)), max(d$yi) + 4 * max(sqrt(d$vi)), length.out = 4001)
  dens <- Reduce(`+`, Map(function(tau, p) {
    w <- 1 / (d$vi + tau^2); m <- sum(w * d$yi) / sum(w); s <- sqrt(1 / sum(w)); p * dnorm(grid, m, s)
  }, taus, pw))
  cdf <- cumsum(dens) / sum(dens)
  q <- function(a) grid[which.min(abs(cdf - a))]
  list(estimate = round(q(0.5), digits), lower = round(q(0.025), digits), upper = round(q(0.975), digits),
       prob_above_zero = round(1 - cdf[which.min(abs(grid))], 3), tau_prior_scale = scale)
}
loo <- function(d, label) {
  if (is.null(d) || nrow(d) < 3) return(NULL)
  fits <- lapply(seq_len(nrow(d)), function(i) re(d[-i, ], paste(label, "without", d$comparison[i])))
  list(min_estimate = min(sapply(fits, `[[`, "estimate")), max_estimate = max(sapply(fits, `[[`, "estimate")),
       max_p = signif(max(sapply(fits, `[[`, "p")), 3), min_p = signif(min(sapply(fits, `[[`, "p")), 3))
}
retention_ok <- function(d) {
  ok <- mapply(function(cmp, wave) {
    t <- comparisons$trial[comparisons$comparison == cmp]
    a <- Filter(function(x) identical(x$comparison, cmp) && identical(x$wave, wave),
                (Filter(function(x) identical(x$trial, t), extraction)[[1]])$attrition %||% list())
    if (!length(a)) return(FALSE)
    a <- a[[1]]
    rt <- num(a$responded_treated) / num(a$randomized_treated); rc <- num(a$responded_control) / num(a$randomized_control)
    if (any(is.na(c(rt, rc)))) return(FALSE)
    (num(a$responded_treated) + num(a$responded_control)) / (num(a$randomized_treated) + num(a$randomized_control)) >= P$min_retention &&
      abs(rt - rc) <= P$max_differential_attrition
  }, d$comparison, d$wave)
  d[ok, ]
}
each_of <- function(d, digits = 3) {
  setNames(lapply(seq_len(nrow(d)), function(i) list(
    trial = d$trial[i], name = comparisons$name[match(d$comparison[i], comparisons$comparison)], wave = d$wave[i], months = d$months[i],
    class = d$class[i], measure = d$measure[i], estimate = round(d$yi[i], digits),
    lower = round(d$yi[i] - 1.96 * sqrt(d$vi[i]), digits), upper = round(d$yi[i] + 1.96 * sqrt(d$vi[i]), digits),
    how = d$how[i], sd_rule = d$sd_rule[i], analysed = d$nt[i] + d$nc[i])), gsub("-", "_", d$comparison))
}
write_result <- function(name, x) write(toJSON(x, auto_unbox = TRUE, digits = NA, pretty = TRUE, null = "null", na = "null"),
                                        file.path(out_dir, paste0(name, ".json")))

recurring <- comparisons$comparison[comparisons$kind == "recurring"]
lump <- comparisons$comparison[comparisons$kind == "lump_sum"]
employed <- function(cmp, wave) class_at(cmp, wave, "employed_any")
employed_ft <- function(cmp, wave) { r <- class_at(cmp, wave, "employed_any"); if (is.null(r)) class_at(cmp, wave, "employed_full_time") else r }

# ---------------------------------------------------------------------------------------------------
# Primary analyses.

d1 <- collect(recurring, closest_during, mh_at)
d2 <- collect(recurring, closest_during, employed)
# P3: the change in the mental-health effect between the earliest and the latest wave during payments, in
# comparisons with both, at least `fade_min_months` apart; the two waves' estimates are treated as independent.
d3 <- do.call(rbind, Filter(Negate(is.null), lapply(recurring, function(cmp) {
  a <- earliest_during(cmp, mh_at); b <- latest_during(cmp, mh_at)
  if (is.null(a) || is.null(b) || b$months - a$months < P$fade_min_months) return(NULL)
  data.frame(trial = a$trial, comparison = cmp, wave = paste(a$wave, "to", b$wave), months = b$months - a$months,
             early_months = a$months, late_months = b$months, early = a$yi, late = b$yi, v_early = a$vi, v_late = b$vi,
             yi = b$yi - a$yi, vi = a$vi + b$vi, nt = b$nt, nc = b$nc, class = b$class, measure = b$measure, how = b$how,
             sd_rule = b$sd_rule, stringsAsFactors = FALSE)
})))
if (!is.null(d3)) {
  d3 <- split_shared(d3)
  f <- d3$vi / (d3$v_early + d3$v_late)
  d3$v_early <- d3$v_early * f; d3$v_late <- d3$v_late * f
}

p1 <- re(d1, "P1 mental health during payments, near 12 months")
p2 <- re(d2, "P2 employment during payments, near 12 months")
p3 <- re(d3, "P3 change from earliest to latest wave during payments")
ps <- c(p1$p %||% NA, p2$p %||% NA, p3$p %||% NA)
holm <- p.adjust(ps, "holm")

with_holm <- function(r, i) { if (is.null(r)) return(NULL); r$p_holm <- signif(holm[i], 3); r$p_holm_text <- fmt_p(holm[i]); r }
p1 <- with_holm(p1, 1); p2 <- with_holm(p2, 2); p3 <- with_holm(p3, 3)

sensitivity <- function(d, label, scale) {
  if (is.null(d) || nrow(d) == 0) return(NULL)
  att <- retention_ok(d)
  sd_ok <- d[is.na(d$sd_rule) | d$sd_rule %in% c("wave", "baseline", "same_trial"), ]
  reg <- d[comparisons$registered[match(d$comparison, comparisons$comparison)], ]
  adj <- d[d$kind %in% c("adjusted_difference", "standardized_difference") | d$how == "means", ]
  big <- d[d$nt + d$nc >= P$large_n, ]
  list(latest_during = re(collect(recurring, latest_during, if (identical(label, "P1")) mh_at else employed), paste(label, "latest wave during payments")),
       common_effect = fe(d, paste(label, "common effect")),
       bayes = bayes(d, scale),
       leave_one_out = loo(d, label),
       retention = re(att, paste(label, "retention at least", P$min_retention)),
       printed_sd = if (identical(label, "P1")) re(sd_ok, paste(label, "standard deviations printed by the trial")) else NULL,
       registered = re(reg, paste(label, "trials registered before payments began")),
       large = re(big, paste(label, "comparisons with at least", P$large_n, "analysed")),
       egger = if (nrow(d) >= 3) {
         f <- rma(d$yi, d$vi, mods = ~ sqrt(d$vi), method = "REML", test = "knha", control = list(stepadj = 0.5, maxiter = 1000))
         list(slope = round(as.numeric(f$beta[2]), 3), lower = round(f$ci.lb[2], 3), upper = round(f$ci.ub[2], 3), p = signif(f$pval[2], 3))
       } else NULL)
}

robust <- function(r, s) {
  if (is.null(r) || is.null(s)) return(FALSE)
  isTRUE(r$p_holm < 0.05) && isTRUE((s$common_effect$p %||% 1) < 0.05) && isTRUE((s$leave_one_out$max_p %||% r$p) < 0.05) &&
    isTRUE((s$retention$p %||% 1) < 0.05)
}

s1 <- sensitivity(d1, "P1", P$bayes_tau_scale_smd)
s2 <- sensitivity(d2, "P2", P$bayes_tau_scale_pp)
if (!is.null(p1)) {
  p1 <- c(p1, s1, list(robust = robust(p1, s1), each = each_of(d1)))
  p1$measures <- as.list(table(d1$class))
}
if (!is.null(p2)) {
  p2 <- c(p2, s2, list(robust = robust(p2, s2), each = each_of(d2, 2),
                       largest_reduction_allowed_pp = round(-p2$lower, 2), largest_increase_allowed_pp = round(p2$upper, 2)))
  p2$with_full_time <- re(collect(recurring, closest_during, employed_ft), "P2 with full-time employment where any isn't printed", 2)
}
if (!is.null(p3)) {
  p3$early <- re(transform(d3, yi = early, vi = v_early), "P3 earliest wave")
  p3$late <- re(transform(d3, yi = late, vi = v_late), "P3 latest wave")
  p3$mean_gap_months <- round(mean(d3$months), 1)
  p3$common_effect <- fe(d3, "P3 common effect")
  p3$leave_one_out <- loo(d3, "P3")
  p3$each <- each_of(d3)
}
write_result("R1", p1 %||% list(k = 0))
write_result("R2", p2 %||% list(k = 0))
write_result("R3", p3 %||% list(k = 0))

# ---------------------------------------------------------------------------------------------------
# Secondary analyses.

# R4: mental health in time during payments: every wave during payments, with a random effect for each
# comparison and for each wave within it, and the slope per year since the first payment.
all_waves <- do.call(rbind, lapply(recurring, function(cmp) {
  w <- cmp_waves(cmp); w <- w[w$during & !is.na(w$months), ]
  do.call(rbind, Filter(Negate(is.null), lapply(w$wave, function(x) mh_at(cmp, x))))
}))
r4 <- list(k = 0)
if (!is.null(all_waves) && nrow(all_waves) >= 3 && length(unique(all_waves$comparison)) >= 2) {
  all_waves <- split_shared(all_waves)
  f <- rma.mv(yi, vi, mods = ~ I(months / 12), random = ~ 1 | comparison / wave, data = all_waves, test = "t")
  r4 <- list(k = nrow(all_waves), comparisons = length(unique(all_waves$comparison)),
             intercept = round(as.numeric(f$beta[1]), 3), slope_per_year = round(as.numeric(f$beta[2]), 3),
             slope_lower = round(f$ci.lb[2], 3), slope_upper = round(f$ci.ub[2], 3), p = signif(f$pval[2], 3), p_text = fmt_p(f$pval[2]))
  bands <- list(first_6 = c(0, 6), m7_12 = c(6, 12), m13_24 = c(12, 24), after_24 = c(24, Inf))
  r4$bands <- lapply(bands, function(b) {
    dd <- all_waves[all_waves$months > b[1] & all_waves$months <= b[2], ]
    dd <- do.call(rbind, lapply(split(dd, dd$comparison), function(x) x[which.max(x$months), ]))
    re(dd, paste("R4 band", b[1], "to", b[2], "months"))
  })
}
write_result("R4", r4)

# R5: after payments ended; stress; life satisfaction; hours; full-time employment.
r5 <- list(
  after_payments = re(collect(recurring, first_after, mh_at), "R5 mental health after payments ended"),
  stress = re(collect(recurring, closest_during, function(c, w) class_at(c, w, "stress")), "R5 stress during payments"),
  life_satisfaction = re(collect(recurring, closest_during, function(c, w) class_at(c, w, "life_satisfaction")), "R5 life satisfaction during payments"),
  hours_per_week = re(collect(recurring, closest_during, function(c, w) class_at(c, w, "hours")), "R5 hours per week during payments", 2),
  employed_full_time = re(collect(recurring, closest_during, function(c, w) class_at(c, w, "employed_full_time")), "R5 full-time employment during payments", 2),
  employment_after_payments = re(collect(recurring, first_after, employed), "R5 employment after payments ended", 2))
write_result("R5", r5)

# R6: lump sums.
r6 <- list(mental_health = re(collect(lump, latest_lump, mh_at), "R6 lump sum mental health"),
           employment = re(collect(lump, latest_lump, employed), "R6 lump sum employment", 2))
write_result("R6", r6)

# R7: dose. Among recurring comparisons paid in US dollars, the slope on log2 of the monthly amount.
dose <- function(d, label, digits = 3) {
  if (is.null(d)) return(NULL)
  d$amount <- comparisons$monthly_amount[match(d$comparison, comparisons$comparison)]
  d <- d[comparisons$currency[match(d$comparison, comparisons$comparison)] %in% "USD" & !is.na(d$amount) & d$amount > 0, ]
  if (nrow(d) < 4 || length(unique(d$amount)) < 2) return(list(k = nrow(d)))
  f <- rma(d$yi, d$vi, mods = ~ log2(d$amount), method = "REML", test = "knha", control = list(stepadj = 0.5, maxiter = 1000))
  list(k = f$k, slope_per_doubling = round(as.numeric(f$beta[2]), digits), lower = round(f$ci.lb[2], digits), upper = round(f$ci.ub[2], digits),
       p = signif(f$pval[2], 3), amounts = sort(unique(d$amount)))
}
write_result("R7", list(mental_health = dose(d1, "R7 MH"), employment = dose(d2, "R7 employment", 2)))

# R8: everything with any cash: recurring and lump sums together.
write_result("R8", list(mental_health = re(rbind(collect(recurring, closest_during, mh_at), collect(lump, latest_lump, mh_at)), "R8 all cash mental health"),
                        employment = re(rbind(collect(recurring, closest_during, employed), collect(lump, latest_lump, employed)), "R8 all cash employment", 2)))

# R9: the file drawer.
due <- Filter(function(r) isTRUE(r$due), registry)
unreported <- Filter(function(r) !isTRUE(r$reported), due)
write_result("R9", list(due = length(due), reported = length(due) - length(unreported), unreported = length(unreported),
                        unreported_ids = vapply(unreported, function(r) r$id, ""),
                        unreported_names = vapply(unreported, function(r) r$name %||% r$id, "")))

write.csv(do.call(rbind, analyses), file.path(out_dir, "analyses.csv"), row.names = FALSE)
cat(sprintf("P1 k=%s  P2 k=%s  P3 k=%s\n", p1$k %||% 0, p2$k %||% 0, p3$k %||% 0))