R/runSA.R

In ELAplus: Energy Landscape Analysis for Ecological Communities

#' Stochastic approximation for parameter fitting
#'
#' Estimate species interaction parameters and (optionally) environmental
#' preference parameters by stochastic approximation (SA). The function runs
#' multiple SA optimizations in parallel (when \code{threads > 1}) using
#' \pkg{foreach}.
#'
#' @param ocmat Numeric matrix (0/1). Rows are samples and columns are species.
#' @param enmat Optional numeric matrix. Environmental variables for each sample.
#'   Should be normalized/standardized beforehand.
#' @param rep Integer. Number of parallel runs (replicates) for optimization.
#' @param threads Integer. Number of cores used when running in parallel.
#' @param we Numeric. Weight decay parameter for AdamW.
#' @param totalit Integer. Total number of optimization iterations.
#' @param lambda Numeric. Sparsity penalty parameter (used when \code{sparse = TRUE}).
#' @param intv Integer. Interval of iterations at which intermediate results are recorded.
#' @param runadamW Logical. If \code{TRUE}, use AdamW-style optimization.
#' @param maxlr Numeric. Maximum learning rate.
#' @param sparse Logical. If \code{TRUE}, apply sparsity constraint.
#' @param getall Logical. If \code{TRUE}, return all replicate trajectories.
#' @param reporting Logical. If \code{TRUE}, print progress messages.
#'
#' @return If \code{getall = TRUE}, a list of SA trajectories (length \code{rep}).
#' Otherwise, a list with two elements, \code{c(params, enmat)} is returned.
#' \code{params} is the concatenated matrix of the estimated parameter set:
#' \describe{
#'   \item{h}{Implicit environmental effects on species occurrence}
#'   \item{J}{Species-species interaction matrix}
#'   \item{g}{(If \code{enmat} is not \code{NULL}) Explicit effects of the environmental factors.}
#' }
#'
#'
#' @useDynLib ELAplus, .registration=TRUE
#' @importFrom Rcpp sourceCpp
#' @import foreach
#' @export
runSA <- function(ocmat, enmat=NULL,  rep=128, threads=1, we=0.001, 
                  totalit=1000, lambda=0.01,intv=100, runadamW=TRUE,
                  maxlr = 0.005, sparse=TRUE, getall=FALSE, reporting=TRUE){
    #if(!is.matrix(ocmat)){ocmat <- as.matrix(ocmat)}
    if(!sparse){lambda <- 0}
  
    #if(!runadamW){we <- 0}
  
    if (threads > 1) {
      cluster <- parallel::makeCluster(threads, type = if (.Platform$OS.type == "windows") "PSOCK" else "FORK")
      parallel::clusterCall(cluster, function(x) .libPaths(x), .libPaths())
      parallel::clusterSetRNGStream(cluster)
      doParallel::registerDoParallel(cluster)
      on.exit(parallel::stopCluster(cluster), add = TRUE)
    }

    if(is.null(enmat)){

    	## ============================================== ##
        ## -- without explicit variables

        if(reporting){cat('Start parameter fitting:\n')}
        s <- proc.time()[3]
        if(threads>1){fittingMat <- foreach(i = 1:rep, .packages="ELAplus") %dopar% {
            SAres <- SA(ocData = ocmat, we = we, totalit = totalit, 
                        runadamW = runadamW, sparse = sparse,
                        lambda = lambda, intv = intv, maxlearningrate = maxlr)
            return(SAres)
        }}else{fittingMat <- foreach(i = 1:rep) %do% {
	      if(reporting){cat('.')}
            SAres <- SA(ocData = ocmat, we = we, totalit = totalit, 
                        runadamW = runadamW, sparse = sparse,
                        lambda = lambda, intv = intv,maxlearningrate = maxlr)
            return(SAres)
        }}

        if(reporting){cat(sprintf('SA: elapsed time %.2f sec\n\n', proc.time()[3]-s))}
        ## ============================================== ##
      
        if(getall){
          return(fittingMat)
        }else{
            return(list(saEndpoint(fittingMat,ocmat),enmat))}
          
	    }else{

        ## ============================================== ##
        ## -- with explicit variables
	      #if(!is.matrix(enmat)){enmat <- as.matrix(enmat)}

        if(reporting){cat('Start parameter fitting\n')}
        s <- proc.time()[3]
        if(threads>1){fittingMat <- foreach(i = 1:rep, .packages="ELAplus") %dopar% {
            SAres <- fullSA(ocData = ocmat, envData=as.matrix(enmat), 
                            runadamW = runadamW, sparse = sparse, we = we, 
                            totalit = totalit, lambda = lambda, intv = intv,
                            maxlearningrate = maxlr)
            return(SAres)
        }}else{fittingMat <- foreach(i = 1:rep) %do% {
            if(reporting){cat('.')}
            SAres <- fullSA(ocData = ocmat, envData=as.matrix(enmat),
                            runadamW = runadamW, sparse = sparse, we = we, 
                            totalit = totalit, lambda = lambda, intv = intv,
                            maxlearningrate = maxlr)
            return(SAres)
        }}

        if(reporting){cat(sprintf('SA: elapsed time %.2f sec\n\n', proc.time()[3]-s))}
        if(getall){
          return(fittingMat)
        }else{
          return(list(saEndpoint(fittingMat,ocmat,enmat=enmat),enmat))}
        ## ============================================== ##
	
    }
}


