# The simulated world: physicians, patients, and the monthly course of disease and care.
# Analysis plan v1.0, section 5.1.
#
# Common random numbers. simulate_course() draws the same uniforms for every patient
# in every month, in a fixed order, whatever the patient's state or the policy. So two
# runs on the same population with the same seed differ only through the policy. This
# keeps truths under different strategies paired, and makes the sharp null exact at the
# patient level (tested).

STREAMS <- c("eps", "proxy", "clti", "false_code", "amp", "bypass", "rr", "clti_revasc",
             "nco1", "nco2", "nce", "pvi", "death")

#' Physicians: latent skill k and style s (standard normal, correlated), setting, panel weight.
make_physicians <- function(par, seed) {
  set.seed(seed)
  J <- par$size$n_physicians
  rho <- par$physician$rho_skill_style
  z1 <- rnorm(J); z2 <- rnorm(J)
  s <- z1
  k <- rho * z1 + sqrt(1 - rho^2) * z2
  obl <- rbinom(J, 1, plogis(par$physician$obl_intercept + par$physician$obl_style * s))
  w <- rlnorm(J, 0, par$size$panel_sdlog)
  data.table(phys_id = seq_len(J), skill = k, style = s, obl = obl, weight = w / sum(w))
}

#' Patients at diagnosis: measured baseline, latent traits, and assigned physician.
make_patients <- function(par, n, phys, seed) {
  set.seed(seed)
  pp <- par$patient
  phys_id <- sample(phys$phys_id, n, replace = TRUE, prob = phys$weight)
  age <- pp$age_min + rgamma(n, shape = pp$age_shape, scale = pp$age_scale)
  frailty <- rnorm(n)
  a10 <- (age - 76) / 10
  comorb <- lapply(pp$comorbidity, function(cm)
    rbinom(n, 1, plogis(qlogis(cm$prev) + cm$age_coef * a10 + cm$frailty_coef * frailty)))
  smoking <- rbinom(n, 1, plogis(qlogis(pp$smoking$prev) + pp$smoking$age_coef * a10))
  sv <- par$severity
  z0 <- sv$z0_diabetes * comorb$diabetes + sv$z0_smoking * smoking + sv$z0_ckd * comorb$ckd +
    rnorm(n, 0, sv$z0_sd)
  rate <- sv$r0 + sv$r_diabetes * comorb$diabetes + sv$r_smoking * smoking + rnorm(n, 0, sv$r_sd)
  venous <- rnorm(n)
  year <- sample(pp$year_min:pp$year_max, n, replace = TRUE)
  # what claims show for smoking and frailty (the analyst's versions; plan D17)
  ms <- par$measurement
  smoking_coded <- ifelse(smoking == 1, rbinom(n, 1, ms$smoking_sens), rbinom(n, 1, 1 - ms$smoking_spec))
  fi_noise <- rnorm(n)
  pat <- data.table(id = seq_len(n), phys_id = phys_id, age = age,
                    female = rbinom(n, 1, pp$p_female), black = rbinom(n, 1, pp$p_black),
                    year = year)
  for (nm in names(comorb)) set(pat, j = nm, value = comorb[[nm]])
  pat[, `:=`(smoking_true = smoking, smoking_coded = smoking_coded, frailty = frailty,
             fi_noise = fi_noise, z0 = z0, rate = rate, venous = venous)]
  pat <- merge(pat, phys[, .(phys_id, skill, style, obl)], by = "phys_id", sort = FALSE)
  setorder(pat, id)
  pat[]
}

#' A full population for one replicate (or for the truths).
make_population <- function(par, n = par$size$n_patients, seed = 1L, phys = NULL) {
  if (is.null(phys)) phys <- make_physicians(par, seed)
  pat <- make_patients(par, n, phys, seed + 7919L)
  list(phys = phys, pat = pat)
}

# ---- policies -----------------------------------------------------------------------
# A policy overrides the natural elective-PVI decision among patients eligible in month t
# (alive, no PVI yet, no MALE yet). Non-elective revascularization after CLTI is never
# overridden: it happens after a MALE event, which every strategy allows.
#   natural            observed practice
#   A                  natural in months 0..(w-2); everyone still eligible gets PVI at w-1
#   B1                 no PVI in months 0..(w-1); natural afterwards
#   B2                 no elective PVI ever
#   seq_force / seq_forbid (m)  natural except PVI forced / forbidden at month m

#   custom             pol$decide(t, natural, idx): idx are the eligible patients' row
#                      numbers, so the decision can differ by patient (simulated RCTs)

