# Locate a bundled Stan file
#
# @param model One of `"delta0_fixed"` or `"delta_free"`.
#
# @return Path to the `.stan` file shipped with the package.
# @export
rctme_stan_file <- function(model = c("delta0_fixed", "delta_free")) {
  model <- match.arg(model)
  file <- switch(model,
    delta0_fixed = system.file("stan", "rctme_delta0_fixed.stan", package = "prome"),
    delta_free = system.file("stan", "rctme_delta_free.stan", package = "prome")
  )
  if (!nzchar(file)) stop("Stan file not found.", call. = FALSE)
  file
}

.stan_model_cache <- new.env(parent = emptyenv())

.check_rstan <- function() {
  if (!requireNamespace("rstan", quietly = TRUE)) {
    stop("rstan is required for Stan fitting. Install it first.", call. = FALSE)
  }
}

.get_rstan_model <- function(stan_file) {
  key <- normalizePath(stan_file, winslash = "/", mustWork = TRUE)
  if (!exists(key, envir = .stan_model_cache, inherits = FALSE)) {
    .stan_model_cache[[key]] <- rstan::stan_model(file = stan_file, auto_write = TRUE)
  }
  get(key, envir = .stan_model_cache, inherits = FALSE)
}

# Bayesian Stan fit with `delta0` fixed at 0
#
# Fits the restricted Bayesian latent-outcome model using `rstan` and the
# Stan program bundled with the package.
#
# @inheritParams fit_rctme_em
# @param chains Number of chains.
# @param cores Number of CPU cores for parallel sampling. Defaults to `chains`.
# @param iter_warmup Warmup iterations.
# @param iter_sampling Sampling iterations.
# @param control Optional named list of `rstan` control arguments such as
#   `list(adapt_delta = 0.95, max_treedepth = 12)`.
# @param refresh Stan refresh interval.
#
# @return A list with components `fit`, `summary`, and `stan_file`.
# @export
fit_rctme_stan_delta0_fixed <-
    function(W, T, G,
             chains = 4,
             cores = chains,
             iter_warmup = 1000,
             iter_sampling = 1000,
             control = list(adapt_delta = 0.95, max_treedepth = 12),
             refresh = 0) {
  .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 <- rctme_stan_file("delta0_fixed")
  mod <- .get_rstan_model(stan_file)
  fit <- rstan::sampling(
    object = mod,
    data = data_list,
    chains = chains,
    cores = cores,
    iter = iter_warmup + iter_sampling,
    warmup = iter_warmup,
    refresh = refresh,
    control = control
    )
  .make_stan_result1(fit, call = match.call())
}

# Bayesian Stan fit with free `delta0` and `delta1`
#
# Fits the unrestricted Bayesian latent-outcome model using `rstan` and the
# Stan program bundled with the package.
#
# @inheritParams fit_rctme_stan_delta0_fixed
# @return A list with components `fit`, `summary`, and `stan_file`.
# @export
#
fit_rctme_stan <- function(W, T, G,
                           chains = 4,
                           cores = chains,
                           iter_warmup = 1000,
                           iter_sampling = 1000,
                           control = list(adapt_delta = 0.95, max_treedepth = 12),
                           refresh = 0) {
  .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 <- rctme_stan_file("delta_free")
  mod <- .get_rstan_model(stan_file)
  fit <- rstan::sampling(
    object = mod,
    data = data_list,
    chains = chains,
    cores = cores,
    iter = iter_warmup + iter_sampling,
    warmup = iter_warmup,
    refresh = refresh,
    control = control
  )
  .make_stan_result2(fit, call = match.call())
}