#' Extract endpoint parameters from \link{runSA} output
#'
#' Summarize the final parameter estimates across SA replicates by averaging
#' the last recorded block of parameters.
#'
#' @param fittingMat A list of SA trajectories returned by \code{runSA(getall = TRUE)}.
#' @param ocmat Numeric matrix (0/1). Rows are samples and columns are species.
#' @param enmat Optional environmental parameter matrix used in fitting.
#'
#' @return A numeric matrix of median endpoint parameters across serials. 
#' Rownames correspond to samples and Colnames correspond to species list
#'
#' @export
saEndpoint <- function(fittingMat,ocmat,enmat=NULL){
  
  num_cols <- ncol(ocmat)
  packed_results_end <- lapply(fittingMat, 
                               function(x) x[(nrow(x) - num_cols + 1):nrow(x),])
  params_all <- lapply(packed_results_end, function(x) x[, 1:(ncol(x) - 2)])
  #upd <- lapply(sa, function(x) x[, (ncol(x) - 1):ncol(x)]) 
  
  # median
  if (length(params_all) == 1) {
    params_mat <- params_all[[1]]
  } else {
    arr <- simplify2array(params_all)
    params_mat <- apply(arr, c(1, 2), median, na.rm = TRUE)
  }
  
  # mean
  #params_mat <- params_all[[1]]
  #if (length(params_all)>1){
  #  for (n in 2:length(params_all)) {
  #    params_mat <- params_mat + params_all[[n]]
  #  }
  #
  #  params_mat <- params_mat/length(params_all)
  #}
  rownames(params_mat) <- colnames(ocmat)
  
  if(is.null(enmat)){
    colnames(params_mat) <- c("h",paste('J',colnames(ocmat),sep='.'))
  }
  else{
    encols <- ncol(enmat)
    
    if( ncol(as.matrix(enmat))==1 ){
      gname <- paste('g', colnames(enmat), sep='.')
      ncols <- ncol(ocmat)+2
    }else{
      gname <- paste('g',1:ncol(enmat), sep='.')
      ncols <- ncol(ocmat)+ncol(enmat)+1
    }
    
    colnames(params_mat) <- c('h', gname, paste('J',colnames(ocmat),sep='.'))
  }
  
  return(params_mat)
}


#' Convert SA output to model parameter objects
#'
#' Convert the output of \link{runSA} / \link{saEndpoint} into parameter
#' objects in (extended) maximum pairwise entropy model:
#' \code{h}, \code{J}, and optionally \code{g} when environmental parameter
#' matrix is specified.
#'
#' @param sa A list in the form \code{list(param_matrix, enmat)} where
#'   \code{param_matrix} contains estimated parameters.
#' @param env Optional numeric vector of environmental values (column names correspond
#'    to environmental factors).Default is NULL.
#'
#' @return A list with elements \code{he} (h), \code{je} (J), \code{ge} (g or NULL),
#'   and \code{hge} (combined field).
#'
#'
#' @useDynLib ELAplus, .registration=TRUE
#' @export
sa2params <- function(sa, env=NULL){
  sadim <- dim(sa[[1]])
  na <- rownames(sa[[1]])
  if(sadim[2] > sadim[1]+1){
      if(length(env)==0){env <- rep(0, sadim[2]-sadim[1]-1)} ## !!!
        he <- `names<-`(sa[[1]][, 1], na)
        je <- `dimnames<-`(sa[[1]][, ((sadim[2] - sadim[1]) + 1):sadim[2]],list(na, na))
        ge <- `dimnames<-`(as.matrix(sa[[1]][, 2:(sadim[2] - sadim[1])]),list(na, colnames(sa[[2]])))
        hge <- (he + ge %*% env)[,1]
        #hge <- t(t(envecs %*% t(ge)) + he)
  }else{
        he <- sa[[1]][, 1] %>% 'names<-'(na)
        je <- sa[[1]][, -1] %>% 'dimnames<-'(list(na,na))
        ge <- NULL
        hge <- he
  }
  return(list(he, je, ge, hge))}


