# Estimators (plan sections 7 and 9; D13-D15, D26, D28).
#
# One machinery for every observational design: inverse-probability weights from
# pooled logistic models of elective PVI, then a weighted Aalen-Johansen estimator of
# the cumulative incidence of the outcome (death competing), nonparametric in time.
# Standard errors come from the influence function of the weighted Aalen-Johansen
# estimator with the weights treated as known (usually conservative for IPW; checked
# against a physician-cluster bootstrap, D28). Summed by patient for the naive SE, by
# physician for the cluster-robust SE. Clones and person-trials of one patient are summed
# together, so cloning and stacking are accounted for.

BASE_COVS <- c("age", "female", "black", "year", "diabetes", "ckd", "chf", "copd", "prior_cvd",
               "smoking", "frailty_index", "obl")
VARIANTS <- c("V1_ignore", "V2_cluster", "V3_random", "V4_within")

# ---- inputs -------------------------------------------------------------------------

#' Person-month analysis table: eligibility for elective PVI and the PVI decision.
#'
#' A row is month t for a patient who survived month t's events and death and had no
#' MALE and no PVI before the decision in month t: the risk set of the PVI model.
pvi_risk_set <- function(rec) {
  p <- rec$patients
  extra <- intersect("t_pvi_orig", names(p))
  pm <- merge(rec$pm, p[, c("id", "phys_id", BASE_COVS, "t_pvi", "pvi_elective", "t_male", "t_death", extra),
                        with = FALSE], by = "id")
  pm <- pm[(is.na(t_death) | t_death > t) & (is.na(t_male) | t_male > t) & (is.na(t_pvi) | t_pvi >= t)]
  pm[, pvi := as.integer(!is.na(t_pvi) & t_pvi == t)]
  sev0 <- rec$pm[t == 0, .(id, sev0 = sev_obs)]
  pm <- merge(pm, sev0, by = "id")
  pm[, `:=`(win = as.integer(t < rec$par$size$window), tt = pmin(t, 12))]
  if (length(extra)) pm[, pvi_before := as.integer(!is.na(t_pvi_orig) & t_pvi_orig < t)]
  pm[]
}

# ---- treatment models ------------------------------------------------------------------

#' Fit P(PVI at t | eligible, history) and return predictions for `newdata`.
#'
#' The decision window (months 0..W-1) and the deferred period are fitted as separate
#' models: practice inside and after the window differs, and a pooled model would
#' misspecify the window probabilities that the CCW and sequential weights depend on.
#' variant: V1/V2 no physician terms; V3 physician random intercept (glmer, nAGQ = 0);
#' V4 physician fixed effects (fixest). Under V4 a physician with no PVI in a phase has
#' no estimable intercept for it; the maximum-likelihood probability there is 0.
fit_pvi <- function(data, newdata, variant, extra_covs = NULL) {
  out <- numeric(nrow(newdata))
  for (w in c(1L, 0L)) {
    i_new <- newdata$win == w
    if (!any(i_new)) next
    d <- data[win == w]
    nd <- newdata[i_new]
    tt_levels <- if (w == 1L) "factor(tt)" else "splines::ns(t, df = 4)"
    # deferred period (about 2 million person-months, few PVIs): fit on every PVI and a
    # fixed 10% sample of the other months, with offset log(1/0.1) so the coefficients
    # are those of the full model (case-control sampling of person-months)
    if (w == 0L) d <- sample_controls(d, frac = 0.1)
    # no exposure at all in a phase (the NCE after the window): probability 0
    out[i_new] <- if (sum(d$pvi) == 0) 0 else fit_pvi_phase(d, nd, variant, tt_levels, extra_covs)
  }
  out
}

#' Keep all events and a deterministic fraction of non-events; offset column `off`.
sample_controls <- function(d, frac) {
  keep <- d$pvi == 1L | (((d$id * 7919L + d$t * 104729L) %% 1000L) < frac * 1000)
  d <- d[keep]
  d[, off := log(1 / frac)]   # sampling raises every row's odds by 1/frac
  d
}

