#' Bootstrap a modsem Model
#'
#' A generic interface for parametric and non-parametric bootstrap procedures
#' for structural equation models estimated with the \pkg{modsem} ecosystem.
#' The function dispatches on the class of \code{model}; currently dedicated
#' methods exist for \code{\link{modsem_pi}} (product‑indicator approach) and
#' \code{\link{modsem_da}} (distributional‑analytic approach).
#'
#' @param model A fitted \code{modsem} object, or a function to be bootstrapped
#'   (e.g., \code{\link{modsem}}, \code{\link{modsem_da}} and \code{\link{modsem_pi}})
#' @param FUN  A function that returns the statistic of interest when applied to
#'   a fitted model.  The function must accept a single argument, the model
#'   object, and should ideally return a numeric vector; see Value.
#' @param ...  Additional arguments forwarded to \code{lavaan::bootstrapLavaan}
#'   for \code{\link{modsem_pi}} objects, or \code{\link{modsem_da}} for
#'   \code{\link{modsem_da}} objects.
#'
#' @return Depending on the return type of \code{FUN} either
#'   \describe{
#'     \item{numeric}{A matrix with \code{R} rows (bootstrap replicates) and as
#'       many columns as \code{length(FUN(model))}.}
#'     \item{other}{A list of length \code{R}; each element is the raw output of
#'       \code{FUN}. \strong{NOTE}: Only applies for \code{\link{modsem_da}} objects}
#'   }
#' @seealso \code{\link[lavaan]{bootstrapLavaan}}, \code{\link{modsem_pi}},
#'   \code{\link{modsem_da}}
#' @export
bootstrap_modsem <- function(model = modsem, FUN, ...) {
  UseMethod("bootstrap_modsem")
}


#' @describeIn bootstrap_modsem Bootstrap a \code{modsem_pi} model by delegating
#'   to \code{\link[lavaan]{bootstrapLavaan}}.
#'
#' @details A thin wrapper around \code{lavaan::bootstrapLavaan()} that performs the
#'   necessary book‑keeping so that \code{FUN} receives a fully‑featured
#'   \code{modsem_pi} object—rather than a bare \code{lavaan} fit—at every
#'   iteration.
#'
#' @examples
#'
#' m1 <- '
#'   X =~ x1 + x2
#'   Z =~ z1 + z2
#'   Y =~ y1 + y2
#'
#'   Y ~ X + Z + X:Z
#' '
#'
#' fit_pi <- modsem(m1, oneInt)
#' bootstrap_modsem(fit_pi, FUN = coef, R = 10L)
#'
#' @export
bootstrap_modsem.modsem_pi <- function(model, FUN = "coef", ...) {
  checkWarnRCS(model)

  wrapLavFit <- function(lavfit) {
    model$lavaan       <- lavfit
    model$coefParTable <- lavaan::parameterEstimates(lavfit)
    model$data         <- lavaan::lavInspect(lavfit, what = "data")
    model
  }

  fun.mod <- \(fit, ...) do.call(FUN, list(wrapLavFit(fit), ...))

  lavaan::bootstrapLavaan(extract_lavaan(model), FUN = fun.mod, ...)
}