#'testfittingSA: testrun of SA across given parameter sets using subset of ocmatrix
#'
#' Evaluate SA trajectories by validating the predicted occurrence matrix under
#' estimated parameters at each recorded time point.
#'
#' @param ocmat Numeric matrix (0/1). Observed community matrix.
#' @param sa_results List of SA parameter trajectories (e.g., from \code{runSA(getall=TRUE)}).
#' @param rep Integer. Number of replicates.
#' @param intv Integer. Recording interval used in fitting.
#' @param enmat Optional environmental matrix.
#'
#' @return A numeric matrix of validation scores (rows = replicates, columns = time points).
#'
#' @keywords internal
#' @noRd
testfittingSA <- function(ocmat,sa_results,rep,intv,enmat=NULL){
  nspecies = ncol(ocmat)
  num_timepoints <- nrow(sa_results[[1]]) / nspecies - 1
  timepoints <- sapply(0:num_timepoints, function(x) (x + 1) * intv)
  serials <- 1:rep
  res_serials <- matrix(0, ncol= (num_timepoints+1), nrow=rep,
                       dimnames=list(serials,timepoints))
  for(i in 1:length(sa_results)){
    rowidx <- 1
    timeidx <- 1
    SAparams <- sa_results[[i]]
    while(rowidx < nrow(SAparams)){
      
      params<-sa2params(list(SAparams[rowidx:(rowidx+nspecies-1),
                                                  1:(ncol(SAparams)-2)],enmat))
      res_serials[i,timeidx] <- validateSA(ocmat,params[[1]],params[[2]],
                                           g=params[[3]],enmat=enmat)
      rowidx <- rowidx + nspecies
      timeidx <- timeidx + 1
    }
  }
  return(res_serials)
}

#' Search best hyperparameters for SA fitting (lambda and weight decay) by cross-validation
#'
#' Repeating cross-validation over a grid of hyperparameters for
#' simulated annealing-based fitting, evaluates prediction/fitting errors on
#' held-out samples, and returns candidate best parameter/timepoint combinations.
#'
#' @param ocmat Numeric matrix (0/1). occurrence matrix of the communities.
#' @param enmat Optional environmental parameter matrix.
#' @param cvmode Cross-validation mode. One of `"10fold"`, `"5fold"`, or `"loo"`.
#'   `"loo"` performs leave-one-out cross-validation. Default is `"10fold"`.
#' @param nrepeat Number of repeated cross-validation runs for k-fold CV.
#'   Ignored when `cvmode = "loo"`. Default is `10`.
#' @param we Numeric vector. Candidate weight decay values.
#' @param totalit Integer. Total iterations for SA.
#' @param lmd Numeric vector. Candidate \code{lambda} values.
#' @param maxlrs Numeric vector. Maximum learning rate.
#' @param rep Integer. Number of replicates for \link{runSA}.
#' @param intv Integer. Recording interval.
#' @param threads Integer. Number of cores.
#' @param tol Numeric. Error tolerance for finding the best hyperparameters (default: 0.03)
#' @param runadamW Logical. Use AdamW.
#' @param sparse Logical. Use sparsity penalty.
#' @param reporting Logical. Print progress.
#' @param fastfitting Logical. If \code{TRUE}, returns early iterations within the range
#' specified by the parameter \code{tol} (default: 0.03) of best error.
#' @importFrom utils head
#'
#' @return A list with:
#' \describe{
#'   \item{best_candidates}{A named list of best hyperparameter settings and scores.}
#'   \item{allresults}{All validation results for each grid point.}
#' }
#'
#' @export

