# Illustrative DGP (labeled simulation): binary "amenorrhea" response,
# rising over time, two dose groups. MONOTONE MAR dropout: women NOT
# experiencing the event (yprev = 0) are far more likely to drop out.
set.seed(667)
Nn <- 400; vtimes <- c(3, 6, 9, 12)
bi   <- rnorm(Nn, 0, 1.2)
dose <- rep(c(0, 1), each = Nn / 2)

ipw_long <- expand.grid(id = 1:Nn, t = vtimes) |>
  arrange(id, t) |>
  mutate(dose = dose[id], bi = bi[id],
         eta = -1.6 + 0.16 * t + 0.45 * dose + bi,
         y_full = rbinom(n(), 1, plogis(eta))) |>
  group_by(id) |>
  mutate(occ = row_number(), yprev = lag(y_full)) |>
  ungroup()

# Monotone MAR dropout: depends on PREVIOUS observed response (observed -> MAR).
set.seed(2024)
ipw_long <- ipw_long |>
  group_by(id) |>
  mutate(p_drop = ifelse(occ == 1, 0,
                         plogis(-0.4 - 2.0 * replace(yprev, is.na(yprev), 0))),
         observed = cumsum(rbinom(n(), 1, p_drop)) == 0) |>
  ungroup() |>
  mutate(y = ifelse(observed, y_full, NA))

round(tapply(ipw_long$observed, ipw_long$occ, mean), 3)  # % observed by occasion

# Stacked data: from occasion 2, among those observed at the prior occasion.
stacked <- ipw_long |>
  group_by(id) |>
  mutate(prev_obs = lag(observed, default = TRUE)) |>
  ungroup() |>
  filter(occ >= 2, prev_obs) |>
  mutate(R = as.integer(observed))

# Logistic model for remaining in the study (uses past observed response).
pi_model <- glm(R ~ factor(t) + dose + yprev, data = stacked, family = binomial)

# Predicted pi_ij for each subject-occasion, then cumulative-product weight.
wdat <- ipw_long |>
  group_by(id) |>
  mutate(prev_obs = lag(observed, default = TRUE)) |>
  ungroup()
wdat$pi_hat <- 1
idx <- wdat$occ >= 2 & wdat$prev_obs
wdat$pi_hat[idx] <- predict(pi_model, newdata = wdat[idx, ], type = "response")

wdat <- wdat |>
  group_by(id) |>
  arrange(occ) |>
  mutate(w = 1 / cumprod(ifelse(is.na(pi_hat), 1, pi_hat))) |>
  ungroup()

obs_dat <- wdat |> filter(observed)
summary(obs_dat$w)   # inspect for extreme weights (FLW p. 529 caution)

# Weights require working-independence in standard GEE software (FLW p. 530).
gee_unw <- geeglm(y ~ factor(t) + dose, id = id, data = obs_dat,
                  family = binomial, corstr = "independence")
gee_ipw <- geeglm(y ~ factor(t) + dose, id = id, data = obs_dat,
                  weights = w, family = binomial, corstr = "independence")
# "Truth": the full-data fit we would have gotten with no dropout.
gee_full <- geeglm(y_full ~ factor(t) + dose, id = id, data = ipw_long,
                   family = binomial, corstr = "independence")

data.frame(
  term       = names(coef(gee_full)),
  full_data  = round(coef(gee_full), 3),
  unweighted = round(coef(gee_unw), 3),
  IPW_GEE    = round(coef(gee_ipw), 3),
  row.names  = NULL
)

# Download-fallback load
tlc_file <- "../../data/tlc.csv"
if (!file.exists(tlc_file)) {
  dir.create(dirname(tlc_file), showWarnings = FALSE, recursive = TRUE)
  download.file("https://content.sph.harvard.edu/fitzmaur/ala2e/tlc.csv", tlc_file)
}
tlc_raw <- read.csv(tlc_file, check.names = FALSE, fileEncoding = "UTF-8-BOM")
names(tlc_raw) <- c("id", "group", "y0", "y1", "y4", "y6")
tlc_raw$group <- ifelse(trimws(tlc_raw$group) == "A", "S", "P")  # A = succimer

