# Repo-local MCMC-frequency covariance extensions for BMS.
# --------------------------------------------------------
# This file exposes an explicit local API on top of the vendored upstream
# `BMS` source in `vendor/BMS/`. The installed package remains untouched.
#
# The core idea is simple:
# - the standard BMS MCMC summary stores coefficient-wise first moments and
#   diagonal second moments,
# - that is sufficient for posterior means and posterior SDs,
# - but not sufficient for interaction marginal effects, which need full
#   posterior covariance.
#
# The vendored function `bms_mcmc_cov()` extends the sampler so it also
# accumulates:
#
#   E[beta beta' | Y]
#
# under MCMC frequencies. This file then turns that additional moment into a
# user-facing covariance matrix and marginal-effect calculations.

suppressPackageStartupMessages({
  library(BMS)
})


# Vendored source loader
# ======================
# The vendored package source is sourced into a dedicated environment whose
# parent is the installed BMS namespace. This keeps the local fork explicit
# and isolated while still allowing the vendored functions to reuse upstream
# objects they were written against.
get_vendor_bms_env <- local({
  cached_env <- NULL

  function(root_dir = normalizePath(getwd())) {
    if (!is.null(cached_env)) {
      return(cached_env)
    }

    vendor_env <- new.env(parent = asNamespace("BMS"))
    vendor_files <- c(
      file.path(root_dir, "vendor", "BMS", "R", "aux_inner.R"),
      file.path(root_dir, "vendor", "BMS", "R", "aux_outer.R"),
      file.path(root_dir, "vendor", "BMS", "R", "bma.R")
    )

    for (single_file in vendor_files) {
      source(single_file, local = vendor_env)
    }

    cached_env <<- vendor_env
    vendor_env
  }
})


# Validation helpers
# ==================

bms_mcmc_assert_ready <- function(bmao) {
  if (!BMS::is.bma(bmao)) {
    stop("bms_mcmc_assert_ready() requires an object of class 'bma'.")
  }

  if (is.null(bmao$info$b12mo)) {
    stop(
      "The provided BMS object does not contain MCMC covariance moments. ",
      "Use bms_mcmc_cov(...) to create a covariance-enabled object."
    )
  }

  if (!identical(bmao$info$covariance_mode, "mcmc_frequency")) {
    stop(
      "The provided BMS object is missing covariance_mode = 'mcmc_frequency'. ",
      "Use a repo-local MCMC covariance object."
    )
  }

  if (identical(bmao$arguments$mcmc, "enum")) {
    stop(
      "Enumeration objects are not supported by the MCMC-frequency covariance path. ",
      "The manuscript's BMA uses the birth-death MCMC sampler (mcmc = \"bd\")."
    )
  }

  invisible(TRUE)
}


bms_mcmc_validate_terms <- function(bmao, terms = NULL) {
  available_terms <- bmao$reg.names

  if (is.null(terms)) {
    return(available_terms)
  }

  duplicate_terms <- unique(terms[duplicated(terms)])
  if (length(duplicate_terms) > 0) {
    stop(
      "Requested term vector contains duplicates: ",
      paste(duplicate_terms, collapse = ", ")
    )
  }

  missing_terms <- setdiff(terms, available_terms)
  if (length(missing_terms) > 0) {
    stop(
      "Requested term(s) not found in the BMS object: ",
      paste(missing_terms, collapse = ", ")
    )
  }

  terms
}


# Public estimator wrapper
# ========================

bms_mcmc_cov <- function(...) {
  # run_all.R defines ROOT from its own script path, so the vendored sources are
  # found no matter which working directory the replication is launched from.
  # Fall back to the working directory only when this file is sourced standalone.
  vendor_env <- get_vendor_bms_env(
    if (exists("ROOT", inherits = TRUE)) ROOT else normalizePath(getwd())
  )
  vendor_env$bms_mcmc_cov(...)
}


# Posterior moments and covariance
# ================================