Findbp <- function(ocmat,enmat=NULL,cvmode=c("10fold","5fold","loo"),nrepeat=5,
                   we=c(0.001),maxlrs = c(0.005),
                   totalit=4000,lmd=c(0.0005,0.001,0.0025,0.005,0.0075),
                   rep=8,intv=100,threads=1,tol=0.03,
                   runadamW=TRUE,sparse=TRUE,fastfitting=TRUE,reporting=TRUE){
  min_points <- list()
  n_samples <- nrow(ocmat)
  if(n_samples < 2){
    stop("At least two samples are required.")
  }
  cvmode <- match.arg(cvmode)
  nrepeat <- as.integer(nrepeat)
  if(nrepeat < 1){
    stop("nrepeat must be >= 1.")
  }
  if(!runadamW){we <- c(0)}
  if(!sparse){lmd <- c(0)}
  ## ------------------------------------------------------------
  ## Make repeated cross-validation folds
  ## ------------------------------------------------------------
  split_list <- list()
  if(cvmode == "loo"){
    ## Leave-one-out is deterministic; nrepeat is ignored.
    for(i in seq_len(n_samples)){
      split_list[[length(split_list) + 1]] <- list(
        repeat_id = 1,
        fold_id = i,
        label = paste0("loo_", i),
        train = setdiff(seq_len(n_samples), i),
        test  = i
      )
    }
    split_prefix <- "loo"
  }else{
    k <- ifelse(cvmode == "10fold", 10, 5)
    if(n_samples < k){
      stop("Number of samples must be >= number of folds.")
    }
    for(r in seq_len(nrepeat)){
      ## Randomly assign samples to approximately equal-sized folds
      fold_id <- sample(rep(seq_len(k), length.out=n_samples))
      for(f in seq_len(k)){
        test_idx <- which(fold_id == f)
        train_idx <- setdiff(seq_len(n_samples), test_idx)
        split_list[[length(split_list) + 1]] <- list(
          repeat_id = r,
          fold_id = f,
          label = paste0("rep", r, "_fold", f),
          train = train_idx,
          test  = test_idx
        )
      }
    }
    split_prefix <- paste0(k, "fold")
  }
  n_splits <- length(split_list)
  ## ------------------------------------------------------------
  ## Make task grid:
  ## one task = one hyperparameter setting x one repeat/fold split
  ## ------------------------------------------------------------
  task_grid <- expand.grid(
    lambda = lmd,
    we = we,
    maxlr = maxlrs,
    split_id = seq_len(n_splits),
    stringsAsFactors = FALSE
  )
  n_tasks <- nrow(task_grid)
  ## ------------------------------------------------------------
  ## Required packages
  ## ------------------------------------------------------------
  #if(!requireNamespace("foreach", quietly=TRUE)){
  #  stop("Package 'foreach' is required.")
  #}
  #if(!requireNamespace("magrittr", quietly=TRUE)){
  #  stop("Package 'magrittr' is required.")
  #}
  #suppressPackageStartupMessages({
  #  library(foreach)
  #  library(magrittr)
  #})
  ## ------------------------------------------------------------
  ## Parallel setup
  ## ------------------------------------------------------------
  use_parallel <- 1
  if(use_parallel){
    if(!requireNamespace("doParallel", quietly=TRUE)){
      stop("Package 'doParallel' is required when threads > 1.")
    }
    n_workers <- min(threads, n_tasks)
    cluster <- parallel::makeCluster(n_workers)
    parallel::clusterCall(cluster, function(x) .libPaths(x), .libPaths())
    #parallel::clusterEvalQ(cluster, {
    #  suppressPackageStartupMessages({
    #    library(foreach)
    #    library(magrittr)
    #    library(ELAplus)
    #  })
    #  NULL
    #})
    doParallel::registerDoParallel(cluster)
    on.exit(parallel::stopCluster(cluster), add=TRUE)
  }
  ## ------------------------------------------------------------
  ## Run all tasks
  ## ------------------------------------------------------------
  run_one_task <- function(i){
    l <- task_grid$lambda[i]
    w <- task_grid$we[i]
    maxlr <- task_grid$maxlr[i]
    split_id <- task_grid$split_id[i]
    split <- split_list[[split_id]]
    train_idx <- split$train
    test_idx  <- split$test
    octraining <- ocmat[train_idx,,drop=FALSE]
    octest     <- ocmat[test_idx,,drop=FALSE]
    if(!is.null(enmat)){
      envtraining <- enmat[train_idx,,drop=FALSE]
      envtest     <- enmat[test_idx,,drop=FALSE]
    }else{
      envtraining <- NULL
      envtest     <- NULL
    }
    sa_results <- runSA(octraining, enmat=envtraining,
                        rep=rep, threads=1, we=w,maxlr=maxlr,
                        totalit=totalit, lambda=l, intv=intv,
                        runadamW=runadamW, sparse=sparse,
                        getall=TRUE, reporting=FALSE)
    current_results <- testfittingSA(octest, sa_results,
                                     rep=rep, intv=intv,
                                     enmat=envtest)
    ## Average across SA replicates within this split
    error_vec <- colMeans(current_results, na.rm=TRUE)
    list(
      lambda = l,
      we = w,
      maxlr = maxlr,
      split_id = split_id,
      split_label = split$label,
      repeat_id = split$repeat_id,
      fold_id = split$fold_id,
      error = error_vec
    )
  }
  if(reporting){cat('Start hyper-parameter search:\n')}
  if(reporting){cat(sprintf("There are %d tasks\n", n_tasks))}
  
  s <- proc.time()[3]
  
  task_results <- if(use_parallel){
    foreach(
      i = seq_len(n_tasks),
      .packages = c("foreach", "magrittr", "ELAplus"),
      .export = c("testfittingSA", "validateSA", "sa2params")
    ) %dopar% {
      run_one_task(i)
    }
  }else{
    foreach(
      i = seq_len(n_tasks),
      .packages = c("foreach", "magrittr", "ELAplus"),
      .export = c("testfittingSA", "validateSA", "sa2params")
    ) %do% {
      run_one_task(i)
    }
  }
  
  ## ------------------------------------------------------------
  ## Reconstruct allresults in the same format:
  ## allresults[[lambda_we]] = matrix of CV errors
  ## rows = repeat/fold splits, columns = timepoints
  ## ------------------------------------------------------------
  allresults <- list()
  min_points <- list()
  combo_names <- unique(paste(task_grid$lambda, task_grid$we, task_grid$maxlr, sep="_"))
  for(combo_name in combo_names){
    idx <- which(paste(task_grid$lambda, task_grid$we, task_grid$maxlr, sep="_") == combo_name)
    combo_results <- task_results[idx]
    current_results <- do.call(
      rbind,
      lapply(combo_results, function(x) x$error)
    )
    rownames(current_results) <- sapply(combo_results, function(x) x$split_label)
    allresults[[combo_name]] <- current_results
    mean_results <- apply(current_results, 2, mean, na.rm=TRUE)
    combo_parts <- strsplit(combo_name, "_", fixed=TRUE)[[1]]
    l <- combo_parts[1]
    w <- combo_parts[2]
    maxlr <- combo_parts[3]
    for (m in names(mean_results)){
      min_points[[paste(l, w, maxlr, as.integer(m), sep = "_")]] <- as.numeric(mean_results[m])
    }
  }
  if(reporting){cat(sprintf("Finished %d tasks\n", n_tasks))}
  if(reporting){cat(sprintf('SA: elapsed time %.2f sec\n\n', proc.time()[3]-s))}
  ## ------------------------------------------------------------
  ## Select best parameter/timepoint combinations
  ## ------------------------------------------------------------
  if(fastfitting){
    min_err <- min(unlist(min_points))
    min_points <- min_points[min_points < min_err*(1+tol)]
    sorted_min_points <- min_points[order(sapply(names(min_points),
                                                 function(x) {as.numeric(unlist(strsplit(x,split="_"))[[4]])}))]
  }else{
    sorted_min_points <- min_points[order(sapply(min_points, function(x) x[1]))]
  }
  return(list(head(sorted_min_points,min(length(sorted_min_points),5)),allresults))
}