# FLW-style MAR mechanism: succimer-group missingness at later weeks depends on
# the OBSERVED baseline y0 (MAR); strength grows over time. Baseline fully observed.
set.seed(667)
tlc <- tlc_raw
mar_miss <- function(y0, grp, strength) {
  lp <- ifelse(grp == "S", 3.0 - 0.04 * y0, 3.0 - 0.04 * mean(tlc_raw$y0))
  rbinom(length(y0), 1, pmin(plogis(-lp) * strength, 0.85)) == 1
}
tlc$y1[mar_miss(tlc$y0, tlc$group, 1.0)] <- NA
tlc$y4[mar_miss(tlc$y0, tlc$group, 1.4)] <- NA
tlc$y6[mar_miss(tlc$y0, tlc$group, 1.8)] <- NA

colSums(is.na(tlc[, c("y0", "y1", "y4", "y6")]))   # y0 never missing

tlc_mice <- tlc
tlc_mice$group <- factor(tlc_mice$group)
md.pattern(tlc_mice[, c("y0", "y1", "y4", "y6")], rotate.names = TRUE)

imp_tlc <- mice(tlc_mice[, c("group", "y0", "y1", "y4", "y6")],
                m = 20, method = "pmm", printFlag = FALSE,
                seed = 667, maxit = 10)

# Scientific question: succimer vs placebo change in blood lead, baseline to week 6.
fit_tlc <- with(imp_tlc, lm(I(y6 - y0) ~ group))
pooled_tlc <- pool(fit_tlc)
summary(pooled_tlc)[, c("term", "estimate", "std.error", "statistic", "p.value")]

# Show pool() is not a black box: compute W, B, T by hand for the group effect.
fits_list <- lapply(1:20, function(k) lm(I(y6 - y0) ~ group, data = complete(imp_tlc, k)))
ests <- sapply(fits_list, function(f) coef(f)["groupS"])
ses  <- sapply(fits_list, function(f) summary(f)$coefficients["groupS", "Std. Error"])

m     <- 20
Q_bar <- mean(ests)                 # pooled point estimate
W_bar <- mean(ses^2)                # within-imputation variance
B     <- var(ests)                  # between-imputation variance
T_var <- W_bar + (1 + 1/m) * B      # Rubin total variance
r     <- (1 + 1/m) * B / W_bar      # relative increase in variance
fmi   <- (r + 2 / ((m - 1) + 3)) / (r + 1)  # fraction of missing information

cat("Q_bar (pooled):", round(Q_bar, 3), "\n")

# Define dropout pattern
dat_pm <- dat_mnar |>
  group_by(id) |>
  summarise(
    last_obs_time = max(time[observed]),
    trt = first(trt),
    .groups = "drop"
  ) |>
  mutate(pattern = factor(last_obs_time))

# Merge back
dat_pm_full <- dat_mnar |>
  left_join(dat_pm |> dplyr::select(id, pattern), by = "id") |>
  filter(observed)

table(dat_pm$pattern)

# Fit model within each pattern (for patterns with enough data)
patterns_with_data <- dat_pm_full |>
  group_by(pattern) |>
  summarise(n = n_distinct(id)) |>
  filter(n >= 10) |>
  pull(pattern)

fits_by_pattern <- lapply(patterns_with_data, function(p) {
  dat_p <- dat_pm_full |> filter(pattern == p)
  tryCatch({
    lme(y_obs ~ time * trt, random = ~ 1 | id, data = dat_p, method = "REML")
  }, error = function(e) NULL)
})
names(fits_by_pattern) <- patterns_with_data

