#' Moment-corrected estimating equations for measurement error regression
#'
#' Constructs the moment-corrected estimating equations for a linear regression
#' model with averaged error-prone covariates and error-free covariates.
#'
#' @param y Numeric vector of length \code{n} containing the response.
#' @param Wbar Numeric matrix of dimension \code{n x p} containing averaged
#'   error-prone covariates.
#' @param Zmat Numeric matrix of dimension \code{n x q} containing error-free
#'   covariates (including an intercept if required).
#' @param b_in Numeric vector of length \code{p + q} containing the regression
#'   coefficients at which the estimating equations are evaluated.
#' @param estSigmaU_bar Numeric array of dimension \code{n x p x p} (or numeric
#'   vector in the univariate case) containing measurement error covariance
#'   estimates.
#'
#' @return Numeric vector of length \code{p + q} containing the stacked
#'   estimating equations.
#'
#'@keywords internal

estim_eq_corrected <- function(y, Wbar, Zmat, b_in, estSigmaU_bar) {

  n <- length(y)
  p <- ncol(Wbar)

  beta  <- b_in[1:p]
  gamma <- b_in[-(1:p)]

  r <- y - Wbar %*% beta - Zmat %*% gamma

  ## Score for Z
  SZ <- (t(Zmat) %*% r) / sqrt(n)

  ## Moment-corrected score for W

  # Univariate case
  if (is.null(dim(estSigmaU_bar)) || dim(estSigmaU_bar)[2] == 1) {

    meanSigmaU <- mean(estSigmaU_bar, na.rm = TRUE)

    if (is.na(meanSigmaU)) {
      # no replicate information → naive score
      return(SZ)
    } else {
      SW <- (t(Wbar) %*% r + n * meanSigmaU * beta) / sqrt(n)
    }

  } else {
    # Multivariate case
    meanSigmaU <- apply(estSigmaU_bar, c(2, 3), mean, na.rm = TRUE)

    if (any(is.na(meanSigmaU))) {
      # no replicate information → naive score
      return(SZ)
    } else {
      SW <- (t(Wbar) %*% r + n * meanSigmaU %*% beta) / sqrt(n)
    }
  }

  ## Stack equations
  rbind(SW, SZ)
}