#' @describeIn bootstrap_modsem Parametric or non-parametric bootstrap for
#'   \code{modsem_da} models.
#'
#' @param R number of bootstrap replicates.
#' @param P.max ceiling for the simulated population size.
#' @param type Bootstrap flavour, see Details.
#' @param verbose Should progress information be printed to the console?
#' @param calc.se Should standard errors for each replicate. Defaults to \code{FALSE}.
#' @param optimize Should starting values be re-optimized for each replicate. Defaults to \code{FALSE}.
#' @param algorithm Which algorithm should be used for the LMS approach? Defaults to \code{"EM"}, as
#'   opposed to \code{"EMA"}.
#' @param convergence.abs Absolute convergence criteria. Defaults to 10 times the original criteria.
#' @param convergence.rel Relative convergence criteria. Defaults to 10 times the original criteria.
#'
#' @details The function internally resamples the observed data (non-parametric
#'   case) or simulates from the estimated parameter table (parametric case),
#'   feeds the sample to \code{\link{modsem_da}}, evaluates \code{FUN} on the
#'   refitted object and finally collates the results.
#'
#' @examples
#'
#' m1 <- '
#'   X =~ x1 + x2
#'   Z =~ z1 + z2
#'   Y =~ y1 + y2
#'
#'   Y ~ X + Z + X:Z
#' '
#'
#' \dontrun{
#' fit_lms <- modsem(m1, oneInt, method = "lms")
#' bootstrap_modsem(fit_lms, FUN = coef, R = 10L)
#' }
#' @export
bootstrap_modsem.modsem_da <- function(model,
                                       FUN = "coef",
                                       R = 1000L,
                                       P.max = 1e5,
                                       type = c("nonparametric", "parametric"),
                                       verbose = interactive(),
                                       calc.se = FALSE,
                                       optimize = FALSE,
                                       convergence.abs = 10 * model$args$convergence.abs,
                                       convergence.rel = 10 * model$args$convergence.rel,
                                       algorithm = "EM",
                                       ...) {
  checkWarnRCS(model)

  type <- tolower(type)
  type <- match.arg(type)

  INSPECT  <- modsem_inspect(model, what = "all", is.public = FALSE)
  group    <- INSPECT$group

  if (!is.null(group)) {
    group.label <- INSPECT$group.label
    DATA        <- INSPECT$data

    sampling.weights <- NULL
    cluster          <- NULL
    data             <- NULL

    for (g in group.label) {
      data.g.mat         <- DATA[[g]]
      data.g             <- as.data.frame(data.g.mat)
      cluster.g          <- attr(data.g.mat, "cluster")
      sampling.weights.g <- attr(data.g.mat, "weights")
      data.g[[group]]    <- g

      data             <- rbind(data, data.g)
      cluster          <- c(cluster, cluster.g)
      sampling.weights <- c(sampling.weights, sampling.weights.g)
    }

  } else {
    data.mat         <- INSPECT$data
    data             <- as.data.frame(data.mat)
    cluster          <- attr(data.mat, "cluster")
    sampling.weights <- attr(data.mat, "weights")
  }

  allvars  <- colnames(data)
  ovs      <- INSPECT$ovs
  N        <- NROW(data)
  P        <- min(P.max, N * R)
  parTable <- parameter_estimates(model)

  mod_warnif(P.max <= N, "`P.max` is less than `N`!")

  mod_stopif(!is.null(cluster) && type != "nonparametric",
         "cluster bootstrap is only available with `type=\"nonparametric\"`!")

  if (!is.null(cluster) && !is.null(model$args$cluster))
    data[[model$args$cluster]] <- cluster

  if (!is.null(sampling.weights) && !is.null(model$args$sampling.weights))
    data[[model$args$sampling.weights]] <- sampling.weights

  population <- switch(type,
    parametric    = simulatedGroupsToDf(simulateDataParTable(parTable, N = P, colsOVs = ovs), type = "OV"),
    nonparametric = data,
    mod_msg_stop("Unrecognized type!\n")
  )

  # Use settings from the original model, but allow the user to
  # overrride the defaults
  userArgs <- list(...)
  argList  <- model$args
  argList[names(userArgs)] <- userArgs
 
  # Except these...
  argList$algorithm       <- algorithm
  argList$convergence.abs <- convergence.abs
  argList$convergence.rel <- convergence.rel
  argList$calc.se         <- calc.se
  argList$verbose         <- verbose
  argList$model.syntax    <- model$model$info$group.info$syntax
  argList$cov.syntax      <- model$model$info$group.info$cov.syntax
  argList$method          <- model$method
  argList$sampling.weights.normalization <- "none" # This has already been done

  if (type == "parametric" && !is.null(argList$sampling.weights)) {
    # By desing, the sampling weights aren't needed when simulating
    # parametric data, so we delete the argument.
    argList$sampling.weights <- NULL
  }

  if (!optimize) {
    argList$start        <- model$theta
    argList$optimize     <- FALSE
  } else {
    argList$start        <- NULL
    argList$optimize     <- TRUE
  }

  f0 <- do.call(FUN, args = list(model))
  FUN.VEC <- is.vector(f0) || inherits(f0, "ModsemVector")

  if (FUN.VEC) {
    out <- matrix(NA, nrow = R, ncol = length(f0),
                  dimnames = list(NULL, names(f0)))

  } else out <- vector("list", length = R)

  ERROR <- \(e) {mod_msg_warn(e); NULL}

  for (i in seq_len(R)) {
    printedLines <- utils::capture.output(split = TRUE, {
      if (verbose) printf("Bootstrap %d/%d...\n", i, R)

      sample_i  <- resample(population, n.out = N, cluster = cluster)
      argList_i <- c(argList, list(data = sample_i))

      fit_i <- tryCatch(
        do.call(modsem_da, args = argList_i),
        error = ERROR
      )

      if (is.null(fit_i)) {
        fi <- NA

      } else {
        fi <- tryCatch(do.call(FUN, args = list(fit_i)),
                       error = ERROR)
      }

      if (FUN.VEC) out[i, ] <- fi else out[[i]] <- fi
    })

    nprinted <- length(printedLines)
    if (i < R) eraseConsoleLines(nprinted)
  }

  out
}