# Extract treatment effects at final time
sapply(fits_by_pattern, function(fit) {
  if (is.null(fit)) return(NA)
  coef(summary(fit))["trt", "Value"]
})

# Calculate pattern weights (proportion of subjects in each pattern)
pattern_weights <- dat_pm |>
  count(pattern) |>
  mutate(weight = n / sum(n))

# Extract treatment effects and SEs from each pattern
extract_trt_effect <- function(fit) {
  if (is.null(fit)) return(c(estimate = NA, se = NA))
  coefs <- coef(summary(fit))
  c(estimate = coefs["trt", "Value"],
    se = coefs["trt", "Std.Error"])
}

pattern_effects <- bind_rows(
  lapply(names(fits_by_pattern), function(p) {
    eff <- extract_trt_effect(fits_by_pattern[[p]])
    data.frame(pattern = p, estimate = eff["estimate"], se = eff["se"])
  })
)

# Merge with weights
pattern_combined <- pattern_effects |>
  left_join(pattern_weights, by = "pattern") |>
  filter(!is.na(estimate))

# Weighted average treatment effect (marginal estimate)
marginal_estimate <- sum(pattern_combined$estimate * pattern_combined$weight) /
                     sum(pattern_combined$weight)

# Approximate SE using delta method (simplified)
marginal_se <- sqrt(sum((pattern_combined$weight * pattern_combined$se)^2)) /
               sum(pattern_combined$weight)

cat("Pattern-Mixture Model Results:\n")

# Under CCMV, we use completer model parameters for all dropout patterns
# This means we assume dropouts would have followed the same trajectory as completers

# Get completer pattern (last time point = 4 for 0-indexed times 0-4)
completer_fit <- fits_by_pattern[["4"]]

if (!is.null(completer_fit)) {
  # Extract completer model coefficients
  completer_coefs <- fixef(completer_fit)

  cat("CCMV Identifying Restriction:\n")
  cat("==============================\n")
  cat("Under CCMV, we assume dropouts would have had outcomes\n")
  cat("distributed like completers at unobserved time points.\n\n")
  cat("Completer model coefficients:\n")
  print(round(completer_coefs, 3))

  cat("\nUnder CCMV, marginal treatment effect = completer treatment effect\n")
  cat("Treatment effect:", round(completer_coefs["trt"], 3), "\n")
} else {
  cat("Completer pattern model did not converge.\n")
}

# Selection-model intuition: contrast a naive observed-data fit with the truth.
# In practice a full selection model needs joint ML (e.g. SAS PROC NLMIXED).
#
# Illustrative MNAR DGP from Part I: dropout hazard rises with the current y.
# A MAR-style analysis of the observed data ignores that; the truth model uses
# the complete data (fit_truth was built in the shared setup chunk).
fit_mar <- lme(y_obs ~ time * trt, random = ~ 1 | id,
               data = dat_pm_full, method = "REML")

data.frame(
  Model = c("Truth (complete)", "MAR (observed)"),
  Intercept = round(c(fixef(fit_truth)[1], fixef(fit_mar)[1]), 3),
  Time = round(c(fixef(fit_truth)[2], fixef(fit_mar)[2]), 3),
  Trt = round(c(fixef(fit_truth)[3], fixef(fit_mar)[3]), 3),
  Time_x_Trt = round(c(fixef(fit_truth)[4], fixef(fit_mar)[4]), 3)
)

# Step 1: Estimate the selection (dropout) model from observed data.
# Among subjects still in the study at time t-1, model whether they are
# observed at time t (R_it), using the previous observed response.
dat_dropout <- dat_mnar |>
  group_by(id) |>
  mutate(
    y_prev   = lag(y, 1),                       # previous (observed) Y
    prev_obs = lag(observed, default = TRUE),   # in the study at t-1?
    R        = as.integer(observed)             # observed at t?
  ) |>
  ungroup() |>
  filter(time > 0, prev_obs)                    # at-risk rows after baseline