fit_pvi_phase <- function(d, nd, variant, time_terms, extra_covs = NULL) {
  if (!"off" %in% names(d)) d[, off := 0]
  nd <- copy(nd)[, off := 0]                 # predictions are for the full population
  covs <- c(BASE_COVS, "sev_obs", "sev0", extra_covs)
  if (variant == "V4_within") covs <- setdiff(covs, "obl")
  rhs <- paste(c(covs, time_terms), collapse = " + ")
  if (variant %in% c("V1_ignore", "V2_cluster")) {
    f <- as.formula(paste("~", rhs))
    X <- model.matrix(f, d)
    fit <- fastglm::fastglm(X, d$pvi, family = binomial(), method = 2, offset = d$off)
    b <- coef(fit); b[is.na(b)] <- 0
    Xn <- model.matrix(f, nd)
    return(as.vector(plogis(Xn %*% b)))
  }
  if (variant == "V3_random") {
    fit <- lme4::glmer(as.formula(paste("pvi ~", rhs, "+ offset(off) + (1 | phys_id)")), data = d,
                       family = binomial(), nAGQ = 0, control = lme4::glmerControl(calc.derivs = FALSE))
    return(as.vector(predict(fit, newdata = nd, type = "response", allow.new.levels = TRUE)))
  }
  if (variant == "V4_within") {
    fit <- fixest::feglm(as.formula(paste("pvi ~", rhs, "| phys_id")), data = d, offset = ~off,
                         family = binomial(), notes = FALSE, warn = FALSE)
    pr <- as.vector(predict(fit, newdata = nd, type = "response"))
    pr[is.na(pr)] <- 0
    return(pr)
  }
  stop("unknown variant ", variant)
}

#' Physicians with variation in early (window) PVI: the V4 analysis population.
physicians_with_variation <- function(rec) {
  W <- rec$par$size$window
  x <- rec$patients[, .(e = sum(pvi_elective %in% TRUE & t_pvi < W), n = .N), by = phys_id]
  x[e > 0 & e < n, phys_id]
}

# ---- weighted Aalen-Johansen with influence functions --------------------------------

#' Cumulative incidence at `horizon` from person-month rows.
#'
#' @param d data.table with id, phys_id, arm, t, w (weight for the row's risk set),
#'   ev (1 if the outcome occurs in month t), dth (1 if death without the outcome in t).
#' @return list(F = named vector of cumulative incidence by arm at the horizon,
#'   inf = data.table(id, phys_id, arm, inf) of influence-function contributions summed
#'   by patient, curve = cumulative incidence by arm and month).
aj_cif <- function(d, horizon = 60L) {
  d <- d[t < horizon & w > 0]
  s <- d[, .(R = sum(w), E = sum(w * ev), D = sum(w * dth)), keyby = .(arm, t)]
  s[, `:=`(he = E / R, hd = D / R)]
  s[, h := he + hd]
  s[, S_prev := shift(cumprod(1 - h), fill = 1), by = arm]
  s[, Fcum := cumsum(S_prev * he), by = arm]
  s[, FH := Fcum[.N], by = arm]
  s[, `:=`(a = S_prev - (FH - Fcum) / (1 - h), b = -(FH - Fcum) / (1 - h))]
  s[!is.finite(a), a := 0]; s[!is.finite(b), b := 0]
  d <- s[d, on = .(arm, t), .(id, phys_id, arm, x = w * ((ev - he) * a + (dth - hd) * b) / R)]
  inf <- d[, .(inf = sum(x)), by = .(id, phys_id, arm)]
  FH <- s[, .(FH = FH[1]), by = arm]
  list(F = setNames(FH$FH, FH$arm), inf = inf, curve = s[, .(arm, t, cuminc = Fcum)])
}

#' Risk difference arm1 - arm0 with naive (patient) and cluster-robust (physician) SEs.
rd_from_aj <- function(aj, arm1, arm0) {
  x <- dcast(aj$inf, id + phys_id ~ arm, value.var = "inf", fill = 0)
  diff <- x[[arm1]] - x[[arm0]]
  se_pat <- sqrt(sum(diff^2))
  se_phys <- sqrt(sum(tapply(diff, x$phys_id, sum)^2))
  data.table(est = aj$F[[arm1]] - aj$F[[arm0]], se_patient = se_pat, se_cluster = se_phys,
             risk1 = aj$F[[arm1]], risk0 = aj$F[[arm0]])
}

# outcome and death indicators for a person-month table that has columns t_out, t_death
mark_events <- function(d) {
  d[, ev := as.integer(!is.na(t_out) & t_out == t)]
  d[, dth := as.integer(ev == 0L & !is.na(t_death) & t_death == t)]
  d
}

# person-months from `start` to the outcome, death or the horizon, one row per month
expand_followup <- function(p, start, horizon = 60L) {
  end <- pmin(ifelse(is.na(p$t_out), horizon - 1L, p$t_out),
              ifelse(is.na(p$t_death), horizon - 1L, p$t_death), horizon - 1L)
  keep <- end >= start
  p <- p[keep]; st <- if (length(start) > 1) start[keep] else rep(start, nrow(p)); end <- end[keep]
  nrep <- end - st + 1L
  d <- p[rep(seq_len(nrow(p)), nrep)]
  d[, t := st[rep(seq_along(nrep), nrep)] + sequence(nrep) - 1L]
  mark_events(d)
}

