# Exploratory analyses added after registration (listed in deviations.json), beside the registered analysis in
# analyze.R. None of these is a planned test.
#
#   Rscript code/summarize.R [data dir] [results dir]
#
# 1. P1 with OpenResearch added. Its health paper prints mental-health effects only pooled over its three years of
#    payments (Table 3) and by year only in figures, so the registered rule (a wave near 12 months) leaves it out.
#    Here its pooled Kessler 6 difference is standardized by the control group's standard deviation the same table
#    prints, and added to P1's comparisons.
# 2. P1 by the kind of estimate: a scale's score (or an estimate printed in standard-deviation units) against a share
#    above a cut-off converted by Chinn's formula, each pooled on its own, and the meta-regression coefficient of
#    being a share.
# 3. P1 by place: the United States against Europe (Finland and Ireland).
# 4. P1 among comparisons with at least 500 people analysed, and the rest.
# 5. People analysed in P1 and P2, each control group counted once where two comparisons share it, and P2's
#    surveyed people apart from the Finnish experiment's registers.
# 6. How many trials are eligible, have results, and how many comparisons they have; and each primary analysis's
#    comparisons as a table (results/p1_comparisons.csv, results/p2_comparisons.csv, results/p3_comparisons.csv).
# Writes results/S.json.

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

args <- commandArgs(trailingOnly = TRUE)
data_dir <- if (length(args) >= 1) args[1] else "data"
out_dir <- if (length(args) >= 2) args[2] else "results"
R1 <- fromJSON(file.path(out_dir, "R1.json"), simplifyVector = FALSE)
trials <- fromJSON(file.path(data_dir, "trials.json"), simplifyVector = FALSE)
est <- read.csv(file.path(out_dir, "estimates.csv"), stringsAsFactors = FALSE)
z <- qnorm(0.975)
`%||%` <- function(a, b) if (is.null(a) || length(a) == 0) b else a

each <- R1$each
d <- data.frame(comparison = gsub("_", "-", names(each)), trial = vapply(each, function(x) x$trial, ""),
                yi = vapply(each, function(x) x$estimate, 1),
                vi = vapply(each, function(x) ((x$upper - x$lower) / (2 * z))^2, 1),
                n = vapply(each, function(x) as.numeric(x$analysed %||% NA), 1), stringsAsFactors = FALSE)
# Whether each chosen estimate is a share above a cut-off: from estimates.csv, the row analyze.R picked.
d$binary <- vapply(seq_len(nrow(d)), function(i) {
  x <- each[[i]]
  r <- est[est$comparison == d$comparison[i] & est$wave == x$wave & est$class == x$class & abs(est$yi - x$estimate) < 5e-4, ]
  if (nrow(r)) isTRUE(as.logical(r$binary[1])) else NA
}, TRUE)
country <- setNames(vapply(trials, function(t) t$country %||% "", ""), vapply(trials, function(t) t$trial, ""))
d$europe <- country[d$trial] %in% c("Finland", "Ireland", "Germany")

fit <- function(x, label) {
  if (is.null(x) || nrow(x) < 2) return(NULL)
  f <- rma(x$yi, x$vi, method = "REML", test = "knha", control = list(stepadj = 0.5, maxiter = 1000))
  list(k = f$k, estimate = round(as.numeric(f$beta), 3), lower = round(f$ci.lb, 3), upper = round(f$ci.ub, 3),
       p = signif(f$pval, 3), i2_pct = round(f$I2, 1), label = label)
}
moderator <- function(x, m) {
  f <- rma(x$yi, x$vi, mods = ~ m, method = "REML", test = "knha", control = list(stepadj = 0.5, maxiter = 1000))
  list(coefficient = 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))
}

out <- list()

# 1. OpenResearch's pooled Kessler 6 (Table 3 of NBER working paper 32711): difference -0.052 (SE 0.124), control
#    mean 6.25 (SD 4.51), n 2,985; higher is worse, so the standardized difference is +0.052 / 4.51.
or <- data.frame(comparison = "T01-1", trial = "T01", yi = 0.052 / 4.51, vi = (0.124 / 4.51)^2, n = 2985, binary = FALSE, europe = FALSE)
with_or <- rbind(d, or)
out$openresearch <- list(standardized = round(or$yi, 3), lower = round(or$yi - z * sqrt(or$vi), 3), upper = round(or$yi + z * sqrt(or$vi), 3))
out$p1_with_openresearch <- fit(with_or, "P1 with OpenResearch's pooled estimate")
fe <- rma(with_or$yi, with_or$vi, method = "FE")
out$p1_with_openresearch$common_effect <- list(estimate = round(as.numeric(fe$beta), 3), lower = round(fe$ci.lb, 3), upper = round(fe$ci.ub, 3), p = signif(fe$pval, 3))

