## Locate a bundled Stan file
prome_stan_file <- function(model = c(
  "delta0_fixed",
  "delta0_fixed_alphaX_zero",
  "delta0_fixed_alphaT_zero",
  "delta0_fixed_alphaTX_zero"
)) {
  model <- match.arg(model)
  file <- switch(model,
    delta0_fixed = system.file("stan", "delta0_fixed.stan", package = "prome"),
    delta0_fixed_alphaX_zero = system.file("stan", "delta0_fixed_alphaX_zero.stan",
                                           package = "prome"),
    delta0_fixed_alphaT_zero = system.file("stan", "delta0_fixed_alphaT_zero.stan",
                                           package = "prome"),
    delta0_fixed_alphaTX_zero = system.file("stan", "delta0_fixed_alphaTX_zero.stan",
                                            package = "prome")
  )
  if (!nzchar(file)) stop("Stan file not found.", call. = FALSE)
  file
}

.prome_stan_fit <- function(W, T, G, stan_model_key, pars)
 {
  .validate_inputs(W, T, G)
  .check_rstan()
  data_list <- list(N = length(W), T = as.integer(T), 
                    G = as.integer(G), W = as.numeric(W))
  stan_file <- prome_stan_file(stan_model_key)
  mod <- .get_rstan_model(stan_file)
  
  chains = 4
  cores = chains
  iter_warmup = 1000
  iter_sampling = 1000
  control = list(adapt_delta = 0.95, max_treedepth = 12)
  refresh = 0
  
  fit <- rstan::sampling(
    object = mod,
    data = data_list,
    chains = chains,
    cores = cores,
    iter = iter_warmup + iter_sampling,
    warmup = iter_warmup,
    refresh = refresh,
    control = control
  )
  summ <- as.data.frame(rstan::summary(fit, pars = pars)$summary)
  structure(list(
    fit = fit,
    summary = summ,
    stan_file = stan_file,
    model_key = stan_model_key
  ), class = "stan_test")
}

## Bayesian Stan fit with `delta0` fixed at 0
fit_delta0_fixed <- function(W, T, G) {
  pars <- c("mu0", "tau", "sigmaX", "sigmaW0", "sigmaW1", "sigmaX2", 
            "sigmaW0_2", "sigmaW1_2", "alpha0", "alphaT", "alphaX", "delta1")
  .prome_stan_fit(W, T, G, "delta0_fixed", pars)
}

## Bayesian Stan fit with `alphaX = 0` and `delta0` fixed at 0
fit_delta0_fixed_alphaX_zero <- function(W, T, G) {
  pars <- c("mu0", "tau", "sigmaX", "sigmaW0", "sigmaW1", "sigmaX2", 
            "sigmaW0_2", "sigmaW1_2", "alpha0", "alphaT", "delta1")
  .prome_stan_fit(W, T, G, "delta0_fixed_alphaX_zero", pars)
}

## Bayesian Stan fit with `alphaT = 0` and `delta0` fixed at 0
fit_delta0_fixed_alphaT_zero <- function(W, T, G) {
  pars <- c("mu0", "tau", "sigmaX", "sigmaW0", "sigmaW1", "sigmaX2", 
            "sigmaW0_2", "sigmaW1_2", "alpha0", "alphaX", "delta1")
  .prome_stan_fit(W, T, G, "delta0_fixed_alphaT_zero", pars)
}

## Bayesian Stan fit with `alphaT = 0`, `alphaX = 0`, and `delta0` fixed at 0
fit_delta0_fixed_alphaTX_zero <- function(W, T, G) {
  pars <- c("mu0", "tau", "sigmaX", "sigmaW0", "sigmaW1", "sigmaX2", 
            "sigmaW0_2", "sigmaW1_2", "alpha0", "delta1")
  .prome_stan_fit(W, T, G, "delta0_fixed_alphaTX_zero", pars)
}

.extract_log_lik_matrix <- function(stanfit) {
  arr <- rstan::extract(stanfit, pars = "log_lik", 
                        permuted = FALSE, inc_warmup = FALSE)
  if (length(dim(arr)) == 3L) {
    mat <- matrix(arr, nrow = dim(arr)[1] * dim(arr)[2], ncol = dim(arr)[3])
  } else {
    mat <- arr
  }
  mat
}