# ---- designs --------------------------------------------------------------------------

#' For the within-physician variant, restrict to physicians whose patients include both
#' early PVI and no early PVI (positivity); the estimand changes accordingly (plan 7).
restrict_variant <- function(rec, variant) {
  if (variant != "V4_within") return(rec)
  keep <- physicians_with_variation(rec)
  rec$patients <- rec$patients[phys_id %in% keep]
  rec$pm <- rec$pm[id %in% rec$patients$id]
  rec
}

#' Run one design under one physician variant for several outcomes.
#' Returns one row per outcome (and contrast), with the SE that the variant reports:
#' patient-level for V1, physician-cluster-robust for V2-V4.
run_design <- function(rec, design, variant, outcomes = c(limb = "t_limb", male = "t_male", disease = "t_disease",
                                                           nco1 = "t_nco1", nco2 = "t_nco2")) {
  run_designs(rec, design, variant, outcomes)
}

#' Several designs under one variant, sharing the PVI fit between CCW and sequential
#' trials (the sequential propensity score is the window part of the same model).
run_designs <- function(rec, designs, variant, outcomes = c(limb = "t_limb", male = "t_male", disease = "t_disease",
                                                             nco1 = "t_nco1", nco2 = "t_nco2")) {
  rec <- restrict_variant(rec, variant)
  need <- intersect(designs, c("ccw", "seq"))
  rs <- if (length(need)) pvi_risk_set(rec) else NULL
  phat <- if ("ccw" %in% need) fit_pvi(rs, rs, variant, rec$extra_covs)
          else if ("seq" %in% need) { x <- rep(NA_real_, nrow(rs)); w <- rs$win == 1L
                                      x[w] <- fit_pvi(rs[w], rs[w], variant, rec$extra_covs); x }
  cmp <- if (is.null(rec$comparators)) c("B1", "B2") else rec$comparators
  use_patient <- variant == "V1_ignore"
  rbindlist(lapply(designs, function(design) {
    r <- switch(design,
      naive = rbindlist(lapply(names(outcomes), function(o)
        cbind(contrast = "early_none", est_naive(rec, outcomes[[o]], variant), outcome = o))),
      landmark = rbindlist(lapply(names(outcomes), function(o)
        cbind(contrast = "early_none", est_landmark(rec, outcomes[[o]], variant, extra_covs = rec$extra_covs), outcome = o))),
      landmark_m5 = rbindlist(lapply(names(outcomes), function(o)
        cbind(contrast = "early_none", est_landmark(rec, outcomes[[o]], variant, extra_covs = rec$extra_covs,
                                                    at_landmark = TRUE), outcome = o))),
      ccw = est_ccw(rec, outcomes, variant, rs = copy(rs), phat = phat, comparators = cmp),
      seq = cbind(contrast = "now_not", est_seq(rec, outcomes, variant, rs = copy(rs), phat = phat)))
    r[, `:=`(design = design, variant = variant, n = nrow(rec$patients))]
    r[, se := if (use_patient) se_patient else se_cluster]
    r
  }), fill = TRUE)
}

#' Naive "ever early PVI" design: exposure = PVI in months 0..W-1, follow-up from month 0,
#' baseline propensity score. Immortal time is present by construction.
est_naive <- function(rec, outcome = "t_male", variant = "V1_ignore", horizon = 60L, extra_covs = NULL) {
  W <- rec$par$size$window
  p <- copy(rec$patients)
  p[, `:=`(t_out = get(outcome), arm = ifelse(!is.na(t_pvi) & t_pvi < W, "early", "none"))]
  p <- merge(p, rec$pm[t == 0, .(id, sev_obs)], by = "id")
  p[, pvi := as.integer(arm == "early")]
  ps <- fit_pvi_ps(p, variant, covs = c(BASE_COVS, "sev_obs", extra_covs))   # sev_obs is the month-0 value here
  p[, w := ifelse(arm == "early", 1 / ps, 1 / (1 - ps))]
  d <- expand_followup(p[!is.na(w)], 0L, horizon)
  rd_from_aj(aj_cif(d, horizon), "early", "none")
}