bms_mcmc_posterior_moments <- function(bmao, terms = NULL) {
  bms_mcmc_assert_ready(bmao)
  terms <- bms_mcmc_validate_terms(bmao, terms)

  posterior_mean <- bmao$info$b1mo / bmao$info$cumsumweights
  names(posterior_mean) <- bmao$reg.names

  posterior_second_moment <- bmao$info$b12mo / bmao$info$cumsumweights
  rownames(posterior_second_moment) <- colnames(posterior_second_moment) <- bmao$reg.names

  selected_mean <- posterior_mean[terms]
  selected_second_moment <- posterior_second_moment[terms, terms, drop = FALSE]
  posterior_cov <- selected_second_moment - tcrossprod(as.numeric(selected_mean))
  rownames(posterior_cov) <- colnames(posterior_cov) <- terms

  posterior_variance <- diag(posterior_cov)

  list(
    terms = terms,
    posterior_mean = selected_mean,
    posterior_second_moment = selected_second_moment,
    posterior_cov = posterior_cov,
    posterior_sd = setNames(sqrt(pmax(posterior_variance, 0)), terms)
  )
}


bms_mcmc_vcov <- function(bmao, terms = NULL) {
  bms_mcmc_posterior_moments(bmao, terms = terms)$posterior_cov
}


# Marginal effects
# ================

bms_mcmc_marginal_effects <- function(
  bmao,
  base_term,
  interaction_term,
  conditioning_values,
  conf_level = 0.90
) {
  term_set <- c(base_term, interaction_term)
  posterior_moments <- bms_mcmc_posterior_moments(bmao, terms = term_set)
  covariance_matrix <- posterior_moments$posterior_cov[term_set, term_set, drop = FALSE]

  base_mean <- posterior_moments$posterior_mean[[base_term]]
  interaction_mean <- posterior_moments$posterior_mean[[interaction_term]]
  base_variance <- covariance_matrix[base_term, base_term]
  interaction_variance <- covariance_matrix[interaction_term, interaction_term]
  cross_covariance <- covariance_matrix[base_term, interaction_term]
  critical_value <- stats::qnorm((1 + conf_level) / 2)

  estimate <- base_mean + interaction_mean * conditioning_values
  std_error <- sqrt(
    pmax(
      0,
      base_variance +
        conditioning_values^2 * interaction_variance +
        2 * conditioning_values * cross_covariance
    )
  )

  data.frame(
    conditioning_value = conditioning_values,
    estimate = estimate,
    std_error = std_error,
    conf_low = estimate - critical_value * std_error,
    conf_high = estimate + critical_value * std_error
  )
}


# Validation against native non-exact summaries
# =============================================

bms_mcmc_validate_against_bms <- function(bmao, terms = NULL) {
  bms_mcmc_assert_ready(bmao)
  terms <- bms_mcmc_validate_terms(bmao, terms)

  package_summary <- as.data.frame(
    coef(
      bmao,
      exact = FALSE,
      order.by.pip = FALSE,
      std.coefs = FALSE
    )
  )
  package_summary$term <- rownames(package_summary)
  rownames(package_summary) <- NULL

  helper_moments <- bms_mcmc_posterior_moments(bmao, terms = terms)
  comparison <- data.frame(
    term = terms,
    helper_post_mean = unname(helper_moments$posterior_mean[terms]),
    helper_post_sd = unname(helper_moments$posterior_sd[terms]),
    stringsAsFactors = FALSE
  )

  package_subset <- package_summary[
    package_summary$term %in% terms,
    c("term", "Post Mean", "Post SD")
  ]
  names(package_subset) <- c("term", "bms_post_mean", "bms_post_sd")

  comparison <- merge(comparison, package_subset, by = "term", sort = FALSE)
  comparison$abs_diff_mean <- abs(comparison$helper_post_mean - comparison$bms_post_mean)
  comparison$abs_diff_sd <- abs(comparison$helper_post_sd - comparison$bms_post_sd)
  comparison
}
