## ----include = FALSE----------------------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  fig.width = 7,
  fig.height = 4.5,
  warning = FALSE,
  message = FALSE
)

## ----eval = FALSE-------------------------------------------------------------
# fit <- exnex_surv(
#   Surv(time, event) ~ group,
#   data = d,
#   priors = list(
#     p_mix  = 0.7,
#     v_nex  = 10,
#     a_tau  = 3,
#     b_tau  = 3,
#     v_beta = 10
#   )
# )

## -----------------------------------------------------------------------------
library(exnexSurv)
library(survival)
library(ggplot2)

simulate_basket_data <- function(
  theta,
  sigma = 1.2,
  beta = c(0.8, -0.5),
  group_sizes = c(30, 24, 20, 16, 14, 11, 8, 7, 5),
  target_cens = 0.30,
  seed = 1
) {
  set.seed(seed)
  K <- length(theta)
  group <- rep(seq_len(K), times = group_sizes)
  n <- length(group)
  X1 <- rnorm(n)
  X2 <- rnorm(n)
  eta <- theta[group] + X1 * beta[1] + X2 * beta[2]
  log_time <- eta + rnorm(n, 0, sigma)
  true_time <- exp(log_time)
  censor_fun <- function(c) mean(true_time > c) - target_cens
  censor_time <- uniroot(censor_fun, c(min(true_time), max(true_time)))$root
  data.frame(
    time = pmin(true_time, censor_time),
    event = as.integer(true_time <= censor_time),
    group = factor(group),
    X1 = X1,
    X2 = X2
  )
}

theta_true <- c(1.5, 1.4, 1.6, 1.45, 1.55, 1.5, -0.4, -0.6, -0.5)
d <- simulate_basket_data(theta_true, seed = 1)

# small, quick fit just to inspect resolved_priors
quick <- exnex_surv(
  Surv(time, event) ~ group + X1 + X2,
  data = d,
  priors = list(p_mix = 0.7, v_nex = 10),
  iter = 400, warmup = 200, chains = 1
)
print(quick$resolved_priors)

## ----eval = FALSE-------------------------------------------------------------
# # 1. Impute censored log-event times z_i ~ TN(log Y_i, Inf)(eta_i, sigma^2)
# # 2. Update residual variance sigma^2
# # 3. Update regression coefficients beta
# # 4. Update basket effects theta_j ~ N(m_j, V_j)
# # 5. Update mixture indicators Z_j ~ Bern(p_j)
# # 6. Update exchangeable mean mu and between-basket variance tau^2

## ----eval = FALSE-------------------------------------------------------------
# fit$draws            # posterior draws, data.frame
# fit$data             # time, event, group, X, n, n_groups, n_covariates, cov_names, chain_seeds
# fit$priors           # the priors list you supplied
# fit$resolved_priors  # defaults merged with your overrides
# fit$iter, fit$warmup, fit$chains   # MCMC settings
# fit$blueprint        # hardhat blueprint for the formula / data

## -----------------------------------------------------------------------------
theta_true
mean(d$event)  # realized censoring proportion

## -----------------------------------------------------------------------------
fit <- exnex_surv(
  Surv(time, event) ~ group + X1 + X2,
  data = d,
  iter = 1500, warmup = 750, chains = 2, parallel_chains = 2, seed = 42
)
print(fit, show_trace = FALSE)

## -----------------------------------------------------------------------------
summ <- summary(fit)
th <- summ[grepl("^theta_", summ$parameter), ]
th$true <- theta_true
ggplot(th, aes(x = mean, y = reorder(parameter, true))) +
  geom_vline(xintercept = 0, linetype = 2, colour = "grey50") +
  geom_errorbarh(aes(xmin = q05, xmax = q95), height = 0.25, colour = "grey40") +
  geom_point(aes(x = true), shape = 3, size = 2.2, colour = "#C0392B") +
  geom_point(size = 2.6, colour = "#2C3E50") +
  labs(
    title = "EXNEX posterior basket effects",
    subtitle = "Points = posterior mean; horizontal bars = 90% interval; red crosses = true values",
    x = "theta_j (log-survival location)", y = NULL
  )

## -----------------------------------------------------------------------------
fit_nex <- exnex_surv(
  Surv(time, event) ~ group + X1 + X2,
  data = d,
  priors = list(
    p_mix = c(rep(0.9, 6), rep(1e-6, 3)),   # baskets 7-9 almost surely non-exchangeable
    v_nex = c(rep(1e4, 6), rep(1e-4, 3))    # and tightly prior-ed around m_nex = 0
  ),
  iter = 1000, warmup = 500, chains = 1, seed = 7
)
fit_nex$resolved_priors$p_mix

## -----------------------------------------------------------------------------
fit_ex <- exnex_surv(
  Surv(time, event) ~ group + X1 + X2, data = d,
  priors = list(p_mix = 0.9999999),       # EX: full borrowing
  iter = 1000, warmup = 500, chains = 1, seed = 11
)
fit_np <- exnex_surv(
  Surv(time, event) ~ group + X1 + X2, data = d,
  priors = list(p_mix = 1e-6, v_nex = 1e4),  # No pooling
  iter = 1000, warmup = 500, chains = 1, seed = 13
)

means_of <- function(f) {
  s <- summary(f)
  setNames(s$mean, s$parameter)
}
compare <- data.frame(
  theta = theta_true,
  truth = theta_true,
  exnex = means_of(fit)[grepl("^theta_", names(means_of(fit)))],
  ex    = means_of(fit_ex)[grepl("^theta_", names(means_of(fit_ex)))],
  npool = means_of(fit_np)[grepl("^theta_", names(means_of(fit_np)))]
)
compare