# Fit selection (remaining-in-study) model
dropout_model <- glm(R ~ y_prev + time + trt,
                     data = dat_dropout,
                     family = binomial)

cat("Selection (Dropout) Model:\n")

# The model is for REMAINING in the study (R = 1). The y_prev coefficient
# is the log-OR of staying per 1-unit increase in the previous response.
dropout_coefs <- coef(dropout_model)
or_remain <- exp(dropout_coefs["y_prev"])

cat("Selection Model Interpretation (modeling P(remain)):\n")

# Reshape the illustrative MNAR data to wide for mice
dat_wide <- dat_mnar |>
  dplyr::select(id, time, trt, y_obs) |>
  pivot_wider(names_from = time, values_from = y_obs, names_prefix = "y")

# Impute ONCE under MAR, reuse for every delta (faster and self-consistent)
m_imp <- 10
imp_base <- mice(dat_wide, m = m_imp, method = "pmm", printFlag = FALSE, seed = 123)
imp_long_all <- complete(imp_base, action = "long") |>
  pivot_longer(starts_with("y"), names_to = "tv", values_to = "y") |>
  mutate(time = as.numeric(gsub("y", "", tv)))

# Flag which (id, time) cells were originally missing (so delta hits only imputed)
orig_missing <- dat_wide |>
  pivot_longer(starts_with("y"), names_to = "tv", values_to = "y") |>
  mutate(time = as.numeric(gsub("y", "", tv)), was_missing = is.na(y)) |>
  dplyr::select(id, time, was_missing)

imp_flagged <- imp_long_all |> left_join(orig_missing, by = c("id", "time"))

# Apply a delta offset to imputed cells, refit, and pool with Rubin's rules
analyze_with_delta <- function(delta) {
  imp_flagged |>
    mutate(y_adj = ifelse(was_missing, y + delta, y)) |>
    group_by(.imp) |>
    do(tidy(lm(y_adj ~ time * trt, data = .))) |>
    ungroup() |>
    group_by(term) |>
    summarise(
      # std.error first: it needs the full per-imputation `estimate` vector for the
      # between-imputation variance var(estimate). Collapsing estimate first would
      # make var() see a scalar and return NA (dplyr evaluates summarise() in order).
      std.error = sqrt(mean(std.error^2) + (1 + 1/m_imp) * var(estimate)),
      estimate  = mean(estimate),
      .groups   = "drop"
    ) |>
    mutate(delta = delta)
}

deltas <- c(-2, -1, 0, 1, 2)
sensitivity_results <- bind_rows(lapply(deltas, analyze_with_delta))

sensitivity_results |>
  filter(term == "trt") |>
  dplyr::select(delta, estimate, std.error) |>
  mutate(
    lower = estimate - 1.96 * std.error,
    upper = estimate + 1.96 * std.error
  )

sens_trt <- sensitivity_results |>
  filter(term == "trt") |>
  mutate(
    lower = estimate - 1.96 * std.error,
    upper = estimate + 1.96 * std.error
  )

ggplot(sens_trt, aes(x = delta, y = estimate)) +
  geom_point(size = 3) +
  geom_errorbar(aes(ymin = lower, ymax = upper), width = 0.1) +
  geom_hline(yintercept = 0, linetype = 2, color = "red") +
  geom_hline(yintercept = fixef(fit_truth)["trt"], linetype = 3, color = "blue") +
  labs(
    title = "Sensitivity Analysis: Treatment Effect by Delta",
    subtitle = "Blue dashed = true effect; Red dashed = null",
    x = expression(delta ~ "(adjustment to imputed values)"),
    y = "Estimated Treatment Effect"
  ) +
  theme_minimal()

# Find tipping point (where CI includes 0)
sens_trt |>
  mutate(
    significant = (lower > 0) | (upper < 0),
    direction = ifelse(estimate > 0, "positive", "negative")
  ) |>
  dplyr::select(delta, estimate, lower, upper, significant)