#'plotSAtest: plot the results of SA fitting (Findbp output)
#' 
#' Plot mean validation scores over time for the results returned by
#' \link{Findbp}, including \code{95\%} confidence intervals.
#'
#' @param sa_results Output list from \link{Findbp}.
#' @param ylim Numeric length-2 vector (ymin, ymax). Default: \code{c("auto","auto")}.
#'
#' @return A \pkg{ggplot2} object.
#' @importFrom tidyr gather
#' @importFrom ggplot2 geom_line geom_ribbon scale_y_continuous
#' @details
#' This function uses \pkg{ggplot2} and tidyverse-style data manipulation.
#' Ensure required packages are installed.
#'@export
plotSAtest <- function(sa_results,ylim=c("auto","auto")){
  
  if(is.null(names(sa_results))){
    names(sa_results) <- paste0("set_", seq_along(sa_results))
  }
  
  summary_df <- data.frame()
  
  for(i in seq_along(sa_results)){
    tmp <- as.data.frame(sa_results[[i]])
    tmp$rowname <- rownames(tmp)
    rownames(tmp) <- NULL
    res_serials_melt <- tidyr::pivot_longer(
      tmp,
      cols = setdiff(names(tmp), "rowname"),
      names_to = "timepoint",
      values_to = "value"
    )
    
    res_serials_melt$timepoint <- as.numeric(res_serials_melt$timepoint)
    results_i <- res_serials_melt
    results_i$id <- names(sa_results)[i]
    summary_df <- rbind(summary_df,results_i)
  }
  
  summary_mean <- summary_df %>%
    dplyr::group_by(.data$id,.data$timepoint) %>%
    dplyr::summarise(n=sum(!is.na(.data$value)),
                     mean=mean(.data$value, na.rm=TRUE),
                     sd=stats::sd(.data$value, na.rm=TRUE),
                     .groups="drop") %>%
    dplyr::mutate(sem=ifelse(.data$n > 1, .data$sd / sqrt(.data$n), 0),
                  tcrit=ifelse(.data$n > 1, stats::qt(0.975, .data$n - 1), 0),
                  CI_lower=.data$mean - .data$tcrit * .data$sem,
                  CI_upper=.data$mean + .data$tcrit * .data$sem)
  
  if(ylim[[1]]=="auto"){
    ymin <- min(summary_mean$CI_lower, na.rm=TRUE)
  }else{
    ymin <- as.numeric(ylim[[1]])
  }
  
  if(ylim[[2]]=="auto"){
    ymax <- max(summary_mean$CI_upper, na.rm=TRUE)
  }else{
    ymax <- as.numeric(ylim[[2]])
  }
  
  ggplot(summary_mean, aes(x=.data$timepoint, y=.data$mean, color=.data$id)) +
    geom_line(aes(color=.data$id)) +
    geom_ribbon(aes(ymin=.data$CI_lower, ymax=.data$CI_upper, fill=.data$id),
                color="grey70", alpha=0.4) +
    scale_y_continuous(limits=c(ymin,ymax))
}


