#' Compute heat-bath update probabilities
#'
#' Internal helper for heat-bath sampling. Given a binary state matrix \code{y},
#' parameters \code{h} and \code{J}, and optionally environmental effects \code{g} with an
#' environmental matrix \code{enmat}, this function returns the per-entry update
#' probabilities used by the heat-bath step.
#'
#' @param y Numeric matrix of size \code{nproc x nn} (typically 0/1), current states.
#' @param h Numeric vector of length \code{nn}, implicit environmental coefficients.
#' @param J Numeric matrix of size \code{nn x nn}, interaction matrix.
#' @param g Optional numeric matrix of size \code{nn x ne}, explicit environmental coefficients.
#' @param enmat Optional numeric matrix of size \code{nproc x ne}, environmental values.
#'
#' @return A numeric matrix of update probabilities (same shape as \code{y} after
#'   internal transposition), with values in (0, 1).
#'
#' @noRd
logmon <- function(y, h, J, g = NULL, enmat = NULL) {
  si <- y %*% J
  if (is.null(enmat)) {
    term <- t(sweep(-si, 2, h, FUN = "-"))
  } else {
    gei <- enmat %*% t(g)
    hei <- sweep(gei, 2, h, FUN = "+")
    term <- t(-hei - si)
  }
  result <- 1 / (1 + exp(term))
  return(result)
}

#' Perform one heat-bath update step
#'
#' Updates each row (sample) of a binary state matrix by choosing one random
#' position per row and resampling it according to precomputed heat-bath
#' probabilities.
#' 
#' @param y Numeric matrix of size \code{nproc x nn} (typically 0/1), current states.
#' @param logmo Numeric matrix of probabilities compatible with \code{y}. Each row
#'   corresponds to a row in \code{y}; entries are used to sample 0/1 updates.
#'
#' @return A numeric matrix with the same dimensions as \code{y}, updated states.
#'
#' @export
OnestepHBS <- function(y, logmo) {
  m <- ncol(y)
  
  update_row <- function(row_y, row_logmo) {
    position <- sample.int(m, 1)
    ne <- ifelse(runif(1) < row_logmo[position], 1, 0)
    row_y[position] <- ne
    row_y
  }
  
  updated_y <- t(mapply(update_row, as.data.frame(t(y)), as.data.frame(t(logmo))))
  return(updated_y)
}

#' Heat-bath sampling of binary states
#'
#' Generates samples of binary state vectors using heat-bath updates under
#' an extended maximum pairwise entropy model parameterized 
#' by implicit environmental parameters \code{h} and interaction matrix \code{J}.
#' Optionally includes environmental effects \code{g}, in which case an environmental
#' matrix is generated internally.
#'
#' @param steps Integer. Number of heat-bath iterations to run.
#' @param nproc Integer. Number of parallel samples (rows) to maintain.
#' @param h Numeric vector of length \code{nn}, implicit environmental parameters.
#' @param J Numeric matrix of size \code{nn x nn}, interaction matrix.
#' @param g Optional numeric matrix of size \code{nn x ne}, environmental coefficients.
#'
#' @return A list with components:
#' \describe{
#'   \item{y}{Numeric matrix (\code{nproc x nn}) of sampled states (0/1).}
#'   \item{enmat}{Numeric matrix (\code{nproc x ne}) of environmental values, or \code{NULL}
#'     if \code{g} is \code{NULL}.}
#' }
#'
#' @examples
#' set.seed(1)
#' nn <- 5
#' pars <- hb.paramgen(nn)
#' res <- HeatBath(steps = 10, nproc = 8, h = pars[[1]], J = pars[[2]])
#' str(res)
#'
#' @importFrom stats runif
#' @export
HeatBath <- function(steps, nproc, h, J, g = NULL) {
  nn <- length(h)
  speciesname <- paste0("species.", seq_len(nn))
  samplename <- paste0("sample.", seq_len(nproc))
  
  y <- array(0, dim = c(nproc, nn), dimnames = list(samplename, speciesname))
  
  if (!is.null(g)) {
    ne <- dim(g)[2]
    factorname <- paste0("factor.", seq_len(ne))
    enmat <- matrix(runif(nproc * ne, min = 0, max = 1),
                    nrow = nproc, ncol = ne,
                    dimnames = list(samplename, factorname))
  } else {
    enmat <- NULL
  }
  
  for (i in seq_len(steps)) {
    logmo <- logmon(y, h = h, J = J, g = g, enmat = enmat)
    y <- OnestepHBS(y, t(logmo))
  }
  
  colnames(y) <- speciesname
  return(list(y = y, enmat = enmat))
}

