#############################################################
# _setup.R: helper functions bundled with the Day 14 walkthrough.
#
# Each session's .qmd starts with:
#
#     source("_setup.R")
#
# and then loads only the packages that day needs. Keep this file
# small and dependency-free enough to read quickly. Nothing here
# installs anything or reaches outside the day's own folder.
#
# -----------------------------------------------------------
# THE CACHING PATTERN  (cache_fit)
# -----------------------------------------------------------
# Every expensive fit in these documents is wrapped in cache_fit().
# The point is that the real ame()/ergm()/siena07() call stays visible
# on screen -- there is no hidden second version of the model that
# actually ran. What cache_fit does is:
#
#   1. hash the code you passed it
#   2. if cache/<name>.rda exists AND the hash matches, load it
#   3. otherwise actually fit, save, and report how long it took
#
# So the first run pays for the fit, later runs load it in milliseconds,
# and if you edit the model the cache
# invalidates itself instead of silently handing you yesterday's
# answer.
#
# Why not knitr's `#| cache: true`? Because it only works inside the
# knit loop. When you run a chunk interactively -- which is what we do
# interactively, knitr's cache is not consulted and you wait for the
# full refit. cache_fit works identically under knit and under
# 'Run Chunk'.
#
# Why not eval=FALSE + a silent load()? Because then the code you are
# reading is not the code that produced the output, and the two drift
# apart without anyone noticing.
#
# Delete cache/ to force fresh refits of everything.
#
# NOTE ON THE TWO NAMES. The ERGM (4_ergm) and SAOM (5_saom) sessions
# were adopted as-built and carry an equivalent helper called
# fit_or_load() in their own _helpers.R -- same idiom, precompute once
# and load forever, seeded; it writes .rds and gates refits on
# options(session.refit=). New days use cache_fit() from this file.
# `fit_or_load` is also defined below as a thin wrapper so either
# vocabulary reads the same in any document.
#############################################################

cache_fit <- function(name, expr, dir = "cache", refresh = FALSE) {
  if (!dir.exists(dir)) dir.create(dir, recursive = TRUE)
  path <- file.path(dir, paste0(name, ".rda"))
  code_hash <- substr(digest::digest(deparse(substitute(expr))), 1, 8)

  if (!refresh && file.exists(path)) {
    e <- new.env()
    load(path, envir = e)
    if (identical(e$.hash, code_hash)) {
      message("cache hit: '", name, "'")
      return(e$.value)
    }
    message("cache STALE for '", name, "' (code changed) -- refitting")
  }

  message("fitting '", name, "' ...")
  t0 <- Sys.time()
  .value <- force(expr)
  .hash  <- code_hash
  save(.value, .hash, file = path)
  message("  done in ",
          round(as.numeric(difftime(Sys.time(), t0, units = "secs")), 1),
          "s -> ", path)
  .value
}

# Alias so the fit_or_load() vocabulary of the ERGM/SAOM days reads the
# same in any document. Same cache/ folder, same behaviour.
fit_or_load <- function(name, expr, dir = "cache",
                        refit = getOption("session.refit", FALSE)) {
  cache_fit(name, expr, dir = dir, refresh = refit)
}

#############################################################
# checkpoint() -- portable across every day (lame, ergm, RSiena,
# blockmodels, plain base R). Prints one compact, comparable line so
# two runs can be compared before the next section depends on them.
#
#   checkpoint(density = 0.21, reciprocity = 0.63)
#############################################################
checkpoint <- function(..., label = "CHECKPOINT") {
  vals <- list(...)
  nms  <- names(vals); if (is.null(nms)) nms <- rep("", length(vals))
  body <- paste(mapply(function(n, v) {
    v <- if (is.numeric(v)) format(round(v, 3), nsmall = 0) else as.character(v)
    if (nzchar(n)) paste0(n, " = ", paste(v, collapse = ", "))
    else paste(v, collapse = ", ")
  }, nms, vals), collapse = "   |   ")
  cat(strrep("-", 66), "\n", label, ": ", body, "\n", strrep("-", 66),
      "\n", sep = "")
  invisible(TRUE)
}

#############################################################
# require_pkgs() -- fail loudly and early with an actionable message
# instead of failing later in the walkthrough. Never installs anything.
#############################################################
require_pkgs <- function(pkgs) {
  missing <- pkgs[!vapply(pkgs, requireNamespace, logical(1), quietly = TRUE)]
  if (length(missing)) {
    stop("Missing package(s): ", paste(missing, collapse = ", "),
         "\nRun:  install.packages(c(",
         paste0('"', missing, '"', collapse = ", "), "))",
         call. = FALSE)
  }
  invisible(TRUE)
}

#############################################################
# lame-specific helpers.
#
# These read structures that only a lame/ame() fit carries ($BETA,
# lame::tidy()), so they are used by the modeling days that fit lame --
# SRM + Blockmodels (2_srm_blocks) and Latent Geometry / AME (3_latent)
# -- and not by the descriptive, ERGM, or SAOM days.
#############################################################

# Effective sample size per beta, rounded. We look at this a lot when
# diagnosing an MCMC fit.
ess <- function(fit) round(coda::effectiveSize(coda::mcmc(fit$BETA)))

# lame::tidy(), printed the way we actually read it: point estimate,
# interval, and a flag for whether the interval clears zero.
coef_tbl <- function(fit, digits = 3) {
  t <- lame::tidy(fit)
  data.frame(
    term = t$term,
    est  = round(t$estimate,  digits),
    lo   = round(t$conf.low,  digits),
    hi   = round(t$conf.high, digits),
    clears_zero = ifelse(sign(t$conf.low) == sign(t$conf.high), "yes", ""),
    stringsAsFactors = FALSE
  )
}

# Put two lame fits side by side -- the comparison we make on every
# modeling day: what did adding structure do to the point estimate, and
# what did it do to the width of the interval?
compare_fits <- function(a, b, names = c("a", "b"), digits = 3) {
  ta <- lame::tidy(a); tb <- lame::tidy(b)
  stopifnot(identical(ta$term, tb$term))
  wa <- ta$conf.high - ta$conf.low
  wb <- tb$conf.high - tb$conf.low
  out <- data.frame(
    term        = ta$term,
    est_a       = round(ta$estimate, digits),
    est_b       = round(tb$estimate, digits),
    width_ratio = round(wb / wa, 2),
    moved       = ifelse(sign(ta$conf.low) == sign(ta$conf.high) &
                         sign(tb$conf.low) != sign(tb$conf.high),
                         "interval now includes zero",
                  ifelse(sign(ta$conf.low) != sign(ta$conf.high) &
                         sign(tb$conf.low) == sign(tb$conf.high),
                         "interval now excludes zero", "")),
    stringsAsFactors = FALSE
  )
  names(out)[2:3] <- paste0("est_", names)
  out
}