#' @describeIn bootstrap_modsem Non-parametric bootstrap of \code{modsem} functions
#'
#' @param R number of bootstrap replicates.
#' @param verbose Should progress information be printed to the console?
#' @param data Dataset to be resampled.
#' @param FUN.args Arguments passed to \code{FUN}
#' @param cluster.boot Variable to cluster bootstrapping by
#'
#' @details This is a more general version of \code{boostrap_modsem} for
#'   bootstrapping \code{modsem} functions, not modsem objects.
#'   \code{model} is now a function to be boostrapped, and \code{...}
#'   are now passed to the function (\code{model}), not \code{FUN}. To
#'   pass arguments to \code{FUN} use \code{FUN.args}.
#'
#' @examples
#'
#' tpb <- "
#' # Outer Model (Based on Hagger et al., 2007)
#'   ATT =~ att1 + att2 + att3 + att4 + att5
#'   SN =~ sn1 + sn2
#'   PBC =~ pbc1 + pbc2 + pbc3
#'   INT =~ int1 + int2 + int3
#'   BEH =~ b1 + b2
#'
#' # Inner Model (Based on Steinmetz et al., 2011)
#'   INT ~ ATT + SN + PBC
#'   BEH ~ INT + PBC + INT:PBC
#' "
#'
#' \dontrun{
#' boot <- bootstrap_modsem(model = modsem,
#'                          model.syntax = tpb, data = TPB,
#'                          method = "dblcent", rcs = TRUE,
#'                          rcs.scale.corrected = TRUE,
#'                          FUN = "coef", R = 50L)
#' coef <- apply(boot, MARGIN = 2, FUN = mean, na.rm = TRUE)
#' se   <- apply(boot, MARGIN = 2, FUN = sd, na.rm = TRUE)
#'
#' cat("Parameter Estimates:\n")
#' print(coef)
#'
#' cat("Standard Errors: \n")
#' print(se)
#'
#' }
#' @export
bootstrap_modsem.function <- function(model = modsem,
                                      FUN = "coef",
                                      data,
                                      R = 1000L,
                                      verbose = interactive(),
                                      FUN.args = list(),
                                      cluster.boot = NULL,
                                      ...) {
  MODELFUN <- model # rename for convenience

  out  <- vector("list", length = R)
  args <- list(...)
  data <- as.data.frame(data)
  N    <- NROW(data)

  ERROR <- \(e) {mod_msg_warn(e); NULL}

  if (!is.null(cluster.boot)) {
    mod_stopif(length(cluster.boot) > 1L, "`cluster.boot` must be of length 1!")
    mod_stopif(!cluster.boot %in% colnames(data), "`cluster.boot` must be a column in `data`!")
    cluster.values <- data[[cluster.boot]]

  } else cluster.values <- NULL

  for (i in seq_len(R)) {
    printedLines <- utils::capture.output(split = TRUE, {
      if (verbose) printf("Bootstrap %d/%d...\n", i, R)

      sample_i  <- resample(data, n.out = NROW(data), cluster = cluster.values)
      argList_i <- c(list(data = sample_i), args)

      fit_i <- tryCatch(
        do.call(MODELFUN, args = argList_i),
        error = ERROR
      )

      if (is.null(fit_i)) {
        fi <- NA

      } else {
        fi <- tryCatch(do.call(FUN, args = c(list(fit_i), FUN.args)),
                       error = ERROR)
      }

      out[[i]] <- fi
    })

    nprinted <- length(printedLines)
    if (i < R) eraseConsoleLines(nprinted)
  }

  f1 <- out[[1L]]
  FUN.VEC <- is.vector(f1) || inherits(f1, c("ModsemVector", "lavaan.vector"))

  if (FUN.VEC) {
    MAT <- tryCatch(do.call("rbind", out), error = \(e) NULL)

    if (is.null(MAT)) return(out) else return(MAT)
  }

  out
}


resample <- function(df, n.out = NROW(df), cluster = NULL, replace = TRUE) {
  df.orig <- as.data.frame(df)

  if (is.null(cluster)) {
    idx <- sample(NROW(df.orig), size = n.out, replace = replace)
    return(df.orig[idx, , drop = FALSE])
  }

  mod_stopif(length(cluster) != NROW(df.orig), "Cluster must be of same lenght as data!")

  clusters <- unique(cluster)
  G <- length(clusters)

  clusters.sample <- sample(clusters, size = G, replace = replace)

  purrr::list_rbind(
    lapply(clusters.sample, FUN = \(ci) df.orig[cluster==ci, , drop=FALSE])
  )
}


checkWarnRCS <- function(model) {
  isRCS_Model <- attr(model, "isRCS_Model")

  mod_warnif(!is.null(isRCS_Model) && isRCS_Model,
         paste0("bootstrapping a model with estimated with `rcs=TRUE` directly, will\n",
         "generate naive standard errors! Use `bootstrap_modsem.function()` instead!"))
}
