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

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

simulate_surv_data <- function(
  theta,
  sigma2,
  beta = NULL,
  n_per_group = 60,
  censor_min = 4,
  censor_max = 12,
  seed = NULL
) {
  if (!is.null(seed)) {
    set.seed(seed)
  }

  groups <- rep(seq_along(theta), each = n_per_group)
  age_std <- rnorm(length(groups), mean = 0, sd = 1)
  mean_log_time <- rep(theta, each = n_per_group)

  if (!is.null(beta)) {
    mean_log_time <- mean_log_time + beta * age_std
  }

  log_time <- rnorm(length(groups), mean = mean_log_time, sd = sqrt(sigma2))
  true_time <- exp(log_time)
  censor_time <- runif(length(groups), min = censor_min, max = censor_max)

  data.frame(
    time = pmin(true_time, censor_time),
    event = as.integer(true_time <= censor_time),
    group = factor(groups),
    age_std = age_std
  )
}

sim_data <- simulate_surv_data(
  theta = c(1.2, 1.7, 2.1),
  sigma2 = 0.20,
  beta = -0.35,
  n_per_group = 60,
  seed = 6841
)

mean(sim_data$event)

## -----------------------------------------------------------------------------
fit <- exnex_surv(
  Surv(time, event) ~ group + age_std,
  data = sim_data,
  iter = 3000,
  warmup = 1000,
  chains = 2,
  parallel_chains = 2,
  seed = 6841
)

print(fit, show_trace = FALSE)

## -----------------------------------------------------------------------------
plot(fit, ask = FALSE)

## -----------------------------------------------------------------------------
running_mean_df <- do.call(
  rbind,
  lapply(colnames(fit$draws), function(param) {
    values <- fit$draws[[param]]
    data.frame(
      iteration = seq_along(values),
      running_mean = cumsum(values) / seq_along(values),
      parameter = param
    )
  })
)

ggplot(running_mean_df, aes(iteration, running_mean)) +
  geom_line() +
  facet_wrap(~ parameter, scales = "free_y") +
  labs(
    title = "Running posterior means",
    x = "Post-warmup iteration",
    y = "Running mean"
  )

## -----------------------------------------------------------------------------
ess_simple <- function(x, max_lag = 100) {
  acf_vals <- stats::acf(
    x,
    lag.max = min(max_lag, length(x) - 1),
    plot = FALSE
  )$acf[-1]

  positive_acf <- acf_vals[acf_vals > 0]

  if (length(positive_acf) == 0) {
    return(length(x))
  }

  tau <- 1 + 2 * sum(positive_acf)
  min(length(x), length(x) / tau)
}

diagnostics <- data.frame(
  parameter = colnames(fit$draws),
  mean = colMeans(fit$draws),
  sd = apply(fit$draws, 2, sd),
  lag1_acf = apply(fit$draws, 2, function(x) stats::acf(x, lag.max = 1, plot = FALSE)$acf[2]),
  ess = apply(fit$draws, 2, ess_simple),
  mcse = mapply(function(s, e) s / sqrt(e), apply(fit$draws, 2, sd), apply(fit$draws, 2, ess_simple)),
  row.names = NULL
)

diagnostics

## -----------------------------------------------------------------------------
fit_seeds <- lapply(c(6841, 7313, 8129), function(seed) {
  exnex_surv(
    Surv(time, event) ~ group + age_std,
    data = sim_data,
    iter = 3000,
    warmup = 1000,
    chains = 2,
    parallel_chains = 2,
    seed = seed
  )
})

posterior_means_by_run <- do.call(
  rbind,
  lapply(fit_seeds, function(x) {
    s <- summary(x)
    stats::setNames(s$mean, s$parameter)
  })
)

posterior_means_by_run
apply(posterior_means_by_run, 2, sd)