# a point-treatment propensity score with the same covariates and physician variants
fit_pvi_ps <- function(p, variant, covs = c(BASE_COVS, "sev_obs", "sev0")) {
  if (variant == "V4_within") covs <- setdiff(covs, "obl")
  rhs <- paste(covs, collapse = " + ")
  if (variant %in% c("V1_ignore", "V2_cluster")) {
    X <- model.matrix(as.formula(paste("~", rhs)), p)
    fit <- fastglm::fastglm(X, p$pvi, family = binomial(), method = 2)
    return(as.vector(plogis(X %*% coef(fit))))
  }
  if (variant == "V3_random") {
    fit <- lme4::glmer(as.formula(paste("pvi ~", rhs, "+ (1 | phys_id)")), data = p,
                       family = binomial(), nAGQ = 0, control = lme4::glmerControl(calc.derivs = FALSE))
    return(as.vector(predict(fit, type = "response")))
  }
  fit <- fixest::feglm(as.formula(paste("pvi ~", rhs, "| phys_id")), data = p, family = binomial(),
                       notes = FALSE, warn = FALSE)
  as.vector(predict(fit, newdata = p, type = "response"))
}

#' Landmark at month W (D13): patients alive and MALE-free at the start of month W;
#' exposure = PVI in months 0..W-1; covariates at diagnosis; outcome counted in months
#' W..horizon-1. With at_landmark = TRUE the propensity score also uses severity at month
#' W-1, which a PVI has already changed (supplement, D53).
est_landmark <- function(rec, outcome = "t_male", variant = "V1_ignore", horizon = 60L, extra_covs = NULL,
                         at_landmark = FALSE) {
  W <- rec$par$size$window
  p <- copy(rec$patients)[(is.na(t_death) | t_death >= W) & (is.na(t_male) | t_male >= W)]
  p[, `:=`(t_out = get(outcome), arm = ifelse(!is.na(t_pvi) & t_pvi < W, "early", "none"))]
  p <- p[is.na(t_out) | t_out >= W]
  p <- merge(p, rec$pm[t == W - 1L, .(id, sev_obs)], by = "id")
  p <- merge(p, rec$pm[t == 0, .(id, sev0 = sev_obs)], by = "id")
  p[, pvi := as.integer(arm == "early")]
  ps <- fit_pvi_ps(p, variant, covs = c(BASE_COVS, "sev0", if (at_landmark) "sev_obs", extra_covs))
  p[, w := ifelse(arm == "early", 1 / ps, 1 / (1 - ps))]
  d <- expand_followup(p[!is.na(w)], W, horizon)
  rd_from_aj(aj_cif(d, horizon), "early", "none")
}

#' Clone-censor-weight (D10, D26). Arms: A (PVI by month W-1), B1 (no PVI in months
#' 0..W-1), B2 (no elective PVI ever). A clone is censored after the PVI decision of the
#' month in which it deviates; weights are 1 / P(not censored through the previous month).
est_ccw <- function(rec, outcome = "t_male", variant = "V1_ignore", horizon = 60L,
                    comparators = c("B1", "B2"), rs = NULL, phat = NULL) {
  W <- rec$par$size$window
  if (is.null(rs)) rs <- pvi_risk_set(rec)
  if (is.null(phat)) phat <- fit_pvi(rs, rs, variant, rec$extra_covs)
  rs[, ph := phat]
  rs <- rs[!is.na(ph)]
  # everyone eligible at diagnosis is cloned, including patients with a MALE event or
  # death in month 0 (events in the grace period count in both arms). Clones and weights
  # are built once over follow-up to death or the horizon; each outcome then truncates.
  p <- copy(rec$patients)
  p[, t_out := NA_integer_]
  base <- expand_followup(p, 0L, horizon)
  # P(PVI) for each eligible month; months where a patient is not at risk of PVI have
  # probability 0 of deviating
  base <- rs[, .(id, t, ph)][base, on = .(id, t)]
  base[is.na(ph), ph := 0]
  base[, `:=`(tp = fifelse(is.na(t_pvi), Inf, as.numeric(t_pvi)), el = pvi_elective %in% TRUE)]

  arms <- list()
  # arm A: censored after month W-1 if no PVI by then (and still at risk of PVI then)
  a <- copy(base)[, arm := "A"]
  cens_A <- a[t == W - 1L & ph > 0 & tp > W - 1L, id]          # deviated at month W-1
  a <- a[!(id %in% cens_A & t > W - 1L)]
  pA <- base[t == W - 1L & ph > 0, .(id, pA = ph)]
  a <- pA[a, on = "id"]
  a[, w := ifelse(t > W - 1L & !is.na(pA) & tp == W - 1L, 1 / pA, 1)]
  arms$A <- a[, .(id, phys_id, arm, t, w)]
  for (cmp in comparators) {
    b <- copy(base)[, arm := cmp]
    last <- if (cmp == "B1") W - 1L else horizon
    b <- b[!(el & tp <= last & t > tp)]                         # censored after month tp
    b[, q := ifelse(t <= last, 1 - ph, 1)]
    b[, w := 1 / shift(cumprod(q), fill = 1), by = id]
    arms[[cmp]] <- b[, .(id, phys_id, arm, t, w)]
  }
  clones <- rbindlist(arms)
  ev_times <- rec$patients[, c("id", "t_death", unname(outcome)), with = FALSE]
  rbindlist(lapply(seq_along(outcome), function(k) {
    d <- ev_times[, .(id, t_death, t_out = get(outcome[k]))][clones, on = "id"]
    d <- d[is.na(t_out) | t <= t_out]
    aj <- aj_cif(mark_events(d), horizon)
    rbindlist(lapply(comparators, function(cmp)
      cbind(contrast = paste0("A_", cmp), rd_from_aj(aj, "A", cmp))))[
        , outcome := if (is.null(names(outcome))) outcome[k] else names(outcome)[k]]
  }))
}

