# Calibrate each world's natural course to the Medicare targets (plan section 5.4).
#
# For every world, under the null scenario, solve one intercept per target, cycling
# until all targets hold. Every world gets the same targets, so worlds differ only in
# the threat they switch on, not in base rates.
#
# Usage: Rscript simulations/01_calibrate.R [n] [world ...]
# Writes simulations/config/calibrated.yml (merged over existing entries).

pkgload::load_all("pvisim", quiet = TRUE)
suppressPackageStartupMessages(library(data.table))

args <- commandArgs(TRUE)
n <- if (length(args) && grepl("^[0-9e.]+$", args[1])) as.numeric(args[1]) else 1e5
worlds <- setdiff(args, as.character(args[1]))
if (!length(worlds)) worlds <- names(yaml::read_yaml(default_params_path())$worlds)
H <- 53L   # Sorber 2023 median follow-up, about 4.4 years

# Natural-course statistics that the targets are defined on.
stats_of <- function(par) {
  pop <- make_population(par, n = n, seed = 20261027L)
  ev <- simulate_course(pop, par, policy("natural"), seed = 20261028L, record = FALSE)$events
  early <- ev$pvi_elective %in% TRUE & ev$t_pvi < par$size$window
  late_early <- !is.na(ev$t_rr) & ev$t_rr < H
  late_not <- !is.na(ev$t_pvi) & ev$t_pvi >= par$size$window & ev$t_pvi < H
  list(early = mean(early),
       late_not = mean(late_not[!early]),
       late_early = mean(late_early[early]),
       clti_not = mean((!is.na(ev$t_clti_code) & ev$t_clti_code < H)[!early]),
       death = cum_inc(ev$t_death))
}

# target name -> parameter path and value (plan 5.4)
targets <- list(
  early      = list(path = c("treatment", "window", "intercept"), value = 0.032),
  late_not   = list(path = c("treatment", "deferred", "intercept"), value = 0.036),
  late_early = list(path = c("outcomes", "repeat_revasc", "intercept"), value = 0.225),
  clti_not   = list(path = c("outcomes", "clti", "intercept"), value = 0.078),
  death      = list(path = c("outcomes", "death", "intercept"), value = 0.35)
)

set_path <- function(lst, path, value) {
  if (length(path) == 1) { lst[[path]] <- value; return(lst) }
  lst[[path[1]]] <- set_path(if (is.null(lst[[path[1]]])) list() else lst[[path[1]]], path[-1], value)
  lst
}
get_path <- function(lst, path) Reduce(function(x, p) x[[p]], path, lst)

# Each target depends mainly on its own intercept, and for rare monthly events
# logit(rate) moves about one-for-one with the intercept. So update all intercepts
# together: b <- b - damp * (logit(stat) - logit(target)), until every target is within
# `tol` on the logit scale. Common random numbers make each evaluation deterministic.
calibrate_world <- function(world, max_iter = 25, tol = 0.02, damp = 0.9) {
  base <- load_params(world, "null", calibrated = FALSE)
  for (tg in names(base$calibration_targets)) targets[[tg]]$value <- base$calibration_targets[[tg]]
  ov <- list()
  for (tg in names(targets)) ov <- set_path(ov, targets[[tg]]$path, get_path(base, targets[[tg]]$path))
  for (it in seq_len(max_iter)) {
    s <- stats_of(load_params(world, "null", overrides = ov, calibrated = FALSE))
    gap <- sapply(names(targets), function(tg) qlogis(s[[tg]]) - qlogis(targets[[tg]]$value))
    cat(sprintf("%s iter %2d: %s\n", world, it,
                paste(sprintf("%s %.3f", names(s), unlist(s)), collapse = "  ")))
    if (all(abs(gap) < tol)) break
    for (tg in names(targets)) {
      b <- get_path(ov, targets[[tg]]$path)
      ov <- set_path(ov, targets[[tg]]$path, round(b - damp * gap[[tg]], 4))
    }
  }
  if (any(abs(gap) >= tol)) warning(world, ": not converged; largest logit gap ", round(max(abs(gap)), 3))
  ov
}

out_path <- "simulations/config/calibrated.yml"
cal <- if (file.exists(out_path)) yaml::read_yaml(out_path) else list()
for (w in worlds) {
  t0 <- Sys.time()
  cal[[w]] <- calibrate_world(w)
  cat(sprintf("%s done in %.1f min\n", w, as.numeric(Sys.time() - t0, units = "mins")))
  yaml::write_yaml(cal, out_path)
}