# 2. Scale scores against shares above a cut-off.
out$scores <- fit(d[d$binary %in% FALSE, ], "P1, scale scores")
out$shares <- fit(d[d$binary %in% TRUE, ], "P1, shares above a cut-off")
ok <- !is.na(d$binary)
out$share_coefficient <- moderator(d[ok, ], as.numeric(d$binary[ok]))

# 3. United States against Europe.
out$united_states <- fit(d[!d$europe, ], "P1, United States")
out$europe <- fit(d[d$europe, ], "P1, Europe")
out$europe_coefficient <- moderator(d, as.numeric(d$europe))

# 4. Comparisons with at least 500 people analysed, and the rest.
out$at_least_500 <- fit(d[!is.na(d$n) & d$n >= 500, ], "P1, at least 500 analysed")
out$under_500 <- fit(d[!is.na(d$n) & d$n < 500, ], "P1, under 500 analysed")
out$at_least_500$trials <- sort(unique(d$trial[!is.na(d$n) & d$n >= 500]))

# 5. People analysed.
people <- function(R) {
  rows <- do.call(rbind, lapply(names(R$each), function(k) {
    x <- R$each[[k]]; cmp <- gsub("_", "-", k)
    r <- est[est$comparison == cmp & est$wave == x$wave & est$class == x$class & abs(est$yi - x$estimate) < 5e-4, ][1, ]
    data.frame(trial = x$trial, nt = r$nt, nc = r$nc)
  }))
  sum(rows$nt, na.rm = TRUE) + sum(tapply(rows$nc, rows$trial, max, na.rm = TRUE))
}
R2 <- fromJSON(file.path(out_dir, "R2.json"), simplifyVector = FALSE)
out$people <- list(p1 = people(R1), p2 = people(R2))
out$people$p2_registers <- sum(vapply(R2$each, function(x) if (identical(x$trial, "T52")) as.numeric(x$analysed) else 0, 1))
out$people$p2_surveys <- out$people$p2 - out$people$p2_registers

# 6. Trials, and the tables of comparisons.
el <- Filter(function(t) isTRUE(t$eligible) && length(t$comparisons) > 0, trials)
out$trials <- list(grouped = length(trials), eligible = length(el),
                   with_results = sum(vapply(el, function(t) length(t$result_reports) > 0, TRUE)),
                   comparisons = sum(vapply(el, function(t) length(t$comparisons), 1L)),
                   recurring = sum(vapply(el, function(t) sum(vapply(t$comparisons, function(c) identical(c$kind, "recurring"), TRUE)), 1L)),
                   not_eligible = sum(vapply(trials, function(t) isFALSE(t$eligible) && !identical(t$decided_by, "screeners"), TRUE)))
name_of <- setNames(vapply(trials, function(t) t$name, ""), vapply(trials, function(t) t$trial, ""))
table_of <- function(R, file, digits) {
  rows <- do.call(rbind, lapply(names(R$each), function(k) {
    x <- R$each[[k]]
    data.frame(comparison = gsub("_", "-", k), trial = x$trial, name = name_of[[x$trial]], wave = x$wave, months = x$months %||% NA,
               class = x$class %||% NA, measure = x$measure %||% NA, estimate = x$estimate, lower = x$lower, upper = x$upper,
               analysed = x$analysed %||% NA, how = x$how %||% NA, stringsAsFactors = FALSE)
  }))
  write.csv(rows[order(rows$trial, rows$comparison), ], file.path(out_dir, file), row.names = FALSE)
}
table_of(R1, "p1_comparisons.csv")
table_of(R2, "p2_comparisons.csv")
table_of(fromJSON(file.path(out_dir, "R3.json"), simplifyVector = FALSE), "p3_comparisons.csv")

write(toJSON(out, auto_unbox = TRUE, digits = NA, pretty = TRUE, null = "null", na = "null"), file.path(out_dir, "S.json"))
cat("summaries written\n")