#' Sequential trials, ITT (D14): a trial at each month m = 0..W-1 among patients eligible
#' for PVI at m; treatment = PVI at m; propensity score from covariates at m; outcome
#' counted in months m+1..horizon-1; trial RDs pooled with weights = number eligible.
est_seq <- function(rec, outcome = "t_male", variant = "V1_ignore", horizon = 60L, rs = NULL,
                    phat = NULL) {
  W <- rec$par$size$window
  if (is.null(rs)) rs <- pvi_risk_set(rec)
  if (!is.null(phat)) rs[, ps := phat]
  trials <- rs[t < W]
  # the trial-baseline propensity score is the window model of the PVI fit
  if (is.null(phat)) trials[, ps := fit_pvi(trials, trials, variant, rec$extra_covs)]
  trials <- trials[!is.na(ps)]
  trials[, `:=`(m = t, arm = ifelse(pvi == 1L, "pvi", "no"),
                w = ifelse(pvi == 1L, 1 / ps, 1 / (1 - ps)))]
  trials <- trials[, .(id, phys_id, m, arm, w, t_death)]
  rbindlist(lapply(seq_along(outcome), function(k) {
    tr_k <- rec$patients[, .(id, t_out = get(outcome[k]))][trials, on = "id"][is.na(t_out) | t_out > m]
    res <- list(); infs <- list()
    for (mm in 0:(W - 1)) {
      tr <- tr_k[m == mm]
      d <- expand_followup(tr[, .(id, phys_id, arm, w, t_out, t_death)], mm + 1L, horizon)
      aj <- aj_cif(d, horizon)
      x <- dcast(aj$inf, id + phys_id ~ arm, value.var = "inf", fill = 0)
      res[[mm + 1L]] <- data.table(m = mm, n = nrow(tr), est = aj$F[["pvi"]] - aj$F[["no"]])
      infs[[mm + 1L]] <- data.table(id = x$id, phys_id = x$phys_id, m = mm, inf = x$pvi - x$no)
    }
    res <- rbindlist(res); infs <- rbindlist(infs)
    res[, wt := n / sum(n)]
    infs <- res[, .(m, wt)][infs, on = "m"]
    pi <- infs[, .(inf = sum(wt * inf)), by = .(id, phys_id)]
    data.table(est = sum(res$wt * res$est), se_patient = sqrt(sum(pi$inf^2)),
               se_cluster = sqrt(sum(tapply(pi$inf, pi$phys_id, sum)^2)), risk1 = NA_real_, risk0 = NA_real_,
               outcome = if (is.null(names(outcome))) outcome[k] else names(outcome)[k])
  }))
}

#' Negative-control exposure (D32, D36): the early low-value procedure (NCE, months
#' 0..W-1) takes the place of early PVI as the exposure; early-PVI status is adjusted
#' for (time-varying "PVI before t" in CCW, "PVI in the window" at the landmark).
#' Only the landmark and CCW (A vs B1) designs are run.
nce_records <- function(rec, adjust_pvi = TRUE) {
  W <- rec$par$size$window
  p <- copy(rec$patients)
  p[, `:=`(t_pvi_orig = t_pvi, pvi_before = as.integer(!is.na(t_pvi) & t_pvi < W))]
  p[, `:=`(t_pvi = t_nce, pvi_elective = !is.na(t_nce))]
  rec$patients <- p
  rec$extra_covs <- if (adjust_pvi) "pvi_before" else NULL
  rec$comparators <- "B1"
  rec
}