#' @title Multiplier Bootstrap
#'
#' @description This function take an influence function and use the
#'  multiplier bootstrap to compute standard errors and critical values for
#'  uniform confidence bands.
#'
#' @param inf_func an influence function
#' @param did_preprocessed A `dp` object obtained after preprocess
#' @param use_parallel Boolean of whether or not to use parallel processing in the multiplier
#'  bootstrap, default is \code{use_parallel=FALSE}
#' @param cores the number of cores to use with parallel processing, default is \code{cores=1}
#'
#' @return A list with following elements
#' \item{bres}{results from each bootstrap iteration.}
#' \item{V}{variance matrix.}
#' \item{se}{standard errors.}
#' \item{crit_val}{a critical value for computing uniform confidence bands.}
#'
#' @keywords internal
mboot <- function(inf_func, did_preprocessed, use_parallel = FALSE, cores = 1) {

  # setup needed variables
  data <- did_preprocessed$preprocessed_data
  cluster <- did_preprocessed$cluster
  # Reuse the per-unit cluster vector that att_gt()/compute_aggregation() aligned to the
  # influence-function rows, when available. Falls back to deriving it from the data.
  cluster_vector_in  <- did_preprocessed$cluster_vector
  use_cluster_vector <- !is.null(cluster_vector_in) &&
    length(cluster_vector_in) == NROW(inf_func)
  biters <- did_preprocessed$nboot
  alpha <- did_preprocessed$alpha
  panel <- did_preprocessed$panel
  allow_unbalanced_panel <- did_preprocessed$allow_unbalanced_panel
  tlist <- sort(unique(data$period))

  # Get data for cluster extraction
  # For balanced panels: use first period (all IDs present)
  # For unbalanced panels: get one row per unique ID (not all IDs in first period)
  if (panel && !allow_unbalanced_panel) {
    # Balanced panel: filter to first period works since all IDs are present
    dta <- data[period == tlist[1]]
  } else {
    # Unbalanced panel: need one row per unique ID to match influence function dimensions
    dta <- data[, .SD[1], by = id]
  }

  # Make sure inf_func is matrix because we need this for computing n below
  inf_func <- as.matrix(inf_func)

  # set correct number of units
  n <- nrow(inf_func)

  # multiplier bootstrap
  n_clusters <- n
  if (length(cluster)==0) {
    bres <- sqrt(n) * run_multiplier_bootstrap(inf_func, biters, use_parallel, cores)
  } else {
    # Cluster-robust multiplier bootstrap, Callaway & Sant'Anna (2021, Remark 10):
    # cluster-specific multipliers on the influence function aggregated to cluster SUMS.
    if (use_cluster_vector) {
      cluster_vec <- cluster_vector_in
    } else {
      unique_clusters_df <- unique(dta[, c("id", "cluster")])
      cluster_vec <- unique_clusters_df$cluster
    }
    n_clusters      <- length(unique(cluster_vec))
    cluster_sum_if  <- rowsum(inf_func, cluster_vec, reorder = TRUE)   # SUMS, not means
    bres <- sqrt(n_clusters) * run_multiplier_bootstrap(cluster_sum_if, biters, use_parallel, cores)
  }


  if (isTRUE(class(bres) == "numeric")) bres <- as.matrix(bres)

  # Non-degenerate dimensions
  ndg.dim <- (!is.na(colSums(bres))) & (base::colSums(bres^2) > sqrt(.Machine$double.eps)*10)
  bres <- as.matrix(bres[ , ndg.dim])

  # bootstrap variance matrix (this matrix can be defective because of degenerate cases)
  V <- cov(bres)
  # bootstrap standard error
  bSigma <- apply(bres, 2,
                  function(b) (quantile(b, .75, type=1, na.rm = T) -
                                 quantile(b, .25, type=1, na.rm = T))/(qnorm(.75) - qnorm(.25)))

  bSigma[bSigma <= sqrt(.Machine$double.eps)*10] <- NA

  # critical value for uniform confidence band
  bT <- base::suppressWarnings(apply(bres, 1, function(b) max(abs(b / bSigma), na.rm = T)))
  bT <- bT[is.finite(bT)]
  crit_val <- quantile(bT, 1-alpha, type=1, na.rm = T) # uniform critical value
  se <- rep(NA, length(ndg.dim))
  # With cluster SUMS this is the cluster-robust SE (1/n)*sqrt(sum_c S_c^2); without
  # clustering n_clusters == n and it reduces to the i.i.d. bSigma/sqrt(n). (Equivalent
  # to the old bSigma/sqrt(n_clusters) only when clusters are equal-sized.)
  se[ndg.dim] <- as.numeric(bSigma) * sqrt(n_clusters) / n

  return(list(bres = bres, V = V, se = se, bT= bT, unif_crit_val = crit_val))
}

run_multiplier_bootstrap <- function(inf_func, biters, use_parallel = FALSE, cores = 1) {
  ngroups = ceiling(biters/cores)
  chunks = rep(ngroups, cores)
  # Round down so you end up with the right number of biters
  chunks[1] = chunks[1] + biters - sum(chunks)

  n <- nrow(inf_func)
  parallel.function <- function(biters) {
    BMisc::multiplier_bootstrap(inf_func, biters)
  }
  # From tests, this is about where it becomes worth it to parallelize
  if(n > 2500 & use_parallel == TRUE & cores > 1) {
    results = parallel::mclapply(
      chunks,
      FUN = parallel.function,
      mc.cores = cores
    )
    results = do.call(rbind, results)
  } else {
    results = BMisc::multiplier_bootstrap(inf_func, biters)
  }
  return(results)
}