#' Randomly generate model parameters for heat-bath sampling
#'
#' Generates random implicit environmental parameters \code{h} and a interaction matrix \code{J}. 
#' Optionally generates explicit environmental coefficients \code{g}. 
#' The \code{j_connectivity} parameter controls the fraction of non-zero off-diagonal interactions.
#'
#' @param nn Integer. Number of units (species, taxa, or members).
#' @param ne Integer. Number of environmental factors. If \code{ne > 0}, \code{g} is generated.
#' @param j_connectivity Numeric in [0, 1]. Fraction of possible pairwise
#'   interactions set to non-zero.
#' @importFrom stats setNames
#' @return A list containing at least:
#' \describe{
#'   \item{h}{Named numeric vector of length \code{nn}. Implicit environmental parameters}
#'   \item{J}{Named numeric matrix (\code{nn x nn}), symmetric. Interaction matrix}
#' }
#' If \code{ne > 0}, also includes:
#' \describe{
#'   \item{g}{Named numeric matrix (\code{nn x ne}). Explicit environmental parameters}
#' }
#'
#'
#' @export
hb.paramgen <- function(nn, ne = 0, j_connectivity = 1) {
  speciesname <- paste0("species.", seq_len(nn))
  
  h.act <- setNames(runif(nn, min = -2, max = 2), speciesname)
  
  rnp <- function() {
    tuples <- t(combn(nn, 2))
    rnp_list <- lapply(seq_len(nrow(tuples)), function(i) {
      if (runif(1) < 0.5) rev(tuples[i, ]) else tuples[i, ]
    })
    do.call(rbind, rnp_list)
  }
  
  generate_matrix <- function(rnp_list) {
    nonzerolist <- sample.int(nrow(rnp_list), nrow(rnp_list) * j_connectivity)
    ma <- matrix(0, nn, nn)
    for (i in seq_len(nrow(rnp_list))) {
      index <- rnp_list[i, ]
      ma[index[1], index[2]] <- ifelse(i %in% nonzerolist, runif(1, min = -2, max = 2), 0)
    }
    ma
  }
  
  ma <- generate_matrix(rnp())
  J <- `dimnames<-`(ma + t(ma), list(speciesname, speciesname))
  
  if (ne > 0) {
    gname <- paste0("g.", seq_len(ne))
    g <- matrix(runif(nn * ne, min = -2, max = 2), nrow = nn, ncol = ne)
    g <- `dimnames<-`(g, list(speciesname, gname))
    return(list(h = h.act, J = J, g = g))
  }
  
  return(list(h = h.act, J = J))
}

#' Gradient descent estimation of model parameters
#'
#' Estimates parameters \code{h} and \code{J} by iterative gradient updates to match
#' first and second moments of the observed binary data under an maximum
#' pairwise entropy model.
#'
#' @param data Numeric matrix (typically 0/1). Rows are samples and columns are species or taxa.
#' @param maxit Integer. Maximum number of gradient descent iterations.
#'
#' @return A list with elements:
#' \describe{
#'   \item{h.est}{Numeric vector of estimated implicit environmental parameters.}
#'   \item{j.est}{Numeric matrix of estimated interaction matrix.}
#' }
#'
#' @details
#' This function calls \code{Energy()} to compute energies for candidate states.
#' Ensure that \code{Energy} is available in the package namespace.
#'
#' @export
GradientDescent <- function(data, maxit) {
  units <- ncol(data)
  e <- 10^-8
  datalength <- nrow(data)
  tuples <- as.matrix(expand.grid(rep(list(0:1), units)))
  
  semp <- colMeans(data)
  ssemp <- t(data) %*% data / datalength * abs(diag(units) - 1)
  
  htt <- rep(0, units)
  jtt <- matrix(0, units, units)
  ttt <- 0
  
  while (ttt < maxit) {
    ht <- htt
    jt <- jtt
    
    e0 <- apply(tuples, 1, Energy, htt, jtt)
    en <- exp(-e0)
    prb <- en / sum(en)
    
    pt <- prb * tuples
    smod <- colSums(pt)
    ssmod <- (t(pt) %*% tuples) * abs(diag(units) - 1)
    
    dh <- 0.1 * ((semp + e) - (smod + e))
    dj <- 0.1 * ((ssemp + e) - (ssmod + e))
    
    htt <- ht + dh
    jtt <- jt + dj
    ttt <- ttt + 1
  }
  
  list(h.est = htt, j.est = jtt)
}
