# Run an experiment grid; one CSV per world x scenario x replicate, so runs are resumable.
# Usage: Rscript simulations/run_experiment.R main [--scale draft] [--cores 10] [--worlds W0,W2]
pkgload::load_all("pvisim", quiet = TRUE)
suppressPackageStartupMessages(library(data.table))
a <- commandArgs(TRUE)
opt <- function(flag, default) { i <- match(flag, a); if (is.na(i)) default else a[i + 1] }
grid_name <- a[1]
scale <- opt("--scale", "draft")
cores <- as.integer(opt("--cores", "10"))
cfg <- yaml::read_yaml("simulations/config/experiments.yml")
g <- cfg[[grid_name]]; sc <- cfg$scales[[scale]]
worlds <- if (is.null(opt("--worlds", NULL))) g$worlds else strsplit(opt("--worlds", NULL), ",")[[1]]
outdir <- file.path("results", grid_name, scale)
jobs <- CJ(world = worlds, scenario = unlist(g$scenarios), rep = seq_len(sc$reps), sorted = FALSE)
jobs[, file := file.path(outdir, world, scenario, sprintf("rep%03d.csv", rep))]
todo <- jobs[!file.exists(file)]
cat(sprintf("%s/%s: %d jobs, %d to do, %d cores\n", grid_name, scale, nrow(jobs), nrow(todo), cores))

one_job <- function(i) {
  j <- todo[i]
  par <- load_params(j$world, j$scenario)
  rec <- simulate_records(par, rep = j$rep, n = sc$n_patients)
  out <- list()
  for (v in g$variants) {
    for (d in g$designs) {
      r <- run_design(rec, d, v)
      if (v == "V1_ignore") {   # same fit, two SEs
        r2 <- copy(r)[, `:=`(variant = "V2_cluster", se = se_cluster)]
        r <- rbind(r, r2)
      }
      out[[length(out) + 1]] <- r[, analysis := "main"]
    }
    if (isTRUE(g$nce)) for (d in c("landmark", "ccw")) {
      r <- run_design(nce_records(rec), d, v, outcomes = c(male = "t_male", disease = "t_disease"))
      if (v == "V1_ignore") r <- rbind(r, copy(r)[, `:=`(variant = "V2_cluster", se = se_cluster)])
      out[[length(out) + 1]] <- r[, analysis := "nce"]
    }
  }
  for (d in c("landmark", "ccw")) {     # unadjusted NCE (supplement, D36), V1/V2 only
    r <- run_design(nce_records(rec, adjust_pvi = FALSE), d, "V1_ignore", outcomes = c(male = "t_male", disease = "t_disease"))
    out[[length(out) + 1]] <- r[, analysis := "nce_unadjusted"]
  }
  out[[length(out) + 1]] <- run_oracle(rec)[, analysis := "main"]   # own estimand of landmark / naive
  for (dz in g$rct) {
    r <- run_rct(par, dz, rep = j$rep)
    out[[length(out) + 1]] <- r[, `:=`(contrast = analysis, analysis = "rct", variant = "none", se = se_patient)]
  }
  res <- rbindlist(out, fill = TRUE)[, `:=`(world = j$world, scenario = j$scenario, rep = j$rep)]
  dir.create(dirname(j$file), recursive = TRUE, showWarnings = FALSE)
  fwrite(res, j$file)
  NULL
}
t0 <- Sys.time()
res <- parallel::mclapply(seq_len(nrow(todo)), function(i) tryCatch(one_job(i), error = function(e)
  sprintf("%s/%s/%d: %s", todo$world[i], todo$scenario[i], todo$rep[i], conditionMessage(e))),
  mc.cores = cores, mc.preschedule = FALSE)
err <- unlist(res[!vapply(res, is.null, TRUE)])
cat(sprintf("done in %.1f h; %d errors\n", as.numeric(Sys.time() - t0, units = "hours"), length(err)))
dir.create(outdir, recursive = TRUE, showWarnings = FALSE)
if (length(err)) { writeLines(err, file.path(outdir, sprintf("errors_%s.txt", format(Sys.time(), "%Y%m%d_%H%M")))); cat(head(err, 3), sep = "\n") }