#'validateSA: to check the predictability of SA results by computing ocmatrix using
#'the estimated parameter set.
#'
#' @param ocmat Numeric matrix (0/1). Occurrence matrix.
#' @param h Numeric vector. Implicit environmental effects on species occurrence.
#' @param j Numeric matrix. Species-interaction matrix.
#' @param g Optional numeric matrix. Explicit environmental effects on species occurrence.
#' @param enmat Optional environmental parameters matrix.
#'
#' @return A numeric score (negative mean log of prediction error).
#'
#' @noRd
validateSA <- function(ocmat,h,j,g=NULL,enmat=NULL){
  if(is.null(enmat)){
    pp <- abs(1 - ocmat - 1 / (1 + exp(t(-h - t(ocmat %*% j)))))
  }
  else{
    pp <- abs(1 - ocmat - 1 / (1 + exp(-t(t(enmat %*% t(g)) + h) - ocmat %*% j)))
    #pp <- abs(1 - ocmat - 1 / (1 + exp(t(t(-hge - (ocmat %*% J))))))
  }
  return( -mean(log(pp[pp > 0])))
}

#' Parameter fitting with bootstrap sampling of datasets
#'
#' Perform \link{runSA} algorithm and estimate parameters using bootstrap resampled
#' data set (\code{ocmat} and \code{enmat}). 
#' The basic configuration is same as \link{runSA}.
#'
#' @param ocmat Numeric matrix (0/1). Rows are samples and columns are species.
#' @param enmat Optional numeric matrix. Environmental variables for each sample.
#'   Should be normalized/standardized beforehand.
#' @param rep Integer. Number of parallel runs (replicates) for optimization.
#' @param threads Integer. Number of cores used when running in parallel.
#' @param we Numeric. Weight decay parameter for AdamW.
#' @param totalit Integer. Total number of optimization iterations.
#' @param lambda Numeric. Sparsity penalty parameter (used when \code{sparse = TRUE}).
#' @param runadamW Logical. If \code{TRUE}, use AdamW-style optimization.
#' @param maxlr Numeric. Maximum learning rate.
#' @param sparse Logical. If \code{TRUE}, apply sparsity constraint.
#' @param bootstrap.it number of bootstrap trials (default: 1000).
#' @param reporting Logical. If \code{TRUE}, print progress messages.
#'
#' @return a list of SA parameters estimated with bootstrap resampling (length \code{bootstrap.it}).
#'
#' @useDynLib ELAplus, .registration=TRUE
#' @importFrom Rcpp sourceCpp
#' @import foreach
#' @export
bootstrap_SA <- function(ocmat, enmat=NULL, bootstrap.it=2000, rep=128, threads=1, we=0.001, 
                         totalit=1000, lambda=0.01, runadamW=TRUE,maxlr = 0.005, sparse=TRUE, reporting=TRUE){
  if(reporting){cat('Start parameter fitting:\n')}
  s <- proc.time()[3]
  #if(!is.matrix(ocmat)){ocmat <- as.matrix(ocmat)}
  if(!sparse){lambda <- 0}
  
  #if(!runadamW){we <- 0}
  
  bootstrap_res <- vector("list", bootstrap.it)
  n_samples <- nrow(ocmat)
  
  for (i in 1:bootstrap.it){
    # bootstrap sampling
    if(reporting){
      if(stats::runif(1) > max((bootstrap.it - 100)/bootstrap.it, 0))
      cat('.')
      }
    randidxs <- sample(n_samples, size = n_samples, replace = TRUE)
    ocrand <- ocmat[randidxs, , drop = FALSE]
    
    if(!is.null(enmat)){
      envrand <- enmat[randidxs, , drop = FALSE]
    }
    else{envrand=NULL}
    
    sa_results <- runSA(ocrand, enmat=envrand, rep=rep, threads=threads, we=we, 
                        totalit=totalit, lambda=lambda,intv=100, runadamW=runadamW, 
                        sparse=sparse,maxlr=maxlr,getall=FALSE, reporting=FALSE)
    
    bootstrap_res[[i]] <- sa_results[[1]]
  }
  if(reporting){cat(sprintf('\nSA: elapsed time %.2f sec\n\n', proc.time()[3]-s))}
  return(bootstrap_res)
}