.waic_from_loglik <- function(log_lik_mat) {
    lppd <- sum(log(colMeans(exp(log_lik_mat))))
    tmp <- apply(log_lik_mat, 2, var)
    p_waic <- sum(tmp)
    elpd_waic <- lppd - p_waic
    waic <- -2 * elpd_waic
    list(lppd = lppd, p_waic = p_waic, elpd_waic = elpd_waic, waic = waic)
}

## Compare four fixed-`delta0` Stan bias models and compute model probabilities
##
## Fits four nested response models for the `delta0 = 0` setting: null
## (`alphaT = 0`, `alphaX = 0`), design-only (`alphaT` free, `alphaX = 0`),
## outcome-only (`alphaT = 0`, `alphaX` free), and full (`alphaT` and
## `alphaX` free). Model probabilities are computed either by bridge sampling
## when the `bridgesampling` package is available, or by normalized WAIC weights
## as an approximation.
TestStan <- function(W, T, G, method = c("auto", "bridge", "waic")) {
  .validate_inputs(W, T, G)
  .check_rstan()
  method <- match.arg(method)
  model_priors <- rep(0.25, 4)
  model_priors <- as.numeric(model_priors) / sum(model_priors)
  names(model_priors) <- c("null", "design_only", "outcome_only", "full")

  fit_args <- list(W = W, T = T, G = G)

  fits <- list(
    null = do.call(fit_delta0_fixed_alphaTX_zero, fit_args),
    design_only = do.call(fit_delta0_fixed_alphaX_zero, fit_args),
    outcome_only = do.call(fit_delta0_fixed_alphaT_zero, fit_args),
    full = do.call(fit_delta0_fixed, fit_args)
  )

  use_bridge <- FALSE
  if (method == "bridge") {
    if (!requireNamespace("bridgesampling", quietly = TRUE)) {
      stop("bridgesampling is required for method = 'bridge'.", call. = FALSE)
    }
    use_bridge <- TRUE
  } else if (method == "auto" && requireNamespace("bridgesampling", quietly = TRUE)) {
    use_bridge <- TRUE
  }

  details <- list()
  if (use_bridge) {
    bridges <- lapply(fits, function(x) bridgesampling::bridge_sampler(x$fit, silent = TRUE))
    logml <- vapply(bridges, function(b) b$logml, numeric(1))
    log_post_unnorm <- logml + log(model_priors[names(logml)])
    log_post_norm <- log_post_unnorm - max(log_post_unnorm)
    probs <- exp(log_post_norm)
    probs <- probs / sum(probs)
    details$log_marginal_lik <- logml
    details$bridge_objects <- bridges
    method_used <- "bridge"
  } else {
    waic_list <- lapply(fits, function(x) .waic_from_loglik(.extract_log_lik_matrix(x$fit)))
    elpd <- vapply(waic_list, function(w) w$elpd_waic, numeric(1))
    log_post_unnorm <- elpd + log(model_priors[names(elpd)])
    log_post_norm <- log_post_unnorm - max(log_post_unnorm)
    probs <- exp(log_post_norm)
    probs <- probs / sum(probs)
    details$waic <- waic_list
    details$elpd_waic <- elpd
    method_used <- "waic"
  }

  summary_full <- fits$full$summary
  out <- list(
    method = method_used,
    model_priors = model_priors,
    model_probabilities = probs,
    full_summary = summary_full,
    details = details,
    fits = fits
  )
  class(out) <- "stanTest"
  out
}

#' @export
print.stanTest <- function(x, ...) {
  cat("Method:", x$method, "\n")
  probs <- round(x$model_probabilities, 4)
  print(probs)
  if (!is.null(x$full_summary) && all(c("mean", "2.5%", "97.5%") %in% colnames(x$full_summary))) {
    keep <- intersect(c("alphaT", "alphaX"), rownames(x$full_summary))
    if (length(keep)) {
      cat("Full-model posterior summary for bias coefficients:\n")
      print(x$full_summary[keep, c("mean", "sd", "2.5%", "50%", "97.5%"), drop = FALSE])
    }
  }
  invisible(x)
}
