# 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, if (is.null(sc$dir)) scale else sc$dir)
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()
  with_v2 <- function(r) if (r$variant[1] == "V1_ignore")   # same fit, two SEs
    rbind(r, copy(r)[, `:=`(variant = "V2_cluster", se = se_cluster)]) else r
  nce_out <- c(limb = "t_limb", male = "t_male", disease = "t_disease")
  for (v in g$variants) {
    if (j$world %in% unlist(g$skip_variant[[v]])) next
    des <- unlist(g$designs)
    if (v != "V1_ignore") des <- setdiff(des, unlist(g$v1_only_designs))
    out[[length(out) + 1]] <- with_v2(run_designs(rec, des, v))[, analysis := "main"]
    if (isTRUE(g$nce) && (is.null(g$nce_variants) || v %in% unlist(g$nce_variants)))
      out[[length(out) + 1]] <- with_v2(run_designs(nce_records(rec), c("landmark", "ccw"), v, nce_out))[, analysis := "nce"]
  }
  # unadjusted NCE (supplement, D36), V1/V2 only
  out[[length(out) + 1]] <- with_v2(run_designs(nce_records(rec, adjust_pvi = FALSE), c("landmark", "ccw"),
                                                "V1_ignore", nce_out))[, 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") }