policy <- function(type = "natural", m = NA_integer_, decide = NULL) {
  stopifnot(type %in% c("natural", "A", "B1", "B2", "seq_force", "seq_forbid", "custom"))
  list(type = type, m = m, decide = decide)
}

apply_policy <- function(pol, t, natural, window, idx = NULL) {
  if (pol$type == "custom") return(pol$decide(t, natural, idx))
  switch(pol$type,
    natural = natural,
    A = if (t < window - 1) natural else if (t == window - 1) rep(TRUE, length(natural)) else natural,
    B1 = if (t < window) rep(FALSE, length(natural)) else natural,
    B2 = rep(FALSE, length(natural)),
    seq_force = if (t == pol$m) rep(TRUE, length(natural)) else natural,
    seq_forbid = if (t == pol$m) rep(FALSE, length(natural)) else natural)
}

# ---- the monthly course -------------------------------------------------------------

#' Simulate months 0..(months-1) for every patient under a policy.
#'
#' Within a month the order is: severity update; outcome events (CLTI, coding, amputation,
#' bypass, repeat revascularization, revascularization after CLTI, negative controls);
#' death; then the elective PVI decision among survivors still eligible. A MALE event
#' and death in the same month count the MALE event first.
#'
#' @return list(pat = patient table with event times, z = N x months matrix of Z_t
#'   (record = TRUE only), counts = N x months matrix of proxy counts, pvi_month = vector).
simulate_course <- function(pop, par, pol = policy("natural"), seed = 1L, record = TRUE) {
  pat <- pop$pat
  N <- nrow(pat)
  M <- par$size$months
  W <- par$size$window
  tr <- par$treatment; oc <- par$outcomes; nc <- par$negative_controls; ef <- par$effect
  a10 <- (pat$age - 76) / 10

  # time-fixed parts of each linear predictor
  lp_win <- tr$window$intercept + tr$window$frailty * pat$frailty + tr$window$age10 * a10 +
    tr$window$diabetes * pat$diabetes + tr$window$ckd * pat$ckd + tr$window$smoking * pat$smoking_true +
    tr$window$style * pat$style + tr$window$obl * pat$obl + tr$window$year * (pat$year - par$patient$year_min)
  lp_def <- tr$deferred$intercept + tr$deferred$style * pat$style + tr$deferred$obl * pat$obl
  d <- oc$death
  lp_death0 <- d$intercept + d$chf * pat$chf + d$ckd * pat$ckd + d$copd * pat$copd + d$frailty * pat$frailty
  cl <- oc$clti
  lp_clti0 <- cl$intercept + cl$diabetes * pat$diabetes + cl$ckd * pat$ckd + cl$smoking * pat$smoking_true +
    cl$frailty * pat$frailty
  lp_false <- oc$clti_false_code$intercept + oc$clti_false_code$style * pat$style
  am <- oc$amputation
  lp_amp0 <- am$intercept + am$diabetes * pat$diabetes + am$ckd * pat$ckd
  rr <- oc$repeat_revasc
  lp_rr0 <- rr$intercept + rr$skill * pat$skill + rr$style * pat$style
  lp_nco1 <- nc$nco1$intercept + nc$nco1$frailty * pat$frailty + nc$nco1$smoking * pat$smoking_true +
    nc$nco1$age10 * a10
  nco1_z <- if (is.null(nc$nco1$z)) 0 else nc$nco1$z
  lp_nco2 <- nc$nco2$intercept + nc$nco2$style * pat$style + nc$nco2$venous * pat$venous
  lp_nce <- nc$nce$intercept + nc$nce$style * pat$style + nc$nce$frailty * pat$frailty
  kappa_post <- ef$kappa * exp(-ef$skill_kappa * pat$skill)

  # state
  Z <- pat$z0
  Zn <- pat$z0     # severity on the natural course without any PVI (relief or kappa): drives NCO-1 (D53)
  alive <- rep(TRUE, N); treated <- rep(FALSE, N); clti_true <- rep(FALSE, N)
  clti_code <- rep(FALSE, N); male <- rep(FALSE, N)
  t_pvi <- t_clti_true <- t_clti_code <- t_amp <- t_bypass <- t_rr <- t_male <- t_death <-
    t_nco1 <- t_nco2 <- t_nce <- rep(NA_integer_, N)
  pvi_elective <- rep(NA, N)
  n_proc <- integer(N)
  zmat <- if (record) matrix(NA_real_, N, M) else NULL
  cmat <- if (record) matrix(NA_integer_, N, M) else NULL
  ns <- length(STREAMS)

  set.seed(seed)
  for (t in 0:(M - 1)) {
    U <- matrix(runif(N * ns), N, ns, dimnames = list(NULL, STREAMS))
    a <- alive
    if (t > 0) {
      eps <- par$severity$eps_sd * qnorm(U[a, "eps"])
      Z[a] <- Z[a] + pat$rate[a] + kappa_post[a] * treated[a] + eps
      Zn[a] <- Zn[a] + pat$rate[a] + eps
    }
    if (record) {
      zmat[a, t + 1] <- Z[a]
      cmat[a, t + 1] <- qpois(U[a, "proxy"], exp(par$measurement$proxy_intercept + par$measurement$proxy_z * Z[a]))
    }

    # true CLTI, and its coding (true onset always coded; false codes by physician style)
    new_clti <- a & !clti_true & U[, "clti"] < plogis(lp_clti0 + cl$z * Z)
    clti_true[new_clti] <- TRUE; t_clti_true[new_clti] <- t
    new_code <- a & !clti_code & (new_clti | U[, "false_code"] < plogis(lp_false))
    clti_code[new_code] <- TRUE; t_clti_code[new_code] <- t

    new_amp <- a & is.na(t_amp) & U[, "amp"] < plogis(lp_amp0 + am$clti * clti_true + am$z * Z)
    t_amp[new_amp] <- t
    new_byp <- a & is.na(t_bypass) & U[, "bypass"] < plogis(oc$bypass$intercept + oc$bypass$z * Z)
    t_bypass[new_byp] <- t
    n_proc[new_byp] <- n_proc[new_byp] + 1L

    # repeat revascularization: only after the index PVI (from the month after it)
    new_rr <- a & treated & !is.na(t_pvi) & t > t_pvi & U[, "rr"] < plogis(lp_rr0 + rr$z * Z)
    t_rr[new_rr & is.na(t_rr)] <- t
    n_proc[new_rr] <- n_proc[new_rr] + 1L

    # revascularization after a CLTI diagnosis (not elective; allowed by every strategy)
    cr <- a & clti_code & U[, "clti_revasc"] < plogis(oc$clti_revasc$intercept)
    first <- cr & !treated
    rep_ <- cr & treated & !new_rr
    t_rr[rep_ & is.na(t_rr)] <- t
    n_proc[cr & !new_rr] <- n_proc[cr & !new_rr] + 1L
    if (any(first)) {
      treated[first] <- TRUE; t_pvi[first] <- t; pvi_elective[first] <- FALSE
      Z[first] <- Z[first] - ef$relief
    }

    # negative controls (never depend on PVI)
    x <- a & is.na(t_nco1) & U[, "nco1"] < plogis(lp_nco1 + nco1_z * Zn); t_nco1[x] <- t
    x <- a & is.na(t_nco2) & U[, "nco2"] < plogis(lp_nco2); t_nco2[x] <- t
    if (t < W) { x <- a & is.na(t_nce) & U[, "nce"] < plogis(lp_nce); t_nce[x] <- t }

    # MALE: coded CLTI, amputation, bypass, or repeat revascularization
    new_male <- a & !male & (clti_code | !is.na(t_amp) | !is.na(t_bypass) | !is.na(t_rr))
    male[new_male] <- TRUE; t_male[new_male] <- t

    # death
    dies <- a & U[, "death"] < plogis(lp_death0 + d$age10 * a10 + d$age10 * t / 120 + d$z * Z + d$clti * clti_true)
    alive[dies] <- FALSE; t_death[dies] <- t

    # elective PVI among survivors still eligible
    elig <- alive & !treated & !male
    if (any(elig)) {
      lp <- if (t < W) lp_win[elig] + tr$window$z * Z[elig] else lp_def[elig] + tr$deferred$z * Z[elig]
      natural <- U[elig, "pvi"] < plogis(lp)
      go <- apply_policy(pol, t, natural, W, which(elig))
      idx <- which(elig)[go]
      if (length(idx)) {
        treated[idx] <- TRUE; t_pvi[idx] <- t; pvi_elective[idx] <- TRUE
        n_proc[idx] <- n_proc[idx] + 1L
        Z[idx] <- Z[idx] - ef$relief
      }
    }
  }

  # disease composite (D34): coded CLTI or major amputation
  t_disease <- pmin(t_clti_code, t_amp, na.rm = TRUE)
  out <- data.table(id = pat$id, t_pvi, pvi_elective, t_clti_true, t_clti_code, t_amp, t_bypass,
                    t_rr, t_male, t_disease, t_death, t_nco1, t_nco2, t_nce, n_proc)
  list(events = out, z = zmat, counts = cmat, policy = pol)
}

#' Cumulative incidence by month `h` (event in months 0..h-1, before death or in the same month).
cum_inc <- function(t_event, h = 60L) mean(!is.na(t_event) & t_event < h)