# Power of the two primary tests, from the design coding alone: each trial's group, setting, selection,
# and size, with no outcome data. Simulates effects under stated truths and counts how often each test
# rejects at a two-sided 5% level.
#
#   Rscript code/power.R [data dir] [output file]

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

args <- commandArgs(trailingOnly = TRUE)
data_dir <- if (length(args) >= 1) args[1] else "data"
out_file <- if (length(args) >= 2) args[2] else "power.json"
P <- fromJSON(file.path(data_dir, "parameters.json"), simplifyVector = FALSE)
trials <- fromJSON(file.path(data_dir, "trials.json"), simplifyVector = FALSE)
`%||%` <- function(a, b) if (is.null(a) || length(a) == 0) b else a

# Trials that planned a depression measure on days 1 to 3, the primary window.
early_planned <- function(t) {
  lo <- P$windows$early[[1]]; hi <- P$windows$early[[2]]
  any(vapply(t$design$depression_measures %||% list(), function(m) any(vapply(m$days %||% list(), function(x) is.numeric(x) && x >= lo && x <= hi, TRUE)), TRUE))
}
d <- do.call(rbind, lapply(Filter(function(t) isTRUE(t$design$eligible) && (t$design$conscious_dosing %||% "") %in% c("none", "some") && early_planned(t), trials), function(t) {
  data.frame(trial = t$trial, masked = t$design$conscious_dosing == "none", surgery = identical(t$design$setting, "surgery"),
             selected = identical(t$design$depression_selection, "selected"), n = as.numeric(t$size %||% NA), stringsAsFactors = FALSE)
}))
d$n[is.na(d$n)] <- median(d$n, na.rm = TRUE)

set.seed(P$seed)
sims <- P$power$simulations
simulate <- function(g_masked, g_awake, tau) {
  hit1 <- 0; hit2 <- 0
  m <- d[d$masked, ]
  s <- d[d$surgery, ]
  for (i in seq_len(sims)) {
    g <- function(rows) {
      truth <- ifelse(rows$masked, g_masked, g_awake) + rnorm(nrow(rows), 0, tau)
      vi <- 4 / rows$n + truth^2 / (2 * rows$n)
      list(yi = truth + rnorm(nrow(rows), 0, sqrt(vi)), vi = vi)
    }
    a <- g(m)
    f1 <- tryCatch(rma(a$yi, a$vi, method = "REML", test = "knha"), error = function(e) NULL)
    if (!is.null(f1) && f1$pval < 0.05) hit1 <- hit1 + 1
    b <- g(s)
    f2 <- tryCatch(rma(b$yi, b$vi, mods = ~ as.numeric(!s$masked) + as.numeric(s$selected), method = "REML", test = "knha"), error = function(e) NULL)
    if (!is.null(f2) && f2$pval[2] < 0.05) hit2 <- hit2 + 1
  }
  c(P1 = hit1 / sims, P2 = hit2 / sims)
}

out <- list(
  trials_masked = sum(d$masked), trials_masked_participants = sum(d$n[d$masked]),
  trials_surgery_masked = sum(d$surgery & d$masked), trials_surgery_awake = sum(d$surgery & !d$masked),
  participants_surgery = sum(d$n[d$surgery]), simulations = sims, scenarios = list()
)
for (sc in P$power$scenarios) {
  for (tau in unlist(P$power$tau)) {
    p <- simulate(sc$masked, sc$awake, tau)
    out$scenarios[[length(out$scenarios) + 1]] <- list(masked_g = sc$masked, awake_g = sc$awake, tau = tau, power_P1 = unname(p["P1"]), power_P2 = unname(p["P2"]))
  }
}
write(toJSON(out, auto_unbox = TRUE, pretty = TRUE, digits = NA), out_file)
cat("power written to", out_file, "\n")