#' Check variability in the stochastic approximation algorithm. 
#'
#' Perform bootstrap resampling of the estimated paramters using \link{runSA} with \code{getall=TRUE}.
#' This bootstrap sampling reflects the intrinsic variability of the stochastic approximation algorithm 
#' and does not correspond to standard confidence intervals based on resampled data.
#' Use \link{bootstrap_SA} to assess finite-sample estimation error. 
#'
#' @param fittingMat output of runSA function (need all run results in serials) 
#' @param ocmat Numeric matrix (0/1). Rows are samples and columns are species.
#' @param enmat Optional numeric matrix. Environmental variables for each sample.
#'   Should be normalized/standardized beforehand.
#' @param resamplesize total resampling size (default: length(fittingMat)), necessary when \code{type="quantile"}.
#' @param iterations number of bootstrap iterations (default: 1000).
#'
#' @return a list of SA parameters estimated with bootstrap resampling (length \code{iterations}).
#' @noRd
compute_SA_variance <- function(fittingMat,ocmat,enmat=NULL,resamplesize=NULL,iterations = 1000){
  if(is.null(resamplesize)){resamplesize<-length(fittingMat)}
  rand_sample_idx <- lapply(1:iterations, function(i) {sample(1:length(fittingMat),size=resamplesize,replace=TRUE)})
  rand_sa_params <- vector("list", iterations)
  for (i in 1:iterations){
    rand_serial <- fittingMat[rand_sample_idx[[i]]]
    rand_sa_params[[i]] <- saEndpoint(rand_serial,ocmat,enmat=enmat)
  }
  return(rand_sa_params)
}