mmrm coverage - 97.45%

Files
Source

#' Methods for `mmrm_tmb` Objects
#'
#' @description `r lifecycle::badge("stable")`
#'
#' @param object (`mmrm_tmb`)\cr the fitted MMRM object.
#' @param x (`mmrm_tmb`)\cr same as `object`.
#' @param formula (`mmrm_tmb`)\cr same as `object`.
#' @param complete (`flag`)\cr whether to include potential non-estimable
#'   coefficients.
#' @param ... mostly not used;
#'   Exception is `model.matrix()` passing `...` to the default method.
#' @return Depends on the method, see Functions.
#'
#' @name mmrm_tmb_methods
#'
#' @seealso [`mmrm_methods`], [`mmrm_tidiers`] for additional methods.
#'
#' @examples
#' formula <- FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT | USUBJID)
#' object <- fit_mmrm(formula, fev_data, weights = rep(1, nrow(fev_data)))
NULL

#' @describeIn mmrm_tmb_methods obtains the estimated coefficients.
#' @importFrom stats coef
#' @exportS3Method
#' @examples
#' # Estimated coefficients:
#' coef(object)
coef.mmrm_tmb <- function(object, complete = TRUE, ...) {
  assert_flag(complete)
  nm <- if (complete) "beta_est_complete" else "beta_est"
  component(object, name = nm)
}

#' @describeIn mmrm_tmb_methods obtains the fitted values.
#' @importFrom stats fitted
#' @exportS3Method
#' @examples
#' # Fitted values:
#' fitted(object)
fitted.mmrm_tmb <- function(object, ...) {
  fitted_col <- component(object, "x_matrix") %*% component(object, "beta_est")
  fitted_col[, 1L, drop = TRUE]
}

#' @describeIn mmrm_tmb_methods predict conditional means for new data;
#'  optionally with standard errors and confidence or prediction intervals.
#'  Returns a vector of predictions if `se.fit == FALSE` and
#'  `interval == "none"`; otherwise it returns a data.frame with multiple
#'  columns and one row per input data row.
#'
#' @note If the response in the formula is a complicated expression, e.g. with multiple
#'   variables, and not all variables are provided in `newdata`, a warning is issued.
#'   In this case it is recommended to provide all variables in the `newdata` data set,
#'   and manually set the response variables to `NA` as needed for obtaining unconditional
#'   predictions e.g.
#'
#' @param newdata (`data.frame`)\cr optional new data, otherwise data from `object` is used.
#' @param se.fit (`flag`)\cr indicator if standard errors are required.
#' @param interval (`string`)\cr type of interval calculation. Can be abbreviated.
#' @param level (`number`)\cr tolerance/confidence level.
#' @param nsim (`count`)\cr number of simulations to use.
#' @param conditional (`flag`)\cr indicator if the prediction is conditional on the observation or not.
#'
#' @importFrom stats predict
#' @exportS3Method
#'
#' @examples
#' predict(object, newdata = fev_data)
predict.mmrm_tmb <- function(
  object,
  newdata,
  se.fit = FALSE, # nolint
  interval = c("none", "confidence", "prediction"),
  level = 0.95,
  nsim = 1000L,
  conditional = FALSE,
  ...
) {
  if (missing(newdata)) {
    newdata <- object$data
  }
  assert_data_frame(newdata)
  orig_row_names <- row.names(newdata)
  assert_flag(se.fit)
  assert_number(level, lower = 0, upper = 1)
  assert_count(nsim, positive = TRUE)
  assert_flag(conditional)
  interval <- match.arg(interval)
  formula_parts <- object$formula_parts
  # Check if the response is a simple variable name vs a complex expression
  is_simple_response_var <- is(object$formula_parts$formula[[2]], "name")
  missing_resp_vars <- setdiff(formula_parts$response_var, names(newdata))
  has_missing_resp_vars <- length(missing_resp_vars) > 0
  if (!is_simple_response_var && has_missing_resp_vars) {
    warning(
      "complicated response expression in formula and not all variables in left-hand side are provided, ",
      "using NA to replace the missing variables; ",
      "this could potentially lead to unexpected behavior."
    )
  }
  # Make sure we have at least completely NA filled response variables in newdata.
  for (var in missing_resp_vars) {
    newdata[[var]] <- NA_real_
  }
  if (any(object$tmb_data$x_cols_aliased)) {
    warning(
      "In fitted object there are co-linear variables and therefore dropped terms, ",
      "and this could lead to incorrect prediction on new data."
    )
  }
  colnames <- names(Filter(isFALSE, object$tmb_data$x_cols_aliased))
  if (!conditional && interval %in% c("none", "confidence")) {
    # model.matrix always return a complete matrix (no NA allowed)
    x_mat <- stats::model.matrix(object, data = newdata, use_response = FALSE)[,
      colnames,
      drop = FALSE
    ]
    x_mat_full <- matrix(
      NA,
      nrow = nrow(newdata),
      ncol = ncol(x_mat),
      dimnames = list(row.names(newdata), colnames(x_mat))
    )
    x_mat_full[row.names(x_mat), ] <- x_mat
    predictions <- (x_mat_full %*% component(object, "beta_est"))[, 1]
    predictions_raw <- stats::setNames(
      rep(NA_real_, nrow(newdata)),
      row.names(newdata)
    )
    predictions_raw[names(predictions)] <- predictions
    if (identical(interval, "none")) {
      return(predictions_raw)
    }
    se <- switch(
      interval,
      # can be NA if there are aliased cols
      "confidence" = diag(
        x_mat_full %*% component(object, "beta_vcov") %*% t(x_mat_full)
      ),
      "none" = NA_real_
    )
    res <- cbind(
      fit = predictions,
      se = se,
      lwr = predictions - stats::qnorm(1 - level / 2) * se,
      upr = predictions + stats::qnorm(1 - level / 2) * se
    )
    if (!se.fit) {
      res <- res[, setdiff(colnames(res), "se")]
    }
    res_raw <- matrix(
      NA_real_,
      ncol = ncol(res),
      nrow = nrow(newdata),
      dimnames = list(row.names(newdata), colnames(res))
    )
    res_raw[row.names(res), ] <- res
    return(res_raw)
  }
  tmb_data <- h_mmrm_tmb_data(
    formula_parts,
    newdata,
    weights = rep(1, nrow(newdata)),
    reml = TRUE,
    singular = "keep",
    drop_visit_levels = FALSE,
    allow_na_response = TRUE,
    drop_levels = FALSE,
    xlev = component(object, "xlev"),
    contrasts = component(object, "contrasts")
  )
  tmb_data$x_matrix <- tmb_data$x_matrix[, colnames, drop = FALSE]
  predictions <- h_get_prediction(
    tmb_data,
    object$theta_est,
    object$beta_est,
    component(object, "beta_vcov")
  )$prediction
  res <- cbind(fit = rep(NA_real_, nrow(newdata)))
  new_order <- match(row.names(tmb_data$full_frame), orig_row_names)
  res[new_order, "fit"] <- predictions[, "fit"]
  se <- switch(
    interval,
    "confidence" = sqrt(predictions[, "conf_var"]),
    "prediction" = sqrt(h_get_prediction_variance(object, nsim, tmb_data)),
    "none" = NULL
  )
  if (interval != "none") {
    res <- cbind(
      res,
      se = NA_real_
    )
    res[new_order, "se"] <- se
    alpha <- 1 - level
    z <- stats::qnorm(1 - alpha / 2) * res[, "se"]
    res <- cbind(
      res,
      lwr = res[, "fit"] - z,
      upr = res[, "fit"] + z
    )
    if (!se.fit) {
      res <- res[, setdiff(colnames(res), "se")]
    }
  }
  # Use original names.
  row.names(res) <- orig_row_names
  if (ncol(res) == 1) {
    res <- res[, "fit"]
  }
  return(res)
}

#' Get Prediction
#'
#' @description Get predictions with given `data`, `theta`, `beta`, `beta_vcov`.
#'
#' @details See `predict` function in `predict.cpp` which is called internally.
#'
#' @param tmb_data (`mmrm_tmb_data`)\cr object.
#' @param theta (`numeric`)\cr theta value.
#' @param beta (`numeric`)\cr beta value.
#' @param beta_vcov (`matrix`)\cr beta_vcov matrix.
#'
#' @return List with:
#' - `prediction`: Matrix with columns `fit`, `conf_var`, and `var`.
#' - `covariance`: List with subject specific covariance matrices.
#' - `index`: List of zero-based subject indices.
#'
#' @keywords internal
h_get_prediction <- function(tmb_data, theta, beta, beta_vcov) {
  assert_class(tmb_data, "mmrm_tmb_data")
  assert_numeric(theta)
  n_beta <- ncol(tmb_data$x_matrix)
  assert_numeric(beta, finite = TRUE, any.missing = FALSE, len = n_beta)
  assert_matrix(
    beta_vcov,
    mode = "numeric",
    any.missing = FALSE,
    nrows = n_beta,
    ncols = n_beta
  )
  .Call(`_mmrm_predict`, PACKAGE = "mmrm", tmb_data, theta, beta, beta_vcov)
}

#' Get Prediction Variance
#'
#' @description Get prediction variance with given fit, `tmb_data` with the Monte Carlo sampling method.
#'
#' @param object (`mmrm_tmb`)\cr the fitted MMRM.
#' @param nsim (`count`)\cr number of samples.
#' @param tmb_data (`mmrm_tmb_data`)\cr object.
#'
#' @keywords internal
h_get_prediction_variance <- function(object, nsim, tmb_data) {
  assert_class(object, "mmrm_tmb")
  assert_class(tmb_data, "mmrm_tmb_data")
  assert_count(nsim, positive = TRUE)
  theta_chol <- chol(eval(object$theta_vcov))
  n_theta <- length(object$theta_est)
  res <- replicate(nsim, {
    z <- stats::rnorm(n = n_theta)
    theta_sample <- object$theta_est + z %*% theta_chol
    cond_beta_results <- object$tmb_object$report(theta_sample)
    beta_mean <- cond_beta_results$beta
    beta_cov <- cond_beta_results$beta_vcov
    h_get_prediction(tmb_data, theta_sample, beta_mean, beta_cov)$prediction
  })
  mean_of_var <- rowMeans(res[, "var", ])
  var_of_mean <- apply(res[, "fit", ], 1, stats::var)
  mean_of_var + var_of_mean
}

#' @describeIn mmrm_tmb_methods obtains the model frame.
#' @param data (`data.frame`)\cr object in which to construct the frame.
#' @param include (`character`)\cr names of variable types to include in the model frame creation.
#'   Must be `NULL` or one or more of `c("subject_var", "visit_var", "group_var", "response_var")`.
#' @param exclude (`character`)\cr names of variable types to exclude after the model frame creation.
#'   Same choices as for `include`.
#' @param full (`flag`)\cr indicator whether to return full model frame (deprecated).
#' @param na.action (`string`)\cr na action.
#' @importFrom stats model.frame
#' @exportS3Method
#'
#' @details
#' `include` argument controls the variables the model frame will include upon creation.
#' Possible options are "response_var", "subject_var", "visit_var" and "group_var", representing the
#' response variable, subject variable, visit variable or group variable.
#' By default all four variable types are included, which means that the original model frame will be used.
#' If only a subset of variable types is requested, the model frame will be constructed freshly, and
#' if there are more complete rows (without `NA`s) in the required variables than in the original model frame,
#' those rows will be included as well.
#'
#' The `exclude` argument can be used to exclude specific variable types after the model frame creation.
#' By default no variable types are excluded. The `exclude` argument is useful when the original
#' model frame should be used as basis (with same incomplete rows omitted), but some variable types should be removed
#' afterwards.
#'
#' `character` values in new data will always be factorized according to the data in the fit
#' to avoid mismatched in levels or issues in `model.matrix`.
#'
#' @examples
#' # Model frame:
#' model.frame(object)
#' model.frame(object, include = "subject_var")
model.frame.mmrm_tmb <- function(
  formula,
  data,
  include = c("subject_var", "visit_var", "group_var", "response_var"),
  exclude = NULL,
  full,
  na.action = "na.omit", # nolint
  ...
) {
  # nolint
  # Construct updated formula and data arguments.
  lst_formula_and_data <-
    h_construct_model_frame_inputs(
      formula = formula,
      data = data,
      include = include,
      full = full
    )
  # Only if include is default (full) and also data is missing, and also na.action is na.omit we will
  # use the model frame from the tmb_data.
  include_choice <- c("subject_var", "visit_var", "group_var", "response_var")
  if (
    missing(data) &&
      setequal(include, include_choice) &&
      identical(h_get_na_action(na.action), stats::na.omit)
  ) {
    ret <- formula$tmb_data$full_frame
    # Remove weights column.
    ret[, "(weights)"] <- NULL
    ret
  } else {
    # Construct data frame.
    ret <-
      stats::model.frame(
        formula = lst_formula_and_data$formula,
        data = lst_formula_and_data$data,
        na.action = na.action,
        xlev = component(formula, "xlev"),
        ...
      )
  }
  # Exclude requested variable types.
  if (!is.null(exclude)) {
    assert_subset(
      exclude,
      c("subject_var", "visit_var", "group_var", "response_var")
    )
    vars_to_exclude <- unlist(formula$formula_parts[exclude])
    # Need to preserve the attributes and drop terms according
    # to the excluded variables. Otherwise downstream emmeans
    # will not work correctly.
    ret_attrs <- attributes(ret)
    ret_terms <- ret_attrs$terms
    term_labels <- labels(ret_terms)
    drop_inds <- which(term_labels %in% vars_to_exclude)
    ret_attrs$terms <- stats::drop.terms(
      ret_terms,
      drop_inds,
      keep.response = !("response_var" %in% exclude)
    )
    ret <- ret[, setdiff(colnames(ret), vars_to_exclude), drop = FALSE]
    attributes(ret) <- c(
      attributes(ret),
      ret_attrs[c("terms", "na.action")]
    )
  }
  ret
}


#' Construction of Model Frame Formula and Data Inputs
#'
#' @description
#' Input formulas are converted from `mmrm`-style to a style compatible
#' with default [stats::model.frame()] and [stats::model.matrix()] methods.
#'
#' The full formula is returned so we can construct, for example, the
#' `model.frame()` including all columns as well as the requested subset.
#' The full set is used to identify rows to include in the reduced model frame.
#'
#' @param formula (`mmrm`)\cr fit object.
#' @param data (`data.frame`)\cr optional data frame will be
#'   passed to `model.frame()` or `model.matrix()`.
#' @param include (`character`)\cr specification of variables to include.
#' @param full (`flag`)\cr indicator whether to return full model frame (deprecated).
#'
#' @return A named list with two elements:
#'
#' - `"formula"`: the formula including the columns requested in the `include=` argument.
#' - `"data"`: a data frame including all columns needed in the formula.
#'
#' @keywords internal
h_construct_model_frame_inputs <- function(
  formula,
  data,
  include,
  include_choice = c("subject_var", "visit_var", "group_var", "response_var"),
  full
) {
  if (!missing(full) && identical(full, TRUE)) {
    lifecycle::deprecate_warn("0.3", "model.frame.mmrm_tmb(full)")
    include <- include_choice
  }

  assert_class(formula, classes = "mmrm_tmb")
  assert_subset(include, include_choice)
  if (missing(data)) {
    data <- formula$data
  }
  assert_data_frame(data)

  drop_response <- !"response_var" %in% include
  add_vars <- unlist(formula$formula_parts[setdiff(include, "response_var")])
  new_formula <- h_add_terms(
    formula$formula_parts$model_formula,
    add_vars,
    drop_response
  )

  drop_response_full <- !"response_var" %in% include_choice
  add_vars_full <- unlist(formula$formula_parts[include_choice])
  new_formula_full <-
    h_add_terms(
      formula$formula_parts$model_formula,
      add_vars_full,
      drop_response_full
    )

  # Update data based on the columns in the full formula return.
  all_vars <- all.vars(new_formula_full)
  all_vars <- intersect(colnames(data), all_vars) # Because some variables could be from the environment.
  data <- data[, all_vars, drop = FALSE]

  # Return list with updated formula, data.
  list(
    formula = new_formula,
    data = data
  )
}

#' @describeIn mmrm_tmb_methods obtains the model matrix.
#' @exportS3Method
#' @param use_response (`flag`)\cr whether to use the response for complete rows.
#' @param contrasts.arg (`list`)\cr contrasts to be used.
#'
#' @examples
#' # Model matrix:
#' model.matrix(object)
model.matrix.mmrm_tmb <- function(
  object,
  data,
  use_response = TRUE,
  # nolint start
  contrasts.arg = attr(object$tmb_data$x_matrix, "contrasts"),
  # nolint end
  ...
) {
  # Always return the utilized model matrix if data not provided.
  if (missing(data)) {
    return(object$tmb_data$x_matrix)
  }
  stats::model.matrix(
    h_add_terms(
      object$formula_parts$model_formula,
      NULL,
      drop_response = !use_response
    ),
    data = data,
    contrasts.arg = contrasts.arg,
    xlev = component(object, "xlev"),
    ...
  )
}

#' @describeIn mmrm_tmb_methods obtains the terms object.
#' @importFrom stats model.frame
#' @exportS3Method
#'
#' @examples
#' # terms:
#' terms(object)
#' terms(object, include = "subject_var")
terms.mmrm_tmb <- function(x, include = "response_var", ...) {
  # nolint
  # Construct updated formula and data arguments.
  lst_formula_and_data <-
    h_construct_model_frame_inputs(
      formula = x,
      include = include
    )

  # Use formula method for `terms()` to construct the mmrm terms object.
  stats::terms(
    x = lst_formula_and_data$formula,
    data = lst_formula_and_data$data
  )
}


#' @describeIn mmrm_tmb_methods obtains the attained log likelihood value.
#'   Includes as attributes the number of variance parameters `n_param`, number
#'   of estimated coefficients `n_coef`, degrees of freedom `df`, and number of
#'   subjects `nobs`. The `nobs` attribute is so named so that if this
#'   function's results are passed to `stats::BIC()`, the BIC value will be
#'   calculated correctly. Resulting value is of class
#'   [`logLik`][stats::logLik].
#' @importFrom stats logLik
#' @exportS3Method
#' @examples
#' # Log likelihood given the estimated parameters:
#' logLik(object)
logLik.mmrm_tmb <- function(object, ...) {
  out <- -component(object, "neg_log_lik")

  # Number of variance parameters.
  n_param <- component(object, "n_theta")

  # Number of estimated coefficients.
  n_coef <- length(coef(object, complete = FALSE))

  # Number of degrees of freedom. Note that the number of coefficients is added
  # only if the fit was estimated using ML rather than REML.
  df <- n_param + n_coef * !component(object, "reml")

  n_subjects <- component(object, "n_subjects")

  structure(
    out,
    n_param = n_param,
    n_coef = n_coef,
    df = df,
    nobs = n_subjects,
    class = "logLik"
  )
}

#' @describeIn mmrm_tmb_methods obtains the used formula.
#' @importFrom stats formula
#' @exportS3Method
#' @examples
#' # Formula which was used:
#' formula(object)
formula.mmrm_tmb <- function(x, ...) {
  x$formula_parts$formula
}

#' @describeIn mmrm_tmb_methods obtains the variance-covariance matrix estimate
#'   for the coefficients.
#' @importFrom stats vcov
#' @exportS3Method
#' @examples
#' # Variance-covariance matrix estimate for coefficients:
#' vcov(object)
vcov.mmrm_tmb <- function(object, complete = TRUE, ...) {
  assert_flag(complete)
  nm <- if (complete) "beta_vcov_complete" else "beta_vcov"
  component(object, name = nm)
}

#' @describeIn mmrm_tmb_methods obtains the variance-covariance matrix estimate
#'   for the residuals.
#' @param sigma cannot be used (this parameter does not exist in MMRM).
#' @importFrom nlme VarCorr
#' @export VarCorr
#' @aliases VarCorr
#' @exportS3Method
#' @examples
#' # Variance-covariance matrix estimate for residuals:
#' VarCorr(object)
VarCorr.mmrm_tmb <- function(x, sigma = NA, ...) {
  # nolint
  assert_scalar_na(sigma)

  component(x, name = "varcor")
}

#' @describeIn mmrm_tmb_methods obtains the deviance, which is defined here
#'   as twice the negative log likelihood, which can either be integrated
#'   over the coefficients for REML fits or the usual one for ML fits.
#' @importFrom stats deviance
#' @exportS3Method
#' @examples
#' # REML criterion (twice the negative log likelihood):
#' deviance(object)
deviance.mmrm_tmb <- function(object, ...) {
  2 * component(object, "neg_log_lik")
}

#' @describeIn mmrm_tmb_methods obtains the Akaike Information Criterion,
#'   where the degrees of freedom are the number of variance parameters (`n_theta`).
#'   If `corrected`, then this is multiplied with `m / (m - n_theta - 1)` where
#'   `m` is the number of observations minus the number of coefficients, or
#'   `n_theta + 2` if it is smaller than that \insertCite{hurvich1989regression,burnham1998practical}{mmrm}.
#' @param corrected (`flag`)\cr whether corrected AIC should be calculated.
#' @param k (`number`)\cr the penalty per parameter to be used; default `k = 2`
#'   is the classical AIC.
#' @importFrom stats AIC
#' @exportS3Method
#' @examples
#' # AIC:
#' AIC(object)
#' AIC(object, corrected = TRUE)
#' @references
#' - \insertRef{hurvich1989regression}{mmrm}
#' - \insertRef{burnham1998practical}{mmrm}
AIC.mmrm_tmb <- function(object, corrected = FALSE, ..., k = 2) {
  # nolint
  assert_flag(corrected)
  assert_number(k, lower = 1)

  n_theta <- length(component(object, "theta_est"))
  df <- if (!corrected) {
    n_theta
  } else {
    n_obs <- length(component(object, "y_vector"))
    n_beta <- length(component(object, "beta_est"))
    m <- max(n_theta + 2, n_obs - n_beta)
    n_theta * (m / (m - n_theta - 1))
  }

  2 * component(object, "neg_log_lik") + k * df
}

#' @describeIn mmrm_tmb_methods obtains the Bayesian Information Criterion,
#'   which is using the natural logarithm of the number of subjects for the
#'   penalty parameter `k`.
#' @importFrom stats BIC
#' @exportS3Method
#' @examples
#' # BIC:
#' BIC(object)
BIC.mmrm_tmb <- function(object, ...) {
  # nolint
  k <- log(component(object, "n_subjects"))
  AIC(object, corrected = FALSE, k = k)
}


#' @describeIn mmrm_tmb_methods prints the object.
#' @exportS3Method
print.mmrm_tmb <- function(x, ...) {
  cat("mmrm fit\n\n")

  h_print_call(
    component(x, "call"),
    component(x, "n_obs"),
    component(x, "n_subjects"),
    component(x, "n_timepoints")
  )
  h_print_cov(
    component(x, "cov_type"),
    component(x, "n_theta"),
    component(x, "n_groups")
  )

  cat("Inference:   ")
  cat(ifelse(component(x, "reml"), "REML", "ML"))
  cat("\n")
  cat("Deviance:    ")
  cat(deviance(x))

  cat("\n\nCoefficients: ")
  n_singular_coefs <- sum(component(x, "beta_aliased"))
  if (n_singular_coefs > 0) {
    cat(
      "(",
      n_singular_coefs,
      " not defined because of singularities)",
      sep = ""
    )
  }
  cat("\n")
  print(coef(x, complete = TRUE))

  cat("\nModel Inference Optimization:")

  cat(ifelse(
    component(x, "convergence") == 0,
    "\nConverged",
    "\nFailed to converge"
  ))
  cat(
    " with code",
    component(x, "convergence"),
    "and message:",
    if (is.null(component(x, "conv_message"))) {
      "No message provided."
    } else {
      tolower(component(x, "conv_message"))
    }
  )
  cat("\n")
  invisible(x)
}


#' @describeIn mmrm_tmb_methods to obtain residuals - either unscaled ('response'), 'pearson' or 'normalized'.
#' @param type (`string`)\cr unscaled (`response`), `pearson` or `normalized`. Default is `response`,
#' and this is the only type available for use with models with a spatial covariance structure.
#' @importFrom stats residuals
#' @exportS3Method
#' @examples
#' # residuals:
#' residuals(object, type = "response")
#' residuals(object, type = "pearson")
#' residuals(object, type = "normalized")
#' @references
#' - \insertRef{galecki2013linear}{mmrm}
residuals.mmrm_tmb <- function(
  object,
  type = c("response", "pearson", "normalized"),
  ...
) {
  type <- match.arg(type)
  switch(
    type,
    "response" = h_residuals_response(object),
    "pearson" = h_residuals_pearson(object),
    "normalized" = h_residuals_normalized(object)
  )
}
#' Calculate Pearson Residuals
#'
#' This is used by [residuals.mmrm_tmb()] to calculate Pearson residuals.
#'
#' @param object (`mmrm_tmb`)\cr the fitted MMRM.
#'
#' @return Vector of residuals.
#'
#' @keywords internal
h_residuals_pearson <- function(object) {
  assert_class(object, "mmrm_tmb")
  h_residuals_response(object) * object$tmb_object$report()$diag_cov_inv_sqrt
}

#' Calculate normalized residuals
#'
#' This is used by [residuals.mmrm_tmb()] to calculate normalized / scaled residuals.
#'
#' @param object (`mmrm_tmb`)\cr the fitted MMRM.
#'
#' @return Vector of residuals
#'
#' @keywords internal
h_residuals_normalized <- function(object) {
  assert_class(object, "mmrm_tmb")
  object$tmb_object$report()$epsilonTilde
}
#' Calculate response residuals.
#'
#' This is used by [residuals.mmrm_tmb()] to calculate response residuals.
#'
#' @param object (`mmrm_tmb`)\cr the fitted MMRM.
#'
#' @return Vector of residuals
#'
#' @keywords internal
h_residuals_response <- function(object) {
  assert_class(object, "mmrm_tmb")
  component(object, "y_vector") - unname(fitted(object))
}

#' @describeIn mmrm_tmb_methods simulate responses from a fitted model according
#'   to the simulation `method`, returning a `data.frame` of dimension `[n, m]`
#'   where n is the number of rows in `newdata`,
#'   and m is the number `nsim` of simulated responses.
#'
#' @param seed unused argument from [stats::simulate()].
#' @param method (`string`)\cr simulation method to use. If "conditional",
#'   simulated values are sampled given the estimated covariance matrix of `object`.
#'   If "marginal", the variance of the estimated covariance matrix is taken into account.
#'
#' @importFrom stats simulate
#' @exportS3Method
simulate.mmrm_tmb <- function(
  object,
  nsim = 1,
  seed = NULL,
  newdata,
  ...,
  method = c("conditional", "marginal")
) {
  assert_count(nsim, positive = TRUE)
  assert_null(seed)
  if (missing(newdata)) {
    newdata <- object$data
  }
  assert_data_frame(newdata)
  method <- match.arg(method)

  tmb_data <- h_mmrm_tmb_data(
    object$formula_parts,
    newdata,
    weights = rep(1, nrow(newdata)),
    reml = TRUE,
    singular = "keep",
    drop_visit_levels = FALSE,
    allow_na_response = TRUE,
    drop_levels = FALSE,
    xlev = component(object, "xlev"),
    contrasts = component(object, "contrasts")
  )
  ret <- if (method == "conditional") {
    predict_res <- h_get_prediction(
      tmb_data,
      object$theta_est,
      object$beta_est,
      object$beta_vcov
    )
    as.data.frame(h_get_sim_per_subj(predict_res, tmb_data$n_subjects, nsim))
  } else if (method == "marginal") {
    theta_chol <- t(chol(object$theta_vcov))
    n_theta <- length(object$theta_est)
    as.data.frame(
      sapply(seq_len(nsim), function(x) {
        newtheta <- object$theta_est +
          theta_chol %*% matrix(stats::rnorm(n_theta), ncol = 1)
        # Recalculate betas with sampled thetas.
        hold <- object$tmb_object$report(newtheta)
        # Resample betas given new beta distribution.
        # We first solve L^\top w = D^{-1/2}z_{sample}:
        w_sample <- backsolve(
          r = hold$XtWX_L,
          x = stats::rnorm(length(hold$beta)) / sqrt(hold$XtWX_D),
          upper.tri = FALSE,
          transpose = TRUE
        )
        # Then we add the mean vector, the beta estimate.
        beta_sample <- hold$beta + w_sample
        predict_res <- h_get_prediction(
          tmb_data,
          newtheta,
          beta_sample,
          hold$beta_vcov
        )
        h_get_sim_per_subj(predict_res, tmb_data$n_subjects, 1L)
      })
    )
  }
  orig_row_names <- row.names(newdata)
  new_order <- match(orig_row_names, row.names(tmb_data$full_frame))
  ret[new_order, , drop = FALSE]
}

#' Get simulated values by patient.
#'
#' @param predict_res (`list`)\cr from [h_get_prediction()].
#' @param nsub (`count`)\cr number of subjects.
#' @param nsim (`count`)\cr number of values to simulate.
#'
#' @keywords internal
h_get_sim_per_subj <- function(predict_res, nsub, nsim) {
  assert_list(predict_res)
  assert_count(nsub, positive = TRUE)
  assert_count(nsim, positive = TRUE)

  ret <- matrix(
    predict_res$prediction[, "fit"],
    ncol = nsim,
    nrow = nrow(predict_res$prediction)
  )
  for (i in seq_len(nsub)) {
    # Skip subjects which are not included in predict_res.
    if (length(predict_res$index[[i]]) > 0) {
      # Obtain indices of data.frame belonging to subject i
      # (increment by 1, since indices from cpp are 0-order).
      inds <- predict_res$index[[i]] + 1
      obs <- length(inds)

      # Get relevant covariance matrix for subject i.
      covmat_i <- predict_res$covariance[[i]]
      theta_chol <- t(chol(covmat_i))

      # Simulate epsilon from covariance matrix.
      mus <- ret[inds, , drop = FALSE]
      epsilons <- theta_chol %*% matrix(stats::rnorm(nsim * obs), ncol = nsim)
      ret[inds, ] <- mus + epsilons
    }
  }

  ret
}

#' Processing the Formula for `TMB` Fit
#'
#' @param formula (`formula`)\cr Original formula.
#' @param covariance (`cov_struct`)\cr A covariance structure from which
#'   additional formula parts should be added.
#'
#' @return List of class `mmrm_tmb_formula_parts` with elements:
#'
#' - `formula`: the original input.
#' - `model_formula`: `formula` with the covariance term is removed.
#' - `model_formula`: `formula` with the covariance term removed.
#' - `full_formula`: same as `model_formula` but includes the covariance
#'   structure's subject, visit and (optionally) group variables.
#' - `cov_type`: `string` with covariance term type (e.g. `"us"`).
#' - `is_spatial`: `flag` indicator of whether the covariance structure is
#'   spatial
#' - `visit_var`: `character` with the visit variable name.
#' - `subject_var`: `string` with the subject variable name.
#' - `group_var`: `string` with the group variable name. If no group specified,
#'   this element is `NULL`.
#' - `model_var`: `character` with the variables names of the formula, except `subject_var`.
#' - `response_var`: `string` with the response variable name.
#'
#' @keywords internal
h_mmrm_tmb_formula_parts <- function(
  formula,
  covariance = as.cov_struct(formula, warn_partial = FALSE)
) {
  assert_formula(formula)
  assert_true(identical(length(formula), 3L))

  model_formula <- h_drop_covariance_terms(formula)

  structure(
    list(
      formula = formula,
      model_formula = model_formula,
      full_formula = h_add_covariance_terms(model_formula, covariance),
      cov_type = tmb_cov_type(covariance),
      is_spatial = covariance$type == "sp_exp",
      visit_var = covariance$visits,
      subject_var = covariance$subject,
      group_var = if (length(covariance$group) < 1) NULL else covariance$group,
      model_var = setdiff(all.vars(formula[[3]]), covariance$subject),
      response_var = all.vars(formula[[2]])
    ),
    class = "mmrm_tmb_formula_parts"
  )
}

#' Data for `TMB` Fit
#'
#' @param formula_parts (`mmrm_tmb_formula_parts`)\cr list with formula parts
#'   from [h_mmrm_tmb_formula_parts()].
#' @param data (`data.frame`)\cr which contains variables used in `formula_parts`.
#' @param weights (`vector`)\cr weights to be used in the fitting process.
#' @param reml (`flag`)\cr whether restricted maximum likelihood (REML) estimation is used,
#'   otherwise maximum likelihood (ML) is used.
#' @param singular (`string`)\cr choices of method deal with rank-deficient matrices. "error" to
#'   stop the function return the error, "drop" to drop these columns, and "keep" to keep all the columns.
#' @param drop_visit_levels (`flag`)\cr whether to drop levels for visit variable, if visit variable is a factor.
#' @param allow_na_response (`flag`)\cr whether NA in response is allowed.
#' @param drop_levels (`flag`)\cr whether drop levels for covariates. If not dropped could lead to singular matrix.
#'
#' @return List of class `mmrm_tmb_data` with elements:
#' - `full_frame`: `data.frame` with `n` rows containing all variables needed in the model.
#' - `data`: `data.frame` of input dataset.
#' - `x_matrix`: `matrix` with `n` rows and `p` columns specifying the overall design matrix.
#' - `x_cols_aliased`: `logical` with potentially more than `p` elements indicating which
#'      columns in the original design matrix have been left out to obtain a full rank
#'      `x_matrix`.
#' - `y_vector`: length `n` `numeric` specifying the overall response vector.
#' - `weights_vector`: length `n` `numeric` specifying the weights vector.
#' - `n_visits`: `int` with the number of visits, which is the dimension of the
#'      covariance matrix.
#' - `n_subjects`: `int` with the number of subjects.
#' - `subject_zero_inds`: length `n_subjects` `integer` containing the zero-based start
#'     indices for each subject.
#' - `subject_n_visits`: length `n_subjects` `integer` containing the number of
#'     observed visits for each subjects. So the sum of this vector equals `n`.
#' - `cov_type`: `string` value specifying the covariance type.
#' - `is_spatial_int`: `int` specifying whether the covariance structure is spatial(1) or not(0).
#' - `reml`: `int` specifying whether REML estimation is used (1), otherwise ML (0).
#' - `subject_groups`: `factor` specifying the grouping for each subject.
#' - `n_groups`: `int` with the number of total groups
#'
#' @details Note that the `subject_var` must not be factor but can also be character.
#'   If it is character, then it will be converted to factor internally. Here
#'   the levels will be the unique values, sorted alphabetically and numerically if there
#'   is a common string prefix of numbers in the character elements. For full control
#'   on the order please use a factor.
#'
#' @keywords internal
h_mmrm_tmb_data <- function(
  formula_parts,
  data,
  weights,
  reml,
  singular = c("drop", "error", "keep"),
  drop_visit_levels,
  allow_na_response = FALSE,
  drop_levels = TRUE,
  xlev = NULL,
  contrasts = NULL
) {
  assert_class(formula_parts, "mmrm_tmb_formula_parts")
  assert_data_frame(data)
  assert_numeric(weights, len = nrow(data))
  assert_flag(reml)
  singular <- match.arg(singular)
  assert_flag(drop_visit_levels)

  data <- data.frame(data, weights)
  # Weights is always the last column.
  weights_name <- colnames(data)[ncol(data)]
  # If `y` is allowed to be NA, then first replace y with 1:n, then replace it with original y.
  if (!allow_na_response) {
    h_warn_na_action()
  }

  # Create model frame. This potentially uses variables from the environment if not present in
  # `data`.
  full_frame <- eval(
    bquote(stats::model.frame(
      formula_parts$full_formula,
      data = data,
      weights = .(as.symbol(weights_name)),
      na.action = "na.pass",
      xlev = xlev
    ))
  )

  varname <- formula_parts[c("subject_var", "visit_var", "group_var")]
  assert_names(
    names(full_frame),
    must.include = unlist(varname, use.names = FALSE)
  )
  assert_true(
    is.factor(full_frame[[formula_parts$subject_var]]) ||
      is.character(full_frame[[formula_parts$subject_var]])
  )
  if (is.character(full_frame[[formula_parts$subject_var]])) {
    full_frame[[formula_parts$subject_var]] <- factor(
      full_frame[[formula_parts$subject_var]],
      levels = stringr::str_sort(
        unique(full_frame[[formula_parts$subject_var]]),
        numeric = TRUE
      )
    )
  }
  data_order <- if (formula_parts$is_spatial) {
    order(full_frame[[formula_parts$subject_var]])
  } else {
    subject_visit_data <- full_frame[, c(
      formula_parts$subject_var,
      formula_parts$visit_var
    )]
    is_duplicated <- duplicated(subject_visit_data)
    if (any(is_duplicated)) {
      stop(
        "time points have to be unique for each subject, detected following duplicates in data:\n",
        paste(
          utils::capture.output(print(subject_visit_data[is_duplicated, ])),
          collapse = "\n"
        )
      )
    }
    order(
      full_frame[[formula_parts$subject_var]],
      full_frame[[formula_parts$visit_var]]
    )
  }
  if (identical(formula_parts$is_spatial, FALSE)) {
    h_confirm_large_levels(length(levels(full_frame[[
      formula_parts$visit_var
    ]])))
  }
  full_frame <- full_frame[data_order, ]

  if (drop_levels) {
    full_frame <- h_drop_levels(
      full_frame,
      formula_parts$subject_var,
      formula_parts$visit_var,
      names(xlev)
    )
  }
  has_response <- !identical(attr(attr(full_frame, "terms"), "response"), 0L)
  keep_ind <- if (allow_na_response && has_response) {
    # Note that response is always the first column if there is response.
    stats::complete.cases(full_frame[, -1L, drop = FALSE])
  } else {
    stats::complete.cases(full_frame)
  }
  full_frame <- full_frame[keep_ind, ]
  if (
    drop_visit_levels &&
      !formula_parts$is_spatial &&
      h_extra_levels(full_frame[[formula_parts$visit_var]])
  ) {
    visit_vec <- full_frame[[formula_parts$visit_var]]
    old_levels <- levels(visit_vec)
    full_frame[[formula_parts$visit_var]] <- droplevels(visit_vec)
    new_levels <- levels(full_frame[[formula_parts$visit_var]])
    dropped <- setdiff(old_levels, new_levels)
    message(
      "In ",
      formula_parts$visit_var,
      " there are dropped visits: ",
      toString(dropped),
      ".\n Additional attributes including contrasts are lost.\n",
      "To avoid this behavior, make sure use `drop_visit_levels = FALSE`."
    )
  }
  is_factor_col <- vapply(full_frame, is.factor, FUN.VALUE = TRUE)
  is_factor_col <- intersect(
    names(is_factor_col)[is_factor_col],
    all.vars(formula_parts$model_formula)
  )
  x_matrix <- stats::model.matrix(
    formula_parts$model_formula,
    data = full_frame,
    contrasts.arg = h_default_value(
      contrasts,
      lapply(full_frame[is_factor_col], contrasts)
    )
  )
  x_cols_aliased <- stats::setNames(
    rep(FALSE, ncol(x_matrix)),
    nm = colnames(x_matrix)
  )
  qr_x_mat <- qr(x_matrix)
  x_matrix_complete <- x_matrix
  if (qr_x_mat$rank < ncol(x_matrix)) {
    cols_to_drop <- utils::tail(qr_x_mat$pivot, ncol(x_matrix) - qr_x_mat$rank)
    if (identical(singular, "error")) {
      stop(
        "design matrix only has rank ",
        qr_x_mat$rank,
        " and ",
        length(cols_to_drop),
        " columns (",
        toString(colnames(x_matrix)[cols_to_drop]),
        ") could be dropped",
        " to achieve full rank ",
        ncol(x_matrix),
        " by using `accept_singular = TRUE`"
      )
    } else if (identical(singular, "drop")) {
      contrasts_attr <- attr(x_matrix, "contrasts")
      assign_attr <- attr(x_matrix, "assign")
      x_matrix <- x_matrix[, -cols_to_drop, drop = FALSE]
      x_cols_aliased[cols_to_drop] <- TRUE
      attr(x_matrix, "assign") <- assign_attr[-cols_to_drop]
      attr(x_matrix, "contrasts") <- contrasts_attr
    }
  }
  y_vector <- if (has_response) {
    as.numeric(stats::model.response(full_frame))
  } else {
    rep(NA_real_, nrow(full_frame))
  }
  weights_vector <- as.numeric(stats::model.weights(full_frame))
  n_subjects <- length(unique(full_frame[[formula_parts$subject_var]]))
  subject_zero_inds <- which(
    !duplicated(full_frame[[formula_parts$subject_var]])
  ) -
    1L
  subject_n_visits <- c(utils::tail(subject_zero_inds, -1L), nrow(full_frame)) -
    subject_zero_inds
  # It is possible that `subject_var` is factor with more levels (and this does not affect fit)
  # so no check is needed for `subject_visits`.
  assert_true(all(subject_n_visits > 0))
  if (!is.null(formula_parts$group_var)) {
    assert_factor(full_frame[[formula_parts$group_var]])
    subject_groups <- full_frame[[formula_parts$group_var]][
      subject_zero_inds + 1L
    ]
    n_groups <- nlevels(subject_groups)
  } else {
    subject_groups <- factor(rep(0L, n_subjects))
    n_groups <- 1L
  }
  coordinates <- full_frame[, formula_parts$visit_var, drop = FALSE]
  if (formula_parts$is_spatial) {
    lapply(coordinates, assert_numeric)
    coordinates_matrix <- as.matrix(coordinates)
    n_visits <- max(subject_n_visits)
  } else {
    assert(identical(ncol(coordinates), 1L))
    if (!test_factor(coordinates[[1L]])) {
      stop(
        "Time point variable '",
        formula_parts$visit_var,
        "' must be a factor."
      )
    }
    coordinates_matrix <- as.matrix(as.integer(coordinates[[1L]]) - 1, ncol = 1)
    n_visits <- nlevels(coordinates[[1L]])
    assert_true(all(subject_n_visits <= n_visits))
  }
  structure(
    list(
      full_frame = full_frame,
      data = data,
      x_matrix = x_matrix,
      x_cols_aliased = x_cols_aliased,
      x_matrix_complete = x_matrix_complete,
      coordinates = coordinates_matrix,
      y_vector = y_vector,
      weights_vector = weights_vector,
      n_visits = n_visits,
      n_subjects = n_subjects,
      subject_zero_inds = subject_zero_inds,
      subject_n_visits = subject_n_visits,
      cov_type = formula_parts$cov_type,
      is_spatial_int = as.integer(formula_parts$is_spatial),
      reml = as.integer(reml),
      subject_groups = subject_groups,
      n_groups = n_groups
    ),
    class = "mmrm_tmb_data"
  )
}

#' Start Parameters for `TMB` Fit
#'
#' @param formula_parts (`mmrm_tmb_formula_parts`)\cr produced by
#'  [h_mmrm_tmb_formula_parts()].
#' @param tmb_data (`mmrm_tmb_data`)\cr produced by [h_mmrm_tmb_data()].
#' @param start (`numeric` or `NULL`)\cr optional start values for variance
#'   parameters.
#' @param n_groups (`int`)\cr number of groups.
#' @return List with element `theta` containing the start values for the variance
#'   parameters.
#'
#' @keywords internal
h_mmrm_tmb_parameters <- function(
  formula_parts,
  tmb_data,
  start,
  n_groups = 1L
) {
  assert_class(formula_parts, "mmrm_tmb_formula_parts")
  assert_class(tmb_data, "mmrm_tmb_data")

  m <- tmb_data$n_visits
  start_value0 <- std_start(formula_parts$cov_type, m, n_groups)
  theta_dim <- length(start_value0)
  start_values <- if (is.null(start)) {
    start_value0
  } else if (test_function(start)) {
    do.call(start, utils::modifyList(formula_parts, tmb_data))
  } else {
    start
  }
  assert_numeric(
    start_values,
    len = theta_dim,
    any.missing = FALSE,
    finite = TRUE
  )
  list(theta = start_values)
}

#' Asserting Sane Start Values for `TMB` Fit
#'
#' @param tmb_object (`list`)\cr created with [TMB::MakeADFun()].
#'
#' @return Nothing, only used for assertions.
#'
#' @keywords internal
h_mmrm_tmb_assert_start <- function(tmb_object) {
  assert_list(tmb_object)
  assert_subset(c("fn", "gr", "par"), names(tmb_object))

  if (is.na(tmb_object$fn(tmb_object$par))) {
    stop("negative log-likelihood is NaN at starting parameter values")
  }
  if (any(is.na(tmb_object$gr(tmb_object$par)))) {
    stop("some elements of gradient are NaN at starting parameter values")
  }
}

#' Checking the `TMB` Optimization Result
#'
#' @param tmb_opt (`list`)\cr optimization result.
#' @param mmrm_tmb (`mmrm_tmb`)\cr result from [h_mmrm_tmb_fit()].
#'
#' @return Nothing, only used to generate warnings in case that the model
#' did not converge.
#'
#' @keywords internal
h_mmrm_tmb_check_conv <- function(tmb_opt, mmrm_tmb) {
  assert_list(tmb_opt)
  assert_subset(c("par", "objective", "convergence", "message"), names(tmb_opt))
  assert_class(mmrm_tmb, "mmrm_tmb")

  if (!is.null(tmb_opt$convergence) && tmb_opt$convergence != 0) {
    warning("Model convergence problem: ", tmb_opt$message, ".")
    return()
  }
  theta_vcov <- mmrm_tmb$theta_vcov
  if (is.call(theta_vcov)) {
    message(
      "theta_vcov calculation was disabled and cannot be used for convergence checks"
    )
    return()
  }
  if (is(theta_vcov, "try-error")) {
    warning(
      "Model convergence problem: hessian is singular, theta_vcov not available."
    )
    return()
  }
  if (!all(is.finite(theta_vcov))) {
    warning("Model convergence problem: theta_vcov contains non-finite values.")
    return()
  }
  eigen_vals <- eigen(theta_vcov, only.values = TRUE)$values
  if (mode(eigen_vals) == "complex" || any(eigen_vals <= 0)) {
    # Note: complex eigen values signal that the matrix is not symmetric, therefore not positive definite.
    warning("Model convergence problem: theta_vcov is not positive definite.")
    return()
  }
  qr_rank <- qr(theta_vcov)$rank
  if (qr_rank < ncol(theta_vcov)) {
    warning("Model convergence problem: theta_vcov is numerically singular.")
  }
}

#' Extract covariance matrix from `TMB` report and input data
#'
#' This helper does some simple post-processing to extract covariance matrix or named
#' list of covariance matrices if the fitting is using grouped covariance matrices.
#'
#' @param tmb_report (`list`)\cr report created with [TMB::MakeADFun()] report function.
#' @param tmb_data (`mmrm_tmb_data`)\cr produced by [h_mmrm_tmb_data()].
#' @param visit_var (`character`)\cr character vector of the visit variable
#' @param is_spatial (`flag`)\cr indicator whether the covariance structure is spatial.
#' @return Return a simple covariance matrix if there is no grouping, or a named
#' list of estimated grouped covariance matrices,
#' with its name equal to the group levels.
#'
#' @keywords internal
h_mmrm_tmb_extract_cov <- function(
  tmb_report,
  tmb_data,
  visit_var,
  is_spatial
) {
  d <- dim(tmb_report$covariance_lower_chol)
  visit_names <- if (!is_spatial) {
    levels(tmb_data$full_frame[[visit_var]])
  } else {
    c(0, 1)
  }
  cov <- lapply(
    seq_len(d[1] / d[2]),
    function(i) {
      ret <- tcrossprod(tmb_report$covariance_lower_chol[
        seq(1 + (i - 1) * d[2], i * d[2]),
      ])
      dimnames(ret) <- list(visit_names, visit_names)
      ret
    }
  )
  if (identical(tmb_data$n_groups, 1L)) {
    cov <- cov[[1]]
  } else {
    names(cov) <- levels(tmb_data$subject_groups)
  }
  cov
}

#' Build `TMB` Fit Result List
#'
#' This helper does some simple post-processing of the `TMB` object and
#' optimization results, including setting names, inverting matrices etc.
#'
#' @param tmb_object (`list`)\cr created with [TMB::MakeADFun()].
#' @param tmb_opt (`list`)\cr optimization result.
#' @param formula_parts (`mmrm_tmb_formula_parts`)\cr produced by
#'  [h_mmrm_tmb_formula_parts()].
#' @param tmb_data (`mmrm_tmb_data`)\cr produced by [h_mmrm_tmb_data()].
#' @param disable_theta_vcov (`flag`)\cr whether calculation of `theta_vcov` was disabled.
#'
#' @return List of class `mmrm_tmb` with:
#'   - `cov`: estimated covariance matrix, or named list of estimated group specific covariance matrices.
#'   - `beta_est`: vector of coefficient estimates.
#'   - `beta_vcov`: Variance-covariance matrix for coefficient estimates.
#'   - `beta_vcov_inv_L`: Lower triangular matrix `L` of the inverse variance-covariance matrix decomposition.
#'   - `beta_vcov_inv_D`: vector of diagonal matrix `D` of the inverse variance-covariance matrix decomposition.
#'   - `theta_est`: vector of variance parameter estimates.
#'   - `theta_vcov`: variance-covariance matrix for variance parameter estimates (if not disabled).
#'   - `neg_log_lik`: obtained negative log-likelihood.
#'   - `formula_parts`: input.
#'   - `data`: input.
#'   - `weights`: input.
#'   - `reml`: input as a flag.
#'   - `opt_details`: list with optimization details including convergence code.
#'   - `tmb_object`: original `TMB` object created with [TMB::MakeADFun()].
#'   - `tmb_data`: input.
#'
#' @details Instead of inverting or decomposing `beta_vcov`, it can be more efficient to use its robust
#'   Cholesky decomposition `LDL^T`, therefore we return the corresponding two components `L` and `D`
#'   as well since they have been available on the `C++` side already.
#'
#' @keywords internal
h_mmrm_tmb_fit <- function(
  tmb_object,
  tmb_opt,
  formula_parts,
  tmb_data,
  disable_theta_vcov = FALSE
) {
  assert_list(tmb_object)
  assert_subset(c("fn", "gr", "par", "he"), names(tmb_object))
  assert_list(tmb_opt)
  assert_subset(c("par", "objective", "convergence", "message"), names(tmb_opt))
  assert_class(formula_parts, "mmrm_tmb_formula_parts")
  assert_class(tmb_data, "mmrm_tmb_data")

  tmb_report <- tmb_object$report(par = tmb_opt$par)
  x_matrix_cols <- colnames(tmb_data$x_matrix)
  cov <- h_mmrm_tmb_extract_cov(
    tmb_report,
    tmb_data,
    formula_parts$visit_var,
    formula_parts$is_spatial
  )
  beta_est <- tmb_report$beta
  names(beta_est) <- x_matrix_cols
  beta_vcov <- tmb_report$beta_vcov
  dimnames(beta_vcov) <- list(x_matrix_cols, x_matrix_cols)
  beta_vcov_inv_L <- tmb_report$XtWX_L # nolint
  beta_vcov_inv_D <- tmb_report$XtWX_D # nolint
  theta_est <- tmb_opt$par
  names(theta_est) <- NULL
  theta_vcov <- if (disable_theta_vcov) {
    quote(stop("theta_vcov calculation was disabled"))
  } else {
    try(solve(tmb_object$he(tmb_opt$par)), silent = TRUE)
  }
  opt_details_names <- setdiff(
    names(tmb_opt),
    c("par", "objective")
  )
  structure(
    list(
      cov = cov,
      beta_est = beta_est,
      beta_vcov = beta_vcov,
      beta_vcov_inv_L = beta_vcov_inv_L,
      beta_vcov_inv_D = beta_vcov_inv_D,
      theta_est = theta_est,
      theta_vcov = theta_vcov,
      neg_log_lik = tmb_opt$objective,
      formula_parts = formula_parts,
      data = tmb_data$data,
      weights = tmb_data$weights_vector,
      reml = as.logical(tmb_data$reml),
      opt_details = tmb_opt[opt_details_names],
      tmb_object = tmb_object,
      tmb_data = tmb_data
    ),
    class = "mmrm_tmb"
  )
}

#' Low-Level Fitting Function for MMRM
#'
#' @description `r lifecycle::badge("stable")`
#'
#' This is the low-level function to fit an MMRM. Note that this does not
#' try different optimizers or adds Jacobian information etc. in contrast to
#' [mmrm()].
#'
#' @param formula (`formula`)\cr model formula with exactly one special term
#'   specifying the visits within subjects, see details.
#' @param data (`data.frame`)\cr input data containing the variables used in
#'   `formula`.
#' @param weights (`vector`)\cr input vector containing the weights.
#' @inheritParams h_mmrm_tmb_data
#' @param covariance (`cov_struct`)\cr A covariance structure type definition,
#'   or value that can be coerced to a covariance structure using
#'   [as.cov_struct()]. If no value is provided, a structure is derived from
#'   the provided formula.
#' @param control (`mmrm_control`)\cr list of control options produced by
#'   [mmrm_control()].
#' @inheritParams fit_single_optimizer
#'
#' @return List of class `mmrm_tmb`, see [h_mmrm_tmb_fit()] for details.
#'   In addition, it contains elements `call` and `optimizer`.
#'
#' @details
#' The `formula` typically looks like:
#'
#' `FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT | USUBJID)`
#'
#' which specifies response and covariates as usual, and exactly one special term
#' defines which covariance structure is used and what are the visit and
#' subject variables.
#'
#' Always use only the first optimizer if multiple optimizers are provided.
#'
#' @export
#'
#' @examples
#' formula <- FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT | USUBJID)
#' data <- fev_data
#' system.time(result <- fit_mmrm(formula, data, rep(1, nrow(fev_data))))
fit_mmrm <- function(
  formula,
  data,
  weights,
  reml = TRUE,
  covariance = NULL,
  tmb_data,
  formula_parts,
  control = mmrm_control()
) {
  if (missing(formula_parts) || missing(tmb_data)) {
    covariance <- h_reconcile_cov_struct(formula, covariance)
    formula_parts <- h_mmrm_tmb_formula_parts(formula, covariance)

    assert_class(control, "mmrm_control")
    assert_list(control$optimizers, min.len = 1)
    assert_numeric(weights, any.missing = FALSE)
    assert_true(all(weights > 0))
    tmb_data <- h_mmrm_tmb_data(
      formula_parts,
      data,
      weights,
      reml,
      singular = if (control$accept_singular) "drop" else "error",
      drop_visit_levels = control$drop_visit_levels
    )
  } else {
    assert_class(tmb_data, "mmrm_tmb_data")
    assert_class(formula_parts, "mmrm_tmb_formula_parts")
  }
  tmb_parameters <- h_mmrm_tmb_parameters(
    formula_parts,
    tmb_data,
    start = control$start,
    n_groups = tmb_data$n_groups
  )
  tmb_object <- TMB::MakeADFun(
    data = tmb_data,
    parameters = tmb_parameters,
    hessian = TRUE,
    DLL = "mmrm",
    silent = TRUE
  )
  h_mmrm_tmb_assert_start(tmb_object)
  used_optimizer <- control$optimizers[[1L]]
  used_optimizer_name <- names(control$optimizers)[1L]
  args <- with(
    tmb_object,
    c(
      list(par, fn, gr),
      attr(used_optimizer, "args")
    )
  )
  if (identical(attr(used_optimizer, "use_hessian"), TRUE)) {
    args$hessian <- tmb_object$he
  }
  tmb_opt <- do.call(
    what = used_optimizer,
    args = args
  )
  # Ensure negative log likelihood is stored in `objective` element of list.
  if ("value" %in% names(tmb_opt)) {
    tmb_opt$objective <- tmb_opt$value
    tmb_opt$value <- NULL
  }
  fit <- h_mmrm_tmb_fit(
    tmb_object,
    tmb_opt,
    formula_parts,
    tmb_data,
    disable_theta_vcov = control$disable_theta_vcov
  )
  h_mmrm_tmb_check_conv(tmb_opt, fit)
  fit$call <- match.call()
  fit$call$formula <- formula_parts$formula
  fit$optimizer <- used_optimizer_name
  fit
}

#' Register `mmrm` For Use With `car::Anova`
#'
#' @inheritParams base::requireNamespace
#' @return A logical value indicating whether registration was successful.
#'
#' @keywords internal
car_add_mmrm <- function(quietly = FALSE) {
  if (!requireNamespace("car", quietly = quietly)) {
    return(FALSE)
  }
  envir <- asNamespace("mmrm")
  h_register_s3("car", "Anova", "mmrm", envir)
  TRUE
}

#' Obtain Type 2 Contrast for One Specified Effect
#'
#' This is support function to obtain contrast matrix for type II testing.
#'
#' @param object (`mmrm`)\cr the fitted MMRM.
#' @param effect (`string`) the name of the effect.
#' @param tol (`numeric`) threshold below which values are treated as 0.
#'
#' @return A `matrix` of the contrast.
#'
#' @keywords internal
h_type2_contrast <- function(
  object,
  effect,
  tol = sqrt(.Machine$double.eps)
) {
  assert_class(object, "mmrm")
  assert_string(effect)
  assert_double(tol, finite = TRUE, len = 1L)

  # The complete model matrix including aliased columns is needed. See below.
  mx <- component(object, "x_matrix_complete")
  asg <- attr(mx, "assign")
  formula <- object$formula_parts$model_formula
  tms <- stats::terms(formula)
  fcts <- attr(tms, "factors")[-1L, , drop = FALSE] # Discard the response.

  # We do this because the factors attribute of a terms.object can include 2s.
  fcts[fcts > 1] <- 1

  ods <- attr(tms, "order")
  assert_subset(effect, colnames(fcts))
  idx <- which(effect == colnames(fcts))
  cols <- which(asg == idx)
  xlev <- component(object, "xlev")

  categorical_covars <- intersect(rownames(fcts), names(xlev))

  # An effect contains the model's intercept if the model was not built with an
  # intercept and if the effect is the first one in the model specification to
  # contain a categorical variable.
  effect_contains_intercept <-
    !attr(tms, "intercept") &&
    length(categorical_covars) &&
    effect == h_first_term_containing_categ(fcts, categorical_covars)

  coef_rows <- length(cols) - as.integer(effect_contains_intercept)
  l_mx <-
    matrix(
      0,
      nrow = coef_rows,
      ncol = ncol(mx),
      dimnames = list(NULL, colnames(mx))
    )
  if (coef_rows == 0L) {
    l_mx <- l_mx[, !component(object, "beta_aliased"), drop = FALSE]
    return(l_mx)
  }

  # The effect's submatrix of l_mx shall simply be an identity matrix unless the
  # effect contains the model's intercept. In this case, the submatrix has an
  # extra column that should contain straight -1s.
  l_mx[, cols] <-
    if (effect_contains_intercept) {
      cbind(diag(coef_rows), -1)
    } else {
      diag(coef_rows)
    }

  for (i in setdiff(seq_len(ncol(fcts)), idx)) {
    # This is where replacing the factors attributes' 2s with 1s matters.
    additional_vars <- names(which(fcts[, i] > fcts[, idx]))
    additional_numeric <- !all(additional_vars %in% categorical_covars)
    current_col <- which(asg == i)
    if (
      ods[i] >= ods[idx] && all(fcts[, i] >= fcts[, idx]) && !additional_numeric
    ) {
      # This is where it matters to use the complete design matrix including
      # aliased columns.
      x1 <- mx[, cols, drop = FALSE]
      x0 <- mx[, -c(cols, current_col), drop = FALSE]
      x2 <- mx[, current_col, drop = FALSE]
      m <- diag(nrow(x0)) - h_quad_form_mat(x0, MASS::ginv(crossprod(x0)))
      crossprod_x1_m <- crossprod(x1, m)
      ret <- solve(crossprod_x1_m %*% x1) %*% crossprod_x1_m %*% x2
      sub_mat <- if (effect_contains_intercept) {
        ret[-1, ] - ret[1, ]
      } else {
        ret
      }
      l_mx[, current_col] <- sub_mat
    }
  }
  # Finally, we drop aliased columns. But we needed them to properly calculate
  # the others.
  l_mx <- l_mx[, !component(object, "beta_aliased"), drop = FALSE]
  l_mx[abs(l_mx) < tol] <- 0
  l_mx
}


#' Identify the First Term in a Model to Contain a Categorical Variable
#'
#' This returns the column name of the leftmost column of `factors` containing a
#' nonzero value in a row corresponding to a `categorical` variable.
#'
#' @param factors (matrix)\cr the `factors` attribute of a [`terms.object`],
#'   which is a matrix of 0s, 1s, and 2s.
#' @param categorical (character)\cr a vector of the categorical variables in
#'   the model whose [`terms.object`] is `factors`.
#'
#' @return A `string`: one of the column names of `factors`. If none of the
#'   columns contain a categorical variable, `NULL` is returned.
#'
#' @keywords internal
h_first_term_containing_categ <- function(factors, categorical) {
  factors <- factors[categorical, , drop = FALSE]
  for (term in colnames(factors)) {
    if (any(factors[, term] > 0)) {
      return(term)
    }
  }
}


#' Obtain Type 3 Contrast for All Effects
#'
#' This is support function to obtain contrast matrices for type III testing.
#'
#' @param object (`mmrm`)\cr the fitted MMRM.
#' @param tol (`numeric`) threshold below which values are treated as 0.
#'
#' @return A `list` of contrast matrices, one per effect.
#' @keywords internal
h_type3_contrasts <- function(object, tol = sqrt(.Machine$double.eps)) {
  assert_class(object, "mmrm")
  assert_double(tol, finite = TRUE, len = 1L)

  # First obtain the simple contrasts as if we use all contr.sum only.
  contr_sum_contrasts <- h_contr_sum_type3_contrasts(object)

  # The original design matrix.
  orig_model_matrix <- component(object, "x_matrix")

  # Obtain original contrasts, if there is none or all are contr.sum, return
  # the contr.sum contrasts.
  orig_contrasts <- attr(orig_model_matrix, "contrasts")
  if (
    length(orig_contrasts) == 0L ||
      all(vapply(orig_contrasts, identical, logical(1L), "contr.sum"))
  ) {
    return(contr_sum_contrasts)
  }

  # Compute the modified design matrix using all contr.sum.
  new_contrasts <- orig_contrasts
  new_contrasts[] <- "contr.sum"
  orig_data <- component(object, "full_frame")
  new_model_matrix <-
    stats::model.matrix(
      object,
      data = orig_data,
      contrasts = new_contrasts
    )
  coef_aliased <- component(object, "beta_aliased")
  if (any(coef_aliased)) {
    new_model_matrix <- new_model_matrix[, !coef_aliased, drop = FALSE]
  }

  # Compute the correction matrix and return corrected contrasts,
  # dropping small values.
  type_iii_correction <-
    solve(crossprod(new_model_matrix)) %*%
    crossprod(new_model_matrix, orig_model_matrix)
  result <- lapply(contr_sum_contrasts, `%*%`, type_iii_correction)
  lapply(result, function(mat) {
    mat[abs(mat) < tol] <- 0
    mat
  })
}

#' Conduct type II/III hypothesis testing on the MMRM fit results.
#'
#' @param mod (`mmrm`)\cr the fitted MMRM.
#' @param type (`string`)\cr either `"II"`, `"III"`, `"2"`, or `"3"`, indicating the
#'   type of test to perform.
#' @param test.statistic (`string`)\cr either `"F` or `"Chisq"`, indicating the
#'   kind of test to perform.
#' @param ... arguments passed from other methods.
#' @inheritParams h_type2_contrast
#'
#' @details `Anova()` will return an `anova` object with one row per variable.
#'
#'   If `test.statistic = "F"`, columns will be `Df`(numerator degrees of
#'   freedom), `Res.Df` (denominator degrees of freedom), `F`, and
#'   `Pr(>F)`.
#'
#'   If `test.statistic = "Chisq"`, columns will be `Chisq` (the Chi-squared
#'   test statistic), `Df` (degrees of freedom), and `Pr(>Chisq)` (p-value).
#'
#' @keywords internal
# Please do not load `car` and then create the documentation. The Rd file will
#  be different.
# nolint start
Anova.mmrm <- function(
  # nolint end
  mod,
  type = c("II", "III", "2", "3"),
  tol = sqrt(.Machine$double.eps),
  test.statistic = c("F", "Chisq"), # nolint
  ...
) {
  type <- as.character(type)
  type <- match.arg(type)
  test.statistic <- match.arg(test.statistic) # nolint

  vars <- colnames(attr(terms(mod$formula_parts$model_formula), "factors"))
  if (type %in% c("II", "2")) {
    contrasts <- lapply(vars, h_type2_contrast, object = mod, tol = tol)
    names(contrasts) <- vars
  } else {
    contrasts <- h_type3_contrasts(mod, tol = tol)
    contrasts <- contrasts[intersect(vars, names(contrasts))]
  }

  if (test.statistic == "F") {
    ret <- lapply(contrasts, df_md, object = mod)
    ret_df <- do.call(rbind.data.frame, ret)
    colnames(ret_df) <- c("Df", "Res.Df", "F", "Pr(>F)")
  } else {
    ret <- lapply(contrasts, h_test_md, object = mod, test = "Chisq")
    ret_df <- do.call(rbind.data.frame, ret)
    colnames(ret_df) <- c("Df", "Chisq", "Pr(>Chisq)")
  }

  row.names(ret_df) <- names(contrasts)
  attr(ret_df, "heading") <-
    sprintf(
      "Analysis of Fixed Effect Table (Type %s %s tests)",
      switch(type, "2" = "II", "3" = "III", type),
      test.statistic
    )
  class(ret_df) <- c("anova", "data.frame")

  ret_df
}

#' Obtain Levels Prior and Posterior
#' @param var (`string`) name of the effect.
#' @param additional_vars (`character`) names of additional variables.
#' @param xlev (`list`) named list of character levels.
#' @param factors (`matrix`) the factor matrix.
#' @keywords internal
h_obtain_lvls <- function(var, additional_vars, xlev, factors) {
  assert_string(var)
  assert_character(additional_vars)
  assert_list(xlev, types = "character")
  nms <- names(xlev)
  assert_subset(additional_vars, nms)
  if (var %in% nms) {
    prior_vars <- intersect(nms[seq_len(match(var, nms) - 1)], additional_vars)
    prior_lvls <- vapply(xlev[prior_vars], length, FUN.VALUE = 1L)
    post_vars <- intersect(
      nms[seq(match(var, nms) + 1, length(nms))],
      additional_vars
    )
    post_lvls <- vapply(xlev[post_vars], length, FUN.VALUE = 1L)
    total_lvls <- prod(prior_lvls) * prod(post_lvls)
  } else {
    prior_lvls <- vapply(xlev[additional_vars], length, FUN.VALUE = 1L)
    post_lvls <- 2L
    total_lvls <- prod(prior_lvls)
  }
  list(
    prior = prior_lvls,
    post = post_lvls,
    total = total_lvls
  )
}

#' Test if the First Vector is Subset of the Second Vector
#' @param x (`vector`) the first list.
#' @param y (`vector`) the second list.
#' @keywords internal
h_get_index <- function(x, y) {
  assert_list(x)
  assert_list(y)
  vapply(
    x,
    \(i) {
      r <- vapply(y, \(j) test_subset(j, i), FUN.VALUE = TRUE)
      if (sum(r) == 1L) {
        which(r)
      } else {
        NA_integer_
      }
    },
    FUN.VALUE = 1L
  )
}

#' Construct Preliminary Contrast Matrices for Type III Tests Assuming Sum Contrasts
#'
#' @param object (`mmrm`)\cr the fitted MMRM.
#'
#' @return A `list` of contrast matrices, which are just row-wise term subsets of
#'   an overall identity matrix.
#'
#' @keywords internal
h_contr_sum_type3_contrasts <- function(object) {
  assert_class(object, "mmrm")

  terms <- stats::terms(object)
  has_intercept <- attr(terms, "intercept")
  if (!has_intercept) {
    warning(
      "Type III tests can give misleading results for models without an intercept"
    )
  }
  term_labels <- c(if (has_intercept) "(Intercept)", labels(terms))
  n_coefs <- length(component(object, "beta_est"))
  identity_matrix <- diag(n_coefs)

  x_matrix <- component(object, "x_matrix")
  assign <- attr(x_matrix, "assign")
  assert_integer(assign, lower = 0L, len = ncol(x_matrix))
  map_terms_to_cols <- assign + has_intercept

  terms_per_col <- factor(term_labels[map_terms_to_cols], levels = term_labels)

  split.data.frame(
    identity_matrix,
    terms_per_col,
    drop = TRUE
  )
}

#' Extract Formula Terms used for Covariance Structure Definition
#'
#' @param f (`formula`)\cr a formula from which covariance terms should be
#'   extracted.
#'
#' @return A list of covariance structure expressions found in `f`.
#'
#' @importFrom stats terms
#' @keywords internal
h_extract_covariance_terms <- function(f) {
  specials <- cov_types(c("abbr", "habbr"))
  terms <- stats::terms(formula_rhs(f), specials = specials)
  covariance_terms <- Filter(length, attr(terms, "specials"))
  variables <- attr(terms, "variables")
  lapply(covariance_terms, function(i) variables[[i + 1]])
}

#' Drop Formula Terms used for Covariance Structure Definition
#'
#' @param f (`formula`)\cr a formula from which covariance terms should be
#'   dropped.
#'
#' @return The formula without accepted covariance terms.
#'
#' @details `terms` is used and it will preserve the environment attribute.
#' This ensures the returned formula and the input formula have the same environment.
#' @importFrom stats terms drop.terms
#' @keywords internal
h_drop_covariance_terms <- function(f) {
  specials <- cov_types(c("abbr", "habbr"))

  terms <- stats::terms(f, specials = specials)
  covariance_terms <- Filter(Negate(is.null), attr(terms, "specials"))

  # if no covariance terms were found, return original formula
  if (length(covariance_terms) == 0) {
    return(f)
  }
  if (length(f) != 3) {
    update_str <- "~ . -"
  } else {
    update_str <- ". ~ . -"
  }
  stats::update(
    f,
    stats::as.formula(paste(
      update_str,
      deparse(attr(terms, "variables")[[covariance_terms[[1]] + 1]])
    ))
  )
}

#' Add Individual Covariance Variables As Terms to Formula
#'
#' @param f (`formula`)\cr a formula to which covariance structure terms should
#'   be added.
#' @param covariance (`cov_struct`)\cr a covariance structure object from which
#'   additional variables should be sourced.
#'
#' @return A new formula with included covariance terms.
#'
#' @details [stats::update()] is used to append the covariance structure and the environment
#' attribute will not be changed. This ensures the returned formula and the input formula
#' have the same environment.
#'
#' @keywords internal
h_add_covariance_terms <- function(f, covariance) {
  cov_terms <- with(covariance, c(subject, visits, group))
  cov_terms <- paste(cov_terms, collapse = " + ")
  stats::update(f, stats::as.formula(paste(". ~ . + ", cov_terms)))
}

#' Add Formula Terms with Character
#'
#' Add formula terms from the original formula with character representation.
#'
#' @param f (`formula`)\cr a formula to be updated.
#' @param adds (`character`)\cr representation of elements to be added.
#' @param drop_response (`flag`)\cr whether response should be dropped.
#'
#' @details Elements in `adds` will be added from the formula, while the environment
#' of the formula is unchanged. If `adds` is `NULL` or `character(0)`, the formula is
#' unchanged.
#' @return A new formula with elements in `drops` removed.
#'
#' @keywords internal
h_add_terms <- function(f, adds, drop_response = FALSE) {
  assert_character(adds, null.ok = TRUE)
  if (length(adds) > 0L) {
    add_terms <- stats::as.formula(sprintf(
      ". ~ . + %s",
      paste(adds, collapse = "+")
    ))
    f <- stats::update(f, add_terms)
  }
  if (drop_response && length(f) == 3L) {
    f[[2]] <- NULL
  }
  f
}

#' Methods for `mmrm` Objects
#'
#' @description `r lifecycle::badge("stable")`
#'
#' @param object (`mmrm`)\cr the fitted MMRM including Jacobian and call etc.
#' @param ... not used.
#' @return Depends on the method, see Details and Functions.
#'
#' @details
#' While printing the summary of (`mmrm`)\cr object, the following will be displayed:
#' 1. Formula. The formula used in the model.
#' 2. Data. The data used for analysis, including number of subjects, number of valid observations.
#' 3. Covariance. The covariance structure and number of variance parameters.
#' 4. Method. Restricted maximum likelihood(REML) or maximum likelihood(ML).
#' 5. Model selection criteria. AIC, BIC, log likelihood and deviance.
#' 6. Coefficients. Coefficients of the covariates.
#' 7. Covariance estimate. The covariance estimate(for each group).
#'    1. If the covariance structure is non-spatial, the covariance matrix of all categorical time points available
#'       in data will be displayed.
#'    2. If the covariance structure is spatial, the covariance matrix of two time points with unit distance
#'       will be displayed.
#'
#' `confint` is used to obtain the confidence intervals for the coefficients.
#' Please note that this is different from the confidence interval of difference
#' of least square means from `emmeans`.
#'
#' @name mmrm_methods
#'
#' @seealso [`mmrm_tmb_methods`], [`mmrm_tidiers`] for additional methods.
#'
#' @examples
#' formula <- FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT | USUBJID)
#' object <- mmrm(formula, fev_data)
NULL

#' Coefficients Table for MMRM Fit
#'
#' This is used by [summary.mmrm()] to obtain the coefficients table.
#'
#' @param object (`mmrm`)\cr model fit.
#'
#' @return Matrix with one row per coefficient and columns
#'   `Estimate`, `Std. Error`, `df`, `t value` and `Pr(>|t|)`.
#'
#' @keywords internal
h_coef_table <- function(object) {
  assert_class(object, "mmrm")

  coef_est <- component(object, "beta_est")
  coef_contrasts <- diag(x = rep(1, length(coef_est)))
  rownames(coef_contrasts) <- names(coef_est)
  coef_table <- t(apply(
    coef_contrasts,
    MARGIN = 1L,
    FUN = function(contrast) unlist(df_1d(object, contrast))
  ))
  assert_names(
    colnames(coef_table),
    identical.to = c("est", "se", "df", "t_stat", "p_val")
  )
  colnames(coef_table) <- c(
    "Estimate",
    "Std. Error",
    "df",
    "t value",
    "Pr(>|t|)"
  )

  coef_aliased <- component(object, "beta_aliased")
  if (any(coef_aliased)) {
    names_coef_na <- names(which(coef_aliased))
    coef_na_table <- matrix(
      data = NA,
      nrow = length(names_coef_na),
      ncol = ncol(coef_table),
      dimnames = list(names_coef_na, colnames(coef_table))
    )
    coef_table <- rbind(coef_table, coef_na_table)[names(coef_aliased), ]
  }

  coef_table
}

#' @describeIn mmrm_methods summarizes the MMRM fit results.
#' @exportS3Method
#' @examples
#' # Summary:
#' summary(object)
summary.mmrm <- function(object, ...) {
  aic_list <- list(
    AIC = AIC(object),
    BIC = BIC(object),
    logLik = logLik(object),
    deviance = deviance(object)
  )
  coefficients <- h_coef_table(object)
  call <- stats::getCall(object)
  components <- component(
    object,
    c(
      "cov_type",
      "reml",
      "n_groups",
      "n_theta",
      "n_subjects",
      "n_timepoints",
      "n_obs",
      "beta_vcov",
      "varcor"
    )
  )
  components$method <- object$method
  components$vcov <- object$vcov
  structure(
    c(
      components,
      list(
        coefficients = coefficients,
        n_singular_coefs = sum(component(object, "beta_aliased")),
        aic_list = aic_list,
        call = call
      )
    ),
    class = "summary.mmrm"
  )
}

#' Printing MMRM Function Call
#'
#' This is used in [print.summary.mmrm()].
#'
#' @param call (`call`)\cr original [mmrm()] function call.
#' @param n_obs (`int`)\cr number of observations.
#' @param n_subjects (`int`)\cr number of subjects.
#' @param n_timepoints (`int`)\cr number of timepoints.
#'
#' @keywords internal
h_print_call <- function(call, n_obs, n_subjects, n_timepoints) {
  pass <- 0
  if (!is.null(tmp <- call$formula)) {
    cat("Formula:    ", deparse(tmp), fill = TRUE)
    rhs <- tmp[[2]]
    pass <- nchar(deparse(rhs))
  }
  if (!is.null(call$data)) {
    cat(
      "Data:       ",
      deparse(call$data),
      "(used",
      n_obs,
      "observations from",
      n_subjects,
      "subjects with maximum",
      n_timepoints,
      "timepoints)",
      fill = TRUE
    )
  }
  # Display the expression of weights
  if (!is.null(call$weights)) {
    cat("Weights:    ", deparse(call$weights), fill = TRUE)
  }
}

#' Printing MMRM Covariance Type
#'
#' This is used in [print.summary.mmrm()].
#'
#' @param cov_type (`string`)\cr covariance structure abbreviation.
#' @param n_theta (`count`)\cr number of variance parameters.
#' @param n_groups (`count`)\cr number of groups.
#' @keywords internal
h_print_cov <- function(cov_type, n_theta, n_groups) {
  assert_string(cov_type)
  assert_count(n_theta, positive = TRUE)
  assert_count(n_groups, positive = TRUE)
  cov_definition <- switch(
    cov_type,
    us = "unstructured",
    toep = "Toeplitz",
    toeph = "heterogeneous Toeplitz",
    ar1 = "auto-regressive order one",
    ar1h = "heterogeneous auto-regressive order one",
    ad = "ante-dependence",
    adh = "heterogeneous ante-dependence",
    cs = "compound symmetry",
    csh = "heterogeneous compound symmetry",
    sp_exp = "spatial exponential"
  )

  catstr <- sprintf(
    "Covariance:  %s (%d variance parameters%s)\n",
    cov_definition,
    n_theta,
    ifelse(n_groups == 1L, "", sprintf(" of %d groups", n_groups))
  )
  cat(catstr)
}

#' Printing AIC and other Model Fit Criteria
#'
#' This is used in [print.summary.mmrm()].
#'
#' @param aic_list (`list`)\cr list as part of from [summary.mmrm()].
#' @param digits (`number`)\cr number of decimal places used with [round()].
#'
#' @keywords internal
h_print_aic_list <- function(aic_list, digits = 1) {
  diag_vals <- round(unlist(aic_list), digits)
  diag_vals <- format(diag_vals)
  print(diag_vals, quote = FALSE)
}

#' @describeIn mmrm_methods prints the MMRM fit summary.
#' @exportS3Method
#' @keywords internal
print.summary.mmrm <- function(
  x,
  digits = max(3, getOption("digits") - 3),
  signif.stars = getOption("show.signif.stars"), # nolint
  ...
) {
  cat("mmrm fit\n\n")
  h_print_call(x$call, x$n_obs, x$n_subjects, x$n_timepoints)
  h_print_cov(x$cov_type, x$n_theta, x$n_groups)
  cat("Method:      ", x$method, "\n", sep = "")
  cat("Vcov Method: ", x$vcov, "\n", sep = "")
  cat("Inference:   ")
  cat(ifelse(x$reml, "REML", "ML"))
  cat("\n\n")
  cat("Model selection criteria:\n")
  h_print_aic_list(x$aic_list)
  cat("\n")
  cat("Coefficients: ")
  if (x$n_singular_coefs > 0) {
    cat(
      "(",
      x$n_singular_coefs,
      " not defined because of singularities)",
      sep = ""
    )
  }
  cat("\n")
  stats::printCoefmat(
    x$coefficients,
    zap.ind = 3,
    digits = digits,
    signif.stars = signif.stars
  )
  cat("\n")
  cat("Covariance estimate:\n")
  if (is.list(x$varcor)) {
    for (v in names(x$varcor)) {
      cat(sprintf("Group: %s\n", v))
      print(round(x$varcor[[v]], digits = digits))
    }
  } else {
    print(round(x$varcor, digits = digits))
  }
  cat("\n")
  invisible(x)
}


#' @describeIn mmrm_methods obtain the confidence intervals for the coefficients.
#' @exportS3Method
#' @examples
#' # Confidence Interval:
#' confint(object)
confint.mmrm <- function(object, parm, level = 0.95, ...) {
  cf <- coef(object)
  pnames <- names(cf)
  if (missing(parm)) {
    parm <- pnames
  }
  assert(
    check_subset(parm, pnames),
    check_integerish(parm, lower = 1L, upper = length(cf))
  )
  if (is.numeric(parm)) {
    parm <- pnames[parm]
  }
  assert_number(level, lower = 0, upper = 1)
  a <- (1 - level) / 2
  pct <- paste(
    format(100 * c(a, 1 - a), trim = TRUE, scientific = FALSE, digits = 3),
    "%"
  )
  coef_table <- h_coef_table(object)
  df <- coef_table[parm, "df"]
  ses <- coef_table[parm, "Std. Error"]
  fac <- stats::qt(a, df = df)
  ci <- array(NA, dim = c(length(parm), 2L), dimnames = list(parm, pct))
  sefac <- ses * fac
  ci[] <- cf[parm] + c(sefac, -sefac)
  ci
}


#' Analysis of Variance for `mmrm` Fits
#'
#' If supplied only one model fit, the function will calculate and return the
#' significance of the model terms. If supplied more than one model fit,
#' standard diagnostics will be returned for each model, optionally including
#' likelihood ratio test (LRT) results for adjacent models.
#'
#' @param object (`mmrm`)\cr an `mmrm` model fit.
#' @param ... (`mmrm`)\cr optional `mmrm` model fits. If left empty, the
#'   significance of each term in `object` will be calculated.
#' @param test (`flag`)\cr indicating whether the output should include
#'   likelihood ratio test (LRT) results comparing the model fits to one
#'   another. Defaults to `TRUE`. Ignored if `...` is empty.
#' @param refit (`flag`)\cr indicating whether the models should be refitted
#'   with the dataset consisting of their shared set of observations before
#'   performing diagnostics and testing. This is ignored if the models already
#'   share the same dataset. If `refit = FALSE` and the models have different
#'   underlying data sets, an error will be thrown. Defaults to `FALSE`. Ignored
#'   if `...` is empty.
#'
#' @details When `test = FALSE` (or, when only one model is supplied), this
#'   function will process any `mmrm` fits, related or unrelated.
#'
#'   When supplying multiple models and `test = TRUE`, adjacent pairs of models
#'   are tested sequentially. In other words, the order of the supplied models
#'   matters. Furthermore, there are are multiple requirements for successful
#'   LRT. See the section "Requirements for LRT" below.
#'
#'   # Requirements for LRT
#'
#'   1. Each supplied model fit must have more degrees of freedom than
#'   the preceding model fit.
#'
#'   1. If all supplied models were estimated using maximum likelihood (ML), the
#'   models must have nested covariates in order to undergo LRT. In other words,
#'   the set of covariates for each model must be a subset of the covariates of
#'   the next model. However, if any of the supplied models were estimated using
#'   restricted maximum likelihood (REML), all models must have the same
#'   covariates.
#'
#'   1. The covariance structure of each model must be either (a) the same as
#'   that of the next model or (b) a special case of that of the next model. See
#'   the section "Covariance structure nesting hierarchy" below.
#'
#'   1. All supplied model fits must either already use the same data or be
#'   refitted using `refit = TRUE`, which refits all models to the dataset of
#'   common observations between all models' respective data sets.
#'
#'   # Covariance structure nesting hierarchy
#'
#'   ## Structured nests within unstructured
#'
#'   Tautologically, all covariance structures are special cases of an
#'   unstructured covariance, and a model *with* a covariance structure can be
#'   considered "nested" within an model *without* a covariance structure
#'   (assuming that the covariates are also nested).
#'
#'   ## Homogeneous nests within analogous heterogeneous
#'
#'   All homogeneous covariance structures are nested within their corresponding
#'   heterogeneous counterparts. For instance, the homogeneous Toeplitz
#'   covariance structure is nested within the heterogeneous Toeplitz covariance
#'   structure.
#'
#'   ## Other nested structures
#'
#'   Some different covariance structure types are also nested:
#'
#'   - First-order auto-regressive (`ar1` / `ar1h`) is nested within:
#'     - ante-dependence (`ad` / `adh`)
#'     - Toeplitz (`toep` / `toeph`)
#'   - Compound symmetry (`cs` / `csh`) is nested within Toeplitz (`toep` / `toeph`)
#'
#' @returns A data frame with a row for each supplied model. If `...` is empty,
#'   this will be the the returned value of [h_anova_single_mmrm_model()] with
#'   `object` supplied as its argument. Otherwise, the resulting data frame will
#'   have the following columns:
#'
#'   - `Model`: the sequence number of the model according to the order in which
#'   the models were supplied to this function.
#'
#'   - `refit`: logical, indicating whether or not the model was refitted. If
#'   the `refit` argument was `FALSE`, all values will be `FALSE`.
#'
#'   - `REML`: logical, indicating whether or not the model was fitted using
#'   restricted maximum likelihood (REML) estimation. If `FALSE`, the model was
#'   fitted using maximum likelihood (ML) estimation.
#'
#'   - `n_param`: the number of variance parameters in the model fit, obtained
#'   via [logLik.mmrm_tmb()].
#'
#'   - `n_coef`: the number of estimated coefficients in the model fit, obtained
#'   via [logLik.mmrm_tmb()].
#'
#'   - `df`: degrees of freedom of the model fit, obtained via
#'   [logLik.mmrm_tmb()].
#'
#'   - `AIC`: Akaike's "An Information Criterion" of the model fit, obtained via
#'   [stats::AIC()].
#'
#'   - `BIC`: the Bayesian Information Criterion of the model fit, obtained via
#'   [stats::BIC()].
#'
#'   - `logLik`: the log likelihood of the model fit, obtained via
#'   [logLik.mmrm_tmb()].
#'
#'   - `call`: the [call] that created the model fit, obtained via [component()]
#'   with `name = "call"`, which is passed to [deparse1()]. If the model was
#'   refitted (i.e., if its `refit` entry in this table is `TRUE`), this `call`
#'   will be different from the `call` component of the pre-refitted model fit.
#'
#'   The data frame will have these additional columns inserted before `call` if
#'   `test = TRUE`. Note that since each of these columns describe the results
#'   of a likelihood ratio test (LRT) between the previous row's and current
#'   row's model fits, the first element of each of these columns will be `NA`.
#'
#' - `test`: character, indicating which two model fits were compared. Values
#'  are of the form `{Model - 1} vs {Model}`.
#'
#' - `log_likelihood_ratio`: the logarithm of the likelihood ratio between the two
#'  models being compared.
#'
#' - `p_value`: the p-value of the `log_likelihood_ratio`.
#'
#' @seealso For details on the single model operation of this function, see
#'   [h_anova_single_mmrm_model()]. For details on the generic, see
#'   [stats::anova()].
#'
#' @export
#'
#' @name stats_anova
#'
#' @examples
#' # Create a few model fits, only adding terms from one to the next.
#' # Notice also that each covariance structure is a special case of the one
#' # that follows.
#' fit_sex_ar1 <-
#'   mmrm(FEV1 ~ FEV1_BL + SEX + ARMCD + ar1(AVISIT | USUBJID),
#'     data = fev_data,
#'     reml = FALSE
#'   )
#' fit_sex_race_toeph <-
#'   mmrm(
#'     FEV1 ~ FEV1_BL + SEX + RACE + ARMCD + toeph(AVISIT | USUBJID),
#'     data = fev_data,
#'     reml = FALSE
#'   )
#' fit_interaction_us <-
#'   mmrm(
#'     FEV1 ~ FEV1_BL + SEX * RACE + ARMCD + us(AVISIT | USUBJID),
#'     data = fev_data,
#'     reml = FALSE
#'   )
#'
#' # Single model fit, showing significance of model terms:
#' anova(fit_interaction_us)
#'
#' # Multiple model fits, with diagnostics for each fit and likelihood ratio
#' # testing (LRT) for each adjacent pair. LRT is possible because when the fits
#' # are in this order, their covariates and covariance structures are nested.
#' anova(fit_sex_ar1, fit_sex_race_toeph, fit_interaction_us)
#'
#' # We can only change the order if we forego LRT using test = FALSE.
#' anova(fit_sex_race_toeph, fit_interaction_us, fit_sex_ar1, test = FALSE)
#'
#' # Create a subset of fev_data set with the 4th visit removed.
#' fev_subset <- droplevels(fev_data[fev_data$VISITN < 4, ])
#'
#' # Recreate fit_sex_race_toeph but this time based off fev_subset:
#' fit_sex_race_toeph_sub <-
#'   mmrm(
#'     FEV1 ~ FEV1_BL + SEX + RACE + ARMCD + toeph(AVISIT | USUBJID),
#'     data = fev_subset,
#'     reml = FALSE
#'   )
#'
#' # If a model was created with a different data set, refit = TRUE is needed.
#' anova(fit_sex_ar1, fit_sex_race_toeph_sub, fit_interaction_us, refit = TRUE)
anova.mmrm <- function(object, ..., test = TRUE, refit = FALSE) {
  assert_class(object, "mmrm")

  fits <- list(object, ...)

  if (length(fits) == 1L) {
    out <- h_anova_single_mmrm_model(object)
  } else {
    # Ensure all objects in ... are mmrm fits.
    lapply(fits[-1L], assert_class, classes = "mmrm", .var.name = "...")

    assert_flag(test)
    assert_flag(refit)

    # If the data are not all the same and refit = TRUE, refit all with the
    # largest common data set.
    if (refit && !h_check_fits_all_data_same(fits)) {
      common_data <- h_fits_common_data(fits)

      fit_data <- lapply(fits, h_get_minimal_fit_data)

      nested_in_common_data <-
        vapply(
          fit_data,
          h_check_columns_nested,
          FUN.VALUE = logical(1L),
          data_augmented = common_data
        )

      needs_refit <- !nested_in_common_data

      fits[needs_refit] <-
        lapply(fits[needs_refit], h_refit_mmrm, data = common_data)
    } else {
      needs_refit <- rep_len(FALSE, length(fits))
    }

    log_likelihood_vec <- lapply(fits, logLik)

    out <-
      data.frame(
        Model = seq_along(fits),
        refit = needs_refit,
        REML = vapply(fits, component, logical(1L), "reml"),
        n_param = vapply(log_likelihood_vec, attr, numeric(1L), "n_param"),
        n_coef = vapply(log_likelihood_vec, attr, numeric(1L), "n_coef"),
        df = vapply(log_likelihood_vec, attr, numeric(1L), "df"),
        AIC = vapply(log_likelihood_vec, AIC, numeric(1L)),
        BIC = vapply(log_likelihood_vec, BIC, numeric(1L)),
        logLik = as.numeric(log_likelihood_vec)
      )

    if (test) {
      h_assert_lrt_suitability(fits, refit, dfs = out$df, is_reml = out$REML)

      model_indices_except_last <- out$Model[-length(fits)]
      model_indices_except_first <- out$Model[-1L]

      # Create labels for the pair being compared (e.g., "3 vs 4"). Force the
      # first element to be NA because a pair of models is the previous row's
      # model plus the current row's model.
      out$test <-
        c(
          NA_character_,
          paste(model_indices_except_last, "vs", model_indices_except_first)
        )

      out$log_likelihood_ratio <- c(NA, diff(out$logLik))

      out$p_value <-
        stats::pchisq(
          2 * abs(out$log_likelihood_ratio),
          df = abs(diff(out$df)),
          lower.tail = FALSE
        )
    }

    fit_calls <- lapply(fits, component, "call")
    out$call <- vapply(fit_calls, deparse1, FUN.VALUE = character(1L))
  }

  class(out) <- union("anova.mmrm", class(out))

  out
}

#' Obtain Kenward-Roger Adjustment Components
#'
#' @description Obtains the components needed downstream for the computation of Kenward-Roger degrees of freedom.
#' Used in [mmrm()] fitting if method is "Kenward-Roger".
#'
#' @param tmb_data (`mmrm_tmb_data`)\cr produced by [h_mmrm_tmb_data()].
#' @param theta (`numeric`)\cr theta estimate.
#'
#' @details the function returns a named list, \eqn{P}, \eqn{Q} and \eqn{R}, which corresponds to the
#' paper in 1997. The matrices are stacked in columns so that \eqn{P}, \eqn{Q} and \eqn{R} has the same
#' column number(number of beta parameters). The number of rows, is dependent on
#' the total number of theta and number of groups, if the fit is a grouped mmrm.
#' For \eqn{P} matrix, it is stacked sequentially. For \eqn{Q} and \eqn{R} matrix, it is stacked so
#' that the \eqn{Q_{ij}} and \eqn{R_{ij}} is stacked from \eqn{j} then to \eqn{i}, i.e. \eqn{R_{i1}}, \eqn{R_{i2}}, etc.
#' \eqn{Q} and \eqn{R} only contains intra-group results and inter-group results should be all zero matrices
#' so they are not stacked in the result.
#'
#' @return Named list with elements:
#' - `P`: `matrix` of \eqn{P} component.
#' - `Q`: `matrix` of \eqn{Q} component.
#' - `R`: `matrix` of \eqn{R} component.
#'
#' @keywords internal
h_get_kr_comp <- function(tmb_data, theta) {
  assert_class(tmb_data, "mmrm_tmb_data")
  assert_class(theta, "numeric")
  .Call(`_mmrm_get_pqr`, PACKAGE = "mmrm", tmb_data, theta)
}

#' Calculation of Kenward-Roger Degrees of Freedom for Multi-Dimensional Contrast
#'
#' @description Used in [df_md()] if method is "Kenward-Roger" or "Kenward-Roger-Linear".
#'
#' @inheritParams h_df_md_sat
#' @inherit h_df_md_sat return
#' @keywords internal
h_df_md_kr <- function(object, contrast) {
  assert_class(object, "mmrm")
  assert_matrix(
    contrast,
    mode = "numeric",
    any.missing = FALSE,
    ncols = length(component(object, "beta_est"))
  )
  if (component(object, "reml") != 1) {
    stop("Kenward-Roger is only for REML")
  }
  kr_comp <- object$kr_comp
  w <- component(object, "theta_vcov")
  v_adj <- object$beta_vcov_adj
  df <- h_kr_df(v0 = object$beta_vcov, l = contrast, w = w, p = kr_comp$P)

  h_test_md(object, contrast, df = df$m, f_stat_factor = df$lambda)
}

#' Calculation of Kenward-Roger Degrees of Freedom for One-Dimensional Contrast
#'
#' @description Used in [df_1d()] if method is
#' "Kenward-Roger" or "Kenward-Roger-Linear".
#'
#' @inheritParams h_df_1d_sat
#' @inherit h_df_1d_sat return
#' @keywords internal
h_df_1d_kr <- function(object, contrast) {
  assert_class(object, "mmrm")
  assert_numeric(contrast, len = length(component(object, "beta_est")))
  if (component(object, "reml") != 1) {
    stop("Kenward-Roger is only for REML!")
  }

  df <- h_kr_df(
    v0 = object$beta_vcov,
    l = matrix(contrast, nrow = 1),
    w = component(object, "theta_vcov"),
    p = object$kr_comp$P
  )

  h_test_1d(object, contrast, df$m)
}

#' Obtain the Adjusted Kenward-Roger degrees of freedom
#'
#' @description Obtains the adjusted Kenward-Roger degrees of freedom and F statistic scale parameter.
#' Used in [h_df_md_kr()] or [h_df_1d_kr].
#'
#' @param v0 (`matrix`)\cr unadjusted covariance matrix.
#' @param l (`matrix`)\cr linear combination matrix.
#' @param w (`matrix`)\cr hessian matrix.
#' @param p (`matrix`)\cr P matrix from [h_get_kr_comp()].
#'
#' @return Named list with elements:
#' - `m`: `numeric` degrees of freedom.
#' - `lambda`: `numeric` F statistic scale parameter.
#'
#' @keywords internal
h_kr_df <- function(v0, l, w, p) {
  n_beta <- ncol(v0)
  assert_matrix(v0, ncols = n_beta, nrows = n_beta)
  assert_matrix(l, ncols = n_beta)
  n_theta <- ncol(w)
  assert_matrix(w, ncols = n_theta, nrows = n_theta)
  n_visits <- ncol(p)
  assert_matrix(p, nrows = n_visits * n_theta)
  # see vignettes/kenward.Rmd#279
  slvol <- solve(h_quad_form_mat(l, v0))
  m <- h_quad_form_mat(t(l), slvol)
  nl <- nrow(l)
  mv0 <- m %*% v0
  pl <- lapply(seq_len(nrow(p) / ncol(p)), function(x) {
    ii <- (x - 1) * ncol(p) + 1
    jj <- x * ncol(p)
    p[ii:jj, ]
  })
  mv0pv0 <- lapply(pl, function(x) {
    mv0 %*% x %*% v0
  })
  a1 <- 0
  a2 <- 0
  # see vignettes/kenward.Rmd#283
  for (i in seq_len(length(pl))) {
    for (j in seq_len(length(pl))) {
      a1 <- a1 + w[i, j] * h_tr(mv0pv0[[i]]) * h_tr(mv0pv0[[j]])
      a2 <- a2 + w[i, j] * h_tr(mv0pv0[[i]] %*% mv0pv0[[j]])
    }
  }
  b <- 1 / (2 * nl) * (a1 + 6 * a2)
  e <- 1 + a2 / nl
  e_star <- 1 / (1 - a2 / nl)
  g <- ((nl + 1) * a1 - (nl + 4) * a2) / ((nl + 2) * a2)
  denom <- (3 * nl + 2 - 2 * g)
  c1 <- g / denom
  c2 <- (nl - g) / denom
  c3 <- (nl + 2 - g) / denom
  v_star <- 2 / nl * (1 + c1 * b) / (1 - c2 * b)^2 / (1 - c3 * b)
  rho <- v_star / (2 * e_star^2)
  m <- 4 + (nl + 2) / (nl * rho - 1)
  lambda <- m / (e_star * (m - 2))
  list(m = m, lambda = lambda)
}

#' Obtain the Adjusted Covariance Matrix
#'
#' @description Obtains the Kenward-Roger adjusted covariance matrix for the
#'   coefficient estimates.
#' Used in [mmrm()] fitting if method is "Kenward-Roger" or "Kenward-Roger-Linear".
#'
#' @param v (`matrix`)\cr unadjusted covariance matrix.
#' @param w (`matrix`)\cr hessian matrix.
#' @param p (`matrix`)\cr P matrix from [h_get_kr_comp()].
#' @param q (`matrix`)\cr Q matrix from [h_get_kr_comp()].
#' @param r (`matrix`)\cr R matrix from [h_get_kr_comp()].
#' @param linear (`flag`)\cr whether to use linear Kenward-Roger approximation.
#'
#' @return The matrix of adjusted covariance matrix.
#'
#' @keywords internal
h_var_adj <- function(v, w, p, q, r, linear = FALSE) {
  assert_flag(linear)
  n_beta <- ncol(v)
  assert_matrix(v, nrows = n_beta)
  n_theta <- ncol(w)
  assert_matrix(w, nrows = n_theta)
  n_visits <- ncol(p)
  theta_per_group <- nrow(q) / nrow(p)
  n_groups <- n_theta / theta_per_group
  assert_matrix(p, nrows = n_theta * n_visits)
  assert_matrix(
    q,
    nrows = theta_per_group^2 * n_groups * n_visits,
    ncols = n_visits
  )
  assert_matrix(
    r,
    nrows = theta_per_group^2 * n_groups * n_visits,
    ncols = n_visits
  )
  if (linear) {
    r <- matrix(0, nrow = nrow(r), ncol = ncol(r))
  }

  # see vignettes/kenward.Rmd#131
  ret <- v
  for (i in seq_len(n_theta)) {
    for (j in seq_len(n_theta)) {
      gi <- ceiling(i / theta_per_group)
      gj <- ceiling(j / theta_per_group)
      iid <- (i - 1) * n_beta + 1
      jid <- (j - 1) * n_beta + 1
      ii <- i - (gi - 1) * theta_per_group
      jj <- j - (gi - 1) * theta_per_group
      ijid <- ((ii - 1) * theta_per_group + jj - 1) *
        n_beta +
        (gi - 1) * n_beta * theta_per_group^2 +
        1
      if (gi != gj) {
        ret <- ret +
          2 *
            w[i, j] *
            v %*%
              (-p[iid:(iid + n_beta - 1), ] %*%
                v %*%
                p[jid:(jid + n_beta - 1), ]) %*%
              v
      } else {
        ret <- ret +
          2 *
            w[i, j] *
            v %*%
              (q[ijid:(ijid + n_beta - 1), ] -
                p[iid:(iid + n_beta - 1), ] %*%
                  v %*%
                  p[jid:(jid + n_beta - 1), ] -
                1 / 4 * r[ijid:(ijid + n_beta - 1), ]) %*%
              v
      }
    }
  }
  ret
}

#' Capture all Output
#'
#' This function silences all warnings, errors & messages and instead returns a list
#' containing the results (if it didn't error), as well as the warnings, errors
#' and messages and divergence signals as character vectors.
#'
#' @param expr (`expression`)\cr to be executed.
#' @param remove (`list`)\cr optional list with elements `warnings`, `errors`,
#'   `messages` which can be character vectors, which will be removed from the
#'   results if specified.
#' @param divergence (`list`)\cr optional list similar as `remove`, but these
#'   character vectors will be moved to the `divergence` result and signal
#'   that the fit did not converge.
#'
#' @return
#' A list containing
#'
#' - `result`: The object returned by `expr` or `list()` if an error was thrown.
#' - `warnings`: `NULL` or a character vector if warnings were thrown.
#' - `errors`: `NULL` or a string if an error was thrown.
#' - `messages`: `NULL` or a character vector if messages were produced.
#' - `divergence`: `NULL` or a character vector if divergence messages were caught.
#'
#' @keywords internal
h_record_all_output <- function(expr, remove = list(), divergence = list()) {
  # Note: We don't need to and cannot assert `expr` here.
  assert_list(remove, types = "character")
  assert_list(divergence, types = "character")
  env <- new.env()
  result <- withCallingHandlers(
    withRestarts(
      expr,
      muffleStop = function(e) structure(e$message, class = "try-error")
    ),
    message = function(m) {
      msg_without_newline <- gsub(m$message, pattern = "\n$", replacement = "")
      env$message <- c(env$message, msg_without_newline)
      invokeRestart("muffleMessage")
    },
    warning = function(w) {
      env$warning <- c(env$warning, w$message)
      invokeRestart("muffleWarning")
    },
    error = function(e) {
      env$error <- c(env$error, e$message)
      invokeRestart("muffleStop", e)
    }
  )
  list(
    result = result,
    warnings = setdiff(env$warning, c(remove$warnings, divergence$warnings)),
    errors = setdiff(env$error, c(remove$errors, divergence$errors)),
    messages = setdiff(env$message, c(remove$messages, divergence$messages)),
    divergence = c(
      intersect(env$warning, divergence$warnings),
      intersect(env$error, divergence$errors),
      intersect(env$message, divergence$messages)
    )
  )
}

#' Trace of a Matrix
#'
#' @description Obtain the trace of a matrix if the matrix is diagonal, otherwise raise an error.
#'
#' @param x (`matrix`)\cr square matrix input.
#'
#' @return The trace of the square matrix.
#'
#' @keywords internal
h_tr <- function(x) {
  if (nrow(x) != ncol(x)) {
    stop("x must be square matrix")
  }
  sum(Matrix::diag(x))
}

#' Split Control List
#'
#' @description Split the [mmrm_control()] object according to its optimizers and use additional arguments
#' to replace the elements in the original object.
#'
#' @param control (`mmrm_control`)\cr object.
#' @param ... additional parameters to update the `control` object.
#'
#' @return A `list` of `mmrm_control` entries.
#' @keywords internal
h_split_control <- function(control, ...) {
  assert_class(control, "mmrm_control")
  l <- length(control$optimizers)
  lapply(seq_len(l), function(i) {
    ret <- utils::modifyList(control, list(...))
    ret$optimizers <- control$optimizers[i]
    ret
  })
}

#' Obtain Optimizer according to Optimizer String Value
#'
#' @description This function creates optimizer functions with arguments.
#'
#' @param optimizer (`character`)\cr names of built-in optimizers to try, subset
#'   of "L-BFGS-B", "BFGS", "CG" and "nlminb".
#' @param optimizer_fun (`function` or `list` of `function`)\cr alternatively to `optimizer`,
#'   an optimizer function or a list of optimizer functions can be passed directly here.
#' @param optimizer_args (`list`)\cr additional arguments for `optimizer_fun`.
#' @param optimizer_control (`list`)\cr passed to argument `control` in `optimizer_fun`.
#'
#' @details
#' If you want to use only the built-in optimizers:
#' - `optimizer` is a shortcut to create a list of built-in optimizer functions
#'   passed to `optimizer_fun`.
#' - Allowed are "L-BFGS-B", "BFGS", "CG" (using [stats::optim()] with corresponding method)
#'   and "nlminb" (using [stats::nlminb()]).
#' - Other arguments should go into `optimizer_args`.
#'
#' If you want to use your own optimizer function:
#' - Make sure that there are three arguments: parameter (start value), objective function
#'   and gradient function are sequentially in the function arguments.
#' - If there are other named arguments in front of these, make sure they are correctly
#'   specified through `optimizer_args`.
#' - If the hessian can be used, please make sure its argument name is `hessian` and
#'   please add attribute `use_hessian = TRUE` to the function,
#'   using `attr(fun, "use_hessian) <- TRUE`.
#'
#' @return Named `list` of optimizers created by [h_partial_fun_args()].
#'
#' @keywords internal
h_get_optimizers <- function(
  optimizer = c("L-BFGS-B", "BFGS", "CG", "nlminb"),
  optimizer_fun = h_optimizer_fun(optimizer),
  optimizer_args = list(),
  optimizer_control = list()
) {
  if ("automatic" %in% optimizer) {
    lifecycle::deprecate_warn(
      when = "0.2.0",
      what = I("\"automatic\" optimizer"),
      details = "please just omit optimizer argument"
    )
    optimizer_fun <- h_optimizer_fun()
  }
  assert(
    test_function(optimizer_fun),
    test_list(optimizer_fun, types = "function", names = "unique")
  )
  if (is.function(optimizer_fun)) {
    optimizer_fun <- list(custom_optimizer = optimizer_fun)
  }
  lapply(optimizer_fun, function(x) {
    do.call(
      h_partial_fun_args,
      c(list(fun = x, control = optimizer_control), optimizer_args)
    )
  })
}

#' Obtain Optimizer Function with Character
#' @description Obtain the optimizer function through the character provided.
#' @param optimizer (`character`)\cr vector of optimizers.
#'
#' @return A (`list`)\cr of optimizer functions generated from [h_partial_fun_args()].
#' @keywords internal
h_optimizer_fun <- function(optimizer = c("L-BFGS-B", "BFGS", "CG", "nlminb")) {
  optimizer <- match.arg(optimizer, several.ok = TRUE)
  lapply(stats::setNames(optimizer, optimizer), function(x) {
    switch(
      x,
      "L-BFGS-B" = h_partial_fun_args(fun = stats::optim, method = x),
      "BFGS" = h_partial_fun_args(fun = stats::optim, method = x),
      "CG" = h_partial_fun_args(fun = stats::optim, method = x),
      "nlminb" = h_partial_fun_args(
        fun = stats::nlminb,
        additional_attr = list(use_hessian = TRUE)
      )
    )
  })
}

#' Create Partial Functions
#' @description Creates partial functions with arguments.
#'
#' @param fun (`function`)\cr to be wrapped.
#' @param ... Additional arguments for `fun`.
#' @param additional_attr (`list`)\cr of additional attributes to apply to the result.
#'
#' @details This function add `args` attribute to the original function,
#' and add an extra class `partial` to the function.
#' `args` is the argument for the function, and elements in `...` will override the existing
#' arguments in attribute `args`. `additional_attr` will override the existing attributes.
#'
#' @return Object with S3 class `"partial"`, a `function` with `args` attribute (and possibly more
#' attributes from `additional_attr`).
#' @keywords internal
h_partial_fun_args <- function(fun, ..., additional_attr = list()) {
  assert_function(fun)
  assert_list(additional_attr, names = "unique")
  a_args <- list(...)
  assert_list(a_args, names = "unique")
  args <- attr(fun, "args")
  if (is.null(args)) {
    args <- list()
  }
  do.call(
    structure,
    args = utils::modifyList(
      list(
        .Data = fun,
        args = utils::modifyList(args, a_args),
        class = c("partial", "function")
      ),
      additional_attr
    )
  )
}

#' Obtain Default Covariance Method
#'
#' @description Obtain the default covariance method depending on
#' the degrees of freedom method used.
#'
#' @param method (`string`)\cr degrees of freedom method.
#'
#' @details The default covariance method is different for different degrees of freedom method.
#' For "Satterthwaite" or "Between-Within", "Asymptotic" is returned.
#' For "Kenward-Roger" only, "Kenward-Roger" is returned.
#' For "Residual" only, "Empirical" is returned.
#'
#' @return String of the default covariance method.
#' @keywords internal
h_get_cov_default <- function(
  method = c("Satterthwaite", "Kenward-Roger", "Residual", "Between-Within")
) {
  assert_string(method)
  method <- match.arg(method)
  switch(
    method,
    "Residual" = "Empirical",
    "Satterthwaite" = "Asymptotic",
    "Kenward-Roger" = "Kenward-Roger",
    "Between-Within" = "Asymptotic"
  )
}

#' Complete `character` Vector Names From Values
#'
#' @param x (`character` or `list`)\cr value whose names should be completed
#'   from element values.
#'
#' @return A named vector or list.
#'
#' @keywords internal
fill_names <- function(x) {
  n <- names(x)
  is_unnamed <- if (is.null(n)) rep_len(TRUE, length(x)) else n == ""
  names(x)[is_unnamed] <- x[is_unnamed]
  x
}

#' Drop Items from an Indexible
#'
#' Drop elements from an indexible object (`vector`, `list`, etc.).
#'
#' @param x Any object that can be consumed by [seq_along()] and indexed by a
#'   logical vector of the same length.
#' @param n (`integer`)\cr the number of terms to drop.
#'
#' @return A subset of `x`.
#'
#' @keywords internal
drop_elements <- function(x, n) {
  x[seq_along(x) > n]
}

#' Ask for Confirmation on Large Visit Levels
#'
#' @description Ask the user for confirmation if there are too many visit levels
#' for non-spatial covariance structure in interactive sessions.
#'
#' @param x (`numeric`)\cr number of visit levels.
#'
#' @return Logical value `TRUE`.
#' @keywords internal
h_confirm_large_levels <- function(x) {
  assert_count(x)
  allowed_lvls <- x <= getOption("mmrm.max_visits", 100)
  if (allowed_lvls) {
    return(TRUE)
  }
  if (!interactive()) {
    stop("Visit levels too large!", call. = FALSE)
  }
  proceed <- utils::askYesNo(
    paste(
      "Visit levels is possibly too large.",
      "This requires large memory. Are you sure to continue?",
      collapse = " "
    )
  )
  if (!identical(proceed, TRUE)) {
    stop("Visit levels too large!", call. = FALSE)
  }
  TRUE
}

#' Default Value on NULL
#' Return default value when first argument is NULL.
#'
#' @param x Object.
#' @param y Object.
#'
#' @details If `x` is NULL, returns `y`. Otherwise return `x`.
#'
#' @keywords internal
h_default_value <- function(x, y) {
  if (is.null(x)) {
    y
  } else {
    x
  }
}

#' Warn on na.action
#' @keywords internal
h_warn_na_action <- function() {
  if (!identical(getOption("na.action"), "na.omit")) {
    warning("na.action is always set to `na.omit` for `mmrm` fit!")
  }
}

#' Obtain `na.action` as Function
#' @keywords internal
h_get_na_action <- function(na_action) {
  if (
    is.function(na_action) &&
      identical(methods::formalArgs(na_action), c("object", "..."))
  ) {
    return(na_action)
  }
  if (is.character(na_action) && length(na_action) == 1L) {
    assert_subset(
      na_action,
      c("na.omit", "na.exclude", "na.fail", "na.pass", "na.contiguous")
    )
    get(na_action, mode = "function", pos = "package:stats")
  }
}

#' Validate mmrm Formula
#' @param formula (`formula`)\cr to check.
#'
#' @details In mmrm models, `.` is not allowed as it introduces ambiguity of covariates
#' to be used, so it is not allowed to be in formula.
#'
#' @keywords internal
h_valid_formula <- function(formula) {
  assert_formula(formula)
  if ("." %in% all.vars(formula)) {
    stop("`.` is not allowed in mmrm models!")
  }
}

#' Standard Starting Value
#'
#' @description Obtain standard start values.
#'
#' @param cov_type (`string`)\cr name of the covariance structure.
#' @param n_visits (`int`)\cr number of visits.
#' @param n_groups (`int`)\cr number of groups.
#' @param ... not used.
#'
#' @details
#' `std_start` will try to provide variance parameter from identity matrix.
#' However, for `ar1` and `ar1h` the corresponding values are not ideal because the
#' \eqn{\rho} is usually a positive number thus using 0 as starting value can lead to
#' incorrect optimization result, and we use 0.5 as the initial value of \eqn{\rho}.
#'
#' @return A numeric vector of starting values.
#'
#' @export
std_start <- function(cov_type, n_visits, n_groups, ...) {
  assert_string(cov_type)
  assert_subset(cov_type, cov_types(c("abbr", "habbr")))
  assert_int(n_visits, lower = 1L)
  assert_int(n_groups, lower = 1L)
  start_value <- switch(
    cov_type,
    us = rep(0, n_visits * (n_visits + 1) / 2),
    toep = rep(0, n_visits),
    toeph = rep(0, 2 * n_visits - 1),
    ar1 = c(0, 0.5),
    ar1h = c(rep(0, n_visits), 0.5),
    ad = rep(0, n_visits),
    adh = rep(0, 2 * n_visits - 1),
    cs = rep(0, 2),
    csh = rep(0, n_visits + 1),
    sp_exp = rep(0, 2)
  )
  rep(start_value, n_groups)
}

#' Empirical Starting Value
#'
#' @description Obtain empirical start value for unstructured covariance
#'
#' @param y_vector (`numeric`) response variable.
#' @param x_matrix (`matrix`) design matrix.
#' @param full_frame (`data.frame`) full data frame used for model fitting.
#' @param visit_var (`string`)\cr visit variable.
#' @param subject_var (`string`)\cr subject id variable.
#' @param subject_groups (`factor`)\cr subject group assignment.
#' @param ... not used.
#'
#' @details
#' This `emp_start` only works for unstructured covariance structure.
#' It uses linear regression to first obtain the coefficients and use the residuals
#' to obtain the empirical variance-covariance, and it is then used to obtain the
#' starting values.
#'
#' @return A numeric vector of starting values.
#'
#' @export
emp_start <- function(
  y_vector,
  x_matrix,
  full_frame,
  visit_var,
  subject_var,
  subject_groups,
  ...
) {
  assert_numeric(y_vector)
  assert_matrix(x_matrix)
  assert_data_frame(full_frame)
  assert_string(visit_var)
  assert_string(subject_var)
  assert_factor(full_frame[[visit_var]])
  n_visits <- length(levels(full_frame[[visit_var]]))
  assert_factor(full_frame[[subject_var]])
  subjects <- droplevels(full_frame[[subject_var]])
  n_subjects <- length(levels(subjects))
  fit <- stats::lm.fit(x = x_matrix, y = y_vector)
  res <- rep(NA, n_subjects * n_visits)
  res[
    n_visits *
      as.integer(subjects) -
      n_visits +
      as.integer(full_frame[[visit_var]])
  ] <- residuals(fit)
  res_mat <- matrix(res, ncol = n_visits, nrow = n_subjects, byrow = TRUE)
  emp_covs <- lapply(
    unname(split(seq_len(n_subjects), subject_groups)),
    function(x) {
      stats::cov(res_mat[x, , drop = FALSE], use = "pairwise.complete.obs")
    }
  )
  unlist(lapply(emp_covs, h_get_theta_from_cov))
}

#' Obtain Theta from Covariance Matrix
#'
#' @description Obtain unstructured theta from covariance matrix.
#'
#' @param covariance (`matrix`) of covariance matrix values.
#'
#' @details
#' If the covariance matrix has `NA` in some of the elements, they will be replaced by
#' 0 (non-diagonal) and 1 (diagonal). This ensures that the matrix is positive definite.
#'
#' @return Numeric vector of the theta values.
#' @keywords internal
h_get_theta_from_cov <- function(covariance) {
  assert_matrix(covariance, mode = "numeric", ncols = nrow(covariance))
  covariance[is.na(covariance)] <- 0
  diag(covariance)[diag(covariance) == 0] <- 1
  # empirical is not always positive definite in some special cases of numeric singularity.
  qr_res <- qr(covariance)
  if (qr_res$rank < ncol(covariance)) {
    covariance <- Matrix::nearPD(covariance)$mat
  }
  emp_chol <- t(chol(covariance))
  mat <- t(solve(diag(diag(emp_chol)), emp_chol))
  ret <- c(log(diag(emp_chol)), mat[upper.tri(mat)])
  unname(ret)
}

#' Register S3 Method
#' Register S3 method to a generic.
#'
#' @param pkg (`string`) name of the package name.
#' @param generic (`string`) name of the generic.
#' @param class (`string`) class name the function want to dispatch.
#' @param envir (`environment`) the location the method is defined.
#'
#' @details This function is adapted from `emmeans:::register_s3_method()`.
#'
#' @keywords internal
h_register_s3 <- function(pkg, generic, class, envir = parent.frame()) {
  assert_string(pkg)
  assert_string(generic)
  assert_string(class)
  assert_environment(envir)
  fun <- get(paste0(generic, ".", class), envir = envir)
  if (isNamespaceLoaded(pkg)) {
    registerS3method(generic, class, fun, envir = asNamespace(pkg))
  }
  setHook(packageEvent(pkg, "onLoad"), function(...) {
    registerS3method(generic, class, fun, envir = asNamespace(pkg))
  })
}

#' Check if a Factor Should Drop Levels
#'
#' @param x (`vector`) vector to check.
#'
#' @keywords internal
h_extra_levels <- function(x) {
  is.factor(x) && length(levels(x)) > length(unique(x))
}

#' Drop Levels from Dataset
#' @param data (`data.frame`) data to drop levels.
#' @param subject_var (`character`) subject variable.
#' @param visit_var (`character`) visit variable.
#' @param except (`character`) variables to exclude from dropping.
#' @keywords internal
h_drop_levels <- function(data, subject_var, visit_var, except) {
  assert_data_frame(data)
  assert_character(subject_var)
  assert_character(visit_var)
  assert_character(except, null.ok = TRUE)
  all_cols <- colnames(data)
  to_drop <- vapply(
    data,
    h_extra_levels,
    logical(1L)
  )
  to_drop <- all_cols[to_drop]
  # only drop levels for those not defined in excep and not in visit_var.
  to_drop <- setdiff(to_drop, c(visit_var, except))
  data[to_drop] <- lapply(data[to_drop], droplevels)
  # subject var are always dropped and no message given.
  dropped <- setdiff(to_drop, subject_var)
  if (length(dropped) > 0) {
    message(
      "Some factor levels are dropped due to singular design matrix: ",
      toString(dropped)
    )
  }
  data
}

#' Predicate if the TMB Version Used to Compile the Package is Sufficient
#'
#' @return Flag whether the TMB version is sufficient.
#' @keywords internal
h_tmb_version_sufficient <- function() {
  # Note: There is no version information saved in the dynamic library, but
  # we can check like this:
  tmb_config <- TMB::config(DLL = "mmrm")
  tape_deterministic <- tmb_config$tmbad_deterministic_hash
  !is.null(tape_deterministic)
}

#' Warn if TMB is Configured to Use Non-Deterministic Hash for Tape Optimizer
#'
#' This function checks the TMB configuration for the `tmbad_deterministic_hash` setting
#' If it is set to `FALSE`, a warning is issued indicating that this may lead to
#' unreproducible results.
#'
#' @return No return value, called for side effects.
#' @keywords internal
h_tmb_warn_non_deterministic <- function() {
  if (!h_tmb_version_sufficient()) {
    return()
  }
  tmb_config <- TMB::config(DLL = "mmrm")
  tape_deterministic <- tmb_config$tmbad_deterministic_hash
  if (!tape_deterministic) {
    msg <- paste(
      "TMB is configured to use a non-deterministic hash for its tape optimizer,",
      "and this may lead to unreproducible results.",
      "To disable this behavior, use `TMB::config(tmbad_deterministic_hash = 1)`.",
      sep = "\n"
    )
    warning(msg)
  }
}


#' Sort a Data Frame by All Its Columns in Ascending Order
#'
#' @param data (`data.frame`)\cr a data frame to be sorted.
#'
#' @returns A data frame. The same as the input `data` with columns in the same
#'   order but with rows sorted by the first column, with ties broken by the
#'   second column, and so forth.
#' @keywords internal
h_dataset_sort_all <- function(data) {
  assert_data_frame(data)
  data[do.call(order, unname(data)), , drop = FALSE]
}


#' Predicate Indicating Whether Two Datasets Contain the Same Observations
#'
#' Checks whether two datasets contain the same observations and that the first
#' dataset's columns are a subset of the second dataset's columns.
#'
#' For efficiency, the inspection takes place in this order:
#'
#' 1. `FALSE` is returned early if the datasets do not have the same number of
#' rows.
#'
#' 1. `FALSE` is returned early if the first dataset has a column not in the
#' second dataset.
#'
#' 1. The columns in common are sorted and compared using [all.equal()] with
#' `check.attributes = FALSE`.
#'
#' @param data_basic,data_augmented (`data.frame`)\cr data frames to be
#'   compared.
#'
#' @returns `TRUE` or `FALSE`, indicating whether the or not `data_basic` and
#'   `data_augmented` contain the same observations and that `data_augmented`
#'   contains all the columns in `data_basic`.
#' @seealso [h_check_fits_all_data_same()], which performs this check on the
#'   datasets of adjacent elements of a list of `mmrm` fits.
#' @keywords internal
h_check_columns_nested <- function(data_basic, data_augmented) {
  assert_data_frame(data_basic)
  assert_data_frame(data_augmented)
  if (nrow(data_basic) != nrow(data_augmented)) {
    return(FALSE)
  }

  colnames_basic <- colnames(data_basic)

  if (anyNA(match(colnames_basic, colnames(data_augmented)))) {
    return(FALSE)
  }

  isTRUE(all.equal(
    h_dataset_sort_all(data_basic),
    h_dataset_sort_all(data_augmented[colnames_basic]),
    check.attributes = FALSE
  ))
}


#' Predicate Indicating Whether `mmrm` Fits' Datasets Contain the Same
#' Observations
#'
#' Checks a `list` of `mmrm` fits to see whether all their datasets contain the
#' same observations and that they only increase in columns from one dataset to
#' the next (i.e, columns are nested).
#'
#' For efficiency, the inspection takes place in this order:
#'
#' 1. `FALSE` is returned early if not all datasets have the same number of
#' rows.
#'
#' 1. `FALSE` is returned early if a dataset has a column not in the next
#' dataset.
#'
#' 1. The columns in common among adjacent datasets are sorted and compared
#' using [all.equal()] with `check.attributes = FALSE`.
#'
#' This function is more efficient than running [h_check_columns_nested()] on
#' all adjacent pairs and supplying the results to [all()].
#'
#' @param fits (`list`)\cr list of `mmrm` fits.
#'
#' @returns `TRUE` or `FALSE` indicating whether or not the datasets underlying
#'   the elements of `fits` contain the same observations, and that each
#'   dataset's columns are a subset of the next dataset's columns.
#' @seealso [h_check_columns_nested()], which performs this check on two data
#'   sets.
#' @keywords internal
h_check_fits_all_data_same <- function(fits) {
  assert_list(fits, types = "mmrm", any.missing = FALSE, min.len = 1)
  datasets <- lapply(fits, h_get_minimal_fit_data)

  # Is nrow() the same for all datasets?
  datasets_nrows <- vapply(datasets, nrow, numeric(1L))
  if (length(unique(datasets_nrows)) > 1L) {
    return(FALSE)
  }

  cols <- lapply(datasets, colnames)

  for (i in seq_along(fits)[-1L]) {
    # Ensure previous dataset's cols are a subset of the current dataset's cols
    if (anyNA(match(cols[[i - 1L]], cols[[i]]))) {
      return(FALSE)
    }
  }

  datasets[[1L]] <- h_dataset_sort_all(datasets[[1L]])
  for (i in seq_along(fits)[-1L]) {
    # Pull prev dataset's columns to the front of current dataset. Then sort.
    cols[[i]] <- union(cols[[i - 1L]], cols[[i]])
    datasets[[i]] <- h_dataset_sort_all(datasets[[i]][cols[[i]]])

    # Ensure the common columns are the same
    if (
      !isTRUE(all.equal(
        datasets[[i - 1L]],
        datasets[[i]][cols[[i - 1L]]],
        check.attributes = FALSE
      ))
    ) {
      return(FALSE)
    }
  }

  TRUE
}


#' Combine the Datasets from `mmrm` Fits
#'
#' Take the data columns used in each `mmrm` fit and [merge()] them.
#'
#' All default arguments for [merge()] are used, resulting in a "natural join":
#' the result will only contain the observations found in all datasets.
#'
#' [droplevels()] is applied to the final product to prevent extraneous
#' warnings.
#'
#' @param fits (`list`)\cr list of `mmrm` fits.
#'
#' @returns A data frame combining all the common observations among the
#'   datasets underlying the elements of `fits`.
#' @keywords internal
h_fits_common_data <- function(fits) {
  assert_list(fits, types = "mmrm", any.missing = FALSE, min.len = 1)
  datasets <- lapply(fits, h_get_minimal_fit_data)
  out <- Reduce(merge, datasets)
  out <- droplevels(out)
  out
}


#' Obtain the Minimal Dataset Needed for an `mmrm` Fit
#'
#' Grab the dataset underlying an `mmrm` fit and select only the used columns.
#'
#' Grabs the response variable along with the predictors named in
#' `fit$formula_parts`.
#'
#' @param fit (`mmrm`)\cr a fitted `mmrm` model.
#'
#' @returns A data frame: a subset of the columns the dataset underlying `fit`
#'   (i.e., `fit$data`):
#'
#' - The response column.
#'
#' - The column denoting the visit index.
#'
#' - The column denoting the subject.
#'
#' - The column denoting the subject's grouping (e.g., study arm).
#'
#' - All other predictors not already specified.
#'
#'   Columns that were not used are excluded.
#'
#' @keywords internal
h_get_minimal_fit_data <- function(fit) {
  assert_class(fit, "mmrm")
  covariance_vars <-
    fit[["formula_parts"]][c("visit_var", "subject_var", "group_var")]
  covariance_vars <- Filter(length, covariance_vars)
  covariance_vars <- lapply(covariance_vars, str2lang)
  covariance_vars <- lapply(covariance_vars, all.vars)
  covariance_vars <- unlist(covariance_vars, use.names = FALSE)
  predictors <- c(covariance_vars, fit[["formula_parts"]][["model_var"]])
  predictors <- unique(predictors)
  terms_attr <- attributes(terms(fit))
  response <- all.vars(terms_attr$variables[[terms_attr$response + 1]])
  cols <- c(response, predictors)
  cols <- intersect(cols, colnames(fit[["data"]]))
  fit[["data"]][cols]
}


#' Ensure LRT Is Appropriate for `list` of `mmrm` Fits
#'
#' Throws an error if the degrees of freedom aren't monotonically increasing, if
#' the models aren't nested, or if `refit = FALSE` and the models have different
#' underlying data.
#'
#' @param fits (`list`)\cr list of `mmrm` fits.
#' @param refit (`flag`)\cr `TRUE` or `FALSE` indicating whether or not the user
#'   gave permission to refit models to make all models suitable for LRT.
#' @param dfs (`numeric`)\cr vector of the degrees of freedom for each element
#'   of `fits`.
#' @param is_reml (`logical`)\cr vector indicating whether or not REML was used
#'   for each element of `fits`.
#'
#' @returns `TRUE` if the list of `fits` are suitable for LRT. Otherwise, an
#'   error is thrown.
#'
#' @keywords internal
h_assert_lrt_suitability <- function(fits, refit, dfs, is_reml) {
  assert_list(fits, types = "mmrm", min.len = 2)
  assert_flag(refit)
  assert_numeric(
    dfs,
    lower = 1,
    any.missing = FALSE,
    finite = TRUE,
    len = length(fits)
  )
  assert_logical(is_reml, any.missing = FALSE, len = length(fits))

  if (any(diff(dfs) <= 0)) {
    stop(
      "The degrees of freedom (df) of each model must be less than the df of ",
      "the next model in order to perform likelihood ratio testing (LRT). ",
      "Bypass LRT with test = FALSE.",
      call. = FALSE
    )
  }

  # Do any models use REML?
  any_reml <- any(is_reml)

  # First iterate through and check the models to make sure their covariates
  # and covariance structures are nested
  for (i in seq_along(fits)[-1L]) {
    h_assert_nested_models(fits[[i - 1L]], fits[[i]], any_reml = any_reml)
  }

  # If we didn't refit, throw errors if not all data are the same
  if (!refit && !h_check_fits_all_data_same(fits)) {
    stop(
      "Likelihood ratio testing requires all fits to use the same data.",
      " Not all fits have the same data and refit = FALSE.",
      " Consider setting test = FALSE or refit = TRUE.",
      call. = FALSE
    )
  }

  TRUE
}


#' Ensure Two Models Are Nested
#'
#' Throws an error if `model_basic` isn't nested within `model_augmented`.
#'
#' The following checks are applied in this order, and an error is thrown if any
#' of the conditions are not met:
#'
#' 1. The fits must have the same visit, subject, and grouping variables.
#'
#' 1. The covariates of `model_basic` must be the same as or a subset of the
#' covariates of `model_augmented`.
#'
#' 1. The covariance structure of `model_basic` must be the same as or a special
#' case of the covariance structure of `model_augmented`.
#'
#' Finally, if all these checks were passed, a warning is thrown if the two fits
#' have identical covariates and covariance structures.
#'
#' @param model_basic,model_augmented (`mmrm`)\cr model fits.
#' @param any_reml (`flag`)\cr `TRUE` or `FALSE` indicating whether or not
#'   either model used REML estimation.
#'
#' @returns `TRUE` if `model_basic` is nested within `model_augmented`.
#'   Otherwise, an error is thrown.
#'
#' @keywords internal
h_assert_nested_models <- function(model_basic, model_augmented, any_reml) {
  assert_class(model_basic, "mmrm")
  assert_class(model_augmented, "mmrm")
  assert_flag(any_reml)

  model_basic <- model_basic[["formula_parts"]]
  model_augmented <- model_augmented[["formula_parts"]]

  if (!identical(model_basic[["visit_var"]], model_augmented[["visit_var"]])) {
    stop("Models must all have the same visit variable.")
  }

  if (
    !identical(model_basic[["subject_var"]], model_augmented[["subject_var"]])
  ) {
    stop("All models must have the same subject variable.")
  }

  if (!identical(model_basic[["group_var"]], model_augmented[["group_var"]])) {
    stop("All models must have the same grouping variable.")
  }

  covar_nesting <- h_check_covar_nesting(model_basic, model_augmented)
  if (any_reml && covar_nesting == "nested") {
    stop("If any models use REML, all models' covariates must be the same.")
  }

  cov_struct_nesting <- h_check_cov_struct_nesting(model_basic, model_augmented)

  if (covar_nesting == "identical" && cov_struct_nesting == "identical") {
    warning(
      "Two models in the sequence have identical covariates and ",
      "covariance structures."
    )
  }

  TRUE
}


#' Ensure Two Models' Covariates Are Nested
#'
#' Throws an error if the covariates of `model_basic` aren't the same as or a
#' subset of the covariates of `model_augmented`.
#'
#' For upstream coding brevity, this function accepts the `formula_parts`
#' element of two `mmrm` model fits rather than `mmrm` objects. Such objects are
#' of class `mmrm_tmb_formula_parts`.
#'
#' @param model_basic,model_augmented (`mmrm_tmb_formula_parts`)\cr the
#'   `formula_parts` element of an `mmrm` model fit.
#'
#' @returns `"identical"` if `model_basic` and `model_augmented` have the same
#'   set of covariates. `"nesting"` if the covariates of `model_basic` are a
#'   subset of the covariates of `model_augmented`.
#'
#' @keywords internal
h_check_covar_nesting <- function(model_basic, model_augmented) {
  assert_class(model_basic, "mmrm_tmb_formula_parts")
  assert_class(model_augmented, "mmrm_tmb_formula_parts")

  basic_terms <- terms(model_basic[["model_formula"]])
  aug_terms <- terms(model_augmented[["model_formula"]])

  basic_factors <- attr(basic_terms, "factors")
  aug_factors <- attr(aug_terms, "factors")

  is_interaction_basic <- attr(basic_terms, "order") > 1
  is_interaction_aug <- attr(aug_terms, "order") > 1

  basic_main_effects <- colnames(basic_factors)[!is_interaction_basic]
  aug_main_effects <- colnames(aug_factors)[!is_interaction_aug]

  if (anyNA(match(basic_main_effects, aug_main_effects))) {
    stop(
      "Each model's covariates must be a subset of the next model's ",
      "covariates.",
      call. = FALSE
    )
  }

  if (any(is_interaction_basic)) {
    aug_factors <- aug_factors[rownames(basic_factors), , drop = FALSE]

    for (j_basic in which(is_interaction_basic)) {
      match_found <- FALSE
      for (j_aug in which(is_interaction_aug)) {
        match_found <- all(basic_factors[, j_basic] == aug_factors[, j_aug])
        if (match_found) break
      }
      if (!match_found) {
        stop(
          "Each model's interaction terms must be a subset of the next ",
          "model's terms.",
          call. = FALSE
        )
      }
    }
  }

  if (
    anyNA(match(aug_main_effects, basic_main_effects)) ||
      sum(is_interaction_basic) < sum(is_interaction_aug)
  ) {
    "nested"
  } else {
    "identical"
  }
}


#' Ensure Two Models' Covariance Structures Are Nested
#'
#' Throws an error if the covariance structure of `model_basic` isn't the same
#' as or a special case of the covariance structure of `model_augmented`.
#'
#' The check for "nesting" is a check against a mathematically determined
#' hierarchy of the available [cov_types]: if restricting an aspect of a
#' covariance structure can result in another covariance structure, the latter
#' structure is nested within the former structure.
#'
#' For upstream coding brevity, this function accepts the `formula_parts`
#' element of two `mmrm` model fits rather than `mmrm` objects. Such objects are
#' of class `mmrm_tmb_formula_parts`.
#'
#' @param model_basic,model_augmented (`mmrm_tmb_formula_parts`)\cr the
#'   `formula_parts` element of an `mmrm` model fit.
#'
#' @returns `"identical"` if `model_basic` and `model_augmented` have the same
#'   covariance structure. `"nesting"` if the covariance structure of
#'   `model_basic` is a special case of the covariance structure of
#'   `model_augmented`.
#'
#' @keywords internal
h_check_cov_struct_nesting <- function(model_basic, model_augmented) {
  assert_class(model_basic, "mmrm_tmb_formula_parts")
  assert_class(model_augmented, "mmrm_tmb_formula_parts")

  basic_cov_struct <- model_basic[["cov_type"]]
  aug_cov_struct <- model_augmented[["cov_type"]]

  # Hierarchy, specifying the "parent" covariance structures of each available
  # type in mmrm
  cov_struct_nesting <-
    list(
      ad = c("adh", "us"),
      adh = "us",
      ar1 = c("ad", "adh", "ar1h", "toep", "toeph", "us"),
      ar1h = c("adh", "toeph", "us"),
      cs = c("csh", "toep", "toeph", "us"),
      csh = c("toeph", "us"),
      toep = c("toeph", "us"),
      toeph = "us",
      us = character(0L),
      sp_exp = character(0L)
    )

  if (identical(basic_cov_struct, aug_cov_struct)) {
    "identical"
  } else if (any(cov_struct_nesting[[basic_cov_struct]] == aug_cov_struct)) {
    "nested"
  } else {
    stop(
      "Each model's covariance structure must be either identical to or a ",
      "special case of the next model's covariance structure.",
      call. = FALSE
    )
  }
}


#' Generate a Name Not Already In an Environment nor Its Parents
#'
#' Alters the user-supplied string `x` using [make.names()] with `unique = TRUE`
#' until it is a syntactically valid name not bound to in `env` nor its
#' [parents][parent.env].
#'
#' @param x (`string`)\cr a candidate name.
#' @param env (`environment`)\cr an [environment] whose bindings (and whose
#'   parents' bindings) are checked to ensure that they do not include the
#'   returned value.
#'
#' @returns A string that does not match any of the bindings of `env` nor its
#'   parents.
#'
#' @keywords internal
h_generate_new_name <- function(x, env) {
  assert_string(x, na.ok = TRUE)
  assert_environment(env)

  # Force x to be a syntactically valid name.
  x <- make.names(x, unique = TRUE)

  # As long as we keep finding a binding in env (or its parents) whose name is
  # the same as the last element of x...
  while (exists(x[length(x)], envir = env, inherits = TRUE)) {
    # ...Copy the first element of x and append it to the end.
    # Run make.names(unique = TRUE) again.
    x <- make.names(x[c(seq_along(x), 1L)], unique = TRUE)
  }

  # If we've gotten here, the last element of doesn't exist in env nor its
  # parents
  x[length(x)]
}


#' Refit an `mmrm` Model Using a New Dataset
#'
#' Extract the `call` [component] of `fit` and evaluate it in a new
#' [environment][new.env] that contains the new `data`.
#'
#' This works as follows:
#'
#' 1. A new environment is created whose parent is
#' `environment(fit$formula_parts$full_formula)`.
#'
#' 1. A name is generated using [h_generate_new_name()], and `data` is bound to
#' the new environment using this new name.
#'
#' 1. The [`call`] component of `fit` is extracted and its `data` argument is
#' changed to the new name.
#'
#' 1. The modified call is evaluated in the new environment.
#'
#' @param fit (`mmrm`)\cr an `mmrm` object to be refit.
#' @param data (`data frame`)\cr a data frame upon which `fit` is to be refit.
#'
#' @returns An `mmrm` object with the same terms as `fit` but based on `data`.
#'
#' @keywords internal
h_refit_mmrm <- function(fit, data) {
  assert_class(fit, "mmrm")
  assert_data_frame(data)

  # Grab the environment of the model's formula.
  env <- environment(fit[["formula_parts"]][["full_formula"]])
  assert_environment(env)

  # Grab the model call.
  expr <- component(fit, "call")

  # Generate a name guaranteed not to be in env nor its enclosing frames
  data_name <- h_generate_new_name("data", env)

  # Create a child environment whose parent is the model's environment.
  env <- new.env(parent = env)

  # Bind the new data to the new name in the new, child environment.
  env[[data_name]] <- data

  # Force the data argument of the model call to be the new name.
  expr[["data"]] <- as.name(data_name)

  # Evaluate the updated model call in the new environment.
  fit <- eval(expr, env)

  fit
}


#' Calculate the Significance of Each Term in an `mmrm` Fit.
#'
#' Runs [df_md()] once for each term in an `mmrm` object and returns the results
#' in a data frame.
#'
#' When only one model fit is passed to [anova.mmrm()], the `object` argument of
#' `anova.mmrm()` is passed directly to this function.
#'
#' @param object (`mmrm`)\cr an `mmrm` object.
#'
#' @returns A data frame with a row for each term in `object`, including an
#'   intercept if present. The [row.names] of the data frame identify the terms.
#'   The data frame will contain the following four columns:
#'
#'  - `num_df`: the numerator degrees of freedom.
#'  - `denom_df`: the denominator degrees of freedom.
#'  - `f_stat`: the test statistic on the F-distribution.
#'  - `p_va`: the associated p-value.
#'
#' @keywords internal
h_anova_single_mmrm_model <- function(object) {
  assert_class(object, "mmrm")

  model_terms_attributes <-
    attributes(terms(object[["formula_parts"]][["model_formula"]]))

  term_labels <- c(
    if (model_terms_attributes[["intercept"]]) "(Intercept)",
    model_terms_attributes[["term.labels"]]
  )

  design_matrix <- component(object, "x_matrix")

  # Integer vector with one entry per model coefficient, giving the index of the
  # term that corresponds to the coefficient.
  assign_coef_term <- attr(design_matrix, "assign")

  assign_coef_term <- factor(assign_coef_term, labels = term_labels)

  coef_identity_matrix <- diag(length(assign_coef_term))

  contrast_matrix_per_term <-
    split.data.frame(coef_identity_matrix, f = assign_coef_term)

  df_md_results_per_term <-
    lapply(contrast_matrix_per_term, df_md, object = object)
  df_md_results_per_term <- lapply(df_md_results_per_term, as.data.frame)

  df_md_results_data_frame <- do.call(rbind, df_md_results_per_term)

  df_md_results_data_frame
}

#' Dynamic Registration for Package Interoperability
#'
#' @seealso See `vignette("xtending", package = "emmeans")` for background.
#' @keywords internal
#' @noRd
.onLoad <- function(libname, pkgname) {
  # nolint
  if (!h_tmb_version_sufficient()) {
    msg <- paste(
      "TMB below version 1.9.15 has been used to compile the mmrm package.",
      "Reproducible model fits are not guaranteed.",
      "Please consider recompiling the package with TMB version 1.9.15 or higher."
    )
    warning(msg, call. = FALSE)
  }

  register_on_load(
    "emmeans",
    c("1.6", NA),
    callback = function() emmeans::.emm_register("mmrm", pkgname),
    message = "mmrm() registered as emmeans extension"
  )

  register_on_load(
    "parsnip",
    c("1.1.0", NA),
    callback = parsnip_add_mmrm,
    message = emit_tidymodels_register_msg
  )
  register_on_load(
    "car",
    c("3.1.2", NA),
    callback = car_add_mmrm,
    message = "mmrm() registered as car::Anova extension"
  )
}

#' Helper Function for Registering Functionality With Suggests Packages
#'
#' @inheritParams check_package_version
#'
#' @param callback (`function(...) ANY`)\cr a callback to execute upon package
#'   load. Note that no arguments are passed to this function. Any necessary
#'   data must be provided upon construction.
#'
#' @param message (`NULL` or `string`)\cr an optional message to print after
#'   the callback is executed upon successful registration.
#'
#' @return A logical (invisibly) indicating whether registration was successful.
#'  If not, a onLoad hook was set for the next time the package is loaded.
#'
#' @keywords internal
register_on_load <- function(
  pkg,
  ver = c(NA_character_, NA_character_),
  callback,
  message = NULL
) {
  if (isNamespaceLoaded(pkg) && check_package_version(pkg, ver)) {
    callback()
    if (is.character(message)) {
      packageStartupMessage(message)
    }
    if (is.function(message)) {
      packageStartupMessage(message())
    }
    return(invisible(TRUE))
  }

  setHook(
    packageEvent(pkg, event = "onLoad"),
    action = "append",
    function(...) {
      register_on_load(
        pkg = pkg,
        ver = ver,
        callback = callback,
        message = message
      )
    }
  )

  invisible(FALSE)
}

#' Check Suggested Dependency Against Version Requirements
#'
#' @param pkg (`string`)\cr package name.
#' @param ver (`character`)\cr of length 2 whose elements can be provided to
#'   [numeric_version()], representing a minimum and maximum (inclusive) version
#'   requirement for interoperability. When `NA`, no version requirement is
#'   imposed. Defaults to no version requirement.
#'
#' @return A logical (invisibly) indicating whether the loaded package meets
#'   the version requirements. A warning is emitted otherwise.
#'
#' @keywords internal
check_package_version <- function(pkg, ver = c(NA_character_, NA_character_)) {
  assert_character(ver, len = 2L)
  pkg_ver <- utils::packageVersion(pkg)
  ver <- numeric_version(ver, strict = FALSE)

  warn_version <- function(pkg, pkg_ver, ver) {
    ver_na <- is.na(ver)
    warning(sprintf(
      "Cannot register mmrm for use with %s (v%s). Version %s required.",
      pkg,
      pkg_ver,
      if (!any(ver_na)) {
        sprintf("%s to %s", ver[1], ver[2])
      } else if (ver_na[2]) {
        paste0(">= ", ver[1])
      } else if (ver_na[1]) {
        paste0("<= ", ver[2])
      }
    ))
  }

  if (identical(pkg_ver < ver[1], TRUE) || identical(pkg_ver > ver[2], TRUE)) {
    warn_version(pkg, pkg_ver, ver)
    return(invisible(FALSE))
  }

  invisible(TRUE)
}

#' Format a Message to Emit When Tidymodels is Loaded
#'
#' @return A character message to emit. Either a ansi-formatted cli output if
#'   package 'cli' is available or a plain-text message otherwise.
#'
#' @keywords internal
emit_tidymodels_register_msg <- function() {
  pkg <- utils::packageName()
  ver <- utils::packageVersion(pkg)

  if (isTRUE(getOption("tidymodels.quiet"))) {
    return()
  }

  # if tidymodels is attached, cli packages come as a dependency
  has_cli <- requireNamespace("cli", quietly = TRUE)
  if (has_cli) {
    # unfortunately, cli does not expose many formatting tools for emitting
    # messages (only via conditions to stderr) which can't be suppressed using
    # suppressPackageStartupMessages() so formatting must be done adhoc,
    # similar to how it's done in {tidymodels} R/attach.R
    paste0(
      cli::rule(
        left = cli::style_bold("Model Registration"),
        right = paste(pkg, ver)
      ),
      "\n",
      cli::col_green(cli::symbol$tick),
      " ",
      cli::col_blue("mmrm"),
      "::",
      cli::col_green("mmrm()")
    )
  } else {
    paste0(pkg, "::mmrm() registered for use with tidymodels")
  }
}

#' Determine Within or Between for each Design Matrix Column
#'
#' @description Used in [h_df_bw_calc()] to determine whether a variable
#'   differs only between subjects or also within subjects.
#'
#' @param x_matrix (`matrix`)\cr the design matrix with column names.
#' @param subject_ids (`factor`)\cr the subject IDs.
#'
#' @return Character vector with "intercept", "within" or "between" for each
#'   design matrix column identified via the names of the vector.
#'
#' @keywords internal
h_within_or_between <- function(x_matrix, subject_ids) {
  assert_matrix(x_matrix, col.names = "unique", min.cols = 1L)
  assert_factor(subject_ids, len = nrow(x_matrix))

  n_subjects <- length(unique(subject_ids))
  vapply(
    colnames(x_matrix),
    function(x) {
      if (x == "(Intercept)") {
        "intercept"
      } else {
        n_unique <- nrow(unique(cbind(x_matrix[, x], subject_ids)))
        if (n_unique > n_subjects) "within" else "between"
      }
    },
    character(1L)
  )
}

#' Calculation of Between-Within Degrees of Freedom
#'
#' @description Used in [h_df_1d_bw()] and [h_df_md_bw()].
#'
#' @param object (`mmrm`)\cr the fitted MMRM.
#'
#' @return List with:
#'   - `coefs_between_within` calculated via [h_within_or_between()]
#'   - `ddf_between`
#'   - `ddf_within`
#'
#' @keywords internal
h_df_bw_calc <- function(object) {
  assert_class(object, "mmrm")

  n_subjects <- component(object, "n_subjects")
  n_obs <- component(object, "n_obs")
  x_mat <- component(object, "x_matrix")

  subject_var <- component(object, "subject_var")
  full_frame <- component(object, "full_frame")
  subject_ids <- full_frame[[subject_var]]

  coefs_between_within <- h_within_or_between(x_mat, subject_ids)
  n_coefs_between <- sum(coefs_between_within == "between")
  n_intercept <- sum(coefs_between_within == "intercept")
  n_coefs_within <- sum(coefs_between_within == "within")
  ddf_between <- n_subjects - n_coefs_between - n_intercept
  ddf_within <- n_obs - n_subjects - n_coefs_within

  list(
    coefs_between_within = coefs_between_within,
    ddf_between = ddf_between,
    ddf_within = ddf_within
  )
}

#' Assign Minimum Degrees of Freedom Given Involved Coefficients
#'
#' @description Used in [h_df_1d_bw()] and [h_df_md_bw()].
#'
#' @param bw_calc (`list`)\cr from [h_df_bw_calc()].
#' @param is_coef_involved (`logical`)\cr whether each coefficient is involved
#'   in the contrast.
#'
#' @return The minimum of the degrees of freedom assigned to each involved
#'   coefficient according to its between-within categorization.
#'
#' @keywords internal
h_df_min_bw <- function(bw_calc, is_coef_involved) {
  assert_list(bw_calc)
  assert_names(
    names(bw_calc),
    identical.to = c("coefs_between_within", "ddf_between", "ddf_within")
  )
  assert_logical(is_coef_involved, len = length(bw_calc$coefs_between_within))
  assert_true(sum(is_coef_involved) > 0)

  coef_categories <- bw_calc$coefs_between_within[is_coef_involved]
  coef_dfs <- vapply(
    X = coef_categories,
    FUN = switch,
    intercept = bw_calc$ddf_within,
    between = bw_calc$ddf_between,
    within = bw_calc$ddf_within,
    FUN.VALUE = integer(1)
  )
  min(coef_dfs)
}

#' Calculation of Between-Within Degrees of Freedom for One-Dimensional Contrast
#'
#' @description Used in [df_1d()] if method is "Between-Within".
#'
#' @inheritParams h_df_1d_sat
#' @inherit h_df_1d_sat return
#' @keywords internal
h_df_1d_bw <- function(object, contrast) {
  assert_class(object, "mmrm")
  assert_numeric(contrast, len = length(component(object, "beta_est")))

  bw_calc <- h_df_bw_calc(object)
  is_coef_involved <- contrast != 0
  df <- h_df_min_bw(bw_calc, is_coef_involved)
  h_test_1d(object, contrast, df)
}

#' Calculation of Between-Within Degrees of Freedom for Multi-Dimensional Contrast
#'
#' @description Used in [df_md()] if method is "Between-Within".
#'
#' @inheritParams h_df_md_sat
#' @inherit h_df_md_sat return
#' @keywords internal
h_df_md_bw <- function(object, contrast) {
  assert_class(object, "mmrm")
  assert_matrix(
    contrast,
    mode = "numeric",
    any.missing = FALSE,
    ncols = length(component(object, "beta_est"))
  )

  bw_calc <- h_df_bw_calc(object)
  is_coef_involved <- apply(X = contrast != 0, MARGIN = 2L, FUN = any)
  df <- h_df_min_bw(bw_calc, is_coef_involved)
  h_test_md(object, contrast, df)
}

#' Obtain List of Jacobian Matrix Entries for Covariance Matrix
#'
#' @description Obtain the Jacobian matrices given the covariance function and variance parameters.
#'
#' @param tmb_data (`mmrm_tmb_data`)\cr produced by [h_mmrm_tmb_data()].
#' @param theta_est (`numeric`)\cr variance parameters point estimate.
#' @param beta_vcov (`matrix`)\cr vairance covariance matrix of coefficients.
#'
#' @return List with one element per variance parameter containing a matrix
#'   of the same dimensions as the covariance matrix. The values are the derivatives
#'   with regards to this variance parameter.
#'
#' @keywords internal
h_jac_list <- function(tmb_data, theta_est, beta_vcov) {
  assert_class(tmb_data, "mmrm_tmb_data")
  assert_numeric(theta_est)
  assert_matrix(beta_vcov)
  .Call(`_mmrm_get_jacobian`, PACKAGE = "mmrm", tmb_data, theta_est, beta_vcov)
}

#' Quadratic Form Calculations
#'
#' @description These helpers are mainly for easier readability and slightly better efficiency
#' of the quadratic forms used in the Satterthwaite calculations.
#'
#' @param center (`matrix`)\cr square numeric matrix with the same dimensions as
#'   `x` as the center of the quadratic form.
#'
#' @name h_quad_form
NULL

#' @describeIn h_quad_form calculates the number `vec %*% center %*% t(vec)`
#'   as a numeric (not a matrix).
#'
#' @param vec (`numeric`)\cr interpreted as a row vector.
#'
#' @keywords internal
h_quad_form_vec <- function(vec, center) {
  vec <- as.vector(vec)
  assert_numeric(vec, any.missing = FALSE)
  assert_matrix(
    center,
    mode = "numeric",
    any.missing = FALSE,
    nrows = length(vec),
    ncols = length(vec)
  )

  sum(vec * (center %*% vec))
}

#' @describeIn h_quad_form calculates the quadratic form `mat %*% center %*% t(mat)`
#'   as a matrix, the result is square and has dimensions identical to the number
#'   of rows in `mat`.
#'
#' @param mat (`matrix`)\cr numeric matrix to be multiplied left and right of
#'   `center`, therefore needs to have as many columns as there are rows and columns
#'   in `center`.
#'
#' @keywords internal
h_quad_form_mat <- function(mat, center) {
  assert_matrix(mat, mode = "numeric", any.missing = FALSE, min.cols = 1L)
  assert_matrix(
    center,
    mode = "numeric",
    any.missing = FALSE,
    nrows = ncol(center),
    ncols = ncol(center)
  )

  mat %*% tcrossprod(center, mat)
}

#' Computation of a Gradient Given Jacobian and Contrast Vector
#'
#' @description Computes the gradient of a linear combination of `beta` given the Jacobian matrix and
#' variance parameters.
#'
#' @param jac_list (`list`)\cr Jacobian list produced e.g. by [h_jac_list()].
#' @param contrast (`numeric`)\cr contrast vector, which needs to have the
#'   same number of elements as there are rows and columns in each element of
#'   `jac_list`.
#'
#' @return Numeric vector which contains the quadratic forms of each element of
#'   `jac_list` with the `contrast` vector.
#'
#' @keywords internal
h_gradient <- function(jac_list, contrast) {
  assert_list(jac_list)
  assert_numeric(contrast)

  vapply(
    jac_list,
    h_quad_form_vec,
    vec = contrast,
    numeric(1L)
  )
}

#' Helper for Calculation of Satterthwaite with Empirical Covariance Matrix
#'
#' @description Used in [h_df_1d_sat()] and [h_df_md_sat()] if empirical covariance
#' matrix is used.
#'
#' @param object (`mmrm`)\cr the MMRM fit.
#' @param contrast_matrix (`matrix`)\cr contrast matrix with number of subjects times
#' number of coefficients as the number of columns.
#'
#' @return Adjusted degrees of freedom value.
#' @keywords internal
h_df_1d_sat_empirical <- function(object, contrast_matrix) {
  assert_class(object, "mmrm")
  assert_matrix(
    contrast_matrix,
    mode = "numeric",
    any.missing = FALSE
  )
  g_matrix <- if (
    is.null(object$empirical_g_mat) && !is.null(object$empirical_df_mat)
  ) {
    warning(
      "mmrm fit was obtained with package version < 0.3.15, ",
      "using deprecated calculation of d.f., consider refitting the model"
    )
    h_quad_form_mat(contrast_matrix, object$empirical_df_mat)
  } else if (!is.null(object$empirical_g_mat)) {
    g_times_contrast_transposed <- tcrossprod(
      object$empirical_g_mat,
      contrast_matrix
    )
    crossprod(g_times_contrast_transposed)
  } else {
    stop(
      "neither empirical_df_mat nor empirical_g_mat are available in mmrm fit object"
    )
  }
  h_tr(g_matrix)^2 / sum(g_matrix^2)
}

#' Calculation of Satterthwaite Degrees of Freedom for One-Dimensional Contrast
#'
#' @description Used in [df_1d()] if method is
#' "Satterthwaite".
#'
#' @param object (`mmrm`)\cr the MMRM fit.
#' @param contrast (`numeric`)\cr contrast vector. Note that this should not include
#'   elements for singular coefficient estimates, i.e. only refer to the
#'   actually estimated coefficients.
#'
#' @return List with `est`, `se`, `df`, `t_stat` and `p_val`.
#' @keywords internal
h_df_1d_sat <- function(object, contrast) {
  assert_class(object, "mmrm")
  contrast <- as.numeric(contrast)
  assert_numeric(contrast, len = length(component(object, "beta_est")))

  df <- if (identical(object$vcov, "Asymptotic")) {
    grad <- h_gradient(component(object, "jac_list"), contrast)
    v_num <- 2 * h_quad_form_vec(contrast, component(object, "beta_vcov"))^2
    v_denom <- h_quad_form_vec(grad, component(object, "theta_vcov"))
    v_num / v_denom
  } else if (
    object$vcov %in%
      c("Empirical", "Empirical-Jackknife", "Empirical-Bias-Reduced")
  ) {
    contrast_matrix <- Matrix::.bdiag(rep(
      list(matrix(contrast, nrow = 1)),
      component(object, "n_subjects")
    ))
    contrast_matrix <- as.matrix(contrast_matrix)
    h_df_1d_sat_empirical(object, contrast_matrix)
  }

  h_test_1d(object, contrast, df)
}

#' Calculating Denominator Degrees of Freedom for the Multi-Dimensional Case
#'
#' @description Calculates the degrees of freedom for multi-dimensional contrast.
#'
#' @param t_stat_df (`numeric`)\cr `n` t-statistic derived degrees of freedom.
#'
#' @return Usually the calculation is returning `2 * E / (E - n)` where
#'   `E` is the sum of `t / (t - 2)` over all `t_stat_df` values `t`.
#'
#' @note If the input values are two similar to each other then just the average
#'   of them is returned. If any of the inputs is not larger than 2 then 2 is
#'   returned.
#'
#' @keywords internal
h_md_denom_df <- function(t_stat_df) {
  assert_numeric(
    t_stat_df,
    min.len = 1L,
    lower = .Machine$double.xmin,
    any.missing = FALSE
  )

  if (test_scalar(t_stat_df)) {
    t_stat_df
  } else if (all(abs(diff(t_stat_df)) < sqrt(.Machine$double.eps))) {
    mean(t_stat_df)
  } else if (any(t_stat_df <= 2)) {
    2
  } else {
    e <- sum(t_stat_df / (t_stat_df - 2))
    2 * e / (e - (length(t_stat_df)))
  }
}

#' Creating F-Statistic Results from One-Dimensional Contrast
#'
#' @description Creates multi-dimensional result from one-dimensional contrast from [df_1d()].
#'
#' @param object (`mmrm`)\cr model fit.
#' @param contrast (`numeric`)\cr one-dimensional contrast.
#'
#' @return The one-dimensional degrees of freedom are calculated and then
#'   based on that the p-value is calculated.
#'
#' @keywords internal
h_df_md_from_1d <- function(object, contrast) {
  res_1d <- h_df_1d_sat(object, contrast)
  list(
    num_df = 1,
    denom_df = res_1d$df,
    f_stat = res_1d$t_stat^2,
    p_val = stats::pf(
      q = res_1d$t_stat^2,
      df1 = 1,
      df2 = res_1d$df,
      lower.tail = FALSE
    )
  )
}

#' Calculation of Satterthwaite Degrees of Freedom for Multi-Dimensional Contrast
#'
#' @description Used in [df_md()] if method is "Satterthwaite".
#'
#' @param object (`mmrm`)\cr the MMRM fit.
#' @param contrast (`matrix`)\cr numeric contrast matrix, if given a `numeric`
#'   then this is coerced to a row vector. Note that this should not include
#'   elements for singular coefficient estimates, i.e. only refer to the
#'   actually estimated coefficients.
#'
#' @return List with `num_df`, `denom_df`, `f_stat` and `p_val` (2-sided p-value).
#' @keywords internal
h_df_md_sat <- function(object, contrast) {
  assert_class(object, "mmrm")
  assert_matrix(
    contrast,
    mode = "numeric",
    any.missing = FALSE,
    ncols = length(component(object, "beta_est"))
  )
  # Early return if we are in the one-dimensional case.
  if (identical(nrow(contrast), 1L)) {
    return(h_df_md_from_1d(object, contrast))
  }

  contrast_cov <- h_quad_form_mat(contrast, component(object, "beta_vcov"))
  eigen_cont_cov <- eigen(contrast_cov)
  eigen_cont_cov_vctrs <- eigen_cont_cov$vectors
  eigen_cont_cov_vals <- eigen_cont_cov$values

  eps <- sqrt(.Machine$double.eps)
  tol <- max(eps * eigen_cont_cov_vals[1], 0)
  rank_cont_cov <- sum(eigen_cont_cov_vals > tol)
  assert_number(rank_cont_cov, lower = .Machine$double.xmin)
  rank_seq <- seq_len(rank_cont_cov)
  vctrs_cont_prod <- crossprod(eigen_cont_cov_vctrs, contrast)[
    rank_seq,
    ,
    drop = FALSE
  ]

  # Early return if rank 1.
  if (identical(rank_cont_cov, 1L)) {
    return(h_df_md_from_1d(object, vctrs_cont_prod))
  }

  t_squared_nums <- drop(vctrs_cont_prod %*% object$beta_est)^2
  t_squared_denoms <- eigen_cont_cov_vals[rank_seq]
  t_squared <- t_squared_nums / t_squared_denoms
  f_stat <- sum(t_squared) / rank_cont_cov
  t_stat_df_nums <- 2 * eigen_cont_cov_vals^2
  t_stat_df <- if (identical(object$vcov, "Asymptotic")) {
    grads_vctrs_cont_prod <- lapply(
      rank_seq,
      function(m) {
        h_gradient(
          component(object, "jac_list"),
          contrast = vctrs_cont_prod[m, ]
        )
      }
    )
    t_stat_df_denoms <- vapply(
      grads_vctrs_cont_prod,
      h_quad_form_vec,
      center = component(object, "theta_vcov"),
      numeric(1)
    )
    t_stat_df_nums / t_stat_df_denoms
  } else {
    vapply(
      rank_seq,
      function(m) {
        contrast_matrix <- Matrix::.bdiag(
          rep(
            list(vctrs_cont_prod[m, , drop = FALSE]),
            component(object, "n_subjects")
          )
        )
        contrast_matrix <- as.matrix(contrast_matrix)
        h_df_1d_sat_empirical(object, contrast_matrix)
      },
      FUN.VALUE = 0
    )
  }
  denom_df <- h_md_denom_df(t_stat_df)

  list(
    num_df = rank_cont_cov,
    denom_df = denom_df,
    f_stat = f_stat,
    p_val = stats::pf(
      q = f_stat,
      df1 = rank_cont_cov,
      df2 = denom_df,
      lower.tail = FALSE
    )
  )
}

#' Tidying Methods for `mmrm` Objects
#'
#' @description `r lifecycle::badge("stable")`
#'
#' These methods tidy the estimates from an `mmrm` object into a
#' summary.
#'
#' @param x (`mmrm`)\cr fitted model.
#' @param conf.int (`flag`)\cr if `TRUE` columns for the lower (`conf.low`) and upper bounds
#'   (`conf.high`) of coefficient estimates are included.
#' @param conf.level (`number`)\cr defines the range of the optional confidence internal.
#' @param newdata (`data.frame` or `NULL`)\cr optional new data frame.
#' @param se_fit (`flag`)\cr whether to return standard errors of fit.
#' @param interval (`string`)\cr type of interval calculation.
#' @param type.residuals (`string`)\cr passed on to [residuals.mmrm_tmb()].
#' @param ... only used by `augment()` to pass arguments to the [predict.mmrm_tmb()] method.
#'
#' @name mmrm_tidiers
#' @aliases mmrm_tidiers
#'
#' @seealso [`mmrm_methods`], [`mmrm_tmb_methods`] for additional methods.
#'
#' @examples
#' fit <- mmrm(
#'   formula = FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT | USUBJID),
#'   data = fev_data
#' )
NULL

#' @describeIn mmrm_tidiers derives tidy `tibble` from an `mmrm` object.
#' @exportS3Method
#' @examples
#' # Applying tidy method to return summary table of covariate estimates.
#' fit |> tidy()
#' fit |> tidy(conf.int = TRUE, conf.level = 0.9)
tidy.mmrm <- function(
  x, # nolint
  conf.int = FALSE, # nolint
  conf.level = 0.95, # nolint
  ...
) {
  assert_flag(conf.int)
  assert_number(conf.level, lower = 0, upper = 1)
  tbl <- tibble::as_tibble(summary(x)$coefficients, rownames = "term")
  colnames(tbl) <- c(
    "term",
    "estimate",
    "std.error",
    "df",
    "statistic",
    "p.value"
  )
  coefs <- coef(x)
  if (length(coefs) != nrow(tbl)) {
    coefs <- tibble::enframe(coefs, name = "term", value = "estimate")
    tbl <- merge(coefs, tbl, by = c("term", "estimate"))
  }
  if (conf.int) {
    ci <- h_tbl_confint_terms(x, level = conf.level)
    tbl <- tibble::as_tibble(merge(tbl, ci, by = "term"))
  }
  tbl
}

#' @describeIn mmrm_tidiers derives `glance` `tibble` from an `mmrm` object.
#' @exportS3Method
#' @examples
#' # Applying glance method to return summary table of goodness of fit statistics.
#' fit |> glance()
glance.mmrm <- function(x, ...) {
  # nolint
  tibble::as_tibble(summary(x)$aic_list)
}

#' @describeIn mmrm_tidiers derives `augment` `tibble` from an `mmrm` object.
#' @exportS3Method
#' @examples
#' # Applying augment method to return merged `tibble` of model data, fitted and residuals.
#' fit |> augment()
#' fit |> augment(interval = "confidence")
#' fit |> augment(type.residuals = "pearson")
augment.mmrm <- function(
  x, # nolint
  newdata = NULL,
  interval = c("none", "confidence", "prediction"),
  se_fit = (interval != "none"),
  type.residuals = c("response", "pearson", "normalized"), # nolint
  ...
) {
  type.residuals <- match.arg(type.residuals) # nolint
  resid_df <- NULL
  if (is.null(newdata)) {
    newdata <- stats::get_all_vars(x, data = stats::na.omit(x$data))
    resid_df <- data.frame(
      .rownames = rownames(newdata),
      .resid = unname(residuals(x, type = type.residuals))
    )
  }
  interval <- match.arg(interval)

  tbl <- h_newdata_add_pred(
    x,
    newdata = newdata,
    se_fit = se_fit,
    interval = interval,
    ...
  )
  if (!is.null(resid_df)) {
    tbl <- merge(tbl, resid_df, by = ".rownames")
    tbl$.rownames <- as.numeric(tbl$.rownames)
    tbl <- tbl[order(tbl$.rownames), , drop = FALSE]
  }
  tibble::as_tibble(tbl)
}

#' Extract `tibble` with Confidence Intervals and Term Names
#'
#' This is used in [tidy.mmrm()].
#'
#' @param x (`mmrm`)\cr fit object.
#' @param ... passed to [stats::confint()], hence not used at the moment.
#'
#' @return A `tibble` with `term`, `conf.low`, `conf.high` columns.
#'
#' @keywords internal
h_tbl_confint_terms <- function(x, ...) {
  df <- stats::confint(x, ...)
  tbl <- tibble::as_tibble(df, rownames = "term", .name_repair = "minimal")
  names(tbl) <- c("term", "conf.low", "conf.high")
  tbl
}

#' Add Prediction Results to New Data
#'
#' This is used in [augment.mmrm()].
#'
#' @param x (`mmrm`)\cr fit.
#' @param newdata (`data.frame`)\cr data to predict.
#' @param se_fit (`flag`)\cr whether to return standard error of prediction,
#'   can only be used when `interval` is not "none".
#' @param interval (`string`)\cr type of interval.
#' @param ... passed to [predict.mmrm_tmb()].
#'
#' @return The `newdata` as a `tibble` with additional columns `.fitted`,
#'   `.lower`, `.upper` (if interval is not `none`) and `.se.fit` (if `se_fit`
#'   requested).
#'
#' @keywords internal
h_newdata_add_pred <- function(x, newdata, se_fit, interval, ...) {
  assert_class(x, "mmrm")
  assert_data_frame(newdata)
  assert_flag(se_fit)
  assert_string(interval)
  if (interval == "none") {
    assert_false(se_fit)
  }

  tbl <- h_df_to_tibble(newdata)
  pred_results <- predict(
    x,
    newdata = newdata,
    na.action = stats::na.pass,
    se.fit = se_fit,
    interval = interval,
    ...
  )
  if (interval == "none") {
    assert_numeric(pred_results)
    tbl$.fitted <- unname(pred_results)
  } else {
    assert_matrix(pred_results)
    tbl$.fitted <- unname(pred_results[, "fit"])
    tbl$.lower <- unname(pred_results[, "lwr"])
    tbl$.upper <- unname(pred_results[, "upr"])
  }
  if (se_fit) {
    tbl$.se.fit <- unname(pred_results[, "se"])
  }
  tbl
}

#' Coerce a Data Frame to a `tibble`
#'
#' This is used in [h_newdata_add_pred()].
#'
#' @details This is only a thin wrapper around [tibble::as_tibble()], except
#' giving a useful error message and it checks for `rownames` and adds them
#' as a new column `.rownames` if they are not just a numeric sequence as
#' per the [tibble::has_rownames()] decision.
#'
#' @param data (`data.frame`)\cr what to coerce.
#'
#' @return The `data` as a `tibble`, potentially with a `.rownames` column.
#'
#' @keywords internal
h_df_to_tibble <- function(data) {
  tryCatch(tbl <- tibble::as_tibble(data), error = function(cnd) {
    stop(
      "Could not coerce data to `tibble`. Try explicitly passing a",
      "dataset to either the `data` or `newdata` argument.",
      call. = FALSE
    )
  })
  if (tibble::has_rownames(data)) {
    tbl <- tibble::add_column(tbl, .rownames = rownames(data), .before = TRUE)
  }
  tbl
}

#' Fitting an MMRM with Single Optimizer
#'
#' @description `r lifecycle::badge("stable")`
#'
#' This function helps to fit an MMRM using `TMB` with a single optimizer,
#' while capturing messages and warnings.
#'
#' @inheritParams mmrm
#' @param control (`mmrm_control`)\cr object.
#' @param tmb_data (`mmrm_tmb_data`)\cr object.
#' @param formula_parts (`mmrm_tmb_formula_parts`)\cr object.
#' @param ... Additional arguments to pass to [mmrm_control()].
#'
#' @details
#' `fit_single_optimizer` will fit the `mmrm` model using the `control` provided.
#' If there are multiple optimizers provided in `control`, only the first optimizer
#' will be used.
#' If `tmb_data` and `formula_parts` are both provided, `formula`, `data`, `weights`,
#' `reml`, and `covariance` are ignored.
#'
#' @return The `mmrm_fit` object, with additional attributes containing warnings,
#'   messages, optimizer used and convergence status in addition to the
#'   `mmrm_tmb` contents.
#' @export
#'
#' @examples
#' mod_fit <- fit_single_optimizer(
#'   formula = FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT | USUBJID),
#'   data = fev_data,
#'   weights = rep(1, nrow(fev_data)),
#'   optimizer = "nlminb"
#' )
#' attr(mod_fit, "converged")
fit_single_optimizer <- function(
  formula,
  data,
  weights,
  reml = TRUE,
  covariance = NULL,
  tmb_data,
  formula_parts,
  ...,
  control = mmrm_control(...)
) {
  to_remove <- list(
    # Transient visit to invalid parameters.
    warnings = c("NA/NaN function evaluation")
  )
  as_diverged <- list(
    errors = c(
      "NA/NaN Hessian evaluation",
      "L-BFGS-B needs finite values of 'fn'"
    )
  )
  if (missing(tmb_data) || missing(formula_parts)) {
    h_valid_formula(formula)
    assert_data_frame(data)
    assert_numeric(weights, any.missing = FALSE, lower = .Machine$double.xmin)
    assert_flag(reml)
    assert_class(control, "mmrm_control")
    assert_list(
      control$optimizers,
      names = "unique",
      types = c("function", "partial")
    )
    quiet_fit <- h_record_all_output(
      fit_mmrm(
        formula = formula,
        data = data,
        weights = weights,
        reml = reml,
        covariance = covariance,
        control = control
      ),
      remove = to_remove,
      divergence = as_diverged
    )
  } else {
    assert_class(tmb_data, "mmrm_tmb_data")
    assert_class(formula_parts, "mmrm_tmb_formula_parts")
    quiet_fit <- h_record_all_output(
      fit_mmrm(
        formula_parts = formula_parts,
        tmb_data = tmb_data,
        control = control
      ),
      remove = to_remove,
      divergence = as_diverged
    )
  }
  if (length(quiet_fit$errors)) {
    stop(quiet_fit$errors)
  }
  converged <- (length(quiet_fit$warnings) == 0L) &&
    (length(quiet_fit$divergence) == 0L) &&
    isTRUE(quiet_fit$result$opt_details$convergence == 0)
  structure(
    quiet_fit$result,
    warnings = quiet_fit$warnings,
    messages = quiet_fit$messages,
    divergence = quiet_fit$divergence,
    converged = converged,
    class = c("mmrm_fit", class(quiet_fit$result))
  )
}

#' Summarizing List of Fits
#'
#' @param all_fits (`list` of `mmrm_fit` or `try-error`)\cr list of fits.
#'
#' @return List with `warnings`, `messages`, `log_liks` and `converged` results.
#' @keywords internal
h_summarize_all_fits <- function(all_fits) {
  assert_list(all_fits, types = c("mmrm_fit", "try-error"))
  is_error <- vapply(all_fits, is, logical(1), class2 = "try-error")

  warnings <- messages <- vector(mode = "list", length = length(all_fits))
  warnings[is_error] <- lapply(all_fits[is_error], as.character)
  warnings[!is_error] <- lapply(all_fits[!is_error], attr, which = "warnings")
  messages[!is_error] <- lapply(all_fits[!is_error], attr, which = "messages")
  log_liks <- as.numeric(rep(NA, length.out = length(all_fits)))
  log_liks[!is_error] <- vapply(all_fits[!is_error], stats::logLik, numeric(1L))
  converged <- rep(FALSE, length.out = length(all_fits))
  converged[!is_error] <- vapply(
    all_fits[!is_error],
    attr,
    logical(1),
    which = "converged"
  )

  list(
    warnings = warnings,
    messages = messages,
    log_liks = log_liks,
    converged = converged
  )
}

#' Refitting MMRM with Multiple Optimizers
#'
#' @description `r lifecycle::badge("stable")`
#'
#' @param fit (`mmrm_fit`)\cr original model fit from [fit_single_optimizer()].
#' @param ... Additional arguments passed to [mmrm_control()].
#' @param control (`mmrm_control`)\cr object.
#'
#' @return The best (in terms of log likelihood) fit which converged.
#'
#' @note For Windows, no parallel computations are currently implemented.
#' @export
#'
#' @examples
#' fit <- fit_single_optimizer(
#'   formula = FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT | USUBJID),
#'   data = fev_data,
#'   weights = rep(1, nrow(fev_data)),
#'   optimizer = "nlminb"
#' )
#' best_fit <- refit_multiple_optimizers(fit)
refit_multiple_optimizers <- function(fit, ..., control = mmrm_control(...)) {
  assert_class(fit, "mmrm_fit")
  assert_class(control, "mmrm_control")

  n_cores_used <- ifelse(
    .Platform$OS.type == "windows",
    1L,
    min(
      length(control$optimizers),
      control$n_cores
    )
  )
  controls <- h_split_control(
    control,
    start = fit$theta_est
  )

  # Take the results from old fit as starting values for new fits.
  all_fits <- suppressWarnings(parallel::mcmapply(
    FUN = fit_single_optimizer,
    control = controls,
    MoreArgs = list(
      tmb_data = fit$tmb_data,
      formula_parts = fit$formula_parts
    ),
    mc.cores = n_cores_used,
    mc.silent = TRUE,
    SIMPLIFY = FALSE
  ))
  all_fits <- c(all_fits, list(old_result = fit))

  # Find the results that are ok and return best in terms of log-likelihood.
  all_fits_summary <- h_summarize_all_fits(all_fits)
  is_ok <- all_fits_summary$converged
  if (!any(is_ok)) {
    stop(
      "No optimizer led to a successful model fit. ",
      "Please try to use a different covariance structure or other covariates."
    )
  }
  best_optimizer <- which.max(all_fits_summary$log_liks[is_ok])
  all_fits[[which(is_ok)[best_optimizer]]]
}

#' Control Parameters for Fitting an MMRM
#'
#' @description `r lifecycle::badge("stable")`
#' Fine-grained specification of the MMRM fit details is possible using this
#' control function.
#'
#' @param n_cores (`count`)\cr number of cores to be used.
#' @param method (`string`)\cr adjustment method for degrees of freedom.
#' @param vcov (`string`)\cr coefficients covariance matrix adjustment method.
#' @param start (`NULL`, `numeric` or `function`)\cr optional start values for variance
#'   parameters. See details for more information.
#' @param accept_singular (`flag`)\cr whether singular design matrices are reduced
#'   to full rank automatically and additional coefficient estimates will be missing.
#' @param optimizers (`list`)\cr optimizer specification, created with [h_get_optimizers()].
#'   Please note that optimizers using the Hessian will not be compatible with
#'   `disable_theta_vcov = TRUE` and an error will be raised in that case.
#' @param drop_visit_levels (`flag`)\cr whether to drop levels for visit variable,
#'   if visit variable is a factor, see details.
#' @param disable_theta_vcov (`flag`)\cr whether to disable calculation of
#'   variance-covariance matrix for variance parameters. This can speed up fitting
#'   when there are many variance parameters, see details.
#' @param ... additional arguments passed to [h_get_optimizers()].
#'
#' @details
#  - The `drop_visit_levels` flag will decide whether unobserved visits will be kept for analysis.
#'   For example, if the data only has observations at visits `VIS1`, `VIS3` and `VIS4`, by default
#'   they are treated to be equally spaced, the distance from `VIS1` to `VIS3`, and from `VIS3` to `VIS4`,
#'   are identical. However, you can manually convert this visit into a factor, with
#'   `levels = c("VIS1", "VIS2", "VIS3", "VIS4")`, and also use `drop_visits_levels = FALSE`,
#'   then the distance from `VIS1` to `VIS3` will be double, as `VIS2` is a valid visit.
#'   However, please be cautious because this can lead to convergence failure
#'   when using an unstructured covariance matrix and there are no observations
#'   at the missing visits.
#' - The `method` and `vcov` arguments specify the degrees of freedom and coefficients
#'   covariance matrix adjustment methods, respectively.
#'   - Allowed `vcov` includes: "Asymptotic", "Kenward-Roger", "Kenward-Roger-Linear", "Empirical" (CR0),
#'     "Empirical-Jackknife" (CR3), and "Empirical-Bias-Reduced" (CR2).
#'   - Allowed `method` includes: "Satterthwaite", "Kenward-Roger", "Between-Within" and "Residual".
#'   - If `method` is "Kenward-Roger" then only "Kenward-Roger" or "Kenward-Roger-Linear" are allowed for `vcov`.
#' - The `vcov` argument can be `NULL` to use the default covariance method depending on the `method`
#'   used for degrees of freedom, see the following table:
#'
#'    | `method`  |  Default `vcov`|
#'    |-----------|----------|
#'    |Satterthwaite| Asymptotic|
#'    |Kenward-Roger| Kenward-Roger|
#'    |Residual| Empirical|
#'    |Between-Within| Asymptotic|
#'
#' - Please note that "Kenward-Roger" for "Unstructured" covariance gives different results
#'   compared to SAS; Use "Kenward-Roger-Linear" for `vcov` instead for better matching
#'   of the SAS results.
#'
#' - The argument `start` is used to facilitate the choice of initial values for fitting the model.
#'   If `function` is provided, make sure its parameter is a valid element of `mmrm_tmb_data`
#'   or `mmrm_tmb_formula_parts` and it returns a numeric vector.
#'   By default or if `NULL` is provided, `std_start` will be used.
#'   Other implemented methods include `emp_start`.
#'
#' - Disabling the `theta_vcov` calculation can speed up the fitting process when there are many variance
#'   parameters. However, this has also drawbacks. We can no longer check that the variance parameter estimates
#'   are indeed at a local maximum of the log-likelihood surface, and therefore convergence issues might go unnoticed.
#'   In addition, some covariance matrix adjustment methods (e.g., Kenward-Roger) require `theta_vcov` to be calculated.
#'   These will then raise errors.
#'
#' @return List of class `mmrm_control` with the control parameters.
#' @export
#'
#' @examples
#' mmrm_control(
#'   optimizer_fun = stats::optim,
#'   optimizer_args = list(method = "L-BFGS-B")
#' )
mmrm_control <- function(
  n_cores = 1L,
  method = c("Satterthwaite", "Kenward-Roger", "Residual", "Between-Within"),
  vcov = NULL,
  start = std_start,
  accept_singular = TRUE,
  drop_visit_levels = TRUE,
  disable_theta_vcov = FALSE,
  ...,
  optimizers = h_get_optimizers(...)
) {
  assert_count(n_cores, positive = TRUE)
  assert_character(method)
  if (is.null(start)) {
    start <- std_start
  }
  assert(
    check_function(start, args = "..."),
    check_numeric(start, null.ok = FALSE),
    combine = "or"
  )
  assert_flag(accept_singular)
  assert_flag(drop_visit_levels)
  assert_flag(disable_theta_vcov)
  assert_list(optimizers, names = "unique", types = c("function", "partial"))
  assert_string(vcov, null.ok = TRUE)
  method <- match.arg(method)
  if (is.null(vcov)) {
    vcov <- h_get_cov_default(method)
  }
  assert_subset(
    vcov,
    c(
      "Asymptotic",
      "Empirical",
      "Empirical-Bias-Reduced",
      "Empirical-Jackknife",
      "Kenward-Roger",
      "Kenward-Roger-Linear"
    )
  )
  if (
    xor(
      identical(method, "Kenward-Roger"),
      vcov %in% c("Kenward-Roger", "Kenward-Roger-Linear")
    )
  ) {
    stop(paste(
      "Kenward-Roger degrees of freedom must work together with Kenward-Roger",
      "or Kenward-Roger-Linear covariance!"
    ))
  }
  if (disable_theta_vcov) {
    if (vcov %in% c("Kenward-Roger", "Kenward-Roger-Linear")) {
      stop("Kenward-Roger requires theta_vcov calculation!")
    }
    uses_hessian <- lapply(
      optimizers,
      attr,
      "use_hessian"
    )
    uses_hessian <- vapply(uses_hessian, isTRUE, logical(1))
    if (any(uses_hessian)) {
      stop(
        "Disabling theta_vcov calculation is incompatible with optimizers using Hessian: ",
        toString(names(optimizers[uses_hessian]))
      )
    }
  }
  structure(
    list(
      optimizers = optimizers,
      start = start,
      accept_singular = accept_singular,
      method = method,
      vcov = vcov,
      n_cores = as.integer(n_cores),
      drop_visit_levels = drop_visit_levels,
      disable_theta_vcov = disable_theta_vcov
    ),
    class = "mmrm_control"
  )
}

#' Fit an MMRM
#'
#' @description `r lifecycle::badge("stable")`
#'
#' This is the main function fitting the MMRM.
#'
#' @param formula (`formula`)\cr the model formula, see details.
#' @param data (`data`)\cr the data to be used for the model.
#' @param weights (`vector`)\cr an optional vector of weights to be used in
#'   the fitting process. Should be `NULL` or a numeric vector.
#' @param reml (`flag`)\cr whether restricted maximum likelihood (REML)
#'   estimation is used, otherwise maximum likelihood (ML) is used.
#' @param covariance (`cov_struct`)\cr a covariance structure type definition
#'   as produced with [cov_struct()], or value that can be coerced to a
#'   covariance structure using [as.cov_struct()]. If no value is provided,
#'   a structure is derived from the provided formula.
#' @param control (`mmrm_control`)\cr fine-grained fitting specifications list
#'   created with [mmrm_control()].
#' @param ... arguments passed to [mmrm_control()].
#'
#' @details
#' The `formula` typically looks like:
#' `FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT | USUBJID)`
#' so specifies response and covariates as usual, and exactly one special term
#' defines which covariance structure is used and what are the time point and
#' subject variables. The covariance structures in the formula can be
#' found in [`covariance_types`].
#'
#' The time points have to be unique for each subject. That is,
#' there cannot be time points with multiple observations for any subject.
#' The rationale is that these observations would need to be correlated, but it
#' is not possible within the currently implemented covariance structure framework
#' to do that correctly. Moreover, for non-spatial covariance structures, the time
#' variable must be a factor variable.
#'
#' When optimizer is not set, first the default optimizer
#' (`L-BFGS-B`) is used to fit the model. If that converges, this is returned.
#' If not, the other available optimizers from [h_get_optimizers()],
#' including `BFGS`, `CG` and `nlminb` are
#' tried (in parallel if `n_cores` is set and not on Windows).
#' If none of the optimizers converge, then the function fails. Otherwise
#' the best fit is returned.
#'
#' Note that fine-grained control specifications can either be passed directly
#' to the `mmrm` function, or via the `control` argument for bundling together
#' with the [mmrm_control()] function. Both cannot be used together, since
#' this would delete the arguments passed via `mmrm`.
#'
#' @return An `mmrm` object.
#'
#' @note The `mmrm` object is also an `mmrm_fit` and an `mmrm_tmb` object,
#' therefore corresponding methods also work (see [`mmrm_tmb_methods`]).
#'
#' Additional contents depend on `vcov` (see [mmrm_control()]):
#' - If Kenward-Roger covariance matrix is used, `kr_comp` contains necessary
#' components and `beta_vcov_adj` includes the adjusted coefficients covariance
#' matrix.
#' - If Empirical covariance matrix is used, `beta_vcov_adj` contains the
#' corresponding coefficients covariance matrix estimate. In addition,
#' `empirical_g_mat` contains the empirical g matrix, which is used to calculate
#' the Satterthwaite degrees of freedom. The `score_per_subject` contains the
#' empirical score per subject.
#' - If Asymptotic covariance matrix is used in combination with Satterthwaite
#' d.f. adjustment, the Jacobian information `jac_list` is included.
#'
#' Note that these additional elements might change over time and are to be considered
#' internal implementation details rather than part of the public user interface of
#' the package.
#'
#' Use of the package `emmeans` is supported, see [`emmeans_support`].
#'
#' NA values are always omitted regardless of `na.action` setting.
#'
#' When the number of visit levels is large, it usually requires large memory to create the
#' covariance matrix. By default, the maximum allowed visit levels is 100, and if there are more
#' visit levels, a confirmation is needed if run interactively.
#' You can use `options(mmrm.max_visits = <target>)` to increase the maximum allowed number of visit
#' levels. In non-interactive sessions the confirmation is not raised and will directly give you an error if
#' the number of visit levels exceeds the maximum.
#'
#' @export
#'
#' @examples
#' fit <- mmrm(
#'   formula = FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT | USUBJID),
#'   data = fev_data
#' )
#'
#' # Direct specification of control details:
#' fit <- mmrm(
#'   formula = FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT | USUBJID),
#'   data = fev_data,
#'   weights = fev_data$WEIGHTS,
#'   method = "Kenward-Roger"
#' )
#'
#' # Alternative specification via control argument (but you cannot mix the
#' # two approaches):
#' fit <- mmrm(
#'   formula = FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT | USUBJID),
#'   data = fev_data,
#'   control = mmrm_control(method = "Kenward-Roger")
#' )
mmrm <- function(
  formula,
  data,
  weights = NULL,
  covariance = NULL,
  reml = TRUE,
  control = mmrm_control(...),
  ...
) {
  assert_false(!missing(control) && !missing(...))
  assert_class(control, "mmrm_control")
  assert_list(control$optimizers, min.len = 1)

  if (control$method %in% c("Kenward-Roger", "Kenward-Roger-Linear") && !reml) {
    stop("Kenward-Roger only works for REML")
  }
  h_valid_formula(formula)
  covariance <- h_reconcile_cov_struct(formula, covariance)
  formula_parts <- h_mmrm_tmb_formula_parts(formula, covariance)
  h_tmb_warn_non_deterministic()

  if (!missing(data)) {
    attr(data, which = "dataname") <- toString(match.call()$data)
  } else {
    # na.action set to na.pass to allow data to be full; will be futher trimmed later
    data <- model.frame(formula_parts$full_formula, na.action = "na.pass")
  }

  if (is.null(weights)) {
    weights <- rep(1, nrow(data))
  } else {
    attr(weights, which = "dataname") <- deparse(match.call()$weights)
  }
  tmb_data <- h_mmrm_tmb_data(
    formula_parts,
    data,
    weights,
    reml,
    singular = if (control$accept_singular) "drop" else "error",
    drop_visit_levels = control$drop_visit_levels,
    allow_na_response = FALSE
  )
  fit <- structure("", class = "try-error")
  names_all_optimizers <- names(control$optimizers)
  while (is(fit, "try-error") && length(control$optimizers) > 0) {
    fit <- fit_single_optimizer(
      tmb_data = tmb_data,
      formula_parts = formula_parts,
      control = control
    )
    if (is(fit, "try-error")) {
      warning(paste0(
        "Divergence with optimizer ",
        names(control$optimizers[1L]),
        " due to problems: ",
        toString(attr(fit, "divergence"))
      ))
    }
    control$optimizers <- control$optimizers[-1]
  }
  if (!attr(fit, "converged")) {
    more_optimizers <- length(control$optimizers) >= 1L
    if (more_optimizers) {
      fit <- refit_multiple_optimizers(
        fit = fit,
        control = control
      )
    } else {
      all_problems <- unlist(
        attributes(fit)[c("errors", "warnings")],
        use.names = FALSE
      )
      stop(paste0(
        "Chosen optimizers '",
        toString(names_all_optimizers),
        "' led to problems during model fit:\n",
        paste(
          paste0(seq_along(all_problems), ") ", all_problems),
          collapse = ";\n"
        ),
        "\n",
        "Consider trying multiple or different optimizers."
      ))
    }
  }
  fit_msg <- attr(fit, "messages")
  if (!is.null(fit_msg)) {
    message(paste(fit_msg, collapse = "\n"))
  }
  fit$call <- match.call()
  fit$call$formula <- formula
  fit$method <- control$method
  fit$vcov <- control$vcov
  if (control$vcov %in% c("Kenward-Roger", "Kenward-Roger-Linear")) {
    fit$kr_comp <- h_get_kr_comp(fit$tmb_data, fit$theta_est)
    fit$beta_vcov_adj <- h_var_adj(
      v = fit$beta_vcov,
      w = component(fit, "theta_vcov"),
      p = fit$kr_comp$P,
      q = fit$kr_comp$Q,
      r = fit$kr_comp$R,
      linear = (control$vcov == "Kenward-Roger-Linear")
    )
  } else if (
    control$vcov %in%
      c("Empirical", "Empirical-Bias-Reduced", "Empirical-Jackknife")
  ) {
    empirical_comp <- h_get_empirical(
      fit$tmb_data,
      fit$theta_est,
      fit$beta_est,
      fit$beta_vcov,
      control$vcov
    )
    fit$beta_vcov_adj <- empirical_comp$cov
    fit$empirical_g_mat <- empirical_comp$g_mat
    fit$score_per_subject <- empirical_comp$score_per_subject
    dimnames(fit$beta_vcov_adj) <- dimnames(fit$beta_vcov)
  } else if (identical(control$vcov, "Asymptotic")) {
    # Note that we only need the Jacobian list under Asymptotic covariance method,
    # cf. the Satterthwaite vignette.
    if (identical(fit$method, "Satterthwaite")) {
      fit$jac_list <- h_jac_list(fit$tmb_data, fit$theta_est, fit$beta_vcov)
    }
  } else {
    stop("Unrecognized coefficent variance-covariance method!")
  }

  class(fit) <- c("mmrm", class(fit))
  fit
}

#' Calculation of Degrees of Freedom for One-Dimensional Contrast
#'
#' @description `r lifecycle::badge("stable")`
#' Calculates the estimate, adjusted standard error, degrees of freedom,
#' t statistic and p-value for one-dimensional contrast.
#'
#' @param object (`mmrm`)\cr the MMRM fit.
#' @param contrast (`numeric`)\cr contrast vector. Note that this should not include
#'   elements for singular coefficient estimates, i.e. only refer to the
#'   actually estimated coefficients.
#' @return List with `est`, `se`, `df`, `t_stat` and `p_val`.
#' @export
#'
#' @examples
#' object <- mmrm(
#'   formula = FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT | USUBJID),
#'   data = fev_data
#' )
#' contrast <- numeric(length(object$beta_est))
#' contrast[3] <- 1
#' df_1d(object, contrast)
df_1d <- function(object, contrast) {
  assert_class(object, "mmrm")
  assert_numeric(
    contrast,
    len = length(component(object, "beta_est")),
    any.missing = FALSE
  )
  contrast <- as.vector(contrast)
  if (all(contrast == 0)) {
    return(
      list(
        est = 0,
        se = NA_real_,
        df = NA_real_,
        t_stat = NA_real_,
        p_val = NA_real_
      )
    )
  }
  switch(
    object$method,
    "Satterthwaite" = h_df_1d_sat(object, contrast),
    "Kenward-Roger" = h_df_1d_kr(object, contrast),
    "Residual" = h_df_1d_res(object, contrast),
    "Between-Within" = h_df_1d_bw(object, contrast),
    stop("Unrecognized degrees of freedom method: ", object$method)
  )
}


#' Calculation of Degrees of Freedom for Multi-Dimensional Contrast
#'
#' @description `r lifecycle::badge("stable")`
#' Calculates the estimate, standard error, degrees of freedom,
#' t statistic and p-value for m-dimensional contrast, depending on the method
#' used in [mmrm()].
#'
#' @param object (`mmrm`)\cr the MMRM fit.
#' @param contrast (`matrix`)\cr numeric contrast matrix, if given a `numeric`
#'   then this is coerced to a row vector. Note that this should not include
#'   elements for singular coefficient estimates, i.e. only refer to the
#'   actually estimated coefficients.
#'
#' @return List with `num_df`, `denom_df`, `f_stat` and `p_val` (2-sided p-value).
#' @export
#'
#' @examples
#' object <- mmrm(
#'   formula = FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT | USUBJID),
#'   data = fev_data
#' )
#' contrast <- matrix(data = 0, nrow = 2, ncol = length(object$beta_est))
#' contrast[1, 2] <- contrast[2, 3] <- 1
#' df_md(object, contrast)
df_md <- function(object, contrast) {
  assert_class(object, "mmrm")
  assert_numeric(contrast, any.missing = FALSE)
  if (!is.matrix(contrast)) {
    contrast <- matrix(contrast, ncol = length(contrast))
  }
  assert_matrix(contrast, ncols = length(component(object, "beta_est")))
  if (nrow(contrast) == 0 || all(contrast == 0)) {
    return(
      list(
        num_df = 0,
        denom_df = NA_real_,
        f_stat = NA_real_,
        p_val = NA_real_
      )
    )
  }
  switch(
    object$method,
    "Satterthwaite" = h_df_md_sat(object, contrast),
    "Kenward-Roger" = h_df_md_kr(object, contrast),
    "Residual" = h_df_md_res(object, contrast),
    "Between-Within" = h_df_md_bw(object, contrast),
    stop("Unrecognized degrees of freedom method: ", object$method)
  )
}

#' Creating T-Statistic Test Results For One-Dimensional Contrast
#'
#' @description Creates a list of results for one-dimensional contrasts using
#' a t-test statistic and the given degrees of freedom.
#'
#' @inheritParams df_1d
#' @param df (`number`)\cr degrees of freedom for the one-dimensional contrast.
#'
#' @return List with `est`, `se`, `df`, `t_stat` and `p_val` (2-sided p-value).
#'
#' @keywords internal
h_test_1d <- function(object, contrast, df) {
  assert_class(object, "mmrm")
  assert_numeric(contrast, len = length(component(object, "beta_est")))
  assert_number(df, lower = .Machine$double.xmin)

  est <- sum(contrast * component(object, "beta_est"))
  var <- h_quad_form_vec(contrast, component(object, "beta_vcov"))
  se <- sqrt(var)
  t_stat <- est / se
  p_val <- 2 * stats::pt(q = abs(t_stat), df = df, lower.tail = FALSE)

  list(
    est = est,
    se = se,
    df = df,
    t_stat = t_stat,
    p_val = p_val
  )
}

#' Creating F- or Chi-squared-statistic Test Results For Multi-Dimensional
#' Contrast
#'
#' @description Creates a list of results for multi-dimensional contrasts using
#'   an F-test statistic and the given degrees of freedom or a Chi-squared test
#'   statistic.
#'
#' @inheritParams df_md
#' @param contrast (`matrix`)\cr numeric contrast matrix.
#' @param df (`number`)\cr denominator degrees of freedom for the
#'   multi-dimensional contrast for an F-test. Ignored if `test = "Chisq"`.
#' @param f_stat_factor (`number`)\cr optional scaling factor on top of the
#'   standard F-statistic. Ignored if `test = "Chisq"`.
#' @param test (`string`)\cr either `"F"` or `"Chisq"`, specifying the kind of
#'   test to perform.
#'
#' @return If `test = "F"`, a list with `num_df`, `denom_df`, `f_stat` and
#'   `p_val`. If `test = "Chisq"`, a list with `df`, `chisq_stat`, and `p_val`.
#'   In both cases, p-values are two sided.
#'
#' @keywords internal
h_test_md <- function(
  object,
  contrast,
  df,
  f_stat_factor = 1,
  test = c("F", "Chisq")
) {
  test <- match.arg(test)
  assert_class(object, "mmrm")
  beta <- component(object, "beta_est")
  assert_matrix(contrast, ncols = length(beta))

  if (test == "F") {
    assert_number(df, lower = .Machine$double.xmin)
    assert_number(f_stat_factor, lower = .Machine$double.xmin)
  }

  prec_contrast <- solve(h_quad_form_mat(
    contrast,
    component(object, "beta_vcov")
  ))
  contrast_est <- tcrossprod(beta, contrast)
  chisq_statistic <- as.numeric(h_quad_form_mat(contrast_est, prec_contrast))

  num_df <- nrow(contrast)

  out <-
    if (test == "F") {
      f_statistic <- f_stat_factor / num_df * chisq_statistic

      p_val <- stats::pf(
        q = f_statistic,
        df1 = num_df,
        df2 = df,
        lower.tail = FALSE
      )

      list(
        num_df = num_df,
        denom_df = df,
        f_stat = f_statistic,
        p_val = p_val
      )
    } else {
      p_val <- stats::pchisq(chisq_statistic, df = num_df, lower.tail = FALSE)
      list(df = num_df, chisq_stat = chisq_statistic, p_val = p_val)
    }

  out
}

#' Register `mmrm` For Use With `tidymodels`
#'
#' @inheritParams base::requireNamespace
#' @return A logical value indicating whether registration was successful.
#'
#' @details We can use `parsnip::show_model_info("linear_reg")` to check the
#'   registration with `parsnip` and thus the wider `tidymodels` ecosystem.
#'
#' @keywords internal
parsnip_add_mmrm <- function(quietly = FALSE) {
  if (!requireNamespace("parsnip", quietly = quietly)) {
    return(FALSE)
  }

  parsnip::set_model_engine(
    model = "linear_reg",
    eng = "mmrm",
    mode = "regression"
  )

  parsnip::set_dependency(
    pkg = "mmrm",
    model = "linear_reg",
    eng = "mmrm",
    mode = "regression"
  )

  parsnip::set_encoding(
    model = "linear_reg",
    eng = "mmrm",
    mode = "regression",
    options = list(
      predictor_indicators = "none",
      compute_intercept = FALSE,
      remove_intercept = FALSE,
      allow_sparse_x = TRUE
    )
  )

  parsnip::set_fit(
    model = "linear_reg",
    eng = "mmrm",
    mode = "regression",
    value = list(
      interface = "formula",
      protect = c("formula", "data", "weights"),
      data = c(formula = "formula", data = "data", weights = "weights"),
      func = c(pkg = "mmrm", fun = "mmrm"),
      defaults = list()
    )
  )

  parsnip::set_pred(
    model = "linear_reg",
    eng = "mmrm",
    mode = "regression",
    type = "numeric",
    value = parsnip::pred_value_template(
      # This is boilerplate.
      func = c(fun = "predict"),
      object = quote(object$fit),
      newdata = quote(new_data)
    )
  )

  parsnip::set_pred(
    model = "linear_reg",
    eng = "mmrm",
    mode = "regression",
    # This type allows to pass arguments via `opts` to `parsnip::predict.model_fit`.
    type = "raw",
    value = parsnip::pred_value_template(
      # This is boilerplate.
      func = c(fun = "predict"),
      object = quote(object$fit),
      newdata = quote(new_data)
      # We don't specify additional argument defaults here since otherwise
      # the user is not able to change them (they will be fixed).
    )
  )

  TRUE
}

#' Covariance Type Database
#'
#' An internal constant for covariance type information.
#'
#' @format A data frame with 5 variables and one record per covariance type:
#'
#' \describe{
#'   \item{name}{
#'     The long-form name of the covariance structure type
#'   }
#'   \item{abbr}{
#'     The abbreviated name of the covariance structure type
#'   }
#'   \item{habbr}{
#'     The abbreviated name of the heterogeneous version of a covariance
#'     structure type (The abbreviated name (`abbr`) with a trailing `"h"` if
#'     the structure has a heterogeneous implementation or `NA` otherwise).
#'   }
#'   \item{heterogeneous}{
#'     A logical value indicating whether the covariance structure has a
#'     heterogeneous counterpart.
#'   }
#'   \item{spatial}{
#'     A logical value indicating whether the covariance structure is spatial.
#'   }
#' }
#'
#' @keywords internal
# nolint start
COV_TYPES <- local({
  # nolint end
  type <- function(name, abbr, habbr, heterogeneous, spatial) {
    args <- as.list(match.call()[-1])
    do.call(data.frame, args)
  }

  as.data.frame(
    col.names = names(formals(type)),
    rbind(
      type("unstructured", "us", NA, FALSE, FALSE),
      type("Toeplitz", "toep", "toeph", TRUE, FALSE),
      type("auto-regressive order one", "ar1", "ar1h", TRUE, FALSE),
      type("ante-dependence", "ad", "adh", TRUE, FALSE),
      type("compound symmetry", "cs", "csh", TRUE, FALSE),
      type("spatial exponential", "sp_exp", NA, FALSE, TRUE)
    )
  )
})

#' Covariance Types
#'
#' @description `r lifecycle::badge("stable")`
#'
#' @param form (`character`)\cr covariance structure type name form. One or
#'   more of `"name"`, `"abbr"` (abbreviation), or `"habbr"` (heterogeneous
#'   abbreviation).
#' @param filter (`character`)\cr covariance structure type filter. One or
#'   more of `"heterogeneous"` or `"spatial"`.
#'
#' @return A character vector of accepted covariance structure type names and
#'   abbreviations.
#'
#' @section Abbreviations for Covariance Structures:
#'
#' ## Common Covariance Structures:
#'
#' \tabular{clll}{
#'
#' \strong{Structure}
#' \tab \strong{Description}
#' \tab \strong{Parameters}
#' \tab \strong{\eqn{(i, j)} element}
#' \cr
#'
#' ad
#' \tab Ante-dependence
#' \tab \eqn{m}
#' \tab \eqn{\sigma^{2}\prod_{k=i}^{j-1}\rho_{k}}
#' \cr
#'
#' adh
#' \tab Heterogeneous ante-dependence
#' \tab \eqn{2m-1}
#' \tab \eqn{\sigma_{i}\sigma_{j}\prod_{k=i}^{j-1}\rho_{k}}
#' \cr
#'
#' ar1
#' \tab First-order auto-regressive
#' \tab \eqn{2}
#' \tab \eqn{\sigma^{2}\rho^{\left \vert {i-j} \right \vert}}
#' \cr
#'
#' ar1h
#' \tab Heterogeneous first-order auto-regressive
#' \tab \eqn{m+1}
#' \tab \eqn{\sigma_{i}\sigma_{j}\rho^{\left \vert {i-j} \right \vert}}
#' \cr
#'
#' cs
#' \tab Compound symmetry
#' \tab \eqn{2}
#' \tab \eqn{\sigma^{2}\left[ \rho I(i \neq j)+I(i=j) \right]}
#' \cr
#'
#' csh
#' \tab Heterogeneous compound symmetry
#' \tab \eqn{m+1}
#' \tab \eqn{\sigma_{i}\sigma_{j}\left[ \rho I(i \neq j)+I(i=j) \right]}
#' \cr
#'
#' toep
#' \tab Toeplitz
#' \tab \eqn{m}
#' \tab \eqn{\sigma_{\left \vert {i-j} \right \vert +1}}
#' \cr
#'
#' toeph
#' \tab Heterogeneous Toeplitz
#' \tab \eqn{2m-1}
#' \tab \eqn{\sigma_{i}\sigma_{j}\rho_{\left \vert {i-j} \right \vert}}
#' \cr
#'
#' us
#' \tab Unstructured
#' \tab \eqn{m(m+1)/2}
#' \tab \eqn{\sigma_{ij}}
#'
#' }
#'
#' where \eqn{i} and \eqn{j} denote \eqn{i}-th and \eqn{j}-th time points,
#' respectively, out of total \eqn{m} time points, \eqn{1 \leq i, j \leq m}.
#'
#' @note The **ante-dependence** covariance structure in this package refers to
#' homogeneous ante-dependence, while the ante-dependence covariance structure
#' from SAS `PROC MIXED` refers to heterogeneous ante-dependence and the
#' homogeneous version is not available in SAS.
#'
#' @note For all non-spatial covariance structures, the time variable must
#' be coded as a factor.
#'
#' ## Spatial Covariance structures:
#'
#' \tabular{clll}{
#'
#' \strong{Structure}
#' \tab \strong{Description}
#' \tab \strong{Parameters}
#' \tab \strong{\eqn{(i, j)} element}
#' \cr
#'
#' sp_exp
#' \tab spatial exponential
#' \tab \eqn{2}
#' \tab \eqn{\sigma^{2}\rho^{-d_{ij}}}
#'
#' }
#'
#' where \eqn{d_{ij}} denotes the Euclidean distance between time points
#' \eqn{i} and \eqn{j}.
#'
#' @family covariance types
#' @name covariance_types
#' @export
cov_types <- function(
  form = c("name", "abbr", "habbr"),
  filter = c("heterogeneous", "spatial")
) {
  form <- match.arg(form, several.ok = TRUE)
  filter <- if (missing(filter)) c() else match.arg(filter, several.ok = TRUE)
  df <- COV_TYPES[form][rowSums(!COV_TYPES[filter]) == 0, ]
  Filter(Negate(is.na), unlist(t(df), use.names = FALSE))
}

#' Retrieve Associated Abbreviated Covariance Structure Type Name
#'
#' @param type (`string`)\cr either a full name or abbreviate covariance
#'   structure type name to collapse into an abbreviated type.
#'
#' @return The corresponding abbreviated covariance type name.
#'
#' @keywords internal
cov_type_abbr <- function(type) {
  row <- which(COV_TYPES == type, arr.ind = TRUE)[, 1]
  COV_TYPES$abbr[row]
}

#' Retrieve Associated Full Covariance Structure Type Name
#'
#' @param type (`string`)\cr either a full name or abbreviate covariance
#'   structure type name to convert to a long-form type.
#'
#' @return The corresponding abbreviated covariance type name.
#'
#' @keywords internal
cov_type_name <- function(type) {
  row <- which(COV_TYPES == type, arr.ind = TRUE)[, 1]
  COV_TYPES$name[row]
}

#' Produce A Covariance Identifier Passing to TMB
#'
#' @param cov (`cov_struct`)\cr a covariance structure object.
#'
#' @return A string used for method dispatch when passed to TMB.
#'
#' @keywords internal
tmb_cov_type <- function(cov) {
  paste0(cov$type, if (cov$heterogeneous) "h")
}

#' Define a Covariance Structure
#'
#' @description `r lifecycle::badge("stable")`
#'
#' @param type (`string`)\cr the name of the covariance structure type to use.
#'   For available options, see `cov_types()`. If a type abbreviation is used
#'   that implies heterogeneity (e.g. `cph`) and no value is provided to
#'   `heterogeneous`, then the heterogeneity is derived from the type name.
#' @param visits (`character`)\cr a vector of variable names to use for the
#'   longitudinal terms of the covariance structure. Multiple terms are only
#'   permitted for the `"spatial"` covariance type.
#' @param subject (`string`)\cr the name of the variable that encodes a subject
#'   identifier.
#' @param group (`string`)\cr optionally, the name of the variable that encodes
#'   a grouping variable for subjects.
#' @param heterogeneous (`flag`)\cr
#'
#' @return A `cov_struct` object.
#'
#' @examples
#' cov_struct("csh", "AVISITN", "USUBJID")
#' cov_struct("spatial", c("VISITA", "VISITB"), group = "GRP", subject = "SBJ")
#'
#' @family covariance types
#' @export
cov_struct <- function(
  type = cov_types(),
  visits,
  subject,
  group = character(),
  heterogeneous = FALSE
) {
  # if heterogeneous isn't provided, derive from provided type
  if (missing(heterogeneous)) {
    heterogeneous <- switch(
      type,
      toeph = ,
      ar1h = ,
      adh = ,
      csh = TRUE,
      heterogeneous
    )
  }

  # coerce all type options into abbreviated form
  type <- match.arg(type)
  type <- cov_type_abbr(type)

  x <- structure(
    list(
      type = type,
      heterogeneous = heterogeneous,
      visits = visits,
      subject = subject,
      group = group
    ),
    class = c("cov_struct", "mmrm_cov_struct", "list")
  )

  validate_cov_struct(x)
}

#' Reconcile Possible Covariance Structure Inputs
#'
#' @inheritParams mmrm
#'
#' @return The value `covariance` if it's provided or a covariance structure
#'   derived from the provided `formula` otherwise. An error is raised of both
#'   are provided.
#'
#' @keywords internal
h_reconcile_cov_struct <- function(formula = NULL, covariance = NULL) {
  assert_multi_class(covariance, c("formula", "cov_struct"), null.ok = TRUE)
  assert_formula(formula, null.ok = FALSE)
  if (inherits(covariance, "formula")) {
    covariance <- as.cov_struct(covariance)
  }
  if (!is.null(covariance) && length(h_extract_covariance_terms(formula)) > 0) {
    stop(paste0(
      "Redundant covariance structure definition in `formula` and ",
      "`covariance` arguments"
    ))
  }

  if (!is.null(covariance)) {
    return(covariance)
  }

  as.cov_struct(formula, warn_partial = FALSE)
}

#' Validate Covariance Structure Data
#'
#' Run checks against relational integrity of covariance definition
#'
#' @param x (`cov_struct`)\cr a covariance structure object.
#'
#' @return `x` if successful, or an error is thrown otherwise.
#'
#' @keywords internal
validate_cov_struct <- function(x) {
  checks <- checkmate::makeAssertCollection()

  with(x, {
    assert_character(subject, len = 1, add = checks)
    assert_logical(heterogeneous, len = 1, add = checks)

    if (length(group) > 1 || length(visits) < 1) {
      checks$push(
        "Covariance structure must be of the form `time | (group /) subject`"
      )
    }

    if (!type %in% cov_types(filter = "spatial") && length(visits) > 1) {
      checks$push(paste(
        "Non-spatial covariance structures must have a single longitudinal",
        "variable"
      ))
    }
  })

  reportAssertions(checks)
  x
}

#' Format Covariance Structure Object
#'
#' @param x (`cov_struct`)\cr a covariance structure object.
#' @param ... Additional arguments unused.
#'
#' @return A formatted string for `x`.
#'
#' @export
format.cov_struct <- function(x, ...) {
  sprintf(
    "<covariance structure>\n%s%s:\n\n  %s | %s%s\n",
    if (x$heterogeneous) "heterogeneous " else "",
    cov_type_name(x$type),
    format_symbols(x$visits),
    if (length(x$group) > 0) paste0(format_symbols(x$group), " / ") else "",
    format_symbols(x$subject)
  )
}

#' Print a Covariance Structure Object
#'
#' @param x (`cov_struct`)\cr a covariance structure object.
#' @param ... Additional arguments unused.
#'
#' @return `x` invisibly.
#'
#' @export
print.cov_struct <- function(x, ...) {
  cat(format(x, ...), "\n")
  invisible(x)
}

#' Coerce into a Covariance Structure Definition
#'
#' @description `r lifecycle::badge("stable")`
#'
#' @details
#' A covariance structure can be parsed from a model definition formula or call.
#' Generally, covariance structures defined using non-standard evaluation take
#' the following form:
#'
#' ```
#' type( (visit, )* visit | (group /)? subject )
#' ```
#'
#' For example, formulas may include terms such as
#'
#' ```r
#' us(time | subject)
#' cp(time | group / subject)
#' sp_exp(coord1, coord2 | group / subject)
#' ```
#'
#' Note that only `sp_exp` (spatial) covariance structures may provide multiple
#' coordinates, which identify the Euclidean distance between the time points.
#'
#' @param x an object from which to derive a covariance structure. See object
#'   specific sections for details.
#' @param warn_partial (`flag`)\cr whether to emit a warning when parts of the
#'   formula are disregarded.
#' @param ... additional arguments unused.
#'
#' @return A [cov_struct()] object.
#'
#' @examples
#' # provide a covariance structure as a right-sided formula
#' as.cov_struct(~ csh(visit | group / subject))
#'
#' # when part of a full formula, suppress warnings using `warn_partial = FALSE`
#' as.cov_struct(y ~ x + csh(visit | group / subject), warn_partial = FALSE)
#'
#' @family covariance types
#' @export
# nolint start
as.cov_struct <- function(x, ...) {
  # nolint end
  UseMethod("as.cov_struct")
}

#' @export
as.cov_struct.cov_struct <- function(x, ...) {
  x
}

#' @describeIn as.cov_struct
#' When provided a formula, any specialized functions are assumed to be
#' covariance structure definitions and must follow the form:
#'
#' ```
#' y ~ xs + type( (visit, )* visit | (group /)? subject )
#' ```
#'
#' Any component on the right hand side of a formula is considered when
#' searching for a covariance definition.
#'
#' @export
as.cov_struct.formula <- function(x, warn_partial = TRUE, ...) {
  x_calls <- h_extract_covariance_terms(x)

  if (length(x_calls) < 1) {
    stop(
      "Covariance structure must be specified in formula. ",
      "Possible covariance structures include: ",
      paste0(cov_types(c("abbr", "habbr")), collapse = ", ")
    )
  }

  if (length(x_calls) > 1) {
    cov_struct_types <- as.character(lapply(x_calls, `[[`, 1L))
    stop(
      "Only one covariance structure can be specified. ",
      "Currently specified covariance structures are: ",
      paste0(cov_struct_types, collapse = ", ")
    )
  }

  # flatten into list of infix operators, calls and names/atomics
  x <- flatten_call(x_calls[[1]])
  type <- as.character(x[[1]])
  x <- drop_elements(x, 1)

  # take visits until "|"
  n <- position_symbol(x, "|", nomatch = 0)
  visits <- as.character(utils::head(x, max(n - 1, 0)))
  x <- drop_elements(x, n)

  # take group until "/"
  n <- position_symbol(x, "/", nomatch = 0)
  group <- as.character(utils::head(x, max(n - 1, 0)))
  x <- drop_elements(x, n)

  # remainder is subject
  subject <- as.character(x)

  cov_struct(type = type, visits = visits, group = group, subject = subject)
}

#' Support for `emmeans`
#'
#' @description `r lifecycle::badge("stable")`
#'
#' This package includes methods that allow `mmrm` objects to be used
#' with the `emmeans` package. `emmeans` computes estimated marginal means
#' (also called least-square means) for the coefficients of the MMRM.
#' We can also e.g. obtain differences between groups by applying
#' [`pairs()`][emmeans::pairs.emmGrid()] on the object returned
#' by [emmeans::emmeans()].
#'
#' @examples
#' fit <- mmrm(
#'   formula = FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT | USUBJID),
#'   data = fev_data
#' )
#' if (require(emmeans)) {
#'   emmeans(fit, ~ ARMCD | AVISIT)
#'   pairs(emmeans(fit, ~ ARMCD | AVISIT), reverse = TRUE)
#' }
#' @name emmeans_support
NULL

#' Returns a `data.frame` for `emmeans` Purposes
#'
#' @seealso See [emmeans::recover_data()] for background.
#' @keywords internal
#' @noRd
# nolint start
recover_data.mmrm <- function(object, ...) {
  # nolint end
  fun_call <- stats::getCall(object)
  # subject_var is excluded because it should not contain fixed effect.
  # visit_var is only excluded for spatial covariance structures -
  # for non-spatial covariance structures, emmeans can provide marginal mean
  # by each visit so we need visit_var.
  model_frame <- stats::model.frame(
    object,
    exclude = c(
      "subject_var",
      if (object$formula_parts$is_spatial) "visit_var"
    )
  )
  model_terms <- stats::delete.response(stats::terms(model_frame))
  emmeans::recover_data(
    fun_call,
    trms = model_terms,
    na.action = "na.omit",
    frame = model_frame,
    ...
  )
}

#' Returns a List of Model Details for `emmeans` Purposes
#'
#' @seealso See [emmeans::emm_basis()] for background.
#' @keywords internal
#' @noRd
# nolint start
emm_basis.mmrm <- function(
  # nolint end
  object,
  trms,
  xlev,
  grid,
  ...
) {
  model_frame <- stats::model.frame(
    trms,
    grid,
    na.action = stats::na.pass,
    xlev = xlev
  )
  contrasts <- component(object, "contrasts")
  model_mat <- stats::model.matrix(trms, model_frame, contrasts.arg = contrasts)
  beta_hat <- component(object, "beta_est")
  nbasis <- if (length(beta_hat) < ncol(model_mat)) {
    kept <- match(names(beta_hat), colnames(model_mat))
    beta_hat <- NA * model_mat[1L, ]
    beta_hat[kept] <- component(object, "beta_est")
    orig_model_mat <- stats::model.matrix(
      trms,
      stats::model.frame(
        object,
        exclude = c(
          "subject_var",
          if (object$formula_parts$is_spatial) "visit_var"
        )
      ),
      contrasts.arg = contrasts
    )
    estimability::nonest.basis(orig_model_mat)
  } else {
    estimability::all.estble
  }
  dfargs <- list(object = object)
  dffun <- function(k, dfargs) {
    mmrm::df_md(dfargs$object, contrast = k)$denom_df
  }
  list(
    X = model_mat,
    bhat = beta_hat,
    nbasis = nbasis,
    V = component(object, "beta_vcov"),
    dffun = dffun,
    dfargs = dfargs
  )
}

#' Component Access for `mmrm_tmb` Objects
#'
#' @description `r lifecycle::badge("stable")`
#'
#' @param object (`mmrm_tmb`)\cr the fitted MMRM.
#' @param name (`character`)\cr the component(s) to be retrieved.
#' @return The corresponding component of the object, see details.
#'
#' @details Available `component()` names are as follows:
#' - `call`: low-level function call which generated the model.
#' - `formula`: model formula.
#' - `dataset`: data set name.
#' - `cov_type`: covariance structure type.
#' - `n_theta`: number of parameters.
#' - `n_subjects`: number of subjects.
#' - `n_timepoints`: number of modeled time points.
#' - `n_obs`: total number of observations.
#' - `reml`: was REML used (ML was used if `FALSE`).
#' - `neg_log_lik`: negative log likelihood.
#' - `convergence`: convergence code from optimizer.
#' - `conv_message`: message accompanying the convergence code.
#' - `evaluations`: number of function evaluations for optimization.
#' - `method`: Adjustment method which was used (for `mmrm` objects),
#'      otherwise `NULL` (for `mmrm_tmb` objects).
#' - `beta_vcov`: estimated variance-covariance matrix of coefficients
#'      (excluding aliased coefficients). When Kenward-Roger/Empirical adjusted
#'      coefficients covariance matrix is used, the adjusted covariance matrix is returned (to still obtain the
#'      original asymptotic covariance matrix use `object$beta_vcov`).
#' - `beta_vcov_complete`: estimated variance-covariance matrix including
#'      aliased coefficients with entries set to `NA`.
#' - `varcor`: estimated covariance matrix for residuals. If there are multiple
#'      groups, a named list of estimated covariance matrices for residuals will be
#'      returned. The names are the group levels.
#' - `score_per_subject`: score per subject in empirical covariance.
#'      See the vignette \code{vignette("coef_vcov", package = "mmrm")}.
#' - `theta_est`: estimated variance parameters.
#' - `beta_est`: estimated coefficients (excluding aliased coefficients).
#' - `beta_est_complete`: estimated coefficients including aliased coefficients
#'      set to `NA`.
#' - `beta_aliased`: whether each coefficient was aliased (i.e. cannot be estimated)
#'      or not.
#' - `theta_vcov`:  estimated variance-covariance matrix of variance parameters.
#' - `x_matrix`: design matrix used (excluding aliased columns).
#' - `x_matrix_complete`: design matrix used, including aliased columns.
#' - `xlev`: 	a named list of character vectors giving the full set of levels to be assumed for each factor.
#' - `contrasts`: a list of contrasts used for each factor.
#' - `y_vector`: response vector used.
#' - `jac_list`: Jacobian, see  [h_jac_list()] for details.
#' - `full_frame`: `data.frame` with `n` rows containing all variables needed in the model.
#'
#' @seealso In the `lme4` package there is a similar function `getME()`.
#'
#' @examples
#' fit <- mmrm(
#'   formula = FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT | USUBJID),
#'   data = fev_data
#' )
#' # Get all available components.
#' component(fit)
#' # Get convergence code and message.
#' component(fit, c("convergence", "conv_message"))
#' # Get modeled formula as a string.
#' component(fit, c("formula"))
#'
#' @export
component <- function(
  object,
  name = c(
    "cov_type",
    "subject_var",
    "n_theta",
    "n_subjects",
    "n_timepoints",
    "n_obs",
    "beta_vcov",
    "beta_vcov_complete",
    "varcor",
    "score_per_subject",
    "formula",
    "dataset",
    "n_groups",
    "reml",
    "convergence",
    "evaluations",
    "method",
    "optimizer",
    "conv_message",
    "call",
    "theta_est",
    "beta_est",
    "beta_est_complete",
    "beta_aliased",
    "x_matrix",
    "x_matrix_complete",
    "y_vector",
    "neg_log_lik",
    "jac_list",
    "theta_vcov",
    "full_frame",
    "xlev",
    "contrasts"
  )
) {
  assert_class(object, "mmrm_tmb")
  name <- match.arg(name, several.ok = TRUE)

  list_components <- sapply(
    X = name,
    FUN = switch,
    "call" = object$call,
    # Strings.
    "cov_type" = object$formula_parts$cov_type,
    "subject_var" = object$formula_parts$subject_var,
    "formula" = deparse(object$call$formula),
    "dataset" = object$call$data,
    "reml" = object$reml,
    "conv_message" = object$opt_details$message,
    # Numeric of length 1.
    "convergence" = object$opt_details$convergence,
    "neg_log_lik" = object$neg_log_lik,
    "n_theta" = length(object$theta_est),
    "n_subjects" = object$tmb_data$n_subjects,
    "n_timepoints" = object$tmb_data$n_visits,
    "n_obs" = length(object$tmb_data$y_vector),
    "n_groups" = ifelse(is.list(object$cov), length(object$cov), 1L),
    # Numeric of length > 1.
    "evaluations" = unlist(ifelse(
      is.null(object$opt_details$evaluations),
      list(object$opt_details$counts),
      list(object$opt_details$evaluations)
    )),
    "method" = object$method,
    "optimizer" = object$optimizer,
    "beta_est" = object$beta_est,
    "beta_est_complete" = if (any(object$tmb_data$x_cols_aliased)) {
      stats::setNames(
        object$beta_est[names(object$tmb_data$x_cols_aliased)],
        names(object$tmb_data$x_cols_aliased)
      )
    } else {
      object$beta_est
    },
    "beta_aliased" = object$tmb_data$x_cols_aliased,
    "theta_est" = object$theta_est,
    "y_vector" = object$tmb_data$y_vector,
    "jac_list" = object$jac_list,
    # Matrices.
    "beta_vcov" = if (
      is.null(object$vcov) || identical(object$vcov, "Asymptotic")
    ) {
      object$beta_vcov
    } else {
      object$beta_vcov_adj
    },
    "beta_vcov_complete" = if (any(object$tmb_data$x_cols_aliased)) {
      stats::.vcov.aliased(
        aliased = object$tmb_data$x_cols_aliased,
        vc = component(object, "beta_vcov"),
        complete = TRUE
      )
    } else {
      component(object, "beta_vcov")
    },
    "varcor" = object$cov,
    "score_per_subject" = object$score_per_subject,
    "x_matrix" = object$tmb_data$x_matrix,
    "x_matrix_complete" = object$tmb_data$x_matrix_complete,
    "xlev" = stats::.getXlevels(terms(object), object$tmb_data$full_frame),
    "contrasts" = attr(object$tmb_data$x_matrix, "contrasts"),
    "theta_vcov" = eval(object$theta_vcov),
    "full_frame" = object$tmb_data$full_frame,
    # If not found.
    "..foo.." = stop(sprintf(
      "component '%s' is not available",
      name,
      paste0(class(object), collapse = ", ")
    )),
    simplify = FALSE
  )

  if (length(name) == 1) list_components[[1]] else list_components
}

#' Search For the Position of a Symbol
#'
#' A thin wrapper around [base::Position()] to search through a list of language
#' objects, as produced by [flatten_call()] or [flatten_expr()], for the
#' presence of a specific symbol.
#'
#' @param x (`list` of `language`)\cr a list of language objects in which to
#'   search for a specific symbol.
#' @param sym (`name` or `symbol` or `character`)\cr a symbol to search for in
#'   `x`.
#' @param ... Additional arguments passed to `Position()`.
#'
#' @return The position of the symbol if found, or the `nomatch` value
#'   otherwise.
#'
#' @keywords internal
position_symbol <- function(x, sym, ...) {
  Position(function(i) identical(i, as.symbol(sym)), x, ...)
}

#' Flatten Expressions for Non-standard Evaluation
#'
#' Used primarily to support the parsing of covariance structure definitions
#' from formulas, these functions flatten the syntax tree into a hierarchy-less
#' grammar, allowing for parsing that doesn't abide by R's native operator
#' precedence.
#'
#' Where \code{1 + 2 | 3} in R's syntax tree is \code{(|, (+, 1, 2), 3)},
#' flattening it into its visual order produces \code{(1, +, 2, |, 3)}, which
#' makes for more fluent interpretation of non-standard grammar rules used in
#' formulas.
#'
#' @param call,expr (`language`)\cr a language object to flatten.
#'
#' @return A list of atomic values, symbols, infix operator names and
#'   subexpressions.
#'
#' @name flat_expr
#' @keywords internal
NULL

#' @describeIn flat_expr
#' Flatten a call into a list of names and argument expressions.
#'
#' The call name and all arguments are flattened into the same list, meaning a
#' call of the form \code{sp_exp(a, b, c | d / e)} produces a list of the form
#' \code{(sp_exp, a, b, c, |, d, /, e)}.
#'
#' ```
#' flatten_call(quote(sp_exp(a, b, c | d / e)))
#' ```
#'
#' @keywords internal
flatten_call <- function(call) {
  flattened_args <- unlist(lapply(call[-1], flatten_expr))
  c(flatten_expr(call[[1]]), flattened_args)
}

#' @describeIn flat_expr
#' Flatten nested expressions
#'
#' ```
#' flatten_expr(quote(1 + 2 + 3 | 4))
#' ```
#'
#' @keywords internal
flatten_expr <- function(expr) {
  if (length(expr) > 1 && is_infix(expr[[1]])) {
    op <- list(expr[[1]])
    lhs <- flatten_expr(expr[[2]])
    rhs <- flatten_expr(expr[[3]])
    c(lhs, op, rhs)
  } else {
    list(expr)
  }
}

#' Extract Right-Hand-Side (rhs) from Formula
#'
#' @param f (`formula`)\cr a formula.
#'
#' @return A formula without a response, derived from the right-hand-side of the
#'   formula, `f`.
#'
#' ```
#' formula_rhs(a ~ b + c)
#' formula_rhs(~ b + c)
#' ```
#'
#' @keywords internal
formula_rhs <- function(f) {
  if (length(f) == 2) {
    f
  } else {
    f[-2]
  }
}

#' Test Whether a Symbol is an Infix Operator
#'
#' @param name (`symbol` or `name` or `string`)\cr a possible reference to an
#'   infix operator to check.
#'
#' @return A logical indicating whether the name is the name of an infix
#'   operator.
#'
#' ```
#' is_infix(as.name("|"))
#' is_infix("|")
#' is_infix("c")
#' ```
#'
#' @keywords internal
is_infix <- function(name) {
  "Ops" %in% methods::getGroup(as.character(name), recursive = TRUE)
}

#' Format Symbol Objects
#'
#' For printing, variable names are converted to symbols and deparsed to use R's
#' built-in formatting of variables that may contain spaces or quote characters.
#'
#' @param x (`character`) A vector of variable names.
#'
#' @return A formatted string of comma-separated variables.
#'
#' @keywords internal
format_symbols <- function(x) {
  paste0(
    collapse = ", ",
    lapply(x, function(i) {
      utils::capture.output(as.symbol(i))
    })
  )
}

# Internal functions used for skipping tests or examples.

# Predicate whether currently running R version is under development.
is_r_devel <- function() {
  grepl("devel", R.version$status)
}

# Predicate whether currently running on a Linux operating system.
is_linux <- function() {
  tolower(Sys.info()[["sysname"]]) == "linux"
}

# Get the compiler information. Workaround for older R versions
# where R_compiled_by() is not available.
get_compiler <- function() {
  r_cmd <- file.path(R.home("bin"), "R")
  system2(r_cmd, args = "CMD config CC", stdout = TRUE)
}

# Predicate whether currently using a clang compiler.
is_using_clang <- function() {
  grepl("clang", get_compiler())
}

# Predicate whether an R-devel version is running on Linux Fedora or
# Debian with a clang compiler.
is_r_devel_linux_clang <- function() {
  is_r_devel() &&
    is_linux() &&
    is_using_clang()
}

#' Obtain Empirical/Jackknife/Bias-Reduced Covariance
#'
#' @description Obtain the empirical or Jackknife covariance for \eqn{\beta}.
#' Used in `mmrm` fitting if method is "Empirical", "Empirical-Jackknife" or
#' "Empirical-Bias-Reduced".
#'
#' @param tmb_data (`mmrm_tmb_data`)\cr produced by [h_mmrm_tmb_data()].
#' @param theta (`numeric`)\cr theta estimate.
#' @param beta (`numeric`)\cr beta estimate.
#' @param beta_vcov (`matrix`)\cr covariance of beta estimate.
#' @param type (`string`)\cr type of empirical method, including "Empirical", "Empirical-Jackknife"
#' and "Empirical-Bias-Reduced".
#'
#' @return Named list with elements:
#' - `cov`: `matrix` empirical covariance.
#' - `g_mat`: `matrix` to calculate Satterthwaite degrees of freedom.
#'
#' @note
#' This function used to return `df_mat`, which was equivalent to `crossproduct(g_mat)`. However,
#' executing the cross product in C++ was a costly matrix multiplication, in particular when the number of coefficients
#' and/or the number of subjects was large. Therefore this is now avoided and `g_mat` is returned instead.
#'
#' @keywords internal
h_get_empirical <- function(tmb_data, theta, beta, beta_vcov, type) {
  assert_class(tmb_data, "mmrm_tmb_data")
  assert_numeric(theta)
  n_beta <- ncol(tmb_data$x_matrix)
  assert_numeric(beta, finite = TRUE, any.missing = FALSE, len = n_beta)
  assert_matrix(
    beta_vcov,
    mode = "numeric",
    any.missing = FALSE,
    nrows = n_beta,
    ncols = n_beta
  )
  assert_subset(
    type,
    c("Empirical", "Empirical-Jackknife", "Empirical-Bias-Reduced")
  )
  .Call(
    `_mmrm_get_empirical`,
    PACKAGE = "mmrm",
    tmb_data,
    theta,
    beta,
    beta_vcov,
    type
  )
}

#' Calculation of Residual Degrees of Freedom for One-Dimensional Contrast
#'
#' @description Used in [df_1d()] if method is
#' "Residual".
#'
#' @inheritParams h_df_1d_sat
#' @inherit h_df_1d_sat return
#' @keywords internal
h_df_1d_res <- function(object, contrast) {
  assert_class(object, "mmrm")
  assert_numeric(contrast, len = length(component(object, "beta_est")))

  df <- component(object, "n_obs") - length(component(object, "beta_est"))

  h_test_1d(object, contrast, df)
}

#' Calculation of Residual Degrees of Freedom for Multi-Dimensional Contrast
#'
#' @description Used in [df_md()] if method is "Residual".
#'
#' @inheritParams h_df_md_sat
#' @inherit h_df_md_sat return
#' @keywords internal
h_df_md_res <- function(object, contrast) {
  assert_class(object, "mmrm")
  assert_matrix(
    contrast,
    mode = "numeric",
    any.missing = FALSE,
    ncols = length(component(object, "beta_est"))
  )

  df <- component(object, "n_obs") - length(component(object, "beta_est"))

  h_test_md(object, contrast, df)
}

#ifndef CHOL_CACHE_INCLUDED_
#define CHOL_CACHE_INCLUDED_

#include "covariance.h"
#include "utils.h"

// Base class of spatial and non-spatial Cholesky.
template <class Type>
struct lower_chol_base {
  virtual ~lower_chol_base() {}
  virtual matrix<Type> get_chol(std::vector<int> visits, matrix<Type> dist) = 0;
  virtual matrix<Type> get_sigma(std::vector<int> visits, matrix<Type> dist) = 0;
  virtual matrix<Type> get_sigma_inverse(std::vector<int> visits, matrix<Type> dist) = 0;
};
// Struct to obtain Cholesky for non-spatial.
template <class Type>
struct lower_chol_nonspatial: virtual lower_chol_base<Type> {
  std::map<std::vector<int>, matrix<Type>> chols;
  std::map<std::vector<int>, matrix<Type>> sigmas;
  std::map<std::vector<int>, matrix<Type>> sigmas_inv;
  std::string cov_type;
  int n_visits;
  std::vector<int> full_visit;
  int n_theta;
  vector<Type> theta;
  matrix<Type> chol_full;
  matrix<Type> sigma_full;
  lower_chol_nonspatial() {
    // This default constructor is needed because the use of `[]` in map.
  }
  // Constructor from theta, n_visits and cov_type, and cache full_visits values.
  lower_chol_nonspatial(vector<Type> theta, int n_visits, std::string cov_type): cov_type(cov_type), n_visits(n_visits), full_visit(std::vector<int>(n_visits)) {
    this->theta = theta;
    std::iota(std::begin(this->full_visit), std::end(this->full_visit), 0);
    this->n_theta = theta.size();
    this->chol_full = get_covariance_lower_chol(this->theta, this->n_visits, this->cov_type);
    this->chols[full_visit]  = this->chol_full;
    this->sigma_full = tcrossprod(this->chol_full, true);
  }
  matrix<Type> get_chol(std::vector<int> visits, matrix<Type> dist) {
    auto target = this->chols.find(visits);
     if (target != this->chols.end()) {
      return target->second;
    } else {
      matrix<Type> cov_i = this->get_sigma(visits, dist);
      Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > cov_i_chol(cov_i);
      matrix<Type> Li = cov_i_chol.matrixL();
      this->chols[visits] = Li;
      return Li;
    }
  }
  matrix<Type> get_sigma(std::vector<int> visits, matrix<Type> dist) {
    auto target = this->sigmas.find(visits);
    if (target != this->sigmas.end()) {
      return target->second;
    } else {
      matrix<Type> ret = subset_matrix<matrix<Type>, vector<int>>(sigma_full, visits, visits);
      this->sigmas[visits] = ret;
      return ret;
    }
  }
  matrix<Type> get_sigma_inverse(std::vector<int> visits, matrix<Type> dist) {
    auto target = this->sigmas_inv.find(visits);
    if (target != this->sigmas_inv.end()) {
      return target->second;
    } else {
      matrix<Type> ret = this->get_sigma(visits, dist).inverse();
      this->sigmas_inv[visits] = ret;
      return ret;
    }
  }
};


// Struct to obtain Cholesky for spatial exponential.
template <class Type>
struct lower_chol_spatial: virtual lower_chol_base<Type> {
  vector<Type> theta;
  std::string cov_type;
  lower_chol_spatial() {
    // This default constructor is needed because the use of `[]` in map.
  }
  // Constructor from theta. For now the cholesky does not need to be cached.
  lower_chol_spatial(vector<Type> theta, std::string cov_type): theta(theta), cov_type(cov_type) {
  }
  matrix<Type> get_chol(std::vector<int> visits, matrix<Type> dist) {
    return get_spatial_covariance_lower_chol(this->theta, dist, this->cov_type);
  }
  matrix<Type> get_sigma(std::vector<int> visits, matrix<Type> dist) {
    return tcrossprod(this->get_chol(visits, dist), true);
  }
  matrix<Type> get_sigma_inverse(std::vector<int> visits, matrix<Type> dist) {
    return this->get_sigma(visits, dist).inverse();
  }
};

template <class T, class Base, class D1, class D2>
struct cache_obj {
  std::map<int, std::shared_ptr<Base>> cache;
  int n_groups;
  bool is_spatial;
  int n_visits;
  cache_obj(vector<T> theta, int n_groups, bool is_spatial, std::string cov_type, int n_visits): n_groups(n_groups), is_spatial(is_spatial), n_visits(n_visits) {
    // Get number of variance parameters for one group.
    int theta_one_group_size = theta.size() / n_groups;
    for (int r = 0; r < n_groups; r++) {
      // Use unique pointers here to better manage resource.
      if (is_spatial) {
        this->cache[r] = std::make_shared<D1>(theta.segment(r * theta_one_group_size, theta_one_group_size), cov_type);
      } else {
        this->cache[r] = std::make_shared<D2>(theta.segment(r * theta_one_group_size, theta_one_group_size), n_visits, cov_type);
      }
    }
  }
};

template <class Type>
struct chol_cache_groups: cache_obj<Type, lower_chol_base<Type>, lower_chol_spatial<Type>, lower_chol_nonspatial<Type>> {
  chol_cache_groups(vector<Type> theta, int n_groups, bool is_spatial, std::string cov_type, int n_visits): cache_obj<Type, lower_chol_base<Type>, lower_chol_spatial<Type>, lower_chol_nonspatial<Type>>(theta, n_groups, is_spatial, cov_type, n_visits) {

  }
  // Return covariance lower Cholesky factor from lower_chol_base objects.
  // For non-spatial return for full visits, for spatial return on two points that the distance is 1.
  matrix<Type> get_default_chol() {
    std::vector<int> visit(this->n_visits);
    std::iota(std::begin(visit), std::end(visit), 0);
    matrix<Type> dist(2, 2);
    dist << 0, 1, 1, 0;
    int dim = this->is_spatial?2:this->n_visits;
    matrix<Type> covariance_lower_chol = matrix<Type>::Zero(dim * this->n_groups, dim);
    for (int r = 0; r < this->n_groups; r++) {
      covariance_lower_chol.block(r * dim, 0, dim, dim) = this->cache[r]->get_chol(visit, dist);
    }
    return covariance_lower_chol;
  }
};

#endif

#ifndef COV_INCLUDED_
#define COV_INCLUDED_

#include "utils.h"

// Unstructured covariance:
// Cholesky factor.
template <class T>
matrix<T> get_unstructured(const vector<T>& theta, int n_visits) {
  vector<T> sd_values = exp(theta.head(n_visits));
  vector<T> lower_tri_chol_values = theta.tail(theta.size() - n_visits);
  matrix<T> covariance_lower_chol = matrix<T>::Zero(n_visits, n_visits);
  int k = 0;
  for(int i = 0; i < n_visits; i++) {
    covariance_lower_chol(i, i) = sd_values(i);
    for(int j = 0; j < i; j++){
      covariance_lower_chol(i, j) = sd_values(i) * lower_tri_chol_values(k++);
    }
  }
  return covariance_lower_chol;
}

// Ante-dependence:

// Correlation function.
template <class T>
struct corr_fun_ante_dependence : generic_corr_fun<T> {
  using generic_corr_fun<T>::generic_corr_fun;
  const T operator() (int i, int j) const {
    return this->corr_values.segment(j, i - j).prod();
  }
};
// Homogeneous Ante-dependence Cholesky factor.
template <class T>
matrix<T> get_ante_dependence(const vector<T>& theta, int n_visits) {
  T const_sd = exp(theta(0));
  corr_fun_ante_dependence<T> fun(theta.tail(n_visits - 1));
  matrix<T> ad_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
  return const_sd * ad_cor_mat_chol;
}
// Heterogeneous Ante-dependence Cholesky factor.
template <class T>
matrix<T> get_ante_dependence_heterogeneous(const vector<T>& theta, int n_visits) {
  vector<T> sd_values = exp(theta.head(n_visits));
  corr_fun_ante_dependence<T> fun(theta.tail(n_visits - 1));
  return get_heterogeneous_cov(sd_values, fun);
}

// Toeplitz:

// Correlation function.
template <class T>
struct corr_fun_toeplitz : generic_corr_fun<T> {
  using generic_corr_fun<T>::generic_corr_fun;
  const T operator() (int i, int j) const {
    int index = (i - j) - 1;  // Note: We need to start at 0.
    return this->corr_values(index);
  }
};
// Homogeneous Toeplitz Cholesky factor.
template <class T>
matrix<T> get_toeplitz(const vector<T>& theta, int n_visits) {
  T const_sd = exp(theta(0));
  corr_fun_toeplitz<T> fun(theta.tail(n_visits - 1));
  matrix<T> toep_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
  return const_sd * toep_cor_mat_chol;
}
// Heterogeneous Toeplitz Cholesky factor.
template <class T>
matrix<T> get_toeplitz_heterogeneous(const vector<T>& theta, int n_visits) {
  vector<T> sd_values = exp(theta.head(n_visits));
  corr_fun_toeplitz<T> fun(theta.tail(n_visits - 1));
  return get_heterogeneous_cov(sd_values, fun);
}

// Autoregressive:

// Correlation function.
template <class T>
struct corr_fun_autoregressive : generic_corr_fun<T> {
  using generic_corr_fun<T>::generic_corr_fun;
  const T operator() (int i, int j) const {
    T diff = T((i - j) * 1.0);
    return pow(this->corr_values(0), diff);  // rho^{|i-j|}
  }
};
// Homogeneous autoregressive Cholesky factor.
template <class T>
matrix<T> get_auto_regressive(const vector<T>& theta, int n_visits) {
  T const_sd = exp(theta(0));
  corr_fun_autoregressive<T> fun(theta.tail(1));
  matrix<T> ar1_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
  return const_sd * ar1_cor_mat_chol;
}
// Heterogeneous autoregressive Cholesky factor.
template <class T>
matrix<T> get_auto_regressive_heterogeneous(const vector<T>& theta, int n_visits) {
  vector<T> sd_values = exp(theta.head(n_visits));
  corr_fun_autoregressive<T> fun(theta.tail(1));
  return get_heterogeneous_cov(sd_values, fun);
}

// Compound symmetry:

// Correlation function.
template <class T>
struct corr_fun_compound_symmetry {
  const vector<T> corr_values;
  // Custom constructor here because we need the special mapping from theta to rho.
  corr_fun_compound_symmetry(const vector<T>& theta, int n_visits) :
    corr_values(map_to_cs_cor(theta, n_visits)) {}
  const T operator() (int i, int j) const {
    return this->corr_values(0);  // rho (constant)
  }
};
// Homogeneous compound symmetry Cholesky factor.
template <class T>
matrix<T> get_compound_symmetry(const vector<T>& theta, int n_visits) {
  T const_sd = exp(theta(0));
  corr_fun_compound_symmetry<T> fun(theta.tail(1), n_visits);
  matrix<T> cs_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
  return const_sd * cs_cor_mat_chol;
}
// Heterogeneous compound symmetry Cholesky factor.
template <class T>
matrix<T> get_compound_symmetry_heterogeneous(const vector<T>& theta, int n_visits) {
  vector<T> sd_values = exp(theta.head(n_visits));
  corr_fun_compound_symmetry<T> fun(theta.tail(1), n_visits);
  return get_heterogeneous_cov(sd_values, fun);
}

// Spatial Exponential Cholesky factor.
template <class T>
matrix<T> get_spatial_exponential(const vector<T>& theta, const matrix<T>& distance) {
  T const_sd = exp(theta(0));
  T rho = invlogit(theta(1));
  matrix<T> expdist = exp(distance.array() * log(rho));
  matrix<T> result = expdist * const_sd;
  Eigen::LLT<Eigen::Matrix<T,Eigen::Dynamic,Eigen::Dynamic> > cov_i_chol(result);
  return cov_i_chol.matrixL();
}

// Creates a new correlation object dynamically.
template <class T>
matrix<T> get_covariance_lower_chol(const vector<T>& theta, int n_visits, std::string cov_type) {
  matrix<T> result;

  if (cov_type == "us") {
    result = get_unstructured<T>(theta, n_visits);
  } else if (cov_type == "toep") {
    result = get_toeplitz<T>(theta, n_visits);
  } else if (cov_type == "toeph") {
    result = get_toeplitz_heterogeneous<T>(theta, n_visits);
  } else if (cov_type == "ar1") {
    result = get_auto_regressive<T>(theta, n_visits);
  } else if (cov_type == "ar1h") {
    result = get_auto_regressive_heterogeneous<T>(theta, n_visits);
  } else if (cov_type == "ad") {
    result = get_ante_dependence<T>(theta, n_visits);
  } else if (cov_type == "adh") {
    result = get_ante_dependence_heterogeneous<T>(theta, n_visits);
  } else if (cov_type == "cs") {
    result = get_compound_symmetry<T>(theta, n_visits);
  } else if (cov_type == "csh") {
    result = get_compound_symmetry_heterogeneous<T>(theta, n_visits);
  } else {
    Rf_error("%s", ("Unknown covariance type '" + cov_type + "'.").c_str());
  }

  return result;
}

// Creates a new spatial covariance cholesky.
template <class T>
matrix<T> get_spatial_covariance_lower_chol(const vector<T>& theta, const matrix<T>& distance, std::string cov_type) {
  matrix<T> result;
  if (cov_type == "sp_exp") {
    result = get_spatial_exponential<T>(theta, distance);
  } else {
    Rf_error("%s", ("Unknown spatial covariance type '" + cov_type + "'.").c_str());
  }
  return result;
}

#endif

#ifndef DERIVATIVE_INCLUDED_
#define DERIVATIVE_INCLUDED_

#include "chol_cache.h"

using namespace Rcpp;
using std::string;
// Struct chol to obtain the cholesky factor given theta.
// The reason to have it is that we need a functor that need only theta to
// obtain the derivatives from autodiff.
// Only non-spatial covariance structure here.
struct chol {
  int dim_cov_mat;
  string cov_type;
  chol(int dim, string cov): dim_cov_mat(dim), cov_type(cov) {};
  template <class T>
    vector<T> operator() (vector<T> &theta) {
      return get_covariance_lower_chol(theta, this->dim_cov_mat, this->cov_type).vec();
    }
};
// Struct chol_jacobian that has jacobian of the cholesky factor given theta.
// The reason to have it is that we need hessian so we use jacobian twice.
struct chol_jacobian {
  int dim_cov_mat;
  string cov_type;
  chol mychol;
  chol_jacobian(int dim, string cov): dim_cov_mat(dim), cov_type(cov), mychol(dim, cov) {};
  template<class T>
    vector<T> operator() (vector<T> &theta) {
      return autodiff::jacobian(this->mychol, theta).vec();
    }
};

// Template function to obtain derivatives from visits, cov_type and theta.
// Basically this is calculating the derivatives for the sigma
// from the derivatives for the cholesky factor.
template <class Type>
std::map<std::string, matrix<Type>> derivatives(int n_visits, std::string cov_type, vector<Type> theta) {
  std::map<std::string, matrix<Type>> ret;
  chol chol_obj(n_visits, cov_type);
  chol_jacobian chol_jac_obj(n_visits, cov_type);
  matrix<Type> l = chol_obj(theta).matrix();
  l.resize(n_visits, n_visits);
  vector<Type> chol_d1_vec = autodiff::jacobian(chol_obj, theta).vec(); // chol_d1_vec is (dim * dim * l_theta)
  vector<Type> chol_d2_vec = autodiff::jacobian(chol_jac_obj, theta).vec(); // chol_d2_vec is (dim * dim * l_theta * l_theta)
  matrix<Type> ret_d1 = matrix<Type>(n_visits * theta.size(), n_visits);
  matrix<Type> ret_d2 = matrix<Type>(n_visits * theta.size() * theta.size(), n_visits);
  int n_visits_sq = n_visits * n_visits;
  for (int i = 0; i < theta.size(); i++) {
    matrix<Type> ld1 = chol_d1_vec.segment(i * n_visits_sq, n_visits_sq).matrix();
    ld1.resize(n_visits, n_visits);
    matrix<Type> ld1_lt = ld1 * l.transpose();
    auto sigma_d1_i = ld1_lt + ld1_lt.transpose();
    ret_d1.block(i * n_visits, 0, n_visits, n_visits) = sigma_d1_i;
    for (int j = 0; j < theta.size(); j++) {
      matrix<Type> ld2 = chol_d2_vec.segment( (j * theta.size() + i) * n_visits_sq, n_visits_sq).matrix();
      matrix<Type> ld1_j = chol_d1_vec.segment(j * n_visits_sq, n_visits_sq).matrix();
      ld2.resize(n_visits, n_visits);
      ld1_j.resize(n_visits, n_visits);
      auto ld2_lt = ld2 * l.transpose();
      auto ld1_ld1j = ld1 * ld1_j.transpose();
      auto sigma_d2_ij = ld2_lt + ld2_lt.transpose() + ld1_ld1j + ld1_ld1j.transpose();
      ret_d2.block((i * theta.size() + j) * n_visits, 0, n_visits, n_visits) = sigma_d2_ij;
    }
  }
  ret["derivative1"] = ret_d1;
  ret["derivative2"] = ret_d2;
  return ret;
}
// Base class of spatial and non-spatial derivatives.
template <class Type>
struct derivatives_base: virtual lower_chol_base<Type> {
  virtual matrix<Type> get_inverse_chol(std::vector<int> visits, matrix<Type> dist) = 0;
  virtual matrix<Type> get_sigma_derivative1(std::vector<int> visits, matrix<Type> dist) = 0;
  virtual matrix<Type> get_sigma_derivative2(std::vector<int> visits, matrix<Type> dist) = 0;
  virtual matrix<Type> get_inverse_derivative(std::vector<int> visits, matrix<Type> dist) = 0;
  // Create virtual destructor to avoid the default desctructor being called.
  virtual ~derivatives_base() {};
};

// Struct derivatives_nonspatial is created to get the derivatives with cache.
// The main reason to have it is that we nearly always have duplicated visits
// and the inverse of a matrix is calculation expensive. In addition, we can save
// the resource needed for select matrix calculations.
template <class Type>
struct derivatives_nonspatial: public lower_chol_nonspatial<Type>, virtual derivatives_base<Type> {
  std::map<std::vector<int>, matrix<Type>> inverse_chol_cache;
  std::map<std::vector<int>, matrix<Type>> sigmad1_cache;
  std::map<std::vector<int>, matrix<Type>> sigmad2_cache;
  std::map<std::vector<int>, matrix<Type>> sigma_inverse_d1_cache;
  derivatives_nonspatial() {
    // This default constructor is needed because the use of `[]` in map.
  }
  // Constructor from theta, n_visits and cov_type, and cache full_visits values.
  derivatives_nonspatial(vector<Type> theta, int n_visits, std::string cov_type): lower_chol_nonspatial<Type>(theta, n_visits, cov_type) {
    std::map<std::string, tmbutils::matrix<Type>> allret = derivatives<Type>(this->n_visits, this->cov_type, this->theta);
    matrix<Type> sigma_d1 = allret["derivative1"];
    matrix<Type> sigma_d2 = allret["derivative2"];
    this->sigmad1_cache[this->full_visit] = sigma_d1;
    this->sigmad2_cache[this->full_visit] = sigma_d2;
  }
  // Cache and return the first order derivatives using select matrix.
  matrix<Type> get_sigma_derivative1(std::vector<int> visits, matrix<Type> dist) override {
    auto target = this->sigmad1_cache.find(visits);
     if (target != this->sigmad1_cache.end()) {
      return target->second;
    } else {
      int n_visits_i = visits.size();
      matrix<Type> ret = matrix<Type>(this->n_theta * n_visits_i, n_visits_i);
      for (int i = 0; i < this->n_theta; i++) {
        ret.block(i  * n_visits_i, 0, n_visits_i, n_visits_i) = subset_matrix<matrix<Type>, vector<int>>(this->sigmad1_cache[this->full_visit].block(i  * this->n_visits, 0, this->n_visits, this->n_visits), visits, visits);
      }
      this->sigmad1_cache[visits] = ret;
      return ret;
    }
  }
  // Cache and return the second order derivatives using select matrix.
  matrix<Type> get_sigma_derivative2(std::vector<int> visits, matrix<Type> dist) override {
    auto target = this->sigmad2_cache.find(visits);
     if (target != this->sigmad2_cache.end()) {
      return target->second;
    } else {
      int n_visits_i = visits.size();
      int theta_sq = this->n_theta * this->n_theta;
      matrix<Type> ret = matrix<Type>(theta_sq * n_visits_i, n_visits_i);
      for (int i = 0; i < theta_sq; i++) {
        ret.block(i  * n_visits_i, 0, n_visits_i, n_visits_i) = subset_matrix<matrix<Type>, vector<int>>(this->sigmad2_cache[this->full_visit].block(i  * this->n_visits, 0, this->n_visits, this->n_visits), visits, visits);
      }
      this->sigmad2_cache[visits] = ret;
      return ret;
    }
  }
  // Cache and return the lower cholesky factor of inverse of sigma using select matrix.
  matrix<Type> get_inverse_chol(std::vector<int> visits, matrix<Type> dist) override {
    auto target = this->inverse_chol_cache.find(visits);
     if (target != this->inverse_chol_cache.end()) {
      return target->second;
    } else {
      matrix<Type> sigmainv = this->get_sigma_inverse(visits, dist);
      Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > sigma_inv_chol(sigmainv);
      matrix<Type> Li = sigma_inv_chol.matrixL();
      this->inverse_chol_cache[visits] = Li;
      return Li;
    }
  }
  // Cache and return the first order derivatives of inverse of sigma using select matrix.
  matrix<Type> get_inverse_derivative(std::vector<int> visits, matrix<Type> dist) override {
    auto target = this->sigma_inverse_d1_cache.find(visits);
     if (target != this->sigma_inverse_d1_cache.end()) {
      return target->second;
    } else {
      auto sigma_d1 = this->get_sigma_derivative1(visits, dist);
      matrix<Type> sigma_inv_d1(sigma_d1.rows(), sigma_d1.cols());
      int n_visits_i = visits.size();
      auto sigma_inv = this->get_sigma_inverse(visits, dist);
      for (int r = 0; r < this->n_theta; r++) {
        sigma_inv_d1.block(r * n_visits_i, 0, n_visits_i, n_visits_i) = - sigma_inv * sigma_d1.block(r * n_visits_i, 0, n_visits_i, n_visits_i) *sigma_inv;
      }
      this->sigma_inverse_d1_cache[visits] = sigma_inv_d1;
      return sigma_inv_d1;
    }
  }
};

// derivatives_sp_exp struct is created to obtain the exact derivatives of spatial exponential
// covariance structure, and its inverse.
// No caching is used because the distance can be hardly the same for spatial covariance
// structures.
template <class Type>
struct derivatives_sp_exp: public lower_chol_spatial<Type>, virtual derivatives_base<Type> {
  Type const_sd;
  Type rho;
  Type logrho;
  derivatives_sp_exp() {
    // This default constructor is needed because the use of `[]` in maps.
  }
  // Initialize the theta values; the reason to have theta is that for a fit, the theta
  // is the same for all subjects, while the distance between each visits for each subject
  // can be different.
  derivatives_sp_exp(vector<Type> theta, std::string cov_type): lower_chol_spatial<Type>(theta, cov_type) ,const_sd(exp(theta(0))), rho(invlogit(theta(1))) {
    this->logrho = log(this->rho);
  }
  // Obtain first order derivatives
  matrix<Type> get_sigma_derivative1(std::vector<int> visits, matrix<Type> dist) override {
    matrix<Type> ret(2 * dist.rows(), dist.cols());
    // partial sigma / partial theta_1 = sigma.
    auto sigma = this->get_sigma(visits, dist);
    ret.block(0, 0, dist.rows(), dist.cols()) = sigma;
    ret.block(dist.rows(), 0, dist.rows(), dist.cols()) = sigma.array() * dist.array() * (1 - this->rho);
    return ret;
  }
  // Obtain second order derivatives.
  matrix<Type> get_sigma_derivative2(std::vector<int> visits, matrix<Type> dist) override {
    matrix<Type> ret(4 * dist.rows(), dist.cols());
    auto sigma = this->get_sigma(visits, dist);
    ret.block(0, 0, dist.rows(), dist.cols()) = sigma;
    Type rho_r = 1 - this->rho;
    auto dtheta1dtheta2 = sigma.array() * dist.array() * rho_r;
    ret.block(dist.rows(), 0, dist.rows(), dist.cols()) =  dtheta1dtheta2;
    ret.block(dist.rows() * 2, 0, dist.rows(), dist.cols()) = dtheta1dtheta2;
    matrix<Type> dtheta2s = dtheta1dtheta2 * (dist.array() * rho_r - this->rho);
    ret.block(dist.rows() * 3, 0, dist.rows(), dist.cols()) = dtheta2s;
    return ret;
  }
  // Obtain the lower cholesky factor of inverse of sigma using select matrix.
  matrix<Type> get_inverse_chol(std::vector<int> visits, matrix<Type> dist) override {
    auto sigmainv = this->get_sigma_inverse(visits, dist);
    Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > sigma_inv_chol(sigmainv);
    matrix<Type> Li = sigma_inv_chol.matrixL();
    return Li;
  }
  // Obtain first order derivatives for inverse of sigma.
  matrix<Type> get_inverse_derivative(std::vector<int> visits, matrix<Type> dist) override {
    matrix<Type> sigma_inv_d1 = matrix<Type>::Zero(2 * dist.rows(), dist.cols());
    auto sigma_inv = this->get_sigma_inverse(visits, dist);
    auto sigma_d1 = this->get_sigma_derivative1(visits, dist);
    for (int r = 0; r < 2; r++) {
      sigma_inv_d1.block(r * dist.rows(), 0, dist.rows(), dist.cols()) = - sigma_inv * sigma_d1.block(r * dist.rows(), 0, dist.rows(), dist.cols()) *sigma_inv;
    }
    return sigma_inv_d1;
  }
};

#endif

#include "derivatives.h"

using namespace Rcpp;
using std::string;
// Obtain the empirical given beta, beta_vcov, theta.
//
// Note: This function previously (version < 0.3.15) returned `df_mat` which was the crossproduct of the `g` matrix.
// This was removed because this matrix can be very large if the number of subjects is large and/or the number of
// coefficients is large. Instead, now the `g_mat` element includes the `g` matrix.
List get_empirical(List mmrm_data, NumericVector theta, NumericVector beta, NumericMatrix beta_vcov, string type) {
  NumericMatrix x = mmrm_data["x_matrix"];
  matrix<double> x_matrix = as_num_matrix_tmb(x);
  NumericVector y = mmrm_data["y_vector"];
  matrix<double> beta_vcov_matrix = as_num_matrix_tmb(beta_vcov);
  IntegerVector subject_zero_inds = mmrm_data["subject_zero_inds"];
  int n_subjects = mmrm_data["n_subjects"];
  int n_observations = x_matrix.rows();
  IntegerVector subject_n_visits = mmrm_data["subject_n_visits"];
  int n_visits = mmrm_data["n_visits"];
  String cov_type = mmrm_data["cov_type"];
  int is_spatial_int = mmrm_data["is_spatial_int"];
  bool is_spatial = is_spatial_int == 1;
  int n_groups = mmrm_data["n_groups"];
  IntegerVector subject_groups = mmrm_data["subject_groups"];
  NumericVector weights_vector = mmrm_data["weights_vector"];
  NumericMatrix coordinates = mmrm_data["coordinates"];
  matrix<double> coords = as_num_matrix_tmb(coordinates);
  matrix<double> beta_m = as_num_vector_tmb(beta).matrix();
  vector<double> theta_v = as_num_vector_tmb(theta);
  matrix<double> fitted = x_matrix * beta_m;
  matrix<double> y_matrix = as_num_vector_tmb(y).matrix();
  matrix<double> residual = y_matrix - fitted;
  vector<double> G_sqrt = as_num_vector_tmb(sqrt(weights_vector));
  int p = x.cols();

  matrix<double> score_per_subject = matrix<double>::Zero(n_subjects, p);

  // Use map to hold these base class pointers (can also work for child class objects).
  auto derivatives_by_group = cache_obj<double, derivatives_base<double>, derivatives_sp_exp<double>, derivatives_nonspatial<double>>(theta_v, n_groups, is_spatial, cov_type, n_visits);
  matrix<double> meat = matrix<double>::Zero(p, p);
  matrix<double> xt_g_simga_inv_chol = matrix<double>::Zero(p, n_observations);
  matrix<double> ax = matrix<double>::Zero(n_observations, p);
  for (int i = 0; i < n_subjects; i++) {
    int start_i = subject_zero_inds[i];
    int n_visits_i = subject_n_visits[i];
    std::vector<int> visit_i(n_visits_i);
    matrix<double> dist_i(n_visits_i, n_visits_i);
    if (!is_spatial) {
      for (int i = 0; i < n_visits_i; i++) {
        visit_i[i] = int(coordinates(i + start_i, 0));
      }
    } else {
      dist_i = euclidean(matrix<double>(coords.block(start_i, 0, n_visits_i, coordinates.cols())));
    }
    int subject_group_i = subject_groups[i] - 1;
    matrix<double> sigma_inv_chol = derivatives_by_group.cache[subject_group_i]->get_inverse_chol(visit_i, dist_i);
    matrix<double> Xi = x_matrix.block(start_i, 0, n_visits_i, x_matrix.cols());
    matrix<double> residual_i = residual.block(start_i, 0, n_visits_i, 1);
    matrix<double> gi_sqrt_root = G_sqrt.segment(start_i, n_visits_i).matrix().asDiagonal();
    matrix<double> gi_simga_inv_chol = gi_sqrt_root * sigma_inv_chol;
    matrix<double> xt_gi_simga_inv_chol = Xi.transpose() * gi_simga_inv_chol;
    matrix<double> ai = matrix<double>::Identity(n_visits_i, n_visits_i);
    if (type != "Empirical") {
      ai = ai - xt_gi_simga_inv_chol.transpose() * beta_vcov_matrix * xt_gi_simga_inv_chol;
    }
    if (type == "Empirical-Jackknife") {
      ai = ai.inverse();
    } else if(type == "Empirical-Bias-Reduced") {
      ai = pseudoInverseSqrt(ai);
    }
    matrix<double> xta = xt_gi_simga_inv_chol * ai;
    matrix<double> z = xta * gi_simga_inv_chol.transpose() * residual_i;
    meat = meat + z * z.transpose();
    xt_g_simga_inv_chol.block(0, start_i, p, n_visits_i) = xt_gi_simga_inv_chol;
    ax.block(start_i, 0, n_visits_i, p) = xta.transpose();

    score_per_subject.row(i) = z.transpose();
  }
  matrix<double> h = xt_g_simga_inv_chol.transpose() * beta_vcov_matrix * xt_g_simga_inv_chol;
  matrix<double> imh = matrix<double>::Identity(n_observations, n_observations) - h;
  matrix<double> ax_xtx =  ax * beta_vcov_matrix;
  matrix<double> g = matrix<double>::Zero(n_observations, p * n_subjects);
  for (int i = 0; i < n_subjects; i++) {
    int start_i = subject_zero_inds[i];
    int n_visits_i = subject_n_visits[i];
    g.block(0, i * p, n_observations, p) = imh.block(0, start_i, n_observations, n_visits_i) * ax_xtx.block(start_i, 0, n_visits_i, p);
  }

  // beta_vcov already take gi into consideration;
  matrix<double> ret = beta_vcov_matrix * meat * beta_vcov_matrix;

  return List::create(
    Named("score_per_subject") = as_num_matrix_rcpp(score_per_subject),
    Named("cov") = as_num_matrix_rcpp(ret),
    Named("g_mat") = as_num_matrix_rcpp(g)
  );
}

#ifndef UTILS_INCLUDED_
#define UTILS_INCLUDED_
#include <Rcpp.h>
#define INCLUDE_RCPP
#include "tmb_includes.h"

#define as_num_matrix_tmb as_matrix<matrix<double>, NumericMatrix>
#define as_num_matrix_rcpp as_matrix<NumericMatrix, matrix<double>>
#define as_num_vector_tmb as_vector<vector<double>, NumericVector>
#define as_num_vector_rcpp as_vector<NumericVector, vector<double>>

// Obtain submatrix from index

template <typename T1, typename T2>
T1 subset_matrix(T1 input, T2 index1, T2 index2) {
  #if EIGEN_VERSION_AT_LEAST(3,4,0)
    T1 ret = input(index1, index2);
  #else
    T1 ret(index1.size(), index2.size());
    for (decltype(index1.size()) i = 0; i < index1.size(); i++) {
      for (decltype(index2.size()) j = 0; j < index2.size(); j++) {
        ret(i, j) = input(index1[i], index2[j]);
      }
    }
  #endif
  return ret;
}

template <typename T1, typename T2>
T1 subset_matrix(T1 input, T2 index1) {
  #if EIGEN_VERSION_AT_LEAST(3,4,0)
    T1 ret = input(index1, Eigen::all);
  #else
    T1 ret(index1.size(), input.cols());
    for (decltype(index1.size()) i = 0; i < index1.size(); i++) {
      for (int j = 0; j < input.cols(); j++) {
        ret(i, j) = input(index1[i], j);
      }
    }
  #endif
  return ret;
}


// Conversion from Rcpp vector/matrix to eigen vector/matrix
template <typename T1, typename T2>
T1 as_vector(T2 input) {
  T1 ret(input.size());
  for (int i = 0; i < input.size(); i++) {
    ret(i) = input(i);
  }
  return ret;
}

template <typename T1, typename T2>
T1 as_matrix(T2 input) {
  T1 ret(input.rows(), input.cols());
  for (int i = 0; i < input.rows(); i++) {
    for (int j = 0; j < input.cols(); j++) {
      ret(i,j) = input(i,j);
    }
  }
  return ret;
}

template <typename T>
T segment(T input, int start, int n) {
  T ret(n);
  for (int i = 0, j = start; i < n; i++, j++) {
    ret(i) = input(j);
  }
  return ret;
}

// Calculate tcrossprod(lower_chol) = lower_chol * t(lower_chol).
// If complete, then adds the upper triangular part to the result as well.
// By default only the lower triangular part is populated, as this should be
// sufficient for downstream use of the result in most cases.
template <class Type>
matrix<Type> tcrossprod(const matrix<Type>& lower_chol, bool complete = false) {
  int n = lower_chol.rows();
  matrix<Type> result = matrix<Type>::Zero(n, n);
  result.template selfadjointView<Eigen::Lower>().rankUpdate(lower_chol);
  if (complete) {
    result.template triangularView<Eigen::Upper>() = result.transpose();
  }
  return result;
}

// Calculate crossprod(x) = t(x) * x.
// Only the lower triangular part is populated, as this should be
// sufficient for downstream use of the result in most cases.
// Note that x does not need to be symmetric or square.
template <class Type>
matrix<Type> crossprod(const matrix<Type>& x) {
  int n = x.cols();
  matrix<Type> result = matrix<Type>::Zero(n, n);
  result.template selfadjointView<Eigen::Lower>().rankUpdate(x.transpose());
  return result;
}

// Mapping from real values to correlation parameters in (-1, 1).
template <class T>
vector<T> map_to_cor(const vector<T>& theta) {
  return theta / sqrt(T(1.0) + theta * theta);
}

// Mapping from real values to correlation parameters in (-1/(m-1), 1) for compound symmetry.
template <class T>
vector<T> map_to_cs_cor(const vector<T>& theta, int n_visits) {
  T a = T(1.0) / (T(n_visits) - T(1.0));
  return invlogit(theta) * (T(1.0) + a) - a;
}

// Generic correlation function class containing and initializing correlation
// values from variance parameters theta.
template <class T>
struct generic_corr_fun {
  const vector<T> corr_values;

  generic_corr_fun(const vector<T>& theta) :
    corr_values(map_to_cor(theta)) {}
};

// Correlation function based Cholesky factor of correlation matrix.
// This is used directly for homogeneous covariance matrices.
template <class T, template<class> class F>
matrix<T> get_corr_mat_chol(int n_visits, const F<T>& corr_fun) {
  matrix<T> correlation(n_visits, n_visits);
  correlation.setIdentity();
  for(int i = 0; i < n_visits; i++) {
    for(int j = 0; j < i; j++){
      correlation(i, j) = corr_fun(i, j);
    }
  }
  Eigen::LLT<Eigen::Matrix<T,Eigen::Dynamic,Eigen::Dynamic> > correlation_chol(correlation);
  matrix<T> L = correlation_chol.matrixL();
  return L;
}

// Heterogeneous covariance matrix calculation given vector of standard deviations (sd_values)
// and a correlation function (corr_fun).
template <class T, template<class> class F>
matrix<T> get_heterogeneous_cov(const vector<T>& sd_values, const F<T>& corr_fun) {
  matrix<T> correlation_chol = get_corr_mat_chol(sd_values.size(), corr_fun);
  Eigen::DiagonalMatrix<T,Eigen::Dynamic,Eigen::Dynamic> D = sd_values.matrix().asDiagonal();
  matrix<T> result = D * correlation_chol;
  return result;
}

// Obtain the Euclidean distance
template <class T>
matrix<T> euclidean(const matrix<T>& coordinates) {
  matrix<T> result(coordinates.rows(), coordinates.rows());
  for (int i = 0; i < coordinates.rows(); i++) {
    result(i, i) = 0;
    for (int j = 0; j < i; j ++) {
      vector<T> diff = coordinates.row(i) - coordinates.row(j);
      T d = sqrt((diff * diff).sum());
      result(i, j) = d;
      result(j, i) = d;
    }
  }
  return result;
}

// Element wise power function of a matrix
template <class T>
Eigen::Matrix<T, -1, -1> cpow(const Eigen::Matrix<T, -1, -1> & input, double p) {
  Eigen::Matrix<T, -1, -1> ret = Eigen::Matrix<T, -1, -1>(input.rows(), input.cols());
  for (int i = 0; i < ret.rows(); i ++) {
    for (int j = 0; j < ret.cols(); j++) {
      ret(i, j) = std::pow(input(i, j), p);
    }
  }
  return ret;
}

// Calculate the square root of the pseudo inverse of a matrix
// adapted from the method for calculating the pseudo-Inverse as recommended by the Eigen developers
template<typename T>
matrix<T> pseudoInverseSqrt(const matrix<T> &input, double epsilon = std::numeric_limits<double>::epsilon()) {
  Eigen::Matrix<T, -1, -1> eigen_mat = as_matrix<Eigen::Matrix<T, -1, -1>, matrix<T>>(input);
	Eigen::JacobiSVD< Eigen::Matrix<T, -1, -1> > svd(eigen_mat ,Eigen::ComputeFullU | Eigen::ComputeFullV);
	double tolerance = epsilon * std::max(input.cols(), input.rows()) *svd.singularValues().array().abs()(0);
  auto singular_vals = Matrix<T,-1,-1>((svd.singularValues().array() > tolerance).select(svd.singularValues().array().inverse(), 0).matrix());
	Eigen::Matrix<T, -1, -1> ret_eigen = svd.matrixV() *  cpow(singular_vals, 0.5).asDiagonal() * svd.matrixU().adjoint();
  return as_matrix<matrix<T>, Eigen::Matrix<T, -1, -1>>(ret_eigen);
}

#endif

#include <RcppEigen.h>
#include "utils.h"

using namespace Rcpp;

#ifdef RCPP_USE_GLOBAL_ROSTREAM
Rcpp::Rostream<true>&  Rcpp::Rcout = Rcpp::Rcpp_cout_get();
Rcpp::Rostream<false>& Rcpp::Rcerr = Rcpp::Rcpp_cerr_get();
#endif

List get_pqr(List mmrm_fit, NumericVector theta);
RcppExport SEXP _mmrm_get_pqr(SEXP mmrm_fit_SEXP, SEXP theta_SEXP) {
BEGIN_RCPP
    Rcpp::RObject rcpp_result_gen;
    Rcpp::RNGScope rcpp_rngScope_gen;
    Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
    Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
    rcpp_result_gen = Rcpp::wrap(get_pqr(mmrm_fit, theta));
    return rcpp_result_gen;
END_RCPP
}

List get_jacobian(List mmrm_fit, NumericVector theta, NumericMatrix beta_vcov);
RcppExport SEXP _mmrm_get_jacobian(SEXP mmrm_fit_SEXP, SEXP theta_SEXP, SEXP beta_vcov_SEXP) {
BEGIN_RCPP
    Rcpp::RObject rcpp_result_gen;
    Rcpp::RNGScope rcpp_rngScope_gen;
    Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
    Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
    Rcpp::traits::input_parameter< NumericMatrix >::type beta_vcov(beta_vcov_SEXP);
    rcpp_result_gen = Rcpp::wrap(get_jacobian(mmrm_fit, theta, beta_vcov));
    return rcpp_result_gen;
END_RCPP
}

List get_empirical(List mmrm_fit, NumericVector theta, NumericVector beta, NumericMatrix beta_vcov, std::string type);
RcppExport SEXP _mmrm_get_empirical(SEXP mmrm_fit_SEXP, SEXP theta_SEXP, SEXP beta_SEXP, SEXP beta_vcov_SEXP, SEXP type_SEXP) {
BEGIN_RCPP
    Rcpp::RObject rcpp_result_gen;
    Rcpp::RNGScope rcpp_rngScope_gen;
    Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
    Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
    Rcpp::traits::input_parameter< NumericVector >::type beta(beta_SEXP);
    Rcpp::traits::input_parameter< NumericMatrix >::type beta_vcov(beta_vcov_SEXP);
    Rcpp::traits::input_parameter< std::string >::type type(type_SEXP);
    rcpp_result_gen = Rcpp::wrap(get_empirical(mmrm_fit, theta, beta, beta_vcov, type));
    return rcpp_result_gen;
END_RCPP
}

List predict(List mmrm_fit, NumericVector theta, NumericVector beta, NumericMatrix beta_vcov);
RcppExport SEXP _mmrm_predict(SEXP mmrm_fit_SEXP, SEXP theta_SEXP, SEXP beta_SEXP, SEXP beta_vcov_SEXP) {
BEGIN_RCPP
    Rcpp::RObject rcpp_result_gen;
    Rcpp::RNGScope rcpp_rngScope_gen;
    Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
    Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
    Rcpp::traits::input_parameter< NumericVector >::type beta(beta_SEXP);
    Rcpp::traits::input_parameter< NumericMatrix >::type beta_vcov(beta_vcov_SEXP);
    rcpp_result_gen = Rcpp::wrap(predict(mmrm_fit, theta, beta, beta_vcov));
    return rcpp_result_gen;
END_RCPP
}


RcppExport SEXP run_testthat_tests(SEXP);

static const R_CallMethodDef CallEntries[] = {
    {"_mmrm_get_pqr", (DL_FUNC) &_mmrm_get_pqr, 2},
    {"_mmrm_get_jacobian", (DL_FUNC) &_mmrm_get_jacobian, 3},
    {"_mmrm_get_empirical", (DL_FUNC) &_mmrm_get_empirical, 5},
    {"_mmrm_predict", (DL_FUNC) &_mmrm_predict, 4},
    {"run_testthat_tests", (DL_FUNC) &run_testthat_tests, 1},
    TMB_CALLDEFS,
    {NULL, NULL, 0}
};

RcppExport void R_init_mmrm(DllInfo *dll) {
    R_registerRoutines(dll, NULL, CallEntries, NULL, NULL);
    R_useDynamicSymbols(dll, FALSE);
    #ifdef TMB_CCALLABLES
    TMB_CCALLABLES("mmrm");
    #endif
}

#include "derivatives.h"

using namespace Rcpp;
using std::string;

// Obtain Jacobian from a mmrm fit, given theta.
List get_jacobian(List mmrm_fit, NumericVector theta, NumericMatrix beta_vcov) {
  NumericMatrix x = mmrm_fit["x_matrix"];
  matrix<double> x_matrix = as_num_matrix_tmb(x);
  IntegerVector subject_zero_inds = mmrm_fit["subject_zero_inds"];
  int n_subjects = mmrm_fit["n_subjects"];
  IntegerVector subject_n_visits = mmrm_fit["subject_n_visits"];
  int n_visits = mmrm_fit["n_visits"];
  String cov_type = mmrm_fit["cov_type"];
  int is_spatial_int = mmrm_fit["is_spatial_int"];
  bool is_spatial = is_spatial_int == 1;
  int n_groups = mmrm_fit["n_groups"];
  IntegerVector subject_groups = mmrm_fit["subject_groups"];
  NumericVector weights_vector = mmrm_fit["weights_vector"];
  NumericMatrix coordinates = mmrm_fit["coordinates"];
  matrix<double> coords = as_num_matrix_tmb(coordinates);
  matrix<double> beta_vcov_m = as_num_matrix_tmb(beta_vcov);
  vector<double> theta_v = as_num_vector_tmb(theta);
  vector<double> G_sqrt = as_num_vector_tmb(sqrt(weights_vector));
  int n_theta = theta.size();
  int theta_size_per_group = n_theta / n_groups;
  int p = x.cols();
  matrix<double> P = matrix<double>::Zero(p * n_theta, p);
  // Use map to hold these base class pointers (can also work for child class objects).
  auto derivatives_by_group = cache_obj<double, derivatives_base<double>, derivatives_sp_exp<double>, derivatives_nonspatial<double>>(theta_v, n_groups, is_spatial, cov_type, n_visits);
  for (int i = 0; i < n_subjects; i++) {
    int start_i = subject_zero_inds[i];
    int n_visits_i = subject_n_visits[i];
    std::vector<int> visit_i(n_visits_i);
    matrix<double> dist_i(n_visits_i, n_visits_i);
    if (!is_spatial) {
      for (int i = 0; i < n_visits_i; i++) {
        visit_i[i] = int(coordinates(i + start_i, 0));
      }
    } else {
      dist_i = euclidean(matrix<double>(coords.block(start_i, 0, n_visits_i, coordinates.cols())));
    }
    int subject_group_i = subject_groups[i] - 1;
    matrix<double> sigma_inv_d1 = derivatives_by_group.cache[subject_group_i]->get_inverse_derivative(visit_i, dist_i);

    matrix<double> Xi = x_matrix.block(start_i, 0, n_visits_i, x_matrix.cols());
    auto gi_sqrt_root = G_sqrt.segment(start_i, n_visits_i).matrix().asDiagonal();
    for (int r = 0; r < theta_size_per_group; r ++) {
      auto Pi = Xi.transpose() * gi_sqrt_root * sigma_inv_d1.block(r * n_visits_i, 0, n_visits_i, n_visits_i) * gi_sqrt_root * Xi;
      P.block(r * p + theta_size_per_group * subject_group_i * p, 0, p, p) += Pi;
    }
  }
  if (Rcpp::any(Rcpp::is_infinite(as_num_matrix_rcpp(P)))) {
    stop("Jacobian is not finite. The model can be over-parameterized.");
  }
  auto ret = List::create();
  for (int i = 0; i < n_theta; i++) {
    // the P is derivative of (XWX), the covariance is (XWX)^{-1}.
    ret.push_back(as_num_matrix_rcpp(-beta_vcov_m * P.block(i * p, 0, p, p) * beta_vcov_m));
  }
  return ret;
}

#include "derivatives.h"

using namespace Rcpp;
using std::string;
// Obtain P,Q,R element from a mmrm fit, given theta.
List get_pqr(List mmrm_fit, NumericVector theta) {
  NumericMatrix x = mmrm_fit["x_matrix"];
  matrix<double> x_matrix = as_num_matrix_tmb(x);
  IntegerVector subject_zero_inds = mmrm_fit["subject_zero_inds"];
  int n_subjects = mmrm_fit["n_subjects"];
  IntegerVector subject_n_visits = mmrm_fit["subject_n_visits"];
  int n_visits = mmrm_fit["n_visits"];
  String cov_type = mmrm_fit["cov_type"];
  int is_spatial_int = mmrm_fit["is_spatial_int"];
  bool is_spatial = is_spatial_int == 1;
  int n_groups = mmrm_fit["n_groups"];
  IntegerVector subject_groups = mmrm_fit["subject_groups"];
  NumericVector weights_vector = mmrm_fit["weights_vector"];
  NumericMatrix coordinates = mmrm_fit["coordinates"];
  matrix<double> coords = as_num_matrix_tmb(coordinates);
  vector<double> theta_v = as_num_vector_tmb(theta);
  vector<double> G_sqrt = as_num_vector_tmb(sqrt(weights_vector));
  int n_theta = theta.size();
  int theta_size_per_group = n_theta / n_groups;
  int p = x.cols();
  matrix<double> P = matrix<double>::Zero(p * n_theta, p);
  matrix<double> Q = matrix<double>::Zero(p * theta_size_per_group * n_theta, p);
  matrix<double> R = matrix<double>::Zero(p * theta_size_per_group * n_theta, p);
  // Use map to hold these base class pointers (can also work for child class objects).
  auto derivatives_by_group = cache_obj<double, derivatives_base<double>, derivatives_sp_exp<double>, derivatives_nonspatial<double>>(theta_v, n_groups, is_spatial, cov_type, n_visits);
  for (int i = 0; i < n_subjects; i++) {
    int start_i = subject_zero_inds[i];
    int n_visits_i = subject_n_visits[i];
    std::vector<int> visit_i(n_visits_i);
    matrix<double> dist_i(n_visits_i, n_visits_i);
    if (!is_spatial) {
      for (int i = 0; i < n_visits_i; i++) {
        visit_i[i] = int(coordinates(i + start_i, 0));
      }
    } else {
      dist_i = euclidean(matrix<double>(coords.block(start_i, 0, n_visits_i, coordinates.cols())));
    }
    int subject_group_i = subject_groups[i] - 1;
    matrix<double> sigma_inv, sigma_d1, sigma_d2, sigma, sigma_inv_d1;

    sigma_inv = derivatives_by_group.cache[subject_group_i]->get_sigma_inverse(visit_i, dist_i);
    sigma_d1 = derivatives_by_group.cache[subject_group_i]->get_sigma_derivative1(visit_i, dist_i);
    sigma_d2 = derivatives_by_group.cache[subject_group_i]->get_sigma_derivative2(visit_i, dist_i);
    sigma = derivatives_by_group.cache[subject_group_i]->get_sigma(visit_i, dist_i);
    sigma_inv_d1 = derivatives_by_group.cache[subject_group_i]->get_inverse_derivative(visit_i, dist_i);

    matrix<double> Xi = x_matrix.block(start_i, 0, n_visits_i, x_matrix.cols());
    auto gi_sqrt_root = G_sqrt.segment(start_i, n_visits_i).matrix().asDiagonal();
    for (int r = 0; r < theta_size_per_group; r ++) {
      auto Pi = Xi.transpose() * gi_sqrt_root * sigma_inv_d1.block(r * n_visits_i, 0, n_visits_i, n_visits_i) * gi_sqrt_root * Xi;
      P.block(r * p + theta_size_per_group * subject_group_i * p, 0, p, p) += Pi;
      for (int j = 0; j < theta_size_per_group; j++) {
        auto Qij = Xi.transpose() * gi_sqrt_root * sigma_inv_d1.block(r * n_visits_i, 0, n_visits_i, n_visits_i) * sigma * sigma_inv_d1.block(j * n_visits_i, 0, n_visits_i, n_visits_i) * gi_sqrt_root * Xi;
        // switch the order so that in the matrix partial(i) and partial(j) increase j first
        Q.block((r * theta_size_per_group + j + theta_size_per_group * theta_size_per_group * subject_group_i) * p, 0, p, p) += Qij;
        auto Rij = Xi.transpose() * gi_sqrt_root * sigma_inv * sigma_d2.block((j * theta_size_per_group + r) * n_visits_i, 0, n_visits_i, n_visits_i) * sigma_inv * gi_sqrt_root * Xi;
        R.block((r * theta_size_per_group + j + theta_size_per_group * theta_size_per_group * subject_group_i) * p, 0, p, p) += Rij;
      }
    }
  }
  return List::create(
    Named("P") = as_num_matrix_rcpp(P),
    Named("Q") = as_num_matrix_rcpp(Q),
    Named("R") = as_num_matrix_rcpp(R)
  );
}

#include "covariance.h"
#include "chol_cache.h"
// Definition:
//
// Y_i = X_i * beta + epsilon_i, i = 1, ..., n_subjects
// where Y_i = (Y_i1, ..., Y_im) are the observations of subject i over the m
// timepoints,
//
// and for the epsilon_i's :
// epsilon_i ~iid N(0, Sigma) where Sigma is a covariance matrix
// parameterized by a vector theta.
//
// Note: This is a special generalized least squares model
// Y = X * beta + epsilon,
// where we have a block structure for the covariance matrix of the epsilon
// vector.
//
// beta itself is not a parameter for TMB here:
// - For maximum likelihood estimation:
//   Given theta and therefore Sigma, and writing W = Sigma^-1, we can determine
//   the beta optimizing the likelihood via the weighted least squares equation
//   (X^T W X) beta = X^T W Y.
// - For restricted maximum likelihood estimation:
//   Given theta, beta is integrated out from the likelihood. Weighted least
//   squares results are used to calculate integrated log likelihood.

template<class Type>
Type objective_function<Type>::operator() ()
{
  // Read data from R.
  DATA_MATRIX(x_matrix);           // Model matrix (dimension n x p).
  DATA_VECTOR(y_vector);           // Response vector (length n).
  DATA_VECTOR(weights_vector);     // Weights vector (length n).
  DATA_MATRIX(coordinates);        // Coordinates matrix.
  DATA_INTEGER(n_visits);          // Number of visits, which is the dimension of the covariance matrix.
  DATA_INTEGER(n_subjects);        // Number of subjects.
  DATA_IVECTOR(subject_zero_inds); // Starting indices for each subject (0-based) (length n_subjects).
  DATA_IVECTOR(subject_n_visits);  // Number of observed visits for each subject (length n_subjects).
  DATA_STRING(cov_type);           // Covariance type name.
  DATA_INTEGER(is_spatial_int);    // Spatial covariance (1)? Otherwise non-spatial covariance.
  DATA_INTEGER(reml);              // REML (1)? Otherwise ML (0).
  DATA_FACTOR(subject_groups);     // subject groups vector(0-based) (length n_subjects).
  DATA_INTEGER(n_groups);          // number of total groups.
  // Read parameters from R.
  PARAMETER_VECTOR(theta);         // Covariance parameters (length k). Contents depend on covariance type.

  // X^T W X will be calculated incrementally into here.
  matrix<Type> XtWX = matrix<Type>::Zero(x_matrix.cols(), x_matrix.cols());
  // X^T W Y will be calculated incrementally into here.
  matrix<Type> XtWY = matrix<Type>::Zero(x_matrix.cols(), 1);
  // W^T/2 X will be saved into here.
  matrix<Type> x_mat_tilde = matrix<Type>::Zero(x_matrix.rows(), x_matrix.cols());
  // W^T/2 Y will be saved into here.
  vector<Type> y_vec_tilde = vector<Type>::Zero(y_vector.rows());
  // Sum of the log determinant will be incrementally calculated here.
  Type sum_log_det = 0.0;

  // Convert is_spatial_int to bool.
  bool is_spatial = (is_spatial_int == 1);
  // Diagonal of weighted covariance
  vector<Type> diag_cov_inv_sqrt(x_matrix.rows());
  // Cholesky group object
  auto chols_group = chol_cache_groups<Type>(theta, n_groups, is_spatial, cov_type, n_visits);
  // Go through all subjects and calculate quantities initialized above.
  for (int i = 0; i < n_subjects; i++) {
    // Start index and number of visits for this subject.
    int start_i = subject_zero_inds(i);
    int n_visits_i = subject_n_visits(i);
    std::vector<int> visit_i(n_visits_i);
    matrix<Type> dist_i(n_visits_i, n_visits_i);
    if (!is_spatial) {
      for (int j = 0; j < n_visits_i; j++) {
        visit_i[j] = int(asDouble(coordinates(start_i + j, 0)));
      }
    } else {
      dist_i = euclidean(matrix<Type>(coordinates.block(start_i, 0, n_visits_i, coordinates.cols())));
    }
    // Obtain Cholesky factor Li.
    matrix<Type> Li = chols_group.cache[subject_groups[i]]->get_chol(visit_i, dist_i);
    // Calculate weighted Cholesky factor for this subject.
    Eigen::DiagonalMatrix<Type,Eigen::Dynamic,Eigen::Dynamic> Gi_inv_sqrt = weights_vector.segment(start_i, n_visits_i).cwiseInverse().sqrt().matrix().asDiagonal();
    Li = Gi_inv_sqrt * Li;
    // Calculate scaled design matrix and response vector for this subject.
    matrix<Type> Xi = x_matrix.block(start_i, 0, n_visits_i, x_matrix.cols());
    matrix<Type> XiTilde = Li.template triangularView<Eigen::Lower>().solve(Xi);
    matrix<Type> Yi = y_vector.segment(start_i, n_visits_i).matrix();
    matrix<Type> YiTilde = Li.template triangularView<Eigen::Lower>().solve(Yi);

    // Increment quantities.
    matrix<Type> XiTildeCrossprod = crossprod(XiTilde);
    XtWX += XiTildeCrossprod.template triangularView<Eigen::Lower>();
    XtWY += XiTilde.transpose() * YiTilde;
    vector<Type> LiDiag = Li.diagonal();
    sum_log_det += sum(log(LiDiag));
    // Cache the reciprocal of square root of diagonal of covariance
    diag_cov_inv_sqrt.segment(start_i, n_visits_i) = vector<Type>(tcrossprod(Li).diagonal()).rsqrt();
    // Save stuff.
    x_mat_tilde.block(start_i, 0, n_visits_i, x_matrix.cols()) = XiTilde;
    y_vec_tilde.segment(start_i, n_visits_i) = YiTilde.col(0);
  }

  // Solve for beta.
  Eigen::LDLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > XtWX_decomposition(XtWX);
  matrix<Type> beta_mat = XtWX_decomposition.solve(XtWY);
  vector<Type> beta = beta_mat.col(0);

  // Define scaled residuals.
  vector<Type> x_mat_tilde_beta = x_mat_tilde * beta;
  vector<Type> epsilonTilde = y_vec_tilde - x_mat_tilde_beta;

  // Calculate negative log-likelihood.
  Type neg_log_lik;

  // Always extract the D vector since we want to report this below.
  vector<Type> XtWX_D = XtWX_decomposition.vectorD();

  if (reml == 1) {
    // Use restricted maximum likelihood.
    Type XtWX_log_det = XtWX_D.log().sum();
    neg_log_lik = (x_matrix.rows() - x_matrix.cols()) / 2.0 * log(2.0 * M_PI) +
      sum_log_det +
      XtWX_log_det / 2.0 +
      0.5 * (y_vec_tilde * y_vec_tilde).sum() - 0.5 * (x_mat_tilde_beta * x_mat_tilde_beta).sum();
  } else {
    // Use maximum likelihood.
    neg_log_lik = x_matrix.rows() / 2.0 * log(2.0 * M_PI) +
      sum_log_det +
      0.5 * (epsilonTilde * epsilonTilde).sum();
  }

  // Report quantities to R.
  REPORT(beta);

  // We already compute the inverse of XtWX here because we already did the
  // matrix decomposition above.
  matrix<Type> Identity(XtWX.rows(), XtWX.cols());
  Identity.setIdentity();
  matrix<Type> beta_vcov = XtWX_decomposition.solve(Identity);
  REPORT(beta_vcov);

  // Also return the decomposition components L and D.
  matrix<Type> XtWX_L(XtWX.rows(), XtWX.cols());
  XtWX_L = XtWX_decomposition.matrixL();
  REPORT(XtWX_L);
  REPORT(XtWX_D);

  // normalized residual
  REPORT(epsilonTilde);
  // inverse square root of diagonal of covariance
  REPORT(diag_cov_inv_sqrt);
  matrix<Type> covariance_lower_chol = chols_group.get_default_chol();
  REPORT(covariance_lower_chol);

  return neg_log_lik;
}

#include "covariance.h"
#include "chol_cache.h"

using namespace Rcpp;
using std::string;
// Obtain the conditional mean/variance of `y` given `beta`, `beta_vcov`, `theta`.
// Given any `theta`, we can obtain `beta` and `beta_vcov` through the `mrmm` fit, and then
// we can use the provided `theta` to obtain the covariance matrix for the residual,
// and use `beta_vcov` to obtain the covariance matrix for the mean of the fit,
// and use `beta` to obtain the estimate of the mean of the fit.
List predict(List mmrm_data, NumericVector theta, NumericVector beta, NumericMatrix beta_vcov) {
  NumericMatrix x = mmrm_data["x_matrix"];
  NumericVector y = mmrm_data["y_vector"];
  LogicalVector y_na = is_na(y);
  LogicalVector y_vd = ! y_na;
  IntegerVector subject_zero_inds = mmrm_data["subject_zero_inds"];
  IntegerVector subject_n_visits = mmrm_data["subject_n_visits"];
  String cov_type = mmrm_data["cov_type"];
  IntegerVector subject_groups = mmrm_data["subject_groups"];
  NumericMatrix coordinates = mmrm_data["coordinates"];

  matrix<double> x_matrix = as_num_matrix_tmb(x);
  matrix<double> coordinates_m = as_num_matrix_tmb(coordinates);
  matrix<double> beta_vcov_matrix = as_num_matrix_tmb(beta_vcov);
  int n_subjects = mmrm_data["n_subjects"];
  int n_visits = mmrm_data["n_visits"];
  int is_spatial_int = mmrm_data["is_spatial_int"];
  bool is_spatial = is_spatial_int == 1;
  int n_groups = mmrm_data["n_groups"];
  vector<double> beta_v = as_num_vector_tmb(beta);
  vector<double> theta_v = as_num_vector_tmb(theta);
  // Use map to hold these base class pointers (can also work for child class objects).
  auto chols_group = chol_cache_groups<double>(theta_v, n_groups, is_spatial, cov_type, n_visits);
  NumericVector y_pred = clone(y); // Predict value of y; observed use the same value.
  NumericVector var(y.size()); // Variance of y with 0 as default.
  NumericVector conf_var(y.size()); // Confidence interval variance.
  List covariance;
  List index;
  NumericMatrix empty(0, 0);
  // Go through all subjects and calculate quantities initialized above.
  for (int i = 0; i < n_subjects; i++) {
    // Start index and number of visits for this subject.
    int start_i = subject_zero_inds(i);
    int n_visits_i = subject_n_visits(i);
    NumericVector y_i = segment(y, start_i, n_visits_i);
    LogicalVector y_na_i = segment(y_na, start_i, n_visits_i);
    LogicalVector y_valid_i = segment(y_vd, start_i, n_visits_i);
    IntegerVector visit_i(n_visits_i);
    matrix<double> dist_i(n_visits_i, n_visits_i);
    IntegerVector index_zero_i = seq(0, n_visits_i - 1);
    if (!is_spatial) {
      for (int i = 0; i < n_visits_i; i++) {
        visit_i(i) = int(coordinates(i + start_i, 0));
      }
    } else {
      visit_i = seq(start_i, start_i + n_visits_i - 1);
      dist_i = euclidean(matrix<double>(coordinates_m.block(start_i, 0, n_visits_i, coordinates_m.cols())));
    }
    std::vector<int> visit_std = as<std::vector<int>>(visit_i);
    IntegerVector visit_na_vec = visit_i[y_na_i];
    IntegerVector visit_valid_vec = visit_i[y_valid_i];

    IntegerVector index_zero_i_na = index_zero_i[y_na_i];
    IntegerVector index_zero_i_valid = index_zero_i[y_valid_i];

    std::vector<int> visit_na = as<std::vector<int>>(visit_na_vec);
    std::vector<int> visit_non_na = as<std::vector<int>>(visit_valid_vec);
    matrix<double> Xi = x_matrix.block(start_i, 0, n_visits_i, x_matrix.cols());
    // Subject_group starts with 1.
    int subject_group_i = subject_groups(i) - 1;
    matrix<double> sigma_full = chols_group.cache[subject_group_i]->get_sigma(visit_std, dist_i);
    matrix<double> sigma_12 = subset_matrix(sigma_full, index_zero_i_na, index_zero_i_valid);
    matrix<double> sigma_11;
    if (!is_spatial) {
      sigma_11 = chols_group.cache[subject_group_i]->get_sigma(visit_na, dist_i);
    } else {
      sigma_11 = subset_matrix(sigma_full, index_zero_i_na, index_zero_i_na);
    }
    matrix<double> x_na = subset_matrix(Xi, index_zero_i_na);
    matrix<double> x_valid = subset_matrix(Xi, index_zero_i_valid);
    vector<double> y_valid = as_num_vector_tmb(y_i[y_valid_i]);
    IntegerVector na_index = index_zero_i_na + start_i;
    vector<double> y_hat, var_conf, var_y_on_theta;
    if (visit_valid_vec.size() == 0) {
      // No observations with valid y.
      y_hat = x_na * beta_v;
      var_conf = (x_na * beta_vcov_matrix * x_na.transpose()).diagonal();
      var_y_on_theta = var_conf + vector<double>(sigma_full.diagonal());
      covariance.push_back(as_num_matrix_rcpp(sigma_full));
    } else if (visit_na_vec.size() > 0) {
      // There are observations with invalid y.
      matrix<double> sigma_22_inv;
      if (is_spatial) {
        sigma_22_inv = subset_matrix(sigma_full, index_zero_i_valid, index_zero_i_valid).inverse(); // No cache available for spatial covariance.
      } else {
        sigma_22_inv = chols_group.cache[subject_group_i]->get_sigma_inverse(visit_non_na, dist_i); // We have the inverse in cache for non spatial covariance.
      }
      matrix<double> ss = sigma_12 * sigma_22_inv;
      matrix<double> zz = x_na - ss * x_valid;
      y_hat = zz * beta_v + ss * y_valid;
      var_conf = (zz * beta_vcov_matrix * zz.transpose()).diagonal();
      matrix<double> conditional_sigma = sigma_11 - ss * sigma_12.transpose();
      var_y_on_theta = var_conf + vector<double>(conditional_sigma.diagonal());
      covariance.push_back(as_num_matrix_rcpp(conditional_sigma));
    } else if (visit_na_vec.size() == 0) {
      covariance.push_back(empty);
    }
    index.push_back(na_index);
    // Replace the values with fitted values. If no missing value there, the `na_index` will be length 0
    // and the left hand side will hence not be modified.
    y_pred[na_index] = as_num_vector_rcpp(y_hat);
    conf_var[na_index] = as_num_vector_rcpp(var_conf);
    var[na_index] = as_num_vector_rcpp(var_y_on_theta);
  }
  NumericMatrix ret = cbind(y_pred, conf_var, var);
  CharacterVector cnms = {"fit", "conf_var", "var"};
  colnames(ret) = cnms;
  return List::create(
    Named("prediction") = ret,
    Named("covariance") = covariance,
    Named("index") = index
  );
}

#ifndef TESTTHAT_WRAP_H
#define TESTTHAT_WRAP_H
#include <testthat.h>
#include <limits>
#include "utils.h"

// Expect equal: Here use a default epsilon which gives around 1e-4 on
// my computer here.
#define expect_equal(TARGET, CURRENT)                          \
{                                                              \
  double const eps =                                           \
    std::pow(std::numeric_limits<double>::epsilon(), 0.25);    \
                                                               \
  if(std::abs((TARGET)) > eps)                                 \
    expect_true(std::abs((TARGET) - (CURRENT)) /               \
      std::abs((TARGET)) < eps);                               \
  else                                                         \
    expect_true(std::abs((TARGET) - (CURRENT)) < eps);         \
}

#define expect_equal_eps(TARGET, CURRENT, EPS)                 \
{                                                              \
  if(std::abs((TARGET)) > (EPS))                               \
    expect_true(std::abs((TARGET) - (CURRENT)) /               \
      std::abs((TARGET)) < (EPS));                             \
  else                                                         \
    expect_true(std::abs((TARGET) - (CURRENT)) < (EPS));       \
}

template <class T>
void expect_equal_matrix(const T& target, const T& current)
{
  int nrow = target.rows();
  int ncol = target.cols();

  expect_true(nrow == current.rows());
  expect_true(ncol == current.cols());

  for (int i = 0; i < nrow; i++) {
    for (int j = 0; j < ncol; j++) {
      expect_equal(target(i, j), current(i, j));
    }
  }
}

template <class T>
void expect_equal_vector(const T& target, const T& current)
{
  int n = target.size();
  expect_true(n == current.size());

  for (int i = 0; i < n; i++) {
    expect_equal(target(i), current(i));
  }
}

#endif

1		#' Dynamic Registration for Package Interoperability
2		#'
3		#' @seealso See `vignette("xtending", package = "emmeans")` for background.
4		#' @keywords internal
5		#' @noRd
6		.onLoad <- function(libname, pkgname) {
7		# nolint
8	!	if (!h_tmb_version_sufficient()) {
9	!	msg <- paste(
10	!	"TMB below version 1.9.15 has been used to compile the mmrm package.",
11	!	"Reproducible model fits are not guaranteed.",
12	!	"Please consider recompiling the package with TMB version 1.9.15 or higher."
13		)
14	!	warning(msg, call. = FALSE)
15		}
16
17	!	register_on_load(
18	!	"emmeans",
19	!	c("1.6", NA),
20	!	callback = function() emmeans::.emm_register("mmrm", pkgname),
21	!	message = "mmrm() registered as emmeans extension"
22		)
23
24	!	register_on_load(
25	!	"parsnip",
26	!	c("1.1.0", NA),
27	!	callback = parsnip_add_mmrm,
28	!	message = emit_tidymodels_register_msg
29		)
30	!	register_on_load(
31	!	"car",
32	!	c("3.1.2", NA),
33	!	callback = car_add_mmrm,
34	!	message = "mmrm() registered as car::Anova extension"
35		)
36		}
37
38		#' Helper Function for Registering Functionality With Suggests Packages
39		#'
40		#' @inheritParams check_package_version
41		#'
42		#' @param callback (`function(...) ANY`)\cr a callback to execute upon package
43		#' load. Note that no arguments are passed to this function. Any necessary
44		#' data must be provided upon construction.
45		#'
46		#' @param message (`NULL` or `string`)\cr an optional message to print after
47		#' the callback is executed upon successful registration.
48		#'
49		#' @return A logical (invisibly) indicating whether registration was successful.
50		#' If not, a onLoad hook was set for the next time the package is loaded.
51		#'
52		#' @keywords internal
53		register_on_load <- function(
54		pkg,
55		ver = c(NA_character_, NA_character_),
56		callback,
57		message = NULL
58		) {
59	5x	if (isNamespaceLoaded(pkg) && check_package_version(pkg, ver)) {
60	4x	callback()
61	4x	if (is.character(message)) {
62	3x	packageStartupMessage(message)
63		}
64	4x	if (is.function(message)) {
65	1x	packageStartupMessage(message())
66		}
67	4x	return(invisible(TRUE))
68		}
69
70	1x	setHook(
71	1x	packageEvent(pkg, event = "onLoad"),
72	1x	action = "append",
73	1x	function(...) {
74	!	register_on_load(
75	!	pkg = pkg,
76	!	ver = ver,
77	!	callback = callback,
78	!	message = message
79		)
80		}
81		)
82
83	1x	invisible(FALSE)
84		}
85
86		#' Check Suggested Dependency Against Version Requirements
87		#'
88		#' @param pkg (`string`)\cr package name.
89		#' @param ver (`character`)\cr of length 2 whose elements can be provided to
90		#' [numeric_version()], representing a minimum and maximum (inclusive) version
91		#' requirement for interoperability. When `NA`, no version requirement is
92		#' imposed. Defaults to no version requirement.
93		#'
94		#' @return A logical (invisibly) indicating whether the loaded package meets
95		#' the version requirements. A warning is emitted otherwise.
96		#'
97		#' @keywords internal
98		check_package_version <- function(pkg, ver = c(NA_character_, NA_character_)) {
99	8x	assert_character(ver, len = 2L)
100	7x	pkg_ver <- utils::packageVersion(pkg)
101	7x	ver <- numeric_version(ver, strict = FALSE)
102
103	7x	warn_version <- function(pkg, pkg_ver, ver) {
104	2x	ver_na <- is.na(ver)
105	2x	warning(sprintf(
106	2x	"Cannot register mmrm for use with %s (v%s). Version %s required.",
107	2x	pkg,
108	2x	pkg_ver,
109	2x	if (!any(ver_na)) {
110	!	sprintf("%s to %s", ver[1], ver[2])
111	2x	} else if (ver_na[2]) {
112	1x	paste0(">= ", ver[1])
113	2x	} else if (ver_na[1]) {
114	1x	paste0("<= ", ver[2])
115		}
116		))
117		}
118
119	7x	if (identical(pkg_ver < ver[1], TRUE) \|\| identical(pkg_ver > ver[2], TRUE)) {
120	2x	warn_version(pkg, pkg_ver, ver)
121	2x	return(invisible(FALSE))
122		}
123
124	5x	invisible(TRUE)
125		}
126
127		#' Format a Message to Emit When Tidymodels is Loaded
128		#'
129		#' @return A character message to emit. Either a ansi-formatted cli output if
130		#' package 'cli' is available or a plain-text message otherwise.
131		#'
132		#' @keywords internal
133		emit_tidymodels_register_msg <- function() {
134	1x	pkg <- utils::packageName()
135	1x	ver <- utils::packageVersion(pkg)
136
137	1x	if (isTRUE(getOption("tidymodels.quiet"))) {
138	!	return()
139		}
140
141		# if tidymodels is attached, cli packages come as a dependency
142	1x	has_cli <- requireNamespace("cli", quietly = TRUE)
143	1x	if (has_cli) {
144		# unfortunately, cli does not expose many formatting tools for emitting
145		# messages (only via conditions to stderr) which can't be suppressed using
146		# suppressPackageStartupMessages() so formatting must be done adhoc,
147		# similar to how it's done in {tidymodels} R/attach.R
148	1x	paste0(
149	1x	cli::rule(
150	1x	left = cli::style_bold("Model Registration"),
151	1x	right = paste(pkg, ver)
152		),
153	1x	"\n",
154	1x	cli::col_green(cli::symbol$tick),
155		" ",
156	1x	cli::col_blue("mmrm"),
157		"::",
158	1x	cli::col_green("mmrm()")
159		)
160		} else {
161	!	paste0(pkg, "::mmrm() registered for use with tidymodels")
162		}
163		}

1		#' Determine Within or Between for each Design Matrix Column
2		#'
3		#' @description Used in [h_df_bw_calc()] to determine whether a variable
4		#' differs only between subjects or also within subjects.
5		#'
6		#' @param x_matrix (`matrix`)\cr the design matrix with column names.
7		#' @param subject_ids (`factor`)\cr the subject IDs.
8		#'
9		#' @return Character vector with "intercept", "within" or "between" for each
10		#' design matrix column identified via the names of the vector.
11		#'
12		#' @keywords internal
13		h_within_or_between <- function(x_matrix, subject_ids) {
14	19x	assert_matrix(x_matrix, col.names = "unique", min.cols = 1L)
15	19x	assert_factor(subject_ids, len = nrow(x_matrix))
16
17	19x	n_subjects <- length(unique(subject_ids))
18	19x	vapply(
19	19x	colnames(x_matrix),
20	19x	function(x) {
21	112x	if (x == "(Intercept)") {
22	19x	"intercept"
23		} else {
24	93x	n_unique <- nrow(unique(cbind(x_matrix[, x], subject_ids)))
25	43x	if (n_unique > n_subjects) "within" else "between"
26		}
27		},
28	19x	character(1L)
29		)
30		}
31
32		#' Calculation of Between-Within Degrees of Freedom
33		#'
34		#' @description Used in [h_df_1d_bw()] and [h_df_md_bw()].
35		#'
36		#' @param object (`mmrm`)\cr the fitted MMRM.
37		#'
38		#' @return List with:
39		#' - `coefs_between_within` calculated via [h_within_or_between()]
40		#' - `ddf_between`
41		#' - `ddf_within`
42		#'
43		#' @keywords internal
44		h_df_bw_calc <- function(object) {
45	18x	assert_class(object, "mmrm")
46
47	18x	n_subjects <- component(object, "n_subjects")
48	18x	n_obs <- component(object, "n_obs")
49	18x	x_mat <- component(object, "x_matrix")
50
51	18x	subject_var <- component(object, "subject_var")
52	18x	full_frame <- component(object, "full_frame")
53	18x	subject_ids <- full_frame[[subject_var]]
54
55	18x	coefs_between_within <- h_within_or_between(x_mat, subject_ids)
56	18x	n_coefs_between <- sum(coefs_between_within == "between")
57	18x	n_intercept <- sum(coefs_between_within == "intercept")
58	18x	n_coefs_within <- sum(coefs_between_within == "within")
59	18x	ddf_between <- n_subjects - n_coefs_between - n_intercept
60	18x	ddf_within <- n_obs - n_subjects - n_coefs_within
61
62	18x	list(
63	18x	coefs_between_within = coefs_between_within,
64	18x	ddf_between = ddf_between,
65	18x	ddf_within = ddf_within
66		)
67		}
68
69		#' Assign Minimum Degrees of Freedom Given Involved Coefficients
70		#'
71		#' @description Used in [h_df_1d_bw()] and [h_df_md_bw()].
72		#'
73		#' @param bw_calc (`list`)\cr from [h_df_bw_calc()].
74		#' @param is_coef_involved (`logical`)\cr whether each coefficient is involved
75		#' in the contrast.
76		#'
77		#' @return The minimum of the degrees of freedom assigned to each involved
78		#' coefficient according to its between-within categorization.
79		#'
80		#' @keywords internal
81		h_df_min_bw <- function(bw_calc, is_coef_involved) {
82	17x	assert_list(bw_calc)
83	17x	assert_names(
84	17x	names(bw_calc),
85	17x	identical.to = c("coefs_between_within", "ddf_between", "ddf_within")
86		)
87	17x	assert_logical(is_coef_involved, len = length(bw_calc$coefs_between_within))
88	17x	assert_true(sum(is_coef_involved) > 0)
89
90	17x	coef_categories <- bw_calc$coefs_between_within[is_coef_involved]
91	17x	coef_dfs <- vapply(
92	17x	X = coef_categories,
93	17x	FUN = switch,
94	17x	intercept = bw_calc$ddf_within,
95	17x	between = bw_calc$ddf_between,
96	17x	within = bw_calc$ddf_within,
97	17x	FUN.VALUE = integer(1)
98		)
99	17x	min(coef_dfs)
100		}
101
102		#' Calculation of Between-Within Degrees of Freedom for One-Dimensional Contrast
103		#'
104		#' @description Used in [df_1d()] if method is "Between-Within".
105		#'
106		#' @inheritParams h_df_1d_sat
107		#' @inherit h_df_1d_sat return
108		#' @keywords internal
109		h_df_1d_bw <- function(object, contrast) {
110	7x	assert_class(object, "mmrm")
111	7x	assert_numeric(contrast, len = length(component(object, "beta_est")))
112
113	7x	bw_calc <- h_df_bw_calc(object)
114	7x	is_coef_involved <- contrast != 0
115	7x	df <- h_df_min_bw(bw_calc, is_coef_involved)
116	7x	h_test_1d(object, contrast, df)
117		}
118
119		#' Calculation of Between-Within Degrees of Freedom for Multi-Dimensional Contrast
120		#'
121		#' @description Used in [df_md()] if method is "Between-Within".
122		#'
123		#' @inheritParams h_df_md_sat
124		#' @inherit h_df_md_sat return
125		#' @keywords internal
126		h_df_md_bw <- function(object, contrast) {
127	7x	assert_class(object, "mmrm")
128	7x	assert_matrix(
129	7x	contrast,
130	7x	mode = "numeric",
131	7x	any.missing = FALSE,
132	7x	ncols = length(component(object, "beta_est"))
133		)
134
135	7x	bw_calc <- h_df_bw_calc(object)
136	7x	is_coef_involved <- apply(X = contrast != 0, MARGIN = 2L, FUN = any)
137	7x	df <- h_df_min_bw(bw_calc, is_coef_involved)
138	7x	h_test_md(object, contrast, df)
139		}

1		#' Register `mmrm` For Use With `tidymodels`
2		#'
3		#' @inheritParams base::requireNamespace
4		#' @return A logical value indicating whether registration was successful.
5		#'
6		#' @details We can use `parsnip::show_model_info("linear_reg")` to check the
7		#' registration with `parsnip` and thus the wider `tidymodels` ecosystem.
8		#'
9		#' @keywords internal
10		parsnip_add_mmrm <- function(quietly = FALSE) {
11	1x	if (!requireNamespace("parsnip", quietly = quietly)) {
12	!	return(FALSE)
13		}
14
15	1x	parsnip::set_model_engine(
16	1x	model = "linear_reg",
17	1x	eng = "mmrm",
18	1x	mode = "regression"
19		)
20
21	1x	parsnip::set_dependency(
22	1x	pkg = "mmrm",
23	1x	model = "linear_reg",
24	1x	eng = "mmrm",
25	1x	mode = "regression"
26		)
27
28	1x	parsnip::set_encoding(
29	1x	model = "linear_reg",
30	1x	eng = "mmrm",
31	1x	mode = "regression",
32	1x	options = list(
33	1x	predictor_indicators = "none",
34	1x	compute_intercept = FALSE,
35	1x	remove_intercept = FALSE,
36	1x	allow_sparse_x = TRUE
37		)
38		)
39
40	1x	parsnip::set_fit(
41	1x	model = "linear_reg",
42	1x	eng = "mmrm",
43	1x	mode = "regression",
44	1x	value = list(
45	1x	interface = "formula",
46	1x	protect = c("formula", "data", "weights"),
47	1x	data = c(formula = "formula", data = "data", weights = "weights"),
48	1x	func = c(pkg = "mmrm", fun = "mmrm"),
49	1x	defaults = list()
50		)
51		)
52
53	1x	parsnip::set_pred(
54	1x	model = "linear_reg",
55	1x	eng = "mmrm",
56	1x	mode = "regression",
57	1x	type = "numeric",
58	1x	value = parsnip::pred_value_template(
59		# This is boilerplate.
60	1x	func = c(fun = "predict"),
61	1x	object = quote(object$fit),
62	1x	newdata = quote(new_data)
63		)
64		)
65
66	1x	parsnip::set_pred(
67	1x	model = "linear_reg",
68	1x	eng = "mmrm",
69	1x	mode = "regression",
70		# This type allows to pass arguments via `opts` to `parsnip::predict.model_fit`.
71	1x	type = "raw",
72	1x	value = parsnip::pred_value_template(
73		# This is boilerplate.
74	1x	func = c(fun = "predict"),
75	1x	object = quote(object$fit),
76	1x	newdata = quote(new_data)
77		# We don't specify additional argument defaults here since otherwise
78		# the user is not able to change them (they will be fixed).
79		)
80		)
81
82	1x	TRUE
83		}

1		#' Support for `emmeans`
2		#'
3		#' @description `r lifecycle::badge("stable")`
4		#'
5		#' This package includes methods that allow `mmrm` objects to be used
6		#' with the `emmeans` package. `emmeans` computes estimated marginal means
7		#' (also called least-square means) for the coefficients of the MMRM.
8		#' We can also e.g. obtain differences between groups by applying
9		#' [`pairs()`][emmeans::pairs.emmGrid()] on the object returned
10		#' by [emmeans::emmeans()].
11		#'
12		#' @examples
13		#' fit <- mmrm(
14		#' formula = FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT \| USUBJID),
15		#' data = fev_data
16		#' )
17		#' if (require(emmeans)) {
18		#' emmeans(fit, ~ ARMCD \| AVISIT)
19		#' pairs(emmeans(fit, ~ ARMCD \| AVISIT), reverse = TRUE)
20		#' }
21		#' @name emmeans_support
22		NULL
23
24		#' Returns a `data.frame` for `emmeans` Purposes
25		#'
26		#' @seealso See [emmeans::recover_data()] for background.
27		#' @keywords internal
28		#' @noRd
29		# nolint start
30		recover_data.mmrm <- function(object, ...) {
31		# nolint end
32	16x	fun_call <- stats::getCall(object)
33		# subject_var is excluded because it should not contain fixed effect.
34		# visit_var is only excluded for spatial covariance structures -
35		# for non-spatial covariance structures, emmeans can provide marginal mean
36		# by each visit so we need visit_var.
37	16x	model_frame <- stats::model.frame(
38	16x	object,
39	16x	exclude = c(
40	16x	"subject_var",
41	16x	if (object$formula_parts$is_spatial) "visit_var"
42		)
43		)
44	16x	model_terms <- stats::delete.response(stats::terms(model_frame))
45	16x	emmeans::recover_data(
46	16x	fun_call,
47	16x	trms = model_terms,
48	16x	na.action = "na.omit",
49	16x	frame = model_frame,
50		...
51		)
52		}
53
54		#' Returns a List of Model Details for `emmeans` Purposes
55		#'
56		#' @seealso See [emmeans::emm_basis()] for background.
57		#' @keywords internal
58		#' @noRd
59		# nolint start
60		emm_basis.mmrm <- function(
61		# nolint end
62		object,
63		trms,
64		xlev,
65		grid,
66		...
67		) {
68	15x	model_frame <- stats::model.frame(
69	15x	trms,
70	15x	grid,
71	15x	na.action = stats::na.pass,
72	15x	xlev = xlev
73		)
74	15x	contrasts <- component(object, "contrasts")
75	15x	model_mat <- stats::model.matrix(trms, model_frame, contrasts.arg = contrasts)
76	15x	beta_hat <- component(object, "beta_est")
77	15x	nbasis <- if (length(beta_hat) < ncol(model_mat)) {
78	6x	kept <- match(names(beta_hat), colnames(model_mat))
79	6x	beta_hat <- NA * model_mat[1L, ]
80	6x	beta_hat[kept] <- component(object, "beta_est")
81	6x	orig_model_mat <- stats::model.matrix(
82	6x	trms,
83	6x	stats::model.frame(
84	6x	object,
85	6x	exclude = c(
86	6x	"subject_var",
87	6x	if (object$formula_parts$is_spatial) "visit_var"
88		)
89		),
90	6x	contrasts.arg = contrasts
91		)
92	6x	estimability::nonest.basis(orig_model_mat)
93		} else {
94	9x	estimability::all.estble
95		}
96	15x	dfargs <- list(object = object)
97	15x	dffun <- function(k, dfargs) {
98	143x	mmrm::df_md(dfargs$object, contrast = k)$denom_df
99		}
100	15x	list(
101	15x	X = model_mat,
102	15x	bhat = beta_hat,
103	15x	nbasis = nbasis,
104	15x	V = component(object, "beta_vcov"),
105	15x	dffun = dffun,
106	15x	dfargs = dfargs
107		)
108		}

1		#' Component Access for `mmrm_tmb` Objects
2		#'
3		#' @description `r lifecycle::badge("stable")`
4		#'
5		#' @param object (`mmrm_tmb`)\cr the fitted MMRM.
6		#' @param name (`character`)\cr the component(s) to be retrieved.
7		#' @return The corresponding component of the object, see details.
8		#'
9		#' @details Available `component()` names are as follows:
10		#' - `call`: low-level function call which generated the model.
11		#' - `formula`: model formula.
12		#' - `dataset`: data set name.
13		#' - `cov_type`: covariance structure type.
14		#' - `n_theta`: number of parameters.
15		#' - `n_subjects`: number of subjects.
16		#' - `n_timepoints`: number of modeled time points.
17		#' - `n_obs`: total number of observations.
18		#' - `reml`: was REML used (ML was used if `FALSE`).
19		#' - `neg_log_lik`: negative log likelihood.
20		#' - `convergence`: convergence code from optimizer.
21		#' - `conv_message`: message accompanying the convergence code.
22		#' - `evaluations`: number of function evaluations for optimization.
23		#' - `method`: Adjustment method which was used (for `mmrm` objects),
24		#' otherwise `NULL` (for `mmrm_tmb` objects).
25		#' - `beta_vcov`: estimated variance-covariance matrix of coefficients
26		#' (excluding aliased coefficients). When Kenward-Roger/Empirical adjusted
27		#' coefficients covariance matrix is used, the adjusted covariance matrix is returned (to still obtain the
28		#' original asymptotic covariance matrix use `object$beta_vcov`).
29		#' - `beta_vcov_complete`: estimated variance-covariance matrix including
30		#' aliased coefficients with entries set to `NA`.
31		#' - `varcor`: estimated covariance matrix for residuals. If there are multiple
32		#' groups, a named list of estimated covariance matrices for residuals will be
33		#' returned. The names are the group levels.
34		#' - `score_per_subject`: score per subject in empirical covariance.
35		#' See the vignette \code{vignette("coef_vcov", package = "mmrm")}.
36		#' - `theta_est`: estimated variance parameters.
37		#' - `beta_est`: estimated coefficients (excluding aliased coefficients).
38		#' - `beta_est_complete`: estimated coefficients including aliased coefficients
39		#' set to `NA`.
40		#' - `beta_aliased`: whether each coefficient was aliased (i.e. cannot be estimated)
41		#' or not.
42		#' - `theta_vcov`: estimated variance-covariance matrix of variance parameters.
43		#' - `x_matrix`: design matrix used (excluding aliased columns).
44		#' - `x_matrix_complete`: design matrix used, including aliased columns.
45		#' - `xlev`: a named list of character vectors giving the full set of levels to be assumed for each factor.
46		#' - `contrasts`: a list of contrasts used for each factor.
47		#' - `y_vector`: response vector used.
48		#' - `jac_list`: Jacobian, see [h_jac_list()] for details.
49		#' - `full_frame`: `data.frame` with `n` rows containing all variables needed in the model.
50		#'
51		#' @seealso In the `lme4` package there is a similar function `getME()`.
52		#'
53		#' @examples
54		#' fit <- mmrm(
55		#' formula = FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT \| USUBJID),
56		#' data = fev_data
57		#' )
58		#' # Get all available components.
59		#' component(fit)
60		#' # Get convergence code and message.
61		#' component(fit, c("convergence", "conv_message"))
62		#' # Get modeled formula as a string.
63		#' component(fit, c("formula"))
64		#'
65		#' @export
66		component <- function(
67		object,
68		name = c(
69		"cov_type",
70		"subject_var",
71		"n_theta",
72		"n_subjects",
73		"n_timepoints",
74		"n_obs",
75		"beta_vcov",
76		"beta_vcov_complete",
77		"varcor",
78		"score_per_subject",
79		"formula",
80		"dataset",
81		"n_groups",
82		"reml",
83		"convergence",
84		"evaluations",
85		"method",
86		"optimizer",
87		"conv_message",
88		"call",
89		"theta_est",
90		"beta_est",
91		"beta_est_complete",
92		"beta_aliased",
93		"x_matrix",
94		"x_matrix_complete",
95		"y_vector",
96		"neg_log_lik",
97		"jac_list",
98		"theta_vcov",
99		"full_frame",
100		"xlev",
101		"contrasts"
102		)
103		) {
104	6403x	assert_class(object, "mmrm_tmb")
105	6403x	name <- match.arg(name, several.ok = TRUE)
106
107	6403x	list_components <- sapply(
108	6403x	X = name,
109	6403x	FUN = switch,
110	6403x	"call" = object$call,
111		# Strings.
112	6403x	"cov_type" = object$formula_parts$cov_type,
113	6403x	"subject_var" = object$formula_parts$subject_var,
114	6403x	"formula" = deparse(object$call$formula),
115	6403x	"dataset" = object$call$data,
116	6403x	"reml" = object$reml,
117	6403x	"conv_message" = object$opt_details$message,
118		# Numeric of length 1.
119	6403x	"convergence" = object$opt_details$convergence,
120	6403x	"neg_log_lik" = object$neg_log_lik,
121	6403x	"n_theta" = length(object$theta_est),
122	6403x	"n_subjects" = object$tmb_data$n_subjects,
123	6403x	"n_timepoints" = object$tmb_data$n_visits,
124	6403x	"n_obs" = length(object$tmb_data$y_vector),
125	6403x	"n_groups" = ifelse(is.list(object$cov), length(object$cov), 1L),
126		# Numeric of length > 1.
127	6403x	"evaluations" = unlist(ifelse(
128	6403x	is.null(object$opt_details$evaluations),
129	6403x	list(object$opt_details$counts),
130	6403x	list(object$opt_details$evaluations)
131		)),
132	6403x	"method" = object$method,
133	6403x	"optimizer" = object$optimizer,
134	6403x	"beta_est" = object$beta_est,
135	6403x	"beta_est_complete" = if (any(object$tmb_data$x_cols_aliased)) {
136	8x	stats::setNames(
137	8x	object$beta_est[names(object$tmb_data$x_cols_aliased)],
138	8x	names(object$tmb_data$x_cols_aliased)
139		)
140		} else {
141	52x	object$beta_est
142		},
143	6403x	"beta_aliased" = object$tmb_data$x_cols_aliased,
144	6403x	"theta_est" = object$theta_est,
145	6403x	"y_vector" = object$tmb_data$y_vector,
146	6403x	"jac_list" = object$jac_list,
147		# Matrices.
148	6403x	"beta_vcov" = if (
149	6403x	is.null(object$vcov) \|\| identical(object$vcov, "Asymptotic")
150		) {
151	1105x	object$beta_vcov
152		} else {
153	119x	object$beta_vcov_adj
154		},
155	6403x	"beta_vcov_complete" = if (any(object$tmb_data$x_cols_aliased)) {
156	2x	stats::.vcov.aliased(
157	2x	aliased = object$tmb_data$x_cols_aliased,
158	2x	vc = component(object, "beta_vcov"),
159	2x	complete = TRUE
160		)
161		} else {
162	6x	component(object, "beta_vcov")
163		},
164	6403x	"varcor" = object$cov,
165	6403x	"score_per_subject" = object$score_per_subject,
166	6403x	"x_matrix" = object$tmb_data$x_matrix,
167	6403x	"x_matrix_complete" = object$tmb_data$x_matrix_complete,
168	6403x	"xlev" = stats::.getXlevels(terms(object), object$tmb_data$full_frame),
169	6403x	"contrasts" = attr(object$tmb_data$x_matrix, "contrasts"),
170	6403x	"theta_vcov" = eval(object$theta_vcov),
171	6403x	"full_frame" = object$tmb_data$full_frame,
172		# If not found.
173	6403x	"..foo.." = stop(sprintf(
174	6403x	"component '%s' is not available",
175	6403x	name,
176	6403x	paste0(class(object), collapse = ", ")
177		)),
178	6403x	simplify = FALSE
179		)
180
181	23x	if (length(name) == 1) list_components[[1]] else list_components
182		}

1		#' Register `mmrm` For Use With `car::Anova`
2		#'
3		#' @inheritParams base::requireNamespace
4		#' @return A logical value indicating whether registration was successful.
5		#'
6		#' @keywords internal
7		car_add_mmrm <- function(quietly = FALSE) {
8	2x	if (!requireNamespace("car", quietly = quietly)) {
9	!	return(FALSE)
10		}
11	2x	envir <- asNamespace("mmrm")
12	2x	h_register_s3("car", "Anova", "mmrm", envir)
13	2x	TRUE
14		}
15
16		#' Obtain Type 2 Contrast for One Specified Effect
17		#'
18		#' This is support function to obtain contrast matrix for type II testing.
19		#'
20		#' @param object (`mmrm`)\cr the fitted MMRM.
21		#' @param effect (`string`) the name of the effect.
22		#' @param tol (`numeric`) threshold below which values are treated as 0.
23		#'
24		#' @return A `matrix` of the contrast.
25		#'
26		#' @keywords internal
27		h_type2_contrast <- function(
28		object,
29		effect,
30		tol = sqrt(.Machine$double.eps)
31		) {
32	54x	assert_class(object, "mmrm")
33	54x	assert_string(effect)
34	54x	assert_double(tol, finite = TRUE, len = 1L)
35
36		# The complete model matrix including aliased columns is needed. See below.
37	54x	mx <- component(object, "x_matrix_complete")
38	54x	asg <- attr(mx, "assign")
39	54x	formula <- object$formula_parts$model_formula
40	54x	tms <- stats::terms(formula)
41	54x	fcts <- attr(tms, "factors")[-1L, , drop = FALSE] # Discard the response.
42
43		# We do this because the factors attribute of a terms.object can include 2s.
44	54x	fcts[fcts > 1] <- 1
45
46	54x	ods <- attr(tms, "order")
47	54x	assert_subset(effect, colnames(fcts))
48	54x	idx <- which(effect == colnames(fcts))
49	54x	cols <- which(asg == idx)
50	54x	xlev <- component(object, "xlev")
51
52	54x	categorical_covars <- intersect(rownames(fcts), names(xlev))
53
54		# An effect contains the model's intercept if the model was not built with an
55		# intercept and if the effect is the first one in the model specification to
56		# contain a categorical variable.
57	54x	effect_contains_intercept <-
58	54x	!attr(tms, "intercept") &&
59	54x	length(categorical_covars) &&
60	54x	effect == h_first_term_containing_categ(fcts, categorical_covars)
61
62	54x	coef_rows <- length(cols) - as.integer(effect_contains_intercept)
63	54x	l_mx <-
64	54x	matrix(
65	54x	0,
66	54x	nrow = coef_rows,
67	54x	ncol = ncol(mx),
68	54x	dimnames = list(NULL, colnames(mx))
69		)
70	54x	if (coef_rows == 0L) {
71	!	l_mx <- l_mx[, !component(object, "beta_aliased"), drop = FALSE]
72	!	return(l_mx)
73		}
74
75		# The effect's submatrix of l_mx shall simply be an identity matrix unless the
76		# effect contains the model's intercept. In this case, the submatrix has an
77		# extra column that should contain straight -1s.
78	54x	l_mx[, cols] <-
79	54x	if (effect_contains_intercept) {
80	4x	cbind(diag(coef_rows), -1)
81		} else {
82	50x	diag(coef_rows)
83		}
84
85	54x	for (i in setdiff(seq_len(ncol(fcts)), idx)) {
86		# This is where replacing the factors attributes' 2s with 1s matters.
87	165x	additional_vars <- names(which(fcts[, i] > fcts[, idx]))
88	165x	additional_numeric <- !all(additional_vars %in% categorical_covars)
89	165x	current_col <- which(asg == i)
90		if (
91	165x	ods[i] >= ods[idx] && all(fcts[, i] >= fcts[, idx]) && !additional_numeric
92		) {
93		# This is where it matters to use the complete design matrix including
94		# aliased columns.
95	24x	x1 <- mx[, cols, drop = FALSE]
96	24x	x0 <- mx[, -c(cols, current_col), drop = FALSE]
97	24x	x2 <- mx[, current_col, drop = FALSE]
98	24x	m <- diag(nrow(x0)) - h_quad_form_mat(x0, MASS::ginv(crossprod(x0)))
99	24x	crossprod_x1_m <- crossprod(x1, m)
100	24x	ret <- solve(crossprod_x1_m %% x1) %% crossprod_x1_m %*% x2
101	24x	sub_mat <- if (effect_contains_intercept) {
102	2x	ret[-1, ] - ret[1, ]
103		} else {
104	22x	ret
105		}
106	24x	l_mx[, current_col] <- sub_mat
107		}
108		}
109		# Finally, we drop aliased columns. But we needed them to properly calculate
110		# the others.
111	54x	l_mx <- l_mx[, !component(object, "beta_aliased"), drop = FALSE]
112	54x	l_mx[abs(l_mx) < tol] <- 0
113	54x	l_mx
114		}
115
116
117		#' Identify the First Term in a Model to Contain a Categorical Variable
118		#'
119		#' This returns the column name of the leftmost column of `factors` containing a
120		#' nonzero value in a row corresponding to a `categorical` variable.
121		#'
122		#' @param factors (matrix)\cr the `factors` attribute of a [`terms.object`],
123		#' which is a matrix of 0s, 1s, and 2s.
124		#' @param categorical (character)\cr a vector of the categorical variables in
125		#' the model whose [`terms.object`] is `factors`.
126		#'
127		#' @return A `string`: one of the column names of `factors`. If none of the
128		#' columns contain a categorical variable, `NULL` is returned.
129		#'
130		#' @keywords internal
131		h_first_term_containing_categ <- function(factors, categorical) {
132	11x	factors <- factors[categorical, , drop = FALSE]
133	11x	for (term in colnames(factors)) {
134	16x	if (any(factors[, term] > 0)) {
135	11x	return(term)
136		}
137		}
138		}
139
140
141		#' Obtain Type 3 Contrast for All Effects
142		#'
143		#' This is support function to obtain contrast matrices for type III testing.
144		#'
145		#' @param object (`mmrm`)\cr the fitted MMRM.
146		#' @param tol (`numeric`) threshold below which values are treated as 0.
147		#'
148		#' @return A `list` of contrast matrices, one per effect.
149		#' @keywords internal
150		h_type3_contrasts <- function(object, tol = sqrt(.Machine$double.eps)) {
151	27x	assert_class(object, "mmrm")
152	27x	assert_double(tol, finite = TRUE, len = 1L)
153
154		# First obtain the simple contrasts as if we use all contr.sum only.
155	27x	contr_sum_contrasts <- h_contr_sum_type3_contrasts(object)
156
157		# The original design matrix.
158	27x	orig_model_matrix <- component(object, "x_matrix")
159
160		# Obtain original contrasts, if there is none or all are contr.sum, return
161		# the contr.sum contrasts.
162	27x	orig_contrasts <- attr(orig_model_matrix, "contrasts")
163		if (
164	27x	length(orig_contrasts) == 0L \|\|
165	27x	all(vapply(orig_contrasts, identical, logical(1L), "contr.sum"))
166		) {
167	!	return(contr_sum_contrasts)
168		}
169
170		# Compute the modified design matrix using all contr.sum.
171	27x	new_contrasts <- orig_contrasts
172	27x	new_contrasts[] <- "contr.sum"
173	27x	orig_data <- component(object, "full_frame")
174	27x	new_model_matrix <-
175	27x	stats::model.matrix(
176	27x	object,
177	27x	data = orig_data,
178	27x	contrasts = new_contrasts
179		)
180	27x	coef_aliased <- component(object, "beta_aliased")
181	27x	if (any(coef_aliased)) {
182	1x	new_model_matrix <- new_model_matrix[, !coef_aliased, drop = FALSE]
183		}
184
185		# Compute the correction matrix and return corrected contrasts,
186		# dropping small values.
187	27x	type_iii_correction <-
188	27x	solve(crossprod(new_model_matrix)) %*%
189	27x	crossprod(new_model_matrix, orig_model_matrix)
190	27x	result <- lapply(contr_sum_contrasts, `%*%`, type_iii_correction)
191	27x	lapply(result, function(mat) {
192	135x	mat[abs(mat) < tol] <- 0
193	135x	mat
194		})
195		}
196
197		#' Conduct type II/III hypothesis testing on the MMRM fit results.
198		#'
199		#' @param mod (`mmrm`)\cr the fitted MMRM.
200		#' @param type (`string`)\cr either `"II"`, `"III"`, `"2"`, or `"3"`, indicating the
201		#' type of test to perform.
202		#' @param test.statistic (`string`)\cr either `"F` or `"Chisq"`, indicating the
203		#' kind of test to perform.
204		#' @param ... arguments passed from other methods.
205		#' @inheritParams h_type2_contrast
206		#'
207		#' @details `Anova()` will return an `anova` object with one row per variable.
208		#'
209		#' If `test.statistic = "F"`, columns will be `Df`(numerator degrees of
210		#' freedom), `Res.Df` (denominator degrees of freedom), `F`, and
211		#' `Pr(>F)`.
212		#'
213		#' If `test.statistic = "Chisq"`, columns will be `Chisq` (the Chi-squared
214		#' test statistic), `Df` (degrees of freedom), and `Pr(>Chisq)` (p-value).
215		#'
216		#' @keywords internal
217		# Please do not load `car` and then create the documentation. The Rd file will
218		# be different.
219		# nolint start
220		Anova.mmrm <- function(
221		# nolint end
222		mod,
223		type = c("II", "III", "2", "3"),
224		tol = sqrt(.Machine$double.eps),
225		test.statistic = c("F", "Chisq"), # nolint
226		...
227		) {
228	31x	type <- as.character(type)
229	31x	type <- match.arg(type)
230	31x	test.statistic <- match.arg(test.statistic) # nolint
231
232	31x	vars <- colnames(attr(terms(mod$formula_parts$model_formula), "factors"))
233	31x	if (type %in% c("II", "2")) {
234	12x	contrasts <- lapply(vars, h_type2_contrast, object = mod, tol = tol)
235	12x	names(contrasts) <- vars
236		} else {
237	19x	contrasts <- h_type3_contrasts(mod, tol = tol)
238	19x	contrasts <- contrasts[intersect(vars, names(contrasts))]
239		}
240
241	31x	if (test.statistic == "F") {
242	24x	ret <- lapply(contrasts, df_md, object = mod)
243	24x	ret_df <- do.call(rbind.data.frame, ret)
244	24x	colnames(ret_df) <- c("Df", "Res.Df", "F", "Pr(>F)")
245		} else {
246	7x	ret <- lapply(contrasts, h_test_md, object = mod, test = "Chisq")
247	7x	ret_df <- do.call(rbind.data.frame, ret)
248	7x	colnames(ret_df) <- c("Df", "Chisq", "Pr(>Chisq)")
249		}
250
251	31x	row.names(ret_df) <- names(contrasts)
252	31x	attr(ret_df, "heading") <-
253	31x	sprintf(
254	31x	"Analysis of Fixed Effect Table (Type %s %s tests)",
255	31x	switch(type, "2" = "II", "3" = "III", type),
256	31x	test.statistic
257		)
258	31x	class(ret_df) <- c("anova", "data.frame")
259
260	31x	ret_df
261		}
262
263		#' Obtain Levels Prior and Posterior
264		#' @param var (`string`) name of the effect.
265		#' @param additional_vars (`character`) names of additional variables.
266		#' @param xlev (`list`) named list of character levels.
267		#' @param factors (`matrix`) the factor matrix.
268		#' @keywords internal
269		h_obtain_lvls <- function(var, additional_vars, xlev, factors) {
270	2x	assert_string(var)
271	2x	assert_character(additional_vars)
272	2x	assert_list(xlev, types = "character")
273	2x	nms <- names(xlev)
274	2x	assert_subset(additional_vars, nms)
275	2x	if (var %in% nms) {
276	1x	prior_vars <- intersect(nms[seq_len(match(var, nms) - 1)], additional_vars)
277	1x	prior_lvls <- vapply(xlev[prior_vars], length, FUN.VALUE = 1L)
278	1x	post_vars <- intersect(
279	1x	nms[seq(match(var, nms) + 1, length(nms))],
280	1x	additional_vars
281		)
282	1x	post_lvls <- vapply(xlev[post_vars], length, FUN.VALUE = 1L)
283	1x	total_lvls <- prod(prior_lvls) * prod(post_lvls)
284		} else {
285	1x	prior_lvls <- vapply(xlev[additional_vars], length, FUN.VALUE = 1L)
286	1x	post_lvls <- 2L
287	1x	total_lvls <- prod(prior_lvls)
288		}
289	2x	list(
290	2x	prior = prior_lvls,
291	2x	post = post_lvls,
292	2x	total = total_lvls
293		)
294		}
295
296		#' Test if the First Vector is Subset of the Second Vector
297		#' @param x (`vector`) the first list.
298		#' @param y (`vector`) the second list.
299		#' @keywords internal
300		h_get_index <- function(x, y) {
301	2x	assert_list(x)
302	2x	assert_list(y)
303	2x	vapply(
304	2x	x,
305	2x	\(i) {
306	4x	r <- vapply(y, \(j) test_subset(j, i), FUN.VALUE = TRUE)
307	4x	if (sum(r) == 1L) {
308	3x	which(r)
309		} else {
310	2x	NA_integer_
311		}
312		},
313	2x	FUN.VALUE = 1L
314		)
315		}
316
317		#' Construct Preliminary Contrast Matrices for Type III Tests Assuming Sum Contrasts
318		#'
319		#' @param object (`mmrm`)\cr the fitted MMRM.
320		#'
321		#' @return A `list` of contrast matrices, which are just row-wise term subsets of
322		#' an overall identity matrix.
323		#'
324		#' @keywords internal
325		h_contr_sum_type3_contrasts <- function(object) {
326	30x	assert_class(object, "mmrm")
327
328	30x	terms <- stats::terms(object)
329	30x	has_intercept <- attr(terms, "intercept")
330	30x	if (!has_intercept) {
331	5x	warning(
332	5x	"Type III tests can give misleading results for models without an intercept"
333		)
334		}
335	30x	term_labels <- c(if (has_intercept) "(Intercept)", labels(terms))
336	30x	n_coefs <- length(component(object, "beta_est"))
337	30x	identity_matrix <- diag(n_coefs)
338
339	30x	x_matrix <- component(object, "x_matrix")
340	30x	assign <- attr(x_matrix, "assign")
341	30x	assert_integer(assign, lower = 0L, len = ncol(x_matrix))
342	30x	map_terms_to_cols <- assign + has_intercept
343
344	30x	terms_per_col <- factor(term_labels[map_terms_to_cols], levels = term_labels)
345
346	30x	split.data.frame(
347	30x	identity_matrix,
348	30x	terms_per_col,
349	30x	drop = TRUE
350		)
351		}

1		#' Extract Formula Terms used for Covariance Structure Definition
2		#'
3		#' @param f (`formula`)\cr a formula from which covariance terms should be
4		#' extracted.
5		#'
6		#' @return A list of covariance structure expressions found in `f`.
7		#'
8		#' @importFrom stats terms
9		#' @keywords internal
10		h_extract_covariance_terms <- function(f) {
11	327x	specials <- cov_types(c("abbr", "habbr"))
12	327x	terms <- stats::terms(formula_rhs(f), specials = specials)
13	327x	covariance_terms <- Filter(length, attr(terms, "specials"))
14	327x	variables <- attr(terms, "variables")
15	327x	lapply(covariance_terms, function(i) variables[[i + 1]])
16		}
17
18		#' Drop Formula Terms used for Covariance Structure Definition
19		#'
20		#' @param f (`formula`)\cr a formula from which covariance terms should be
21		#' dropped.
22		#'
23		#' @return The formula without accepted covariance terms.
24		#'
25		#' @details `terms` is used and it will preserve the environment attribute.
26		#' This ensures the returned formula and the input formula have the same environment.
27		#' @importFrom stats terms drop.terms
28		#' @keywords internal
29		h_drop_covariance_terms <- function(f) {
30	310x	specials <- cov_types(c("abbr", "habbr"))
31
32	310x	terms <- stats::terms(f, specials = specials)
33	310x	covariance_terms <- Filter(Negate(is.null), attr(terms, "specials"))
34
35		# if no covariance terms were found, return original formula
36	310x	if (length(covariance_terms) == 0) {
37	6x	return(f)
38		}
39	304x	if (length(f) != 3) {
40	1x	update_str <- "~ . -"
41		} else {
42	303x	update_str <- ". ~ . -"
43		}
44	304x	stats::update(
45	304x	f,
46	304x	stats::as.formula(paste(
47	304x	update_str,
48	304x	deparse(attr(terms, "variables")[[covariance_terms[[1]] + 1]])
49		))
50		)
51		}
52
53		#' Add Individual Covariance Variables As Terms to Formula
54		#'
55		#' @param f (`formula`)\cr a formula to which covariance structure terms should
56		#' be added.
57		#' @param covariance (`cov_struct`)\cr a covariance structure object from which
58		#' additional variables should be sourced.
59		#'
60		#' @return A new formula with included covariance terms.
61		#'
62		#' @details [stats::update()] is used to append the covariance structure and the environment
63		#' attribute will not be changed. This ensures the returned formula and the input formula
64		#' have the same environment.
65		#'
66		#' @keywords internal
67		h_add_covariance_terms <- function(f, covariance) {
68	308x	cov_terms <- with(covariance, c(subject, visits, group))
69	302x	cov_terms <- paste(cov_terms, collapse = " + ")
70	302x	stats::update(f, stats::as.formula(paste(". ~ . + ", cov_terms)))
71		}
72
73		#' Add Formula Terms with Character
74		#'
75		#' Add formula terms from the original formula with character representation.
76		#'
77		#' @param f (`formula`)\cr a formula to be updated.
78		#' @param adds (`character`)\cr representation of elements to be added.
79		#' @param drop_response (`flag`)\cr whether response should be dropped.
80		#'
81		#' @details Elements in `adds` will be added from the formula, while the environment
82		#' of the formula is unchanged. If `adds` is `NULL` or `character(0)`, the formula is
83		#' unchanged.
84		#' @return A new formula with elements in `drops` removed.
85		#'
86		#' @keywords internal
87		h_add_terms <- function(f, adds, drop_response = FALSE) {
88	764x	assert_character(adds, null.ok = TRUE)
89	764x	if (length(adds) > 0L) {
90	398x	add_terms <- stats::as.formula(sprintf(
91	398x	". ~ . + %s",
92	398x	paste(adds, collapse = "+")
93		))
94	398x	f <- stats::update(f, add_terms)
95		}
96	764x	if (drop_response && length(f) == 3L) {
97	35x	f[[2]] <- NULL
98		}
99	764x	f
100		}

1		#' Methods for `mmrm` Objects
2		#'
3		#' @description `r lifecycle::badge("stable")`
4		#'
5		#' @param object (`mmrm`)\cr the fitted MMRM including Jacobian and call etc.
6		#' @param ... not used.
7		#' @return Depends on the method, see Details and Functions.
8		#'
9		#' @details
10		#' While printing the summary of (`mmrm`)\cr object, the following will be displayed:
11		#' 1. Formula. The formula used in the model.
12		#' 2. Data. The data used for analysis, including number of subjects, number of valid observations.
13		#' 3. Covariance. The covariance structure and number of variance parameters.
14		#' 4. Method. Restricted maximum likelihood(REML) or maximum likelihood(ML).
15		#' 5. Model selection criteria. AIC, BIC, log likelihood and deviance.
16		#' 6. Coefficients. Coefficients of the covariates.
17		#' 7. Covariance estimate. The covariance estimate(for each group).
18		#' 1. If the covariance structure is non-spatial, the covariance matrix of all categorical time points available
19		#' in data will be displayed.
20		#' 2. If the covariance structure is spatial, the covariance matrix of two time points with unit distance
21		#' will be displayed.
22		#'
23		#' `confint` is used to obtain the confidence intervals for the coefficients.
24		#' Please note that this is different from the confidence interval of difference
25		#' of least square means from `emmeans`.
26		#'
27		#' @name mmrm_methods
28		#'
29		#' @seealso [`mmrm_tmb_methods`], [`mmrm_tidiers`] for additional methods.
30		#'
31		#' @examples
32		#' formula <- FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT \| USUBJID)
33		#' object <- mmrm(formula, fev_data)
34		NULL
35
36		#' Coefficients Table for MMRM Fit
37		#'
38		#' This is used by [summary.mmrm()] to obtain the coefficients table.
39		#'
40		#' @param object (`mmrm`)\cr model fit.
41		#'
42		#' @return Matrix with one row per coefficient and columns
43		#' `Estimate`, `Std. Error`, `df`, `t value` and `Pr(>\|t\|)`.
44		#'
45		#' @keywords internal
46		h_coef_table <- function(object) {
47	40x	assert_class(object, "mmrm")
48
49	40x	coef_est <- component(object, "beta_est")
50	40x	coef_contrasts <- diag(x = rep(1, length(coef_est)))
51	40x	rownames(coef_contrasts) <- names(coef_est)
52	40x	coef_table <- t(apply(
53	40x	coef_contrasts,
54	40x	MARGIN = 1L,
55	40x	FUN = function(contrast) unlist(df_1d(object, contrast))
56		))
57	40x	assert_names(
58	40x	colnames(coef_table),
59	40x	identical.to = c("est", "se", "df", "t_stat", "p_val")
60		)
61	40x	colnames(coef_table) <- c(
62	40x	"Estimate",
63	40x	"Std. Error",
64	40x	"df",
65	40x	"t value",
66	40x	"Pr(>\|t\|)"
67		)
68
69	40x	coef_aliased <- component(object, "beta_aliased")
70	40x	if (any(coef_aliased)) {
71	2x	names_coef_na <- names(which(coef_aliased))
72	2x	coef_na_table <- matrix(
73	2x	data = NA,
74	2x	nrow = length(names_coef_na),
75	2x	ncol = ncol(coef_table),
76	2x	dimnames = list(names_coef_na, colnames(coef_table))
77		)
78	2x	coef_table <- rbind(coef_table, coef_na_table)[names(coef_aliased), ]
79		}
80
81	40x	coef_table
82		}
83
84		#' @describeIn mmrm_methods summarizes the MMRM fit results.
85		#' @exportS3Method
86		#' @examples
87		#' # Summary:
88		#' summary(object)
89		summary.mmrm <- function(object, ...) {
90	20x	aic_list <- list(
91	20x	AIC = AIC(object),
92	20x	BIC = BIC(object),
93	20x	logLik = logLik(object),
94	20x	deviance = deviance(object)
95		)
96	20x	coefficients <- h_coef_table(object)
97	20x	call <- stats::getCall(object)
98	20x	components <- component(
99	20x	object,
100	20x	c(
101	20x	"cov_type",
102	20x	"reml",
103	20x	"n_groups",
104	20x	"n_theta",
105	20x	"n_subjects",
106	20x	"n_timepoints",
107	20x	"n_obs",
108	20x	"beta_vcov",
109	20x	"varcor"
110		)
111		)
112	20x	components$method <- object$method
113	20x	components$vcov <- object$vcov
114	20x	structure(
115	20x	c(
116	20x	components,
117	20x	list(
118	20x	coefficients = coefficients,
119	20x	n_singular_coefs = sum(component(object, "beta_aliased")),
120	20x	aic_list = aic_list,
121	20x	call = call
122		)
123		),
124	20x	class = "summary.mmrm"
125		)
126		}
127
128		#' Printing MMRM Function Call
129		#'
130		#' This is used in [print.summary.mmrm()].
131		#'
132		#' @param call (`call`)\cr original [mmrm()] function call.
133		#' @param n_obs (`int`)\cr number of observations.
134		#' @param n_subjects (`int`)\cr number of subjects.
135		#' @param n_timepoints (`int`)\cr number of timepoints.
136		#'
137		#' @keywords internal
138		h_print_call <- function(call, n_obs, n_subjects, n_timepoints) {
139	10x	pass <- 0
140	10x	if (!is.null(tmp <- call$formula)) {
141	10x	cat("Formula: ", deparse(tmp), fill = TRUE)
142	10x	rhs <- tmp[[2]]
143	10x	pass <- nchar(deparse(rhs))
144		}
145	10x	if (!is.null(call$data)) {
146	10x	cat(
147	10x	"Data: ",
148	10x	deparse(call$data),
149	10x	"(used",
150	10x	n_obs,
151	10x	"observations from",
152	10x	n_subjects,
153	10x	"subjects with maximum",
154	10x	n_timepoints,
155	10x	"timepoints)",
156	10x	fill = TRUE
157		)
158		}
159		# Display the expression of weights
160	10x	if (!is.null(call$weights)) {
161	4x	cat("Weights: ", deparse(call$weights), fill = TRUE)
162		}
163		}
164
165		#' Printing MMRM Covariance Type
166		#'
167		#' This is used in [print.summary.mmrm()].
168		#'
169		#' @param cov_type (`string`)\cr covariance structure abbreviation.
170		#' @param n_theta (`count`)\cr number of variance parameters.
171		#' @param n_groups (`count`)\cr number of groups.
172		#' @keywords internal
173		h_print_cov <- function(cov_type, n_theta, n_groups) {
174	10x	assert_string(cov_type)
175	10x	assert_count(n_theta, positive = TRUE)
176	10x	assert_count(n_groups, positive = TRUE)
177	10x	cov_definition <- switch(
178	10x	cov_type,
179	10x	us = "unstructured",
180	10x	toep = "Toeplitz",
181	10x	toeph = "heterogeneous Toeplitz",
182	10x	ar1 = "auto-regressive order one",
183	10x	ar1h = "heterogeneous auto-regressive order one",
184	10x	ad = "ante-dependence",
185	10x	adh = "heterogeneous ante-dependence",
186	10x	cs = "compound symmetry",
187	10x	csh = "heterogeneous compound symmetry",
188	10x	sp_exp = "spatial exponential"
189		)
190
191	10x	catstr <- sprintf(
192	10x	"Covariance: %s (%d variance parameters%s)\n",
193	10x	cov_definition,
194	10x	n_theta,
195	10x	ifelse(n_groups == 1L, "", sprintf(" of %d groups", n_groups))
196		)
197	10x	cat(catstr)
198		}
199
200		#' Printing AIC and other Model Fit Criteria
201		#'
202		#' This is used in [print.summary.mmrm()].
203		#'
204		#' @param aic_list (`list`)\cr list as part of from [summary.mmrm()].
205		#' @param digits (`number`)\cr number of decimal places used with [round()].
206		#'
207		#' @keywords internal
208		h_print_aic_list <- function(aic_list, digits = 1) {
209	6x	diag_vals <- round(unlist(aic_list), digits)
210	6x	diag_vals <- format(diag_vals)
211	6x	print(diag_vals, quote = FALSE)
212		}
213
214		#' @describeIn mmrm_methods prints the MMRM fit summary.
215		#' @exportS3Method
216		#' @keywords internal
217		print.summary.mmrm <- function(
218		x,
219		digits = max(3, getOption("digits") - 3),
220		signif.stars = getOption("show.signif.stars"), # nolint
221		...
222		) {
223	5x	cat("mmrm fit\n\n")
224	5x	h_print_call(x$call, x$n_obs, x$n_subjects, x$n_timepoints)
225	5x	h_print_cov(x$cov_type, x$n_theta, x$n_groups)
226	5x	cat("Method: ", x$method, "\n", sep = "")
227	5x	cat("Vcov Method: ", x$vcov, "\n", sep = "")
228	5x	cat("Inference: ")
229	5x	cat(ifelse(x$reml, "REML", "ML"))
230	5x	cat("\n\n")
231	5x	cat("Model selection criteria:\n")
232	5x	h_print_aic_list(x$aic_list)
233	5x	cat("\n")
234	5x	cat("Coefficients: ")
235	5x	if (x$n_singular_coefs > 0) {
236	1x	cat(
237		"(",
238	1x	x$n_singular_coefs,
239	1x	" not defined because of singularities)",
240	1x	sep = ""
241		)
242		}
243	5x	cat("\n")
244	5x	stats::printCoefmat(
245	5x	x$coefficients,
246	5x	zap.ind = 3,
247	5x	digits = digits,
248	5x	signif.stars = signif.stars
249		)
250	5x	cat("\n")
251	5x	cat("Covariance estimate:\n")
252	5x	if (is.list(x$varcor)) {
253	1x	for (v in names(x$varcor)) {
254	2x	cat(sprintf("Group: %s\n", v))
255	2x	print(round(x$varcor[[v]], digits = digits))
256		}
257		} else {
258	4x	print(round(x$varcor, digits = digits))
259		}
260	5x	cat("\n")
261	5x	invisible(x)
262		}
263
264
265		#' @describeIn mmrm_methods obtain the confidence intervals for the coefficients.
266		#' @exportS3Method
267		#' @examples
268		#' # Confidence Interval:
269		#' confint(object)
270		confint.mmrm <- function(object, parm, level = 0.95, ...) {
271	20x	cf <- coef(object)
272	20x	pnames <- names(cf)
273	20x	if (missing(parm)) {
274	15x	parm <- pnames
275		}
276	20x	assert(
277	20x	check_subset(parm, pnames),
278	20x	check_integerish(parm, lower = 1L, upper = length(cf))
279		)
280	18x	if (is.numeric(parm)) {
281	2x	parm <- pnames[parm]
282		}
283	18x	assert_number(level, lower = 0, upper = 1)
284	18x	a <- (1 - level) / 2
285	18x	pct <- paste(
286	18x	format(100 * c(a, 1 - a), trim = TRUE, scientific = FALSE, digits = 3),
287		"%"
288		)
289	18x	coef_table <- h_coef_table(object)
290	18x	df <- coef_table[parm, "df"]
291	18x	ses <- coef_table[parm, "Std. Error"]
292	18x	fac <- stats::qt(a, df = df)
293	18x	ci <- array(NA, dim = c(length(parm), 2L), dimnames = list(parm, pct))
294	18x	sefac <- ses * fac
295	18x	ci[] <- cf[parm] + c(sefac, -sefac)
296	18x	ci
297		}
298
299
300		#' Analysis of Variance for `mmrm` Fits
301		#'
302		#' If supplied only one model fit, the function will calculate and return the
303		#' significance of the model terms. If supplied more than one model fit,
304		#' standard diagnostics will be returned for each model, optionally including
305		#' likelihood ratio test (LRT) results for adjacent models.
306		#'
307		#' @param object (`mmrm`)\cr an `mmrm` model fit.
308		#' @param ... (`mmrm`)\cr optional `mmrm` model fits. If left empty, the
309		#' significance of each term in `object` will be calculated.
310		#' @param test (`flag`)\cr indicating whether the output should include
311		#' likelihood ratio test (LRT) results comparing the model fits to one
312		#' another. Defaults to `TRUE`. Ignored if `...` is empty.
313		#' @param refit (`flag`)\cr indicating whether the models should be refitted
314		#' with the dataset consisting of their shared set of observations before
315		#' performing diagnostics and testing. This is ignored if the models already
316		#' share the same dataset. If `refit = FALSE` and the models have different
317		#' underlying data sets, an error will be thrown. Defaults to `FALSE`. Ignored
318		#' if `...` is empty.
319		#'
320		#' @details When `test = FALSE` (or, when only one model is supplied), this
321		#' function will process any `mmrm` fits, related or unrelated.
322		#'
323		#' When supplying multiple models and `test = TRUE`, adjacent pairs of models
324		#' are tested sequentially. In other words, the order of the supplied models
325		#' matters. Furthermore, there are are multiple requirements for successful
326		#' LRT. See the section "Requirements for LRT" below.
327		#'
328		#' # Requirements for LRT
329		#'
330		#' 1. Each supplied model fit must have more degrees of freedom than
331		#' the preceding model fit.
332		#'
333		#' 1. If all supplied models were estimated using maximum likelihood (ML), the
334		#' models must have nested covariates in order to undergo LRT. In other words,
335		#' the set of covariates for each model must be a subset of the covariates of
336		#' the next model. However, if any of the supplied models were estimated using
337		#' restricted maximum likelihood (REML), all models must have the same
338		#' covariates.
339		#'
340		#' 1. The covariance structure of each model must be either (a) the same as
341		#' that of the next model or (b) a special case of that of the next model. See
342		#' the section "Covariance structure nesting hierarchy" below.
343		#'
344		#' 1. All supplied model fits must either already use the same data or be
345		#' refitted using `refit = TRUE`, which refits all models to the dataset of
346		#' common observations between all models' respective data sets.
347		#'
348		#' # Covariance structure nesting hierarchy
349		#'
350		#' ## Structured nests within unstructured
351		#'
352		#' Tautologically, all covariance structures are special cases of an
353		#' unstructured covariance, and a model with a covariance structure can be
354		#' considered "nested" within an model without a covariance structure
355		#' (assuming that the covariates are also nested).
356		#'
357		#' ## Homogeneous nests within analogous heterogeneous
358		#'
359		#' All homogeneous covariance structures are nested within their corresponding
360		#' heterogeneous counterparts. For instance, the homogeneous Toeplitz
361		#' covariance structure is nested within the heterogeneous Toeplitz covariance
362		#' structure.
363		#'
364		#' ## Other nested structures
365		#'
366		#' Some different covariance structure types are also nested:
367		#'
368		#' - First-order auto-regressive (`ar1` / `ar1h`) is nested within:
369		#' - ante-dependence (`ad` / `adh`)
370		#' - Toeplitz (`toep` / `toeph`)
371		#' - Compound symmetry (`cs` / `csh`) is nested within Toeplitz (`toep` / `toeph`)
372		#'
373		#' @returns A data frame with a row for each supplied model. If `...` is empty,
374		#' this will be the the returned value of [h_anova_single_mmrm_model()] with
375		#' `object` supplied as its argument. Otherwise, the resulting data frame will
376		#' have the following columns:
377		#'
378		#' - `Model`: the sequence number of the model according to the order in which
379		#' the models were supplied to this function.
380		#'
381		#' - `refit`: logical, indicating whether or not the model was refitted. If
382		#' the `refit` argument was `FALSE`, all values will be `FALSE`.
383		#'
384		#' - `REML`: logical, indicating whether or not the model was fitted using
385		#' restricted maximum likelihood (REML) estimation. If `FALSE`, the model was
386		#' fitted using maximum likelihood (ML) estimation.
387		#'
388		#' - `n_param`: the number of variance parameters in the model fit, obtained
389		#' via [logLik.mmrm_tmb()].
390		#'
391		#' - `n_coef`: the number of estimated coefficients in the model fit, obtained
392		#' via [logLik.mmrm_tmb()].
393		#'
394		#' - `df`: degrees of freedom of the model fit, obtained via
395		#' [logLik.mmrm_tmb()].
396		#'
397		#' - `AIC`: Akaike's "An Information Criterion" of the model fit, obtained via
398		#' [stats::AIC()].
399		#'
400		#' - `BIC`: the Bayesian Information Criterion of the model fit, obtained via
401		#' [stats::BIC()].
402		#'
403		#' - `logLik`: the log likelihood of the model fit, obtained via
404		#' [logLik.mmrm_tmb()].
405		#'
406		#' - `call`: the [call] that created the model fit, obtained via [component()]
407		#' with `name = "call"`, which is passed to [deparse1()]. If the model was
408		#' refitted (i.e., if its `refit` entry in this table is `TRUE`), this `call`
409		#' will be different from the `call` component of the pre-refitted model fit.
410		#'
411		#' The data frame will have these additional columns inserted before `call` if
412		#' `test = TRUE`. Note that since each of these columns describe the results
413		#' of a likelihood ratio test (LRT) between the previous row's and current
414		#' row's model fits, the first element of each of these columns will be `NA`.
415		#'
416		#' - `test`: character, indicating which two model fits were compared. Values
417		#' are of the form `{Model - 1} vs {Model}`.
418		#'
419		#' - `log_likelihood_ratio`: the logarithm of the likelihood ratio between the two
420		#' models being compared.
421		#'
422		#' - `p_value`: the p-value of the `log_likelihood_ratio`.
423		#'
424		#' @seealso For details on the single model operation of this function, see
425		#' [h_anova_single_mmrm_model()]. For details on the generic, see
426		#' [stats::anova()].
427		#'
428		#' @export
429		#'
430		#' @name stats_anova
431		#'
432		#' @examples
433		#' # Create a few model fits, only adding terms from one to the next.
434		#' # Notice also that each covariance structure is a special case of the one
435		#' # that follows.
436		#' fit_sex_ar1 <-
437		#' mmrm(FEV1 ~ FEV1_BL + SEX + ARMCD + ar1(AVISIT \| USUBJID),
438		#' data = fev_data,
439		#' reml = FALSE
440		#' )
441		#' fit_sex_race_toeph <-
442		#' mmrm(
443		#' FEV1 ~ FEV1_BL + SEX + RACE + ARMCD + toeph(AVISIT \| USUBJID),
444		#' data = fev_data,
445		#' reml = FALSE
446		#' )
447		#' fit_interaction_us <-
448		#' mmrm(
449		#' FEV1 ~ FEV1_BL + SEX * RACE + ARMCD + us(AVISIT \| USUBJID),
450		#' data = fev_data,
451		#' reml = FALSE
452		#' )
453		#'
454		#' # Single model fit, showing significance of model terms:
455		#' anova(fit_interaction_us)
456		#'
457		#' # Multiple model fits, with diagnostics for each fit and likelihood ratio
458		#' # testing (LRT) for each adjacent pair. LRT is possible because when the fits
459		#' # are in this order, their covariates and covariance structures are nested.
460		#' anova(fit_sex_ar1, fit_sex_race_toeph, fit_interaction_us)
461		#'
462		#' # We can only change the order if we forego LRT using test = FALSE.
463		#' anova(fit_sex_race_toeph, fit_interaction_us, fit_sex_ar1, test = FALSE)
464		#'
465		#' # Create a subset of fev_data set with the 4th visit removed.
466		#' fev_subset <- droplevels(fev_data[fev_data$VISITN < 4, ])
467		#'
468		#' # Recreate fit_sex_race_toeph but this time based off fev_subset:
469		#' fit_sex_race_toeph_sub <-
470		#' mmrm(
471		#' FEV1 ~ FEV1_BL + SEX + RACE + ARMCD + toeph(AVISIT \| USUBJID),
472		#' data = fev_subset,
473		#' reml = FALSE
474		#' )
475		#'
476		#' # If a model was created with a different data set, refit = TRUE is needed.
477		#' anova(fit_sex_ar1, fit_sex_race_toeph_sub, fit_interaction_us, refit = TRUE)
478		anova.mmrm <- function(object, ..., test = TRUE, refit = FALSE) {
479	3x	assert_class(object, "mmrm")
480
481	3x	fits <- list(object, ...)
482
483	3x	if (length(fits) == 1L) {
484	1x	out <- h_anova_single_mmrm_model(object)
485		} else {
486		# Ensure all objects in ... are mmrm fits.
487	2x	lapply(fits[-1L], assert_class, classes = "mmrm", .var.name = "...")
488
489	2x	assert_flag(test)
490	2x	assert_flag(refit)
491
492		# If the data are not all the same and refit = TRUE, refit all with the
493		# largest common data set.
494	2x	if (refit && !h_check_fits_all_data_same(fits)) {
495	1x	common_data <- h_fits_common_data(fits)
496
497	1x	fit_data <- lapply(fits, h_get_minimal_fit_data)
498
499	1x	nested_in_common_data <-
500	1x	vapply(
501	1x	fit_data,
502	1x	h_check_columns_nested,
503	1x	FUN.VALUE = logical(1L),
504	1x	data_augmented = common_data
505		)
506
507	1x	needs_refit <- !nested_in_common_data
508
509	1x	fits[needs_refit] <-
510	1x	lapply(fits[needs_refit], h_refit_mmrm, data = common_data)
511		} else {
512	1x	needs_refit <- rep_len(FALSE, length(fits))
513		}
514
515	2x	log_likelihood_vec <- lapply(fits, logLik)
516
517	2x	out <-
518	2x	data.frame(
519	2x	Model = seq_along(fits),
520	2x	refit = needs_refit,
521	2x	REML = vapply(fits, component, logical(1L), "reml"),
522	2x	n_param = vapply(log_likelihood_vec, attr, numeric(1L), "n_param"),
523	2x	n_coef = vapply(log_likelihood_vec, attr, numeric(1L), "n_coef"),
524	2x	df = vapply(log_likelihood_vec, attr, numeric(1L), "df"),
525	2x	AIC = vapply(log_likelihood_vec, AIC, numeric(1L)),
526	2x	BIC = vapply(log_likelihood_vec, BIC, numeric(1L)),
527	2x	logLik = as.numeric(log_likelihood_vec)
528		)
529
530	2x	if (test) {
531	2x	h_assert_lrt_suitability(fits, refit, dfs = out$df, is_reml = out$REML)
532
533	2x	model_indices_except_last <- out$Model[-length(fits)]
534	2x	model_indices_except_first <- out$Model[-1L]
535
536		# Create labels for the pair being compared (e.g., "3 vs 4"). Force the
537		# first element to be NA because a pair of models is the previous row's
538		# model plus the current row's model.
539	2x	out$test <-
540	2x	c(
541	2x	NA_character_,
542	2x	paste(model_indices_except_last, "vs", model_indices_except_first)
543		)
544
545	2x	out$log_likelihood_ratio <- c(NA, diff(out$logLik))
546
547	2x	out$p_value <-
548	2x	stats::pchisq(
549	2x	2 * abs(out$log_likelihood_ratio),
550	2x	df = abs(diff(out$df)),
551	2x	lower.tail = FALSE
552		)
553		}
554
555	2x	fit_calls <- lapply(fits, component, "call")
556	2x	out$call <- vapply(fit_calls, deparse1, FUN.VALUE = character(1L))
557		}
558
559	3x	class(out) <- union("anova.mmrm", class(out))
560
561	3x	out
562		}

1		#' Obtain Kenward-Roger Adjustment Components
2		#'
3		#' @description Obtains the components needed downstream for the computation of Kenward-Roger degrees of freedom.
4		#' Used in [mmrm()] fitting if method is "Kenward-Roger".
5		#'
6		#' @param tmb_data (`mmrm_tmb_data`)\cr produced by [h_mmrm_tmb_data()].
7		#' @param theta (`numeric`)\cr theta estimate.
8		#'
9		#' @details the function returns a named list, \eqn{P}, \eqn{Q} and \eqn{R}, which corresponds to the
10		#' paper in 1997. The matrices are stacked in columns so that \eqn{P}, \eqn{Q} and \eqn{R} has the same
11		#' column number(number of beta parameters). The number of rows, is dependent on
12		#' the total number of theta and number of groups, if the fit is a grouped mmrm.
13		#' For \eqn{P} matrix, it is stacked sequentially. For \eqn{Q} and \eqn{R} matrix, it is stacked so
14		#' that the \eqn{Q_{ij}} and \eqn{R_{ij}} is stacked from \eqn{j} then to \eqn{i}, i.e. \eqn{R_{i1}}, \eqn{R_{i2}}, etc.
15		#' \eqn{Q} and \eqn{R} only contains intra-group results and inter-group results should be all zero matrices
16		#' so they are not stacked in the result.
17		#'
18		#' @return Named list with elements:
19		#' - `P`: `matrix` of \eqn{P} component.
20		#' - `Q`: `matrix` of \eqn{Q} component.
21		#' - `R`: `matrix` of \eqn{R} component.
22		#'
23		#' @keywords internal
24		h_get_kr_comp <- function(tmb_data, theta) {
25	50x	assert_class(tmb_data, "mmrm_tmb_data")
26	50x	assert_class(theta, "numeric")
27	50x	.Call(`_mmrm_get_pqr`, PACKAGE = "mmrm", tmb_data, theta)
28		}
29
30		#' Calculation of Kenward-Roger Degrees of Freedom for Multi-Dimensional Contrast
31		#'
32		#' @description Used in [df_md()] if method is "Kenward-Roger" or "Kenward-Roger-Linear".
33		#'
34		#' @inheritParams h_df_md_sat
35		#' @inherit h_df_md_sat return
36		#' @keywords internal
37		h_df_md_kr <- function(object, contrast) {
38	52x	assert_class(object, "mmrm")
39	52x	assert_matrix(
40	52x	contrast,
41	52x	mode = "numeric",
42	52x	any.missing = FALSE,
43	52x	ncols = length(component(object, "beta_est"))
44		)
45	52x	if (component(object, "reml") != 1) {
46	!	stop("Kenward-Roger is only for REML")
47		}
48	52x	kr_comp <- object$kr_comp
49	52x	w <- component(object, "theta_vcov")
50	52x	v_adj <- object$beta_vcov_adj
51	52x	df <- h_kr_df(v0 = object$beta_vcov, l = contrast, w = w, p = kr_comp$P)
52
53	52x	h_test_md(object, contrast, df = df$m, f_stat_factor = df$lambda)
54		}
55
56		#' Calculation of Kenward-Roger Degrees of Freedom for One-Dimensional Contrast
57		#'
58		#' @description Used in [df_1d()] if method is
59		#' "Kenward-Roger" or "Kenward-Roger-Linear".
60		#'
61		#' @inheritParams h_df_1d_sat
62		#' @inherit h_df_1d_sat return
63		#' @keywords internal
64		h_df_1d_kr <- function(object, contrast) {
65	22x	assert_class(object, "mmrm")
66	22x	assert_numeric(contrast, len = length(component(object, "beta_est")))
67	22x	if (component(object, "reml") != 1) {
68	!	stop("Kenward-Roger is only for REML!")
69		}
70
71	22x	df <- h_kr_df(
72	22x	v0 = object$beta_vcov,
73	22x	l = matrix(contrast, nrow = 1),
74	22x	w = component(object, "theta_vcov"),
75	22x	p = object$kr_comp$P
76		)
77
78	22x	h_test_1d(object, contrast, df$m)
79		}
80
81		#' Obtain the Adjusted Kenward-Roger degrees of freedom
82		#'
83		#' @description Obtains the adjusted Kenward-Roger degrees of freedom and F statistic scale parameter.
84		#' Used in [h_df_md_kr()] or [h_df_1d_kr].
85		#'
86		#' @param v0 (`matrix`)\cr unadjusted covariance matrix.
87		#' @param l (`matrix`)\cr linear combination matrix.
88		#' @param w (`matrix`)\cr hessian matrix.
89		#' @param p (`matrix`)\cr P matrix from [h_get_kr_comp()].
90		#'
91		#' @return Named list with elements:
92		#' - `m`: `numeric` degrees of freedom.
93		#' - `lambda`: `numeric` F statistic scale parameter.
94		#'
95		#' @keywords internal
96		h_kr_df <- function(v0, l, w, p) {
97	75x	n_beta <- ncol(v0)
98	75x	assert_matrix(v0, ncols = n_beta, nrows = n_beta)
99	75x	assert_matrix(l, ncols = n_beta)
100	75x	n_theta <- ncol(w)
101	75x	assert_matrix(w, ncols = n_theta, nrows = n_theta)
102	75x	n_visits <- ncol(p)
103	75x	assert_matrix(p, nrows = n_visits * n_theta)
104		# see vignettes/kenward.Rmd#279
105	75x	slvol <- solve(h_quad_form_mat(l, v0))
106	75x	m <- h_quad_form_mat(t(l), slvol)
107	75x	nl <- nrow(l)
108	75x	mv0 <- m %*% v0
109	75x	pl <- lapply(seq_len(nrow(p) / ncol(p)), function(x) {
110	282x	ii <- (x - 1) * ncol(p) + 1
111	282x	jj <- x * ncol(p)
112	282x	p[ii:jj, ]
113		})
114	75x	mv0pv0 <- lapply(pl, function(x) {
115	282x	mv0 %% x %% v0
116		})
117	75x	a1 <- 0
118	75x	a2 <- 0
119		# see vignettes/kenward.Rmd#283
120	75x	for (i in seq_len(length(pl))) {
121	282x	for (j in seq_len(length(pl))) {
122	1460x	a1 <- a1 + w[i, j] * h_tr(mv0pv0[[i]]) * h_tr(mv0pv0[[j]])
123	1460x	a2 <- a2 + w[i, j] * h_tr(mv0pv0[[i]] %*% mv0pv0[[j]])
124		}
125		}
126	75x	b <- 1 / (2 * nl) * (a1 + 6 * a2)
127	75x	e <- 1 + a2 / nl
128	75x	e_star <- 1 / (1 - a2 / nl)
129	75x	g <- ((nl + 1) * a1 - (nl + 4) * a2) / ((nl + 2) * a2)
130	75x	denom <- (3 * nl + 2 - 2 * g)
131	75x	c1 <- g / denom
132	75x	c2 <- (nl - g) / denom
133	75x	c3 <- (nl + 2 - g) / denom
134	75x	v_star <- 2 / nl * (1 + c1 * b) / (1 - c2 * b)^2 / (1 - c3 * b)
135	75x	rho <- v_star / (2 * e_star^2)
136	75x	m <- 4 + (nl + 2) / (nl * rho - 1)
137	75x	lambda <- m / (e_star * (m - 2))
138	75x	list(m = m, lambda = lambda)
139		}
140
141		#' Obtain the Adjusted Covariance Matrix
142		#'
143		#' @description Obtains the Kenward-Roger adjusted covariance matrix for the
144		#' coefficient estimates.
145		#' Used in [mmrm()] fitting if method is "Kenward-Roger" or "Kenward-Roger-Linear".
146		#'
147		#' @param v (`matrix`)\cr unadjusted covariance matrix.
148		#' @param w (`matrix`)\cr hessian matrix.
149		#' @param p (`matrix`)\cr P matrix from [h_get_kr_comp()].
150		#' @param q (`matrix`)\cr Q matrix from [h_get_kr_comp()].
151		#' @param r (`matrix`)\cr R matrix from [h_get_kr_comp()].
152		#' @param linear (`flag`)\cr whether to use linear Kenward-Roger approximation.
153		#'
154		#' @return The matrix of adjusted covariance matrix.
155		#'
156		#' @keywords internal
157		h_var_adj <- function(v, w, p, q, r, linear = FALSE) {
158	52x	assert_flag(linear)
159	52x	n_beta <- ncol(v)
160	52x	assert_matrix(v, nrows = n_beta)
161	52x	n_theta <- ncol(w)
162	52x	assert_matrix(w, nrows = n_theta)
163	52x	n_visits <- ncol(p)
164	52x	theta_per_group <- nrow(q) / nrow(p)
165	52x	n_groups <- n_theta / theta_per_group
166	52x	assert_matrix(p, nrows = n_theta * n_visits)
167	52x	assert_matrix(
168	52x	q,
169	52x	nrows = theta_per_group^2 * n_groups * n_visits,
170	52x	ncols = n_visits
171		)
172	52x	assert_matrix(
173	52x	r,
174	52x	nrows = theta_per_group^2 * n_groups * n_visits,
175	52x	ncols = n_visits
176		)
177	52x	if (linear) {
178	13x	r <- matrix(0, nrow = nrow(r), ncol = ncol(r))
179		}
180
181		# see vignettes/kenward.Rmd#131
182	52x	ret <- v
183	52x	for (i in seq_len(n_theta)) {
184	280x	for (j in seq_len(n_theta)) {
185	2282x	gi <- ceiling(i / theta_per_group)
186	2282x	gj <- ceiling(j / theta_per_group)
187	2282x	iid <- (i - 1) * n_beta + 1
188	2282x	jid <- (j - 1) * n_beta + 1
189	2282x	ii <- i - (gi - 1) * theta_per_group
190	2282x	jj <- j - (gi - 1) * theta_per_group
191	2282x	ijid <- ((ii - 1) * theta_per_group + jj - 1) *
192	2282x	n_beta +
193	2282x	(gi - 1) * n_beta * theta_per_group^2 +
194	2282x	1
195	2282x	if (gi != gj) {
196	592x	ret <- ret +
197	592x	2 *
198	592x	w[i, j] *
199	592x	v %*%
200	592x	(-p[iid:(iid + n_beta - 1), ] %*%
201	592x	v %*%
202	592x	p[jid:(jid + n_beta - 1), ]) %*%
203	592x	v
204		} else {
205	1690x	ret <- ret +
206	1690x	2 *
207	1690x	w[i, j] *
208	1690x	v %*%
209	1690x	(q[ijid:(ijid + n_beta - 1), ] -
210	1690x	p[iid:(iid + n_beta - 1), ] %*%
211	1690x	v %*%
212	1690x	p[jid:(jid + n_beta - 1), ] -
213	1690x	1 / 4 * r[ijid:(ijid + n_beta - 1), ]) %*%
214	1690x	v
215		}
216		}
217		}
218	52x	ret
219		}

1		#' Obtain List of Jacobian Matrix Entries for Covariance Matrix
2		#'
3		#' @description Obtain the Jacobian matrices given the covariance function and variance parameters.
4		#'
5		#' @param tmb_data (`mmrm_tmb_data`)\cr produced by [h_mmrm_tmb_data()].
6		#' @param theta_est (`numeric`)\cr variance parameters point estimate.
7		#' @param beta_vcov (`matrix`)\cr vairance covariance matrix of coefficients.
8		#'
9		#' @return List with one element per variance parameter containing a matrix
10		#' of the same dimensions as the covariance matrix. The values are the derivatives
11		#' with regards to this variance parameter.
12		#'
13		#' @keywords internal
14		h_jac_list <- function(tmb_data, theta_est, beta_vcov) {
15	106x	assert_class(tmb_data, "mmrm_tmb_data")
16	106x	assert_numeric(theta_est)
17	106x	assert_matrix(beta_vcov)
18	106x	.Call(`_mmrm_get_jacobian`, PACKAGE = "mmrm", tmb_data, theta_est, beta_vcov)
19		}
20
21		#' Quadratic Form Calculations
22		#'
23		#' @description These helpers are mainly for easier readability and slightly better efficiency
24		#' of the quadratic forms used in the Satterthwaite calculations.
25		#'
26		#' @param center (`matrix`)\cr square numeric matrix with the same dimensions as
27		#' `x` as the center of the quadratic form.
28		#'
29		#' @name h_quad_form
30		NULL
31
32		#' @describeIn h_quad_form calculates the number `vec %% center %% t(vec)`
33		#' as a numeric (not a matrix).
34		#'
35		#' @param vec (`numeric`)\cr interpreted as a row vector.
36		#'
37		#' @keywords internal
38		h_quad_form_vec <- function(vec, center) {
39	6432x	vec <- as.vector(vec)
40	6432x	assert_numeric(vec, any.missing = FALSE)
41	6432x	assert_matrix(
42	6432x	center,
43	6432x	mode = "numeric",
44	6432x	any.missing = FALSE,
45	6432x	nrows = length(vec),
46	6432x	ncols = length(vec)
47		)
48
49	6432x	sum(vec * (center %*% vec))
50		}
51
52		#' @describeIn h_quad_form calculates the quadratic form `mat %% center %% t(mat)`
53		#' as a matrix, the result is square and has dimensions identical to the number
54		#' of rows in `mat`.
55		#'
56		#' @param mat (`matrix`)\cr numeric matrix to be multiplied left and right of
57		#' `center`, therefore needs to have as many columns as there are rows and columns
58		#' in `center`.
59		#'
60		#' @keywords internal
61		h_quad_form_mat <- function(mat, center) {
62	408x	assert_matrix(mat, mode = "numeric", any.missing = FALSE, min.cols = 1L)
63	408x	assert_matrix(
64	408x	center,
65	408x	mode = "numeric",
66	408x	any.missing = FALSE,
67	408x	nrows = ncol(center),
68	408x	ncols = ncol(center)
69		)
70
71	408x	mat %*% tcrossprod(center, mat)
72		}
73
74		#' Computation of a Gradient Given Jacobian and Contrast Vector
75		#'
76		#' @description Computes the gradient of a linear combination of `beta` given the Jacobian matrix and
77		#' variance parameters.
78		#'
79		#' @param jac_list (`list`)\cr Jacobian list produced e.g. by [h_jac_list()].
80		#' @param contrast (`numeric`)\cr contrast vector, which needs to have the
81		#' same number of elements as there are rows and columns in each element of
82		#' `jac_list`.
83		#'
84		#' @return Numeric vector which contains the quadratic forms of each element of
85		#' `jac_list` with the `contrast` vector.
86		#'
87		#' @keywords internal
88		h_gradient <- function(jac_list, contrast) {
89	590x	assert_list(jac_list)
90	590x	assert_numeric(contrast)
91
92	590x	vapply(
93	590x	jac_list,
94	590x	h_quad_form_vec,
95	590x	vec = contrast,
96	590x	numeric(1L)
97		)
98		}
99
100		#' Helper for Calculation of Satterthwaite with Empirical Covariance Matrix
101		#'
102		#' @description Used in [h_df_1d_sat()] and [h_df_md_sat()] if empirical covariance
103		#' matrix is used.
104		#'
105		#' @param object (`mmrm`)\cr the MMRM fit.
106		#' @param contrast_matrix (`matrix`)\cr contrast matrix with number of subjects times
107		#' number of coefficients as the number of columns.
108		#'
109		#' @return Adjusted degrees of freedom value.
110		#' @keywords internal
111		h_df_1d_sat_empirical <- function(object, contrast_matrix) {
112	18x	assert_class(object, "mmrm")
113	18x	assert_matrix(
114	18x	contrast_matrix,
115	18x	mode = "numeric",
116	18x	any.missing = FALSE
117		)
118	18x	g_matrix <- if (
119	18x	is.null(object$empirical_g_mat) && !is.null(object$empirical_df_mat)
120		) {
121	1x	warning(
122	1x	"mmrm fit was obtained with package version < 0.3.15, ",
123	1x	"using deprecated calculation of d.f., consider refitting the model"
124		)
125	1x	h_quad_form_mat(contrast_matrix, object$empirical_df_mat)
126	18x	} else if (!is.null(object$empirical_g_mat)) {
127	16x	g_times_contrast_transposed <- tcrossprod(
128	16x	object$empirical_g_mat,
129	16x	contrast_matrix
130		)
131	16x	crossprod(g_times_contrast_transposed)
132		} else {
133	1x	stop(
134	1x	"neither empirical_df_mat nor empirical_g_mat are available in mmrm fit object"
135		)
136		}
137	17x	h_tr(g_matrix)^2 / sum(g_matrix^2)
138		}
139
140		#' Calculation of Satterthwaite Degrees of Freedom for One-Dimensional Contrast
141		#'
142		#' @description Used in [df_1d()] if method is
143		#' "Satterthwaite".
144		#'
145		#' @param object (`mmrm`)\cr the MMRM fit.
146		#' @param contrast (`numeric`)\cr contrast vector. Note that this should not include
147		#' elements for singular coefficient estimates, i.e. only refer to the
148		#' actually estimated coefficients.
149		#'
150		#' @return List with `est`, `se`, `df`, `t_stat` and `p_val`.
151		#' @keywords internal
152		h_df_1d_sat <- function(object, contrast) {
153	484x	assert_class(object, "mmrm")
154	484x	contrast <- as.numeric(contrast)
155	484x	assert_numeric(contrast, len = length(component(object, "beta_est")))
156
157	484x	df <- if (identical(object$vcov, "Asymptotic")) {
158	471x	grad <- h_gradient(component(object, "jac_list"), contrast)
159	471x	v_num <- 2 * h_quad_form_vec(contrast, component(object, "beta_vcov"))^2
160	471x	v_denom <- h_quad_form_vec(grad, component(object, "theta_vcov"))
161	471x	v_num / v_denom
162	484x	} else if (
163	484x	object$vcov %in%
164	484x	c("Empirical", "Empirical-Jackknife", "Empirical-Bias-Reduced")
165		) {
166	13x	contrast_matrix <- Matrix::.bdiag(rep(
167	13x	list(matrix(contrast, nrow = 1)),
168	13x	component(object, "n_subjects")
169		))
170	13x	contrast_matrix <- as.matrix(contrast_matrix)
171	13x	h_df_1d_sat_empirical(object, contrast_matrix)
172		}
173
174	484x	h_test_1d(object, contrast, df)
175		}
176
177		#' Calculating Denominator Degrees of Freedom for the Multi-Dimensional Case
178		#'
179		#' @description Calculates the degrees of freedom for multi-dimensional contrast.
180		#'
181		#' @param t_stat_df (`numeric`)\cr `n` t-statistic derived degrees of freedom.
182		#'
183		#' @return Usually the calculation is returning `2 * E / (E - n)` where
184		#' `E` is the sum of `t / (t - 2)` over all `t_stat_df` values `t`.
185		#'
186		#' @note If the input values are two similar to each other then just the average
187		#' of them is returned. If any of the inputs is not larger than 2 then 2 is
188		#' returned.
189		#'
190		#' @keywords internal
191		h_md_denom_df <- function(t_stat_df) {
192	52x	assert_numeric(
193	52x	t_stat_df,
194	52x	min.len = 1L,
195	52x	lower = .Machine$double.xmin,
196	52x	any.missing = FALSE
197		)
198
199	52x	if (test_scalar(t_stat_df)) {
200	1x	t_stat_df
201	51x	} else if (all(abs(diff(t_stat_df)) < sqrt(.Machine$double.eps))) {
202	1x	mean(t_stat_df)
203	50x	} else if (any(t_stat_df <= 2)) {
204	2x	2
205		} else {
206	48x	e <- sum(t_stat_df / (t_stat_df - 2))
207	48x	2 * e / (e - (length(t_stat_df)))
208		}
209		}
210
211		#' Creating F-Statistic Results from One-Dimensional Contrast
212		#'
213		#' @description Creates multi-dimensional result from one-dimensional contrast from [df_1d()].
214		#'
215		#' @param object (`mmrm`)\cr model fit.
216		#' @param contrast (`numeric`)\cr one-dimensional contrast.
217		#'
218		#' @return The one-dimensional degrees of freedom are calculated and then
219		#' based on that the p-value is calculated.
220		#'
221		#' @keywords internal
222		h_df_md_from_1d <- function(object, contrast) {
223	161x	res_1d <- h_df_1d_sat(object, contrast)
224	161x	list(
225	161x	num_df = 1,
226	161x	denom_df = res_1d$df,
227	161x	f_stat = res_1d$t_stat^2,
228	161x	p_val = stats::pf(
229	161x	q = res_1d$t_stat^2,
230	161x	df1 = 1,
231	161x	df2 = res_1d$df,
232	161x	lower.tail = FALSE
233		)
234		)
235		}
236
237		#' Calculation of Satterthwaite Degrees of Freedom for Multi-Dimensional Contrast
238		#'
239		#' @description Used in [df_md()] if method is "Satterthwaite".
240		#'
241		#' @param object (`mmrm`)\cr the MMRM fit.
242		#' @param contrast (`matrix`)\cr numeric contrast matrix, if given a `numeric`
243		#' then this is coerced to a row vector. Note that this should not include
244		#' elements for singular coefficient estimates, i.e. only refer to the
245		#' actually estimated coefficients.
246		#'
247		#' @return List with `num_df`, `denom_df`, `f_stat` and `p_val` (2-sided p-value).
248		#' @keywords internal
249		h_df_md_sat <- function(object, contrast) {
250	206x	assert_class(object, "mmrm")
251	206x	assert_matrix(
252	206x	contrast,
253	206x	mode = "numeric",
254	206x	any.missing = FALSE,
255	206x	ncols = length(component(object, "beta_est"))
256		)
257		# Early return if we are in the one-dimensional case.
258	206x	if (identical(nrow(contrast), 1L)) {
259	159x	return(h_df_md_from_1d(object, contrast))
260		}
261
262	47x	contrast_cov <- h_quad_form_mat(contrast, component(object, "beta_vcov"))
263	47x	eigen_cont_cov <- eigen(contrast_cov)
264	47x	eigen_cont_cov_vctrs <- eigen_cont_cov$vectors
265	47x	eigen_cont_cov_vals <- eigen_cont_cov$values
266
267	47x	eps <- sqrt(.Machine$double.eps)
268	47x	tol <- max(eps * eigen_cont_cov_vals[1], 0)
269	47x	rank_cont_cov <- sum(eigen_cont_cov_vals > tol)
270	47x	assert_number(rank_cont_cov, lower = .Machine$double.xmin)
271	47x	rank_seq <- seq_len(rank_cont_cov)
272	47x	vctrs_cont_prod <- crossprod(eigen_cont_cov_vctrs, contrast)[
273	47x	rank_seq,
274		,
275	47x	drop = FALSE
276		]
277
278		# Early return if rank 1.
279	47x	if (identical(rank_cont_cov, 1L)) {
280	1x	return(h_df_md_from_1d(object, vctrs_cont_prod))
281		}
282
283	46x	t_squared_nums <- drop(vctrs_cont_prod %*% object$beta_est)^2
284	46x	t_squared_denoms <- eigen_cont_cov_vals[rank_seq]
285	46x	t_squared <- t_squared_nums / t_squared_denoms
286	46x	f_stat <- sum(t_squared) / rank_cont_cov
287	46x	t_stat_df_nums <- 2 * eigen_cont_cov_vals^2
288	46x	t_stat_df <- if (identical(object$vcov, "Asymptotic")) {
289	45x	grads_vctrs_cont_prod <- lapply(
290	45x	rank_seq,
291	45x	function(m) {
292	118x	h_gradient(
293	118x	component(object, "jac_list"),
294	118x	contrast = vctrs_cont_prod[m, ]
295		)
296		}
297		)
298	45x	t_stat_df_denoms <- vapply(
299	45x	grads_vctrs_cont_prod,
300	45x	h_quad_form_vec,
301	45x	center = component(object, "theta_vcov"),
302	45x	numeric(1)
303		)
304	45x	t_stat_df_nums / t_stat_df_denoms
305		} else {
306	1x	vapply(
307	1x	rank_seq,
308	1x	function(m) {
309	2x	contrast_matrix <- Matrix::.bdiag(
310	2x	rep(
311	2x	list(vctrs_cont_prod[m, , drop = FALSE]),
312	2x	component(object, "n_subjects")
313		)
314		)
315	2x	contrast_matrix <- as.matrix(contrast_matrix)
316	2x	h_df_1d_sat_empirical(object, contrast_matrix)
317		},
318	1x	FUN.VALUE = 0
319		)
320		}
321	46x	denom_df <- h_md_denom_df(t_stat_df)
322
323	46x	list(
324	46x	num_df = rank_cont_cov,
325	46x	denom_df = denom_df,
326	46x	f_stat = f_stat,
327	46x	p_val = stats::pf(
328	46x	q = f_stat,
329	46x	df1 = rank_cont_cov,
330	46x	df2 = denom_df,
331	46x	lower.tail = FALSE
332		)
333		)
334		}

1		#' Tidying Methods for `mmrm` Objects
2		#'
3		#' @description `r lifecycle::badge("stable")`
4		#'
5		#' These methods tidy the estimates from an `mmrm` object into a
6		#' summary.
7		#'
8		#' @param x (`mmrm`)\cr fitted model.
9		#' @param conf.int (`flag`)\cr if `TRUE` columns for the lower (`conf.low`) and upper bounds
10		#' (`conf.high`) of coefficient estimates are included.
11		#' @param conf.level (`number`)\cr defines the range of the optional confidence internal.
12		#' @param newdata (`data.frame` or `NULL`)\cr optional new data frame.
13		#' @param se_fit (`flag`)\cr whether to return standard errors of fit.
14		#' @param interval (`string`)\cr type of interval calculation.
15		#' @param type.residuals (`string`)\cr passed on to [residuals.mmrm_tmb()].
16		#' @param ... only used by `augment()` to pass arguments to the [predict.mmrm_tmb()] method.
17		#'
18		#' @name mmrm_tidiers
19		#' @aliases mmrm_tidiers
20		#'
21		#' @seealso [`mmrm_methods`], [`mmrm_tmb_methods`] for additional methods.
22		#'
23		#' @examples
24		#' fit <- mmrm(
25		#' formula = FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT \| USUBJID),
26		#' data = fev_data
27		#' )
28		NULL
29
30		#' @describeIn mmrm_tidiers derives tidy `tibble` from an `mmrm` object.
31		#' @exportS3Method
32		#' @examples
33		#' # Applying tidy method to return summary table of covariate estimates.
34		#' fit \|> tidy()
35		#' fit \|> tidy(conf.int = TRUE, conf.level = 0.9)
36		tidy.mmrm <- function(
37		x, # nolint
38		conf.int = FALSE, # nolint
39		conf.level = 0.95, # nolint
40		...
41		) {
42	5x	assert_flag(conf.int)
43	5x	assert_number(conf.level, lower = 0, upper = 1)
44	5x	tbl <- tibble::as_tibble(summary(x)$coefficients, rownames = "term")
45	5x	colnames(tbl) <- c(
46	5x	"term",
47	5x	"estimate",
48	5x	"std.error",
49	5x	"df",
50	5x	"statistic",
51	5x	"p.value"
52		)
53	5x	coefs <- coef(x)
54	5x	if (length(coefs) != nrow(tbl)) {
55	!	coefs <- tibble::enframe(coefs, name = "term", value = "estimate")
56	!	tbl <- merge(coefs, tbl, by = c("term", "estimate"))
57		}
58	5x	if (conf.int) {
59	4x	ci <- h_tbl_confint_terms(x, level = conf.level)
60	4x	tbl <- tibble::as_tibble(merge(tbl, ci, by = "term"))
61		}
62	5x	tbl
63		}
64
65		#' @describeIn mmrm_tidiers derives `glance` `tibble` from an `mmrm` object.
66		#' @exportS3Method
67		#' @examples
68		#' # Applying glance method to return summary table of goodness of fit statistics.
69		#' fit \|> glance()
70		glance.mmrm <- function(x, ...) {
71		# nolint
72	1x	tibble::as_tibble(summary(x)$aic_list)
73		}
74
75		#' @describeIn mmrm_tidiers derives `augment` `tibble` from an `mmrm` object.
76		#' @exportS3Method
77		#' @examples
78		#' # Applying augment method to return merged `tibble` of model data, fitted and residuals.
79		#' fit \|> augment()
80		#' fit \|> augment(interval = "confidence")
81		#' fit \|> augment(type.residuals = "pearson")
82		augment.mmrm <- function(
83		x, # nolint
84		newdata = NULL,
85		interval = c("none", "confidence", "prediction"),
86		se_fit = (interval != "none"),
87		type.residuals = c("response", "pearson", "normalized"), # nolint
88		...
89		) {
90	9x	type.residuals <- match.arg(type.residuals) # nolint
91	9x	resid_df <- NULL
92	9x	if (is.null(newdata)) {
93	4x	newdata <- stats::get_all_vars(x, data = stats::na.omit(x$data))
94	4x	resid_df <- data.frame(
95	4x	.rownames = rownames(newdata),
96	4x	.resid = unname(residuals(x, type = type.residuals))
97		)
98		}
99	9x	interval <- match.arg(interval)
100
101	9x	tbl <- h_newdata_add_pred(
102	9x	x,
103	9x	newdata = newdata,
104	9x	se_fit = se_fit,
105	9x	interval = interval,
106		...
107		)
108	9x	if (!is.null(resid_df)) {
109	4x	tbl <- merge(tbl, resid_df, by = ".rownames")
110	4x	tbl$.rownames <- as.numeric(tbl$.rownames)
111	4x	tbl <- tbl[order(tbl$.rownames), , drop = FALSE]
112		}
113	9x	tibble::as_tibble(tbl)
114		}
115
116		#' Extract `tibble` with Confidence Intervals and Term Names
117		#'
118		#' This is used in [tidy.mmrm()].
119		#'
120		#' @param x (`mmrm`)\cr fit object.
121		#' @param ... passed to [stats::confint()], hence not used at the moment.
122		#'
123		#' @return A `tibble` with `term`, `conf.low`, `conf.high` columns.
124		#'
125		#' @keywords internal
126		h_tbl_confint_terms <- function(x, ...) {
127	8x	df <- stats::confint(x, ...)
128	8x	tbl <- tibble::as_tibble(df, rownames = "term", .name_repair = "minimal")
129	8x	names(tbl) <- c("term", "conf.low", "conf.high")
130	8x	tbl
131		}
132
133		#' Add Prediction Results to New Data
134		#'
135		#' This is used in [augment.mmrm()].
136		#'
137		#' @param x (`mmrm`)\cr fit.
138		#' @param newdata (`data.frame`)\cr data to predict.
139		#' @param se_fit (`flag`)\cr whether to return standard error of prediction,
140		#' can only be used when `interval` is not "none".
141		#' @param interval (`string`)\cr type of interval.
142		#' @param ... passed to [predict.mmrm_tmb()].
143		#'
144		#' @return The `newdata` as a `tibble` with additional columns `.fitted`,
145		#' `.lower`, `.upper` (if interval is not `none`) and `.se.fit` (if `se_fit`
146		#' requested).
147		#'
148		#' @keywords internal
149		h_newdata_add_pred <- function(x, newdata, se_fit, interval, ...) {
150	13x	assert_class(x, "mmrm")
151	13x	assert_data_frame(newdata)
152	13x	assert_flag(se_fit)
153	13x	assert_string(interval)
154	13x	if (interval == "none") {
155	7x	assert_false(se_fit)
156		}
157
158	12x	tbl <- h_df_to_tibble(newdata)
159	12x	pred_results <- predict(
160	12x	x,
161	12x	newdata = newdata,
162	12x	na.action = stats::na.pass,
163	12x	se.fit = se_fit,
164	12x	interval = interval,
165		...
166		)
167	12x	if (interval == "none") {
168	6x	assert_numeric(pred_results)
169	6x	tbl$.fitted <- unname(pred_results)
170		} else {
171	6x	assert_matrix(pred_results)
172	6x	tbl$.fitted <- unname(pred_results[, "fit"])
173	6x	tbl$.lower <- unname(pred_results[, "lwr"])
174	6x	tbl$.upper <- unname(pred_results[, "upr"])
175		}
176	12x	if (se_fit) {
177	5x	tbl$.se.fit <- unname(pred_results[, "se"])
178		}
179	12x	tbl
180		}
181
182		#' Coerce a Data Frame to a `tibble`
183		#'
184		#' This is used in [h_newdata_add_pred()].
185		#'
186		#' @details This is only a thin wrapper around [tibble::as_tibble()], except
187		#' giving a useful error message and it checks for `rownames` and adds them
188		#' as a new column `.rownames` if they are not just a numeric sequence as
189		#' per the [tibble::has_rownames()] decision.
190		#'
191		#' @param data (`data.frame`)\cr what to coerce.
192		#'
193		#' @return The `data` as a `tibble`, potentially with a `.rownames` column.
194		#'
195		#' @keywords internal
196		h_df_to_tibble <- function(data) {
197	15x	tryCatch(tbl <- tibble::as_tibble(data), error = function(cnd) {
198	1x	stop(
199	1x	"Could not coerce data to `tibble`. Try explicitly passing a",
200	1x	"dataset to either the `data` or `newdata` argument.",
201	1x	call. = FALSE
202		)
203		})
204	14x	if (tibble::has_rownames(data)) {
205	5x	tbl <- tibble::add_column(tbl, .rownames = rownames(data), .before = TRUE)
206		}
207	14x	tbl
208		}

1		#' Calculation of Degrees of Freedom for One-Dimensional Contrast
2		#'
3		#' @description `r lifecycle::badge("stable")`
4		#' Calculates the estimate, adjusted standard error, degrees of freedom,
5		#' t statistic and p-value for one-dimensional contrast.
6		#'
7		#' @param object (`mmrm`)\cr the MMRM fit.
8		#' @param contrast (`numeric`)\cr contrast vector. Note that this should not include
9		#' elements for singular coefficient estimates, i.e. only refer to the
10		#' actually estimated coefficients.
11		#' @return List with `est`, `se`, `df`, `t_stat` and `p_val`.
12		#' @export
13		#'
14		#' @examples
15		#' object <- mmrm(
16		#' formula = FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT \| USUBJID),
17		#' data = fev_data
18		#' )
19		#' contrast <- numeric(length(object$beta_est))
20		#' contrast[3] <- 1
21		#' df_1d(object, contrast)
22		df_1d <- function(object, contrast) {
23	338x	assert_class(object, "mmrm")
24	338x	assert_numeric(
25	338x	contrast,
26	338x	len = length(component(object, "beta_est")),
27	338x	any.missing = FALSE
28		)
29	338x	contrast <- as.vector(contrast)
30	338x	if (all(contrast == 0)) {
31	!	return(
32	!	list(
33	!	est = 0,
34	!	se = NA_real_,
35	!	df = NA_real_,
36	!	t_stat = NA_real_,
37	!	p_val = NA_real_
38		)
39		)
40		}
41		switch(
42	338x	object$method,
43	318x	"Satterthwaite" = h_df_1d_sat(object, contrast),
44	19x	"Kenward-Roger" = h_df_1d_kr(object, contrast),
45	!	"Residual" = h_df_1d_res(object, contrast),
46	1x	"Between-Within" = h_df_1d_bw(object, contrast),
47	!	stop("Unrecognized degrees of freedom method: ", object$method)
48		)
49		}
50
51
52		#' Calculation of Degrees of Freedom for Multi-Dimensional Contrast
53		#'
54		#' @description `r lifecycle::badge("stable")`
55		#' Calculates the estimate, standard error, degrees of freedom,
56		#' t statistic and p-value for m-dimensional contrast, depending on the method
57		#' used in [mmrm()].
58		#'
59		#' @param object (`mmrm`)\cr the MMRM fit.
60		#' @param contrast (`matrix`)\cr numeric contrast matrix, if given a `numeric`
61		#' then this is coerced to a row vector. Note that this should not include
62		#' elements for singular coefficient estimates, i.e. only refer to the
63		#' actually estimated coefficients.
64		#'
65		#' @return List with `num_df`, `denom_df`, `f_stat` and `p_val` (2-sided p-value).
66		#' @export
67		#'
68		#' @examples
69		#' object <- mmrm(
70		#' formula = FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT \| USUBJID),
71		#' data = fev_data
72		#' )
73		#' contrast <- matrix(data = 0, nrow = 2, ncol = length(object$beta_est))
74		#' contrast[1, 2] <- contrast[2, 3] <- 1
75		#' df_md(object, contrast)
76		df_md <- function(object, contrast) {
77	250x	assert_class(object, "mmrm")
78	250x	assert_numeric(contrast, any.missing = FALSE)
79	250x	if (!is.matrix(contrast)) {
80	143x	contrast <- matrix(contrast, ncol = length(contrast))
81		}
82	250x	assert_matrix(contrast, ncols = length(component(object, "beta_est")))
83	250x	if (nrow(contrast) == 0 \|\| all(contrast == 0)) {
84	2x	return(
85	2x	list(
86	2x	num_df = 0,
87	2x	denom_df = NA_real_,
88	2x	f_stat = NA_real_,
89	2x	p_val = NA_real_
90		)
91		)
92		}
93		switch(
94	248x	object$method,
95	199x	"Satterthwaite" = h_df_md_sat(object, contrast),
96	48x	"Kenward-Roger" = h_df_md_kr(object, contrast),
97	!	"Residual" = h_df_md_res(object, contrast),
98	1x	"Between-Within" = h_df_md_bw(object, contrast),
99	!	stop("Unrecognized degrees of freedom method: ", object$method)
100		)
101		}
102
103		#' Creating T-Statistic Test Results For One-Dimensional Contrast
104		#'
105		#' @description Creates a list of results for one-dimensional contrasts using
106		#' a t-test statistic and the given degrees of freedom.
107		#'
108		#' @inheritParams df_1d
109		#' @param df (`number`)\cr degrees of freedom for the one-dimensional contrast.
110		#'
111		#' @return List with `est`, `se`, `df`, `t_stat` and `p_val` (2-sided p-value).
112		#'
113		#' @keywords internal
114		h_test_1d <- function(object, contrast, df) {
115	515x	assert_class(object, "mmrm")
116	515x	assert_numeric(contrast, len = length(component(object, "beta_est")))
117	515x	assert_number(df, lower = .Machine$double.xmin)
118
119	515x	est <- sum(contrast * component(object, "beta_est"))
120	515x	var <- h_quad_form_vec(contrast, component(object, "beta_vcov"))
121	515x	se <- sqrt(var)
122	515x	t_stat <- est / se
123	515x	p_val <- 2 * stats::pt(q = abs(t_stat), df = df, lower.tail = FALSE)
124
125	515x	list(
126	515x	est = est,
127	515x	se = se,
128	515x	df = df,
129	515x	t_stat = t_stat,
130	515x	p_val = p_val
131		)
132		}
133
134		#' Creating F- or Chi-squared-statistic Test Results For Multi-Dimensional
135		#' Contrast
136		#'
137		#' @description Creates a list of results for multi-dimensional contrasts using
138		#' an F-test statistic and the given degrees of freedom or a Chi-squared test
139		#' statistic.
140		#'
141		#' @inheritParams df_md
142		#' @param contrast (`matrix`)\cr numeric contrast matrix.
143		#' @param df (`number`)\cr denominator degrees of freedom for the
144		#' multi-dimensional contrast for an F-test. Ignored if `test = "Chisq"`.
145		#' @param f_stat_factor (`number`)\cr optional scaling factor on top of the
146		#' standard F-statistic. Ignored if `test = "Chisq"`.
147		#' @param test (`string`)\cr either `"F"` or `"Chisq"`, specifying the kind of
148		#' test to perform.
149		#'
150		#' @return If `test = "F"`, a list with `num_df`, `denom_df`, `f_stat` and
151		#' `p_val`. If `test = "Chisq"`, a list with `df`, `chisq_stat`, and `p_val`.
152		#' In both cases, p-values are two sided.
153		#'
154		#' @keywords internal
155		h_test_md <- function(
156		object,
157		contrast,
158		df,
159		f_stat_factor = 1,
160		test = c("F", "Chisq")
161		) {
162	92x	test <- match.arg(test)
163	92x	assert_class(object, "mmrm")
164	92x	beta <- component(object, "beta_est")
165	92x	assert_matrix(contrast, ncols = length(beta))
166
167	92x	if (test == "F") {
168	61x	assert_number(df, lower = .Machine$double.xmin)
169	61x	assert_number(f_stat_factor, lower = .Machine$double.xmin)
170		}
171
172	92x	prec_contrast <- solve(h_quad_form_mat(
173	92x	contrast,
174	92x	component(object, "beta_vcov")
175		))
176	92x	contrast_est <- tcrossprod(beta, contrast)
177	92x	chisq_statistic <- as.numeric(h_quad_form_mat(contrast_est, prec_contrast))
178
179	92x	num_df <- nrow(contrast)
180
181	92x	out <-
182	92x	if (test == "F") {
183	61x	f_statistic <- f_stat_factor / num_df * chisq_statistic
184
185	61x	p_val <- stats::pf(
186	61x	q = f_statistic,
187	61x	df1 = num_df,
188	61x	df2 = df,
189	61x	lower.tail = FALSE
190		)
191
192	61x	list(
193	61x	num_df = num_df,
194	61x	denom_df = df,
195	61x	f_stat = f_statistic,
196	61x	p_val = p_val
197		)
198		} else {
199	31x	p_val <- stats::pchisq(chisq_statistic, df = num_df, lower.tail = FALSE)
200	31x	list(df = num_df, chisq_stat = chisq_statistic, p_val = p_val)
201		}
202
203	92x	out
204		}

1		#' Covariance Type Database
2		#'
3		#' An internal constant for covariance type information.
4		#'
5		#' @format A data frame with 5 variables and one record per covariance type:
6		#'
7		#' \describe{
8		#' \item{name}{
9		#' The long-form name of the covariance structure type
10		#' }
11		#' \item{abbr}{
12		#' The abbreviated name of the covariance structure type
13		#' }
14		#' \item{habbr}{
15		#' The abbreviated name of the heterogeneous version of a covariance
16		#' structure type (The abbreviated name (`abbr`) with a trailing `"h"` if
17		#' the structure has a heterogeneous implementation or `NA` otherwise).
18		#' }
19		#' \item{heterogeneous}{
20		#' A logical value indicating whether the covariance structure has a
21		#' heterogeneous counterpart.
22		#' }
23		#' \item{spatial}{
24		#' A logical value indicating whether the covariance structure is spatial.
25		#' }
26		#' }
27		#'
28		#' @keywords internal
29		# nolint start
30		COV_TYPES <- local({
31		# nolint end
32		type <- function(name, abbr, habbr, heterogeneous, spatial) {
33		args <- as.list(match.call()[-1])
34		do.call(data.frame, args)
35		}
36
37		as.data.frame(
38		col.names = names(formals(type)),
39		rbind(
40		type("unstructured", "us", NA, FALSE, FALSE),
41		type("Toeplitz", "toep", "toeph", TRUE, FALSE),
42		type("auto-regressive order one", "ar1", "ar1h", TRUE, FALSE),
43		type("ante-dependence", "ad", "adh", TRUE, FALSE),
44		type("compound symmetry", "cs", "csh", TRUE, FALSE),
45		type("spatial exponential", "sp_exp", NA, FALSE, TRUE)
46		)
47		)
48		})
49
50		#' Covariance Types
51		#'
52		#' @description `r lifecycle::badge("stable")`
53		#'
54		#' @param form (`character`)\cr covariance structure type name form. One or
55		#' more of `"name"`, `"abbr"` (abbreviation), or `"habbr"` (heterogeneous
56		#' abbreviation).
57		#' @param filter (`character`)\cr covariance structure type filter. One or
58		#' more of `"heterogeneous"` or `"spatial"`.
59		#'
60		#' @return A character vector of accepted covariance structure type names and
61		#' abbreviations.
62		#'
63		#' @section Abbreviations for Covariance Structures:
64		#'
65		#' ## Common Covariance Structures:
66		#'
67		#' \tabular{clll}{
68		#'
69		#' \strong{Structure}
70		#' \tab \strong{Description}
71		#' \tab \strong{Parameters}
72		#' \tab \strong{\eqn{(i, j)} element}
73		#' \cr
74		#'
75		#' ad
76		#' \tab Ante-dependence
77		#' \tab \eqn{m}
78		#' \tab \eqn{\sigma^{2}\prod_{k=i}^{j-1}\rho_{k}}
79		#' \cr
80		#'
81		#' adh
82		#' \tab Heterogeneous ante-dependence
83		#' \tab \eqn{2m-1}
84		#' \tab \eqn{\sigma_{i}\sigma_{j}\prod_{k=i}^{j-1}\rho_{k}}
85		#' \cr
86		#'
87		#' ar1
88		#' \tab First-order auto-regressive
89		#' \tab \eqn{2}
90		#' \tab \eqn{\sigma^{2}\rho^{\left \vert {i-j} \right \vert}}
91		#' \cr
92		#'
93		#' ar1h
94		#' \tab Heterogeneous first-order auto-regressive
95		#' \tab \eqn{m+1}
96		#' \tab \eqn{\sigma_{i}\sigma_{j}\rho^{\left \vert {i-j} \right \vert}}
97		#' \cr
98		#'
99		#' cs
100		#' \tab Compound symmetry
101		#' \tab \eqn{2}
102		#' \tab \eqn{\sigma^{2}\left[ \rho I(i \neq j)+I(i=j) \right]}
103		#' \cr
104		#'
105		#' csh
106		#' \tab Heterogeneous compound symmetry
107		#' \tab \eqn{m+1}
108		#' \tab \eqn{\sigma_{i}\sigma_{j}\left[ \rho I(i \neq j)+I(i=j) \right]}
109		#' \cr
110		#'
111		#' toep
112		#' \tab Toeplitz
113		#' \tab \eqn{m}
114		#' \tab \eqn{\sigma_{\left \vert {i-j} \right \vert +1}}
115		#' \cr
116		#'
117		#' toeph
118		#' \tab Heterogeneous Toeplitz
119		#' \tab \eqn{2m-1}
120		#' \tab \eqn{\sigma_{i}\sigma_{j}\rho_{\left \vert {i-j} \right \vert}}
121		#' \cr
122		#'
123		#' us
124		#' \tab Unstructured
125		#' \tab \eqn{m(m+1)/2}
126		#' \tab \eqn{\sigma_{ij}}
127		#'
128		#' }
129		#'
130		#' where \eqn{i} and \eqn{j} denote \eqn{i}-th and \eqn{j}-th time points,
131		#' respectively, out of total \eqn{m} time points, \eqn{1 \leq i, j \leq m}.
132		#'
133		#' @note The ante-dependence covariance structure in this package refers to
134		#' homogeneous ante-dependence, while the ante-dependence covariance structure
135		#' from SAS `PROC MIXED` refers to heterogeneous ante-dependence and the
136		#' homogeneous version is not available in SAS.
137		#'
138		#' @note For all non-spatial covariance structures, the time variable must
139		#' be coded as a factor.
140		#'
141		#' ## Spatial Covariance structures:
142		#'
143		#' \tabular{clll}{
144		#'
145		#' \strong{Structure}
146		#' \tab \strong{Description}
147		#' \tab \strong{Parameters}
148		#' \tab \strong{\eqn{(i, j)} element}
149		#' \cr
150		#'
151		#' sp_exp
152		#' \tab spatial exponential
153		#' \tab \eqn{2}
154		#' \tab \eqn{\sigma^{2}\rho^{-d_{ij}}}
155		#'
156		#' }
157		#'
158		#' where \eqn{d_{ij}} denotes the Euclidean distance between time points
159		#' \eqn{i} and \eqn{j}.
160		#'
161		#' @family covariance types
162		#' @name covariance_types
163		#' @export
164		cov_types <- function(
165		form = c("name", "abbr", "habbr"),
166		filter = c("heterogeneous", "spatial")
167		) {
168	1870x	form <- match.arg(form, several.ok = TRUE)
169	1870x	filter <- if (missing(filter)) c() else match.arg(filter, several.ok = TRUE)
170	1870x	df <- COV_TYPES[form][rowSums(!COV_TYPES[filter]) == 0, ]
171	1870x	Filter(Negate(is.na), unlist(t(df), use.names = FALSE))
172		}
173
174		#' Retrieve Associated Abbreviated Covariance Structure Type Name
175		#'
176		#' @param type (`string`)\cr either a full name or abbreviate covariance
177		#' structure type name to collapse into an abbreviated type.
178		#'
179		#' @return The corresponding abbreviated covariance type name.
180		#'
181		#' @keywords internal
182		cov_type_abbr <- function(type) {
183	335x	row <- which(COV_TYPES == type, arr.ind = TRUE)[, 1]
184	335x	COV_TYPES$abbr[row]
185		}
186
187		#' Retrieve Associated Full Covariance Structure Type Name
188		#'
189		#' @param type (`string`)\cr either a full name or abbreviate covariance
190		#' structure type name to convert to a long-form type.
191		#'
192		#' @return The corresponding abbreviated covariance type name.
193		#'
194		#' @keywords internal
195		cov_type_name <- function(type) {
196	6x	row <- which(COV_TYPES == type, arr.ind = TRUE)[, 1]
197	6x	COV_TYPES$name[row]
198		}
199
200		#' Produce A Covariance Identifier Passing to TMB
201		#'
202		#' @param cov (`cov_struct`)\cr a covariance structure object.
203		#'
204		#' @return A string used for method dispatch when passed to TMB.
205		#'
206		#' @keywords internal
207		tmb_cov_type <- function(cov) {
208	302x	paste0(cov$type, if (cov$heterogeneous) "h")
209		}
210
211		#' Define a Covariance Structure
212		#'
213		#' @description `r lifecycle::badge("stable")`
214		#'
215		#' @param type (`string`)\cr the name of the covariance structure type to use.
216		#' For available options, see `cov_types()`. If a type abbreviation is used
217		#' that implies heterogeneity (e.g. `cph`) and no value is provided to
218		#' `heterogeneous`, then the heterogeneity is derived from the type name.
219		#' @param visits (`character`)\cr a vector of variable names to use for the
220		#' longitudinal terms of the covariance structure. Multiple terms are only
221		#' permitted for the `"spatial"` covariance type.
222		#' @param subject (`string`)\cr the name of the variable that encodes a subject
223		#' identifier.
224		#' @param group (`string`)\cr optionally, the name of the variable that encodes
225		#' a grouping variable for subjects.
226		#' @param heterogeneous (`flag`)\cr
227		#'
228		#' @return A `cov_struct` object.
229		#'
230		#' @examples
231		#' cov_struct("csh", "AVISITN", "USUBJID")
232		#' cov_struct("spatial", c("VISITA", "VISITB"), group = "GRP", subject = "SBJ")
233		#'
234		#' @family covariance types
235		#' @export
236		cov_struct <- function(
237		type = cov_types(),
238		visits,
239		subject,
240		group = character(),
241		heterogeneous = FALSE
242		) {
243		# if heterogeneous isn't provided, derive from provided type
244	332x	if (missing(heterogeneous)) {
245	330x	heterogeneous <- switch(
246	330x	type,
247	330x	toeph = ,
248	330x	ar1h = ,
249	330x	adh = ,
250	330x	csh = TRUE,
251	330x	heterogeneous
252		)
253		}
254
255		# coerce all type options into abbreviated form
256	332x	type <- match.arg(type)
257	331x	type <- cov_type_abbr(type)
258
259	331x	x <- structure(
260	331x	list(
261	331x	type = type,
262	331x	heterogeneous = heterogeneous,
263	331x	visits = visits,
264	331x	subject = subject,
265	331x	group = group
266		),
267	331x	class = c("cov_struct", "mmrm_cov_struct", "list")
268		)
269
270	331x	validate_cov_struct(x)
271		}
272
273		#' Reconcile Possible Covariance Structure Inputs
274		#'
275		#' @inheritParams mmrm
276		#'
277		#' @return The value `covariance` if it's provided or a covariance structure
278		#' derived from the provided `formula` otherwise. An error is raised of both
279		#' are provided.
280		#'
281		#' @keywords internal
282		h_reconcile_cov_struct <- function(formula = NULL, covariance = NULL) {
283	273x	assert_multi_class(covariance, c("formula", "cov_struct"), null.ok = TRUE)
284	273x	assert_formula(formula, null.ok = FALSE)
285	273x	if (inherits(covariance, "formula")) {
286	4x	covariance <- as.cov_struct(covariance)
287		}
288	273x	if (!is.null(covariance) && length(h_extract_covariance_terms(formula)) > 0) {
289	2x	stop(paste0(
290	2x	"Redundant covariance structure definition in `formula` and ",
291	2x	"`covariance` arguments"
292		))
293		}
294
295	271x	if (!is.null(covariance)) {
296	5x	return(covariance)
297		}
298
299	266x	as.cov_struct(formula, warn_partial = FALSE)
300		}
301
302		#' Validate Covariance Structure Data
303		#'
304		#' Run checks against relational integrity of covariance definition
305		#'
306		#' @param x (`cov_struct`)\cr a covariance structure object.
307		#'
308		#' @return `x` if successful, or an error is thrown otherwise.
309		#'
310		#' @keywords internal
311		validate_cov_struct <- function(x) {
312	331x	checks <- checkmate::makeAssertCollection()
313
314	331x	with(x, {
315	331x	assert_character(subject, len = 1, add = checks)
316	331x	assert_logical(heterogeneous, len = 1, add = checks)
317
318	331x	if (length(group) > 1 \|\| length(visits) < 1) {
319	4x	checks$push(
320	4x	"Covariance structure must be of the form `time \| (group /) subject`"
321		)
322		}
323
324	331x	if (!type %in% cov_types(filter = "spatial") && length(visits) > 1) {
325	2x	checks$push(paste(
326	2x	"Non-spatial covariance structures must have a single longitudinal",
327	2x	"variable"
328		))
329		}
330		})
331
332	331x	reportAssertions(checks)
333	325x	x
334		}
335
336		#' Format Covariance Structure Object
337		#'
338		#' @param x (`cov_struct`)\cr a covariance structure object.
339		#' @param ... Additional arguments unused.
340		#'
341		#' @return A formatted string for `x`.
342		#'
343		#' @export
344		format.cov_struct <- function(x, ...) {
345	3x	sprintf(
346	3x	"<covariance structure>\n%s%s:\n\n %s \| %s%s\n",
347	3x	if (x$heterogeneous) "heterogeneous " else "",
348	3x	cov_type_name(x$type),
349	3x	format_symbols(x$visits),
350	3x	if (length(x$group) > 0) paste0(format_symbols(x$group), " / ") else "",
351	3x	format_symbols(x$subject)
352		)
353		}
354
355		#' Print a Covariance Structure Object
356		#'
357		#' @param x (`cov_struct`)\cr a covariance structure object.
358		#' @param ... Additional arguments unused.
359		#'
360		#' @return `x` invisibly.
361		#'
362		#' @export
363		print.cov_struct <- function(x, ...) {
364	3x	cat(format(x, ...), "\n")
365	3x	invisible(x)
366		}
367
368		#' Coerce into a Covariance Structure Definition
369		#'
370		#' @description `r lifecycle::badge("stable")`
371		#'
372		#' @details
373		#' A covariance structure can be parsed from a model definition formula or call.
374		#' Generally, covariance structures defined using non-standard evaluation take
375		#' the following form:
376		#'
377		#' ```
378		#' type( (visit, )* visit \| (group /)? subject )
379		#' ```
380		#'
381		#' For example, formulas may include terms such as
382		#'
383		#' ```r
384		#' us(time \| subject)
385		#' cp(time \| group / subject)
386		#' sp_exp(coord1, coord2 \| group / subject)
387		#' ```
388		#'
389		#' Note that only `sp_exp` (spatial) covariance structures may provide multiple
390		#' coordinates, which identify the Euclidean distance between the time points.
391		#'
392		#' @param x an object from which to derive a covariance structure. See object
393		#' specific sections for details.
394		#' @param warn_partial (`flag`)\cr whether to emit a warning when parts of the
395		#' formula are disregarded.
396		#' @param ... additional arguments unused.
397		#'
398		#' @return A [cov_struct()] object.
399		#'
400		#' @examples
401		#' # provide a covariance structure as a right-sided formula
402		#' as.cov_struct(~ csh(visit \| group / subject))
403		#'
404		#' # when part of a full formula, suppress warnings using `warn_partial = FALSE`
405		#' as.cov_struct(y ~ x + csh(visit \| group / subject), warn_partial = FALSE)
406		#'
407		#' @family covariance types
408		#' @export
409		# nolint start
410		as.cov_struct <- function(x, ...) {
411		# nolint end
412	314x	UseMethod("as.cov_struct")
413		}
414
415		#' @export
416		as.cov_struct.cov_struct <- function(x, ...) {
417	!	x
418		}
419
420		#' @describeIn as.cov_struct
421		#' When provided a formula, any specialized functions are assumed to be
422		#' covariance structure definitions and must follow the form:
423		#'
424		#' ```
425		#' y ~ xs + type( (visit, )* visit \| (group /)? subject )
426		#' ```
427		#'
428		#' Any component on the right hand side of a formula is considered when
429		#' searching for a covariance definition.
430		#'
431		#' @export
432		as.cov_struct.formula <- function(x, warn_partial = TRUE, ...) {
433	314x	x_calls <- h_extract_covariance_terms(x)
434
435	314x	if (length(x_calls) < 1) {
436	4x	stop(
437	4x	"Covariance structure must be specified in formula. ",
438	4x	"Possible covariance structures include: ",
439	4x	paste0(cov_types(c("abbr", "habbr")), collapse = ", ")
440		)
441		}
442
443	310x	if (length(x_calls) > 1) {
444	1x	cov_struct_types <- as.character(lapply(x_calls, `[[`, 1L))
445	1x	stop(
446	1x	"Only one covariance structure can be specified. ",
447	1x	"Currently specified covariance structures are: ",
448	1x	paste0(cov_struct_types, collapse = ", ")
449		)
450		}
451
452		# flatten into list of infix operators, calls and names/atomics
453	309x	x <- flatten_call(x_calls[[1]])
454	309x	type <- as.character(x[[1]])
455	309x	x <- drop_elements(x, 1)
456
457		# take visits until "\|"
458	309x	n <- position_symbol(x, "\|", nomatch = 0)
459	309x	visits <- as.character(utils::head(x, max(n - 1, 0)))
460	309x	x <- drop_elements(x, n)
461
462		# take group until "/"
463	309x	n <- position_symbol(x, "/", nomatch = 0)
464	309x	group <- as.character(utils::head(x, max(n - 1, 0)))
465	309x	x <- drop_elements(x, n)
466
467		# remainder is subject
468	309x	subject <- as.character(x)
469
470	309x	cov_struct(type = type, visits = visits, group = group, subject = subject)
471		}

1		#' Search For the Position of a Symbol
2		#'
3		#' A thin wrapper around [base::Position()] to search through a list of language
4		#' objects, as produced by [flatten_call()] or [flatten_expr()], for the
5		#' presence of a specific symbol.
6		#'
7		#' @param x (`list` of `language`)\cr a list of language objects in which to
8		#' search for a specific symbol.
9		#' @param sym (`name` or `symbol` or `character`)\cr a symbol to search for in
10		#' `x`.
11		#' @param ... Additional arguments passed to `Position()`.
12		#'
13		#' @return The position of the symbol if found, or the `nomatch` value
14		#' otherwise.
15		#'
16		#' @keywords internal
17		position_symbol <- function(x, sym, ...) {
18	622x	Position(function(i) identical(i, as.symbol(sym)), x, ...)
19		}
20
21		#' Flatten Expressions for Non-standard Evaluation
22		#'
23		#' Used primarily to support the parsing of covariance structure definitions
24		#' from formulas, these functions flatten the syntax tree into a hierarchy-less
25		#' grammar, allowing for parsing that doesn't abide by R's native operator
26		#' precedence.
27		#'
28		#' Where \code{1 + 2 \| 3} in R's syntax tree is \code{(\|, (+, 1, 2), 3)},
29		#' flattening it into its visual order produces \code{(1, +, 2, \|, 3)}, which
30		#' makes for more fluent interpretation of non-standard grammar rules used in
31		#' formulas.
32		#'
33		#' @param call,expr (`language`)\cr a language object to flatten.
34		#'
35		#' @return A list of atomic values, symbols, infix operator names and
36		#' subexpressions.
37		#'
38		#' @name flat_expr
39		#' @keywords internal
40		NULL
41
42		#' @describeIn flat_expr
43		#' Flatten a call into a list of names and argument expressions.
44		#'
45		#' The call name and all arguments are flattened into the same list, meaning a
46		#' call of the form \code{sp_exp(a, b, c \| d / e)} produces a list of the form
47		#' \code{(sp_exp, a, b, c, \|, d, /, e)}.
48		#'
49		#' ```
50		#' flatten_call(quote(sp_exp(a, b, c \| d / e)))
51		#' ```
52		#'
53		#' @keywords internal
54		flatten_call <- function(call) {
55	311x	flattened_args <- unlist(lapply(call[-1], flatten_expr))
56	311x	c(flatten_expr(call[[1]]), flattened_args)
57		}
58
59		#' @describeIn flat_expr
60		#' Flatten nested expressions
61		#'
62		#' ```
63		#' flatten_expr(quote(1 + 2 + 3 \| 4))
64		#' ```
65		#'
66		#' @keywords internal
67		flatten_expr <- function(expr) {
68	1385x	if (length(expr) > 1 && is_infix(expr[[1]])) {
69	371x	op <- list(expr[[1]])
70	371x	lhs <- flatten_expr(expr[[2]])
71	371x	rhs <- flatten_expr(expr[[3]])
72	371x	c(lhs, op, rhs)
73		} else {
74	1014x	list(expr)
75		}
76		}
77
78		#' Extract Right-Hand-Side (rhs) from Formula
79		#'
80		#' @param f (`formula`)\cr a formula.
81		#'
82		#' @return A formula without a response, derived from the right-hand-side of the
83		#' formula, `f`.
84		#'
85		#' ```
86		#' formula_rhs(a ~ b + c)
87		#' formula_rhs(~ b + c)
88		#' ```
89		#'
90		#' @keywords internal
91		formula_rhs <- function(f) {
92	330x	if (length(f) == 2) {
93	9x	f
94		} else {
95	321x	f[-2]
96		}
97		}
98
99		#' Test Whether a Symbol is an Infix Operator
100		#'
101		#' @param name (`symbol` or `name` or `string`)\cr a possible reference to an
102		#' infix operator to check.
103		#'
104		#' @return A logical indicating whether the name is the name of an infix
105		#' operator.
106		#'
107		#' ```
108		#' is_infix(as.name("\|"))
109		#' is_infix("\|")
110		#' is_infix("c")
111		#' ```
112		#'
113		#' @keywords internal
114		is_infix <- function(name) {
115	383x	"Ops" %in% methods::getGroup(as.character(name), recursive = TRUE)
116		}
117
118		#' Format Symbol Objects
119		#'
120		#' For printing, variable names are converted to symbols and deparsed to use R's
121		#' built-in formatting of variables that may contain spaces or quote characters.
122		#'
123		#' @param x (`character`) A vector of variable names.
124		#'
125		#' @return A formatted string of comma-separated variables.
126		#'
127		#' @keywords internal
128		format_symbols <- function(x) {
129	12x	paste0(
130	12x	collapse = ", ",
131	12x	lapply(x, function(i) {
132	16x	utils::capture.output(as.symbol(i))
133		})
134		)
135		}

1		# Internal functions used for skipping tests or examples.
2
3		# Predicate whether currently running R version is under development.
4		is_r_devel <- function() {
5	21x	grepl("devel", R.version$status)
6		}
7
8		# Predicate whether currently running on a Linux operating system.
9		is_linux <- function() {
10	1x	tolower(Sys.info()[["sysname"]]) == "linux"
11		}
12
13		# Get the compiler information. Workaround for older R versions
14		# where R_compiled_by() is not available.
15		get_compiler <- function() {
16	3x	r_cmd <- file.path(R.home("bin"), "R")
17	3x	system2(r_cmd, args = "CMD config CC", stdout = TRUE)
18		}
19
20		# Predicate whether currently using a clang compiler.
21		is_using_clang <- function() {
22	2x	grepl("clang", get_compiler())
23		}
24
25		# Predicate whether an R-devel version is running on Linux Fedora or
26		# Debian with a clang compiler.
27		is_r_devel_linux_clang <- function() {
28	20x	is_r_devel() &&
29	20x	is_linux() &&
30	20x	is_using_clang()
31		}

1		#include "testthat-helpers.h"
2		#include "chol_cache.h"
3
4	1x	context("cholesky cache") {
5	3x	test_that("cached cholesky stores result correctly") {
6	3x	vector<double> theta {{log(1.0), log(2.0), 3.0}};
7	1x	auto chol = lower_chol_nonspatial<double>(theta, 2, "us");
8	1x	matrix<double> chol1_expected(2, 2);
9	!	chol1_expected <<
10	1x	1.0, 0.0,
11	1x	6.0, 2.0;
12	2x	std::vector<int> vis{0, 1};
13	1x	matrix<double> dist;
14	1x	expect_equal_matrix(chol.get_chol(vis, dist), chol1_expected);
15	1x	expect_equal_matrix(chol.chols[vis], chol1_expected);
16
17	1x	matrix<double> simga1_expected(2, 2);
18	!	simga1_expected <<
19	1x	1.0, 6.0,
20	1x	6.0, 40.0;
21	1x	expect_equal_matrix(chol.get_sigma(vis, dist), simga1_expected);
22	1x	expect_equal_matrix(chol.sigmas[vis], simga1_expected);
23
24	1x	matrix<double> simga1_inv = chol.get_sigma_inverse(vis, dist);
25	1x	matrix<double> simga1_inv_expected(2, 2);
26	!	simga1_inv_expected <<
27	1x	10.0, -1.5,
28	1x	-1.5, 0.25;
29	1x	expect_equal_matrix(simga1_inv, simga1_inv_expected);
30	1x	expect_equal_matrix(chol.sigmas_inv[vis], simga1_inv_expected);
31
32	1x	matrix<double> chol2_expect(1, 1);
33	1x	chol2_expect << 1.0;
34	2x	std::vector<int> vis2{0};
35	1x	expect_equal_matrix(chol.get_chol(vis2, dist), chol2_expect);
36	1x	expect_equal_matrix(chol.chols[vis2], chol2_expect);
37
38	1x	matrix<double> sigma2_expect(1, 1);
39	1x	sigma2_expect << 1.0;
40	1x	expect_equal_matrix(chol.get_sigma(vis2, dist), sigma2_expect);

1		#include "testthat-helpers.h"
2		#include "covariance.h"
3
4	1x	context("unstructured") {
5	3x	test_that("get_unstructured produces expected result") {
6	2x	vector<double> theta {{log(1.0), log(2.0), 3.0}};
7	1x	matrix<double> result = get_unstructured(theta, 2);
8	1x	matrix<double> expected(2, 2);
9	!	expected <<
10	1x	1.0, 0.0,
11	1x	6.0, 2.0;
12	1x	expect_equal_matrix(result, expected);
13		}
14		}
15
16	3x	context("ante_dependence") {
17	9x	test_that("corr_fun_ante_dependence works as expected") {
18	2x	vector<double> theta {{log(2.0), 1.0, 2.0}};
19	1x	corr_fun_ante_dependence<double> test_fun(theta);
20		expect_equal(test_fun(1, 0), 0.5696763);
21		expect_equal(test_fun(2, 0), 0.4028220);
22		expect_equal(test_fun(3, 0), 0.3602949);
23		expect_equal(test_fun(2, 1), 0.7071068);
24		expect_equal(test_fun(3, 1), 0.6324555);
25		expect_equal(test_fun(3, 2), 0.8944272);
26		}
27
28	9x	test_that("get_ante_dependence produces expected result") {
29	2x	vector<double> theta {{log(2.0), log(2.0), 1.0, 2.0}};
30	1x	matrix<double> result = get_ante_dependence(theta, 4);
31	1x	matrix<double> expected(4, 4);
32	!	expected <<
33	1x	2.0, 0.0, 0.0, 0.0,
34	1x	1.139353, 1.643738, 0.0, 0.0,
35	1x	0.805644, 1.162299, 1.414214, 0.0,
36	1x	0.720590, 1.039591, 1.264911, 0.894427;
37	1x	expect_equal_matrix(result, expected);
38		}
39
40	9x	test_that("get_ante_dependence_heterogeneous produces expected result") {

1		#include "testthat-helpers.h"
2		#include "derivatives.h"
3		#include <iostream>
4
5	2x	context("cho_jacobian") {
6	6x	test_that("cho_jacobian works as expected") {
7	1x	chol_jacobian chol_jac_obj(2, "ar1");
8	2x	vector<double> theta {{1.0, 1.0}};
9	1x	vector<double> result = chol_jac_obj(theta);
10	1x	vector<double> expected(8);
11		// expected obtained from numDeriv::jacobian and ar1 sigma
12	1x	expected << 2.718282, 1.922116, 0, 1.922116, 0.0, 0.9610578, 0.0, -0.9610578;
13	1x	expect_equal_vector(result, expected);
14		}
15	6x	test_that("cho_jacobian's jacabian using autodiff works as expected") {
16	1x	chol_jacobian chol_jac_obj(2, "ar1");
17	2x	vector<double> theta {{1.0, 1.0}};
18	1x	vector<double> result = autodiff::jacobian(chol_jac_obj,theta).vec();
19	1x	vector<double> expected(16);
20		// expected obtained from two numDeriv::jacobian and ar1 sigma
21	1x	expected << 2.718282, 1.9221164, 0, 1.9221167, 0.0, 0.9610586, 0.0, -0.9610586, 0.0, 0.9610586, 0.0, -0.9610586, 0.0, -1.4415871, 0.0, 0.4805284;
22	1x	expect_equal_vector(result, expected);
23		}
24		}
25
26	1x	context("derivatives_nonspatial struct works as expected") {
27	3x	test_that("derivatives_nonspatial struct correct sigma, inverse and derivatives") {
28	3x	vector<double> theta {{1.0, 1.0}};
29	1x	auto mychol = derivatives_nonspatial<double>(theta, 4, "ar1");
30	2x	std::vector<int> v1 {0, 1, 2};
31	2x	std::vector<int> v_full {0, 1, 2, 3};
32	1x	matrix<double> dist(0, 0);
33	1x	auto full_sigma = mychol.get_sigma(v_full, dist);
34	1x	auto part_sigma = mychol.get_sigma(v1, dist);
35	1x	auto full_inverse = matrix<double>(mychol.get_sigma_inverse(v_full, dist));
36	1x	matrix<double> expected_inverse(4, 4);
37		// expected values from R side solve
38	1x	expected_inverse << 0.2706706, -0.191393, 0, 0, -0.191393, 0.4060058, -0.191393, 0, 0, -0.191393, 0.4060058, -0.191393, 0,0,-0.191393, 0.2706706;
39	1x	expect_equal_matrix(expected_inverse, full_inverse);
40

1		#include "testthat-helpers.h"
2		#include "utils.h"
3
4		using namespace Rcpp;
5
6	1x	context("subset_matrix") {
7	3x	test_that("subset_matrix works as expected") {
8	1x	matrix<double> mat(3, 3);
9	!	mat <<
10	1x	1.0, 0.0, 0.5,
11	1x	6.0, 2.0, 1.0,
12	1x	3.0, 0.1, 0.2;
13	2x	std::vector<int> index {1, 0};
14	1x	matrix<double> result1 = subset_matrix(mat, index, index);
15	1x	matrix<double> exp1(2, 2);
16	!	exp1 <<
17	1x	2.0, 6.0,
18	1x	0.0, 1.0;
19	1x	expect_equal_matrix(result1, exp1);
20
21	1x	matrix<double> result2 = subset_matrix(mat, index);
22
23	1x	matrix<double> exp2(2, 3);
24	!	exp2 <<
25	1x	6.0, 2.0, 1.0,
26	1x	1.0, 0.0, 0.5;
27	1x	expect_equal_matrix(result2, exp2);
28		}
29		}
30
31	2x	context("tcrossprod") {
32	6x	test_that("tcrossprod works as expected with complete") {
33	1x	matrix<double> lower_chol(2, 2);
34	!	lower_chol <<
35	1x	1.0, 0.0,
36	1x	6.0, 2.0;
37	1x	matrix<double> result = tcrossprod(lower_chol, true);
38	1x	matrix<double> expected = lower_chol * lower_chol.transpose();
39	1x	expect_equal_matrix(result, expected);
40		}

1		#' Obtain Empirical/Jackknife/Bias-Reduced Covariance
2		#'
3		#' @description Obtain the empirical or Jackknife covariance for \eqn{\beta}.
4		#' Used in `mmrm` fitting if method is "Empirical", "Empirical-Jackknife" or
5		#' "Empirical-Bias-Reduced".
6		#'
7		#' @param tmb_data (`mmrm_tmb_data`)\cr produced by [h_mmrm_tmb_data()].
8		#' @param theta (`numeric`)\cr theta estimate.
9		#' @param beta (`numeric`)\cr beta estimate.
10		#' @param beta_vcov (`matrix`)\cr covariance of beta estimate.
11		#' @param type (`string`)\cr type of empirical method, including "Empirical", "Empirical-Jackknife"
12		#' and "Empirical-Bias-Reduced".
13		#'
14		#' @return Named list with elements:
15		#' - `cov`: `matrix` empirical covariance.
16		#' - `g_mat`: `matrix` to calculate Satterthwaite degrees of freedom.
17		#'
18		#' @note
19		#' This function used to return `df_mat`, which was equivalent to `crossproduct(g_mat)`. However,
20		#' executing the cross product in C++ was a costly matrix multiplication, in particular when the number of coefficients
21		#' and/or the number of subjects was large. Therefore this is now avoided and `g_mat` is returned instead.
22		#'
23		#' @keywords internal
24		h_get_empirical <- function(tmb_data, theta, beta, beta_vcov, type) {
25	39x	assert_class(tmb_data, "mmrm_tmb_data")
26	39x	assert_numeric(theta)
27	39x	n_beta <- ncol(tmb_data$x_matrix)
28	39x	assert_numeric(beta, finite = TRUE, any.missing = FALSE, len = n_beta)
29	39x	assert_matrix(
30	39x	beta_vcov,
31	39x	mode = "numeric",
32	39x	any.missing = FALSE,
33	39x	nrows = n_beta,
34	39x	ncols = n_beta
35		)
36	39x	assert_subset(
37	39x	type,
38	39x	c("Empirical", "Empirical-Jackknife", "Empirical-Bias-Reduced")
39		)
40	39x	.Call(
41	39x	`_mmrm_get_empirical`,
42	39x	PACKAGE = "mmrm",
43	39x	tmb_data,
44	39x	theta,
45	39x	beta,
46	39x	beta_vcov,
47	39x	type
48		)
49		}

1		#' Calculation of Residual Degrees of Freedom for One-Dimensional Contrast
2		#'
3		#' @description Used in [df_1d()] if method is
4		#' "Residual".
5		#'
6		#' @inheritParams h_df_1d_sat
7		#' @inherit h_df_1d_sat return
8		#' @keywords internal
9		h_df_1d_res <- function(object, contrast) {
10	1x	assert_class(object, "mmrm")
11	1x	assert_numeric(contrast, len = length(component(object, "beta_est")))
12
13	1x	df <- component(object, "n_obs") - length(component(object, "beta_est"))
14
15	1x	h_test_1d(object, contrast, df)
16		}
17
18		#' Calculation of Residual Degrees of Freedom for Multi-Dimensional Contrast
19		#'
20		#' @description Used in [df_md()] if method is "Residual".
21		#'
22		#' @inheritParams h_df_md_sat
23		#' @inherit h_df_md_sat return
24		#' @keywords internal
25		h_df_md_res <- function(object, contrast) {
26	1x	assert_class(object, "mmrm")
27	1x	assert_matrix(
28	1x	contrast,
29	1x	mode = "numeric",
30	1x	any.missing = FALSE,
31	1x	ncols = length(component(object, "beta_est"))
32		)
33
34	1x	df <- component(object, "n_obs") - length(component(object, "beta_est"))
35
36	1x	h_test_md(object, contrast, df)
37		}

1		#ifndef CHOL_CACHE_INCLUDED_
2		#define CHOL_CACHE_INCLUDED_
3
4		#include "covariance.h"
5		#include "utils.h"
6
7		// Base class of spatial and non-spatial Cholesky.
8		template <class Type>
9		struct lower_chol_base {
10	5677x	virtual ~lower_chol_base() {}
11		virtual matrix<Type> get_chol(std::vector<int> visits, matrix<Type> dist) = 0;
12		virtual matrix<Type> get_sigma(std::vector<int> visits, matrix<Type> dist) = 0;
13		virtual matrix<Type> get_sigma_inverse(std::vector<int> visits, matrix<Type> dist) = 0;
14		};
15		// Struct to obtain Cholesky for non-spatial.
16		template <class Type>
17		struct lower_chol_nonspatial: virtual lower_chol_base<Type> {
18		std::map<std::vector<int>, matrix<Type>> chols;
19		std::map<std::vector<int>, matrix<Type>> sigmas;
20		std::map<std::vector<int>, matrix<Type>> sigmas_inv;
21		std::string cov_type;
22		int n_visits;
23		std::vector<int> full_visit;
24		int n_theta;
25		vector<Type> theta;
26		matrix<Type> chol_full;
27		matrix<Type> sigma_full;
28		lower_chol_nonspatial() {
29		// This default constructor is needed because the use of `[]` in map.
30		}
31		// Constructor from theta, n_visits and cov_type, and cache full_visits values.
32	32859x	lower_chol_nonspatial(vector<Type> theta, int n_visits, std::string cov_type): cov_type(cov_type), n_visits(n_visits), full_visit(std::vector<int>(n_visits)) {
33	10953x	this->theta = theta;
34	32859x	std::iota(std::begin(this->full_visit), std::end(this->full_visit), 0);
35	10953x	this->n_theta = theta.size();
36	10953x	this->chol_full = get_covariance_lower_chol(this->theta, this->n_visits, this->cov_type);
37	10949x	this->chols[full_visit] = this->chol_full;
38	10949x	this->sigma_full = tcrossprod(this->chol_full, true);
39		}
40	634309x	matrix<Type> get_chol(std::vector<int> visits, matrix<Type> dist) {
41	634309x	auto target = this->chols.find(visits);
42	634309x	if (target != this->chols.end()) {
43	586359x	return target->second;
44		} else {
45	47950x	matrix<Type> cov_i = this->get_sigma(visits, dist);
46	47950x	Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > cov_i_chol(cov_i);
47	47950x	matrix<Type> Li = cov_i_chol.matrixL();
48	47950x	this->chols[visits] = Li;
49	47950x	return Li;
50		}
51		}
52	311436x	matrix<Type> get_sigma(std::vector<int> visits, matrix<Type> dist) {
53	311436x	auto target = this->sigmas.find(visits);
54	311436x	if (target != this->sigmas.end()) {
55	244836x	return target->second;
56		} else {
57	66600x	matrix<Type> ret = subset_matrix<matrix<Type>, vector<int>>(sigma_full, visits, visits);
58	66600x	this->sigmas[visits] = ret;
59	66600x	return ret;
60		}
61		}
62	105018x	matrix<Type> get_sigma_inverse(std::vector<int> visits, matrix<Type> dist) {
63	105018x	auto target = this->sigmas_inv.find(visits);
64	105018x	if (target != this->sigmas_inv.end()) {
65	91497x	return target->second;
66		} else {
67	13521x	matrix<Type> ret = this->get_sigma(visits, dist).inverse();
68	13521x	this->sigmas_inv[visits] = ret;
69	13521x	return ret;
70		}
71		}
72		};
73
74
75		// Struct to obtain Cholesky for spatial exponential.
76		template <class Type>
77		struct lower_chol_spatial: virtual lower_chol_base<Type> {
78		vector<Type> theta;
79		std::string cov_type;
80		lower_chol_spatial() {
81		// This default constructor is needed because the use of `[]` in map.
82		}
83		// Constructor from theta. For now the cholesky does not need to be cached.
84	188x	lower_chol_spatial(vector<Type> theta, std::string cov_type): theta(theta), cov_type(cov_type) {
85		}
86	23736x	matrix<Type> get_chol(std::vector<int> visits, matrix<Type> dist) {
87	23736x	return get_spatial_covariance_lower_chol(this->theta, dist, this->cov_type);
88		}
89	7890x	matrix<Type> get_sigma(std::vector<int> visits, matrix<Type> dist) {
90	7890x	return tcrossprod(this->get_chol(visits, dist), true);
91		}
92	2956x	matrix<Type> get_sigma_inverse(std::vector<int> visits, matrix<Type> dist) {
93	2956x	return this->get_sigma(visits, dist).inverse();
94		}
95		};
96
97		template <class T, class Base, class D1, class D2>
98		struct cache_obj {
99		std::map<int, std::shared_ptr<Base>> cache;
100		int n_groups;
101		bool is_spatial;
102		int n_visits;
103	10225x	cache_obj(vector<T> theta, int n_groups, bool is_spatial, std::string cov_type, int n_visits): n_groups(n_groups), is_spatial(is_spatial), n_visits(n_visits) {
104		// Get number of variance parameters for one group.
105	10225x	int theta_one_group_size = theta.size() / n_groups;
106	21358x	for (int r = 0; r < n_groups; r++) {
107		// Use unique pointers here to better manage resource.
108	11137x	if (is_spatial) {
109	187x	this->cache[r] = std::make_shared<D1>(theta.segment(r * theta_one_group_size, theta_one_group_size), cov_type);
110		} else {
111	10950x	this->cache[r] = std::make_shared<D2>(theta.segment(r * theta_one_group_size, theta_one_group_size), n_visits, cov_type);
112		}
113		}
114		}
115		};
116
117		template <class Type>
118		struct chol_cache_groups: cache_obj<Type, lower_chol_base<Type>, lower_chol_spatial<Type>, lower_chol_nonspatial<Type>> {
119	10030x	chol_cache_groups(vector<Type> theta, int n_groups, bool is_spatial, std::string cov_type, int n_visits): cache_obj<Type, lower_chol_base<Type>, lower_chol_spatial<Type>, lower_chol_nonspatial<Type>>(theta, n_groups, is_spatial, cov_type, n_visits) {
120
121		}
122		// Return covariance lower Cholesky factor from lower_chol_base objects.
123		// For non-spatial return for full visits, for spatial return on two points that the distance is 1.
124	6616x	matrix<Type> get_default_chol() {
125	13232x	std::vector<int> visit(this->n_visits);
126	6616x	std::iota(std::begin(visit), std::end(visit), 0);
127	6616x	matrix<Type> dist(2, 2);
128	6616x	dist << 0, 1, 1, 0;
129	6616x	int dim = this->is_spatial?2:this->n_visits;
130	6616x	matrix<Type> covariance_lower_chol = matrix<Type>::Zero(dim * this->n_groups, dim);
131	13916x	for (int r = 0; r < this->n_groups; r++) {
132	7300x	covariance_lower_chol.block(r * dim, 0, dim, dim) = this->cache[r]->get_chol(visit, dist);
133		}
134	13232x	return covariance_lower_chol;
135		}
136		};
137
138		#endif

1		#ifndef COV_INCLUDED_
2		#define COV_INCLUDED_
3
4		#include "utils.h"
5
6		// Unstructured covariance:
7		// Cholesky factor.
8		template <class T>
9	9064x	matrix<T> get_unstructured(const vector<T>& theta, int n_visits) {
10	9064x	vector<T> sd_values = exp(theta.head(n_visits));
11	9064x	vector<T> lower_tri_chol_values = theta.tail(theta.size() - n_visits);
12	9064x	matrix<T> covariance_lower_chol = matrix<T>::Zero(n_visits, n_visits);
13	9064x	int k = 0;
14	45204x	for(int i = 0; i < n_visits; i++) {
15	36140x	covariance_lower_chol(i, i) = sd_values(i);
16	90272x	for(int j = 0; j < i; j++){
17	54132x	covariance_lower_chol(i, j) = sd_values(i) * lower_tri_chol_values(k++);
18		}
19		}
20	18128x	return covariance_lower_chol;
21		}
22
23		// Ante-dependence:
24
25		// Correlation function.
26		template <class T>
27		struct corr_fun_ante_dependence : generic_corr_fun<T> {
28		using generic_corr_fun<T>::generic_corr_fun;
29	2256x	const T operator() (int i, int j) const {
30	2256x	return this->corr_values.segment(j, i - j).prod();
31		}
32		};
33		// Homogeneous Ante-dependence Cholesky factor.
34		template <class T>
35	158x	matrix<T> get_ante_dependence(const vector<T>& theta, int n_visits) {
36	158x	T const_sd = exp(theta(0));
37	158x	corr_fun_ante_dependence<T> fun(theta.tail(n_visits - 1));
38	158x	matrix<T> ad_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
39	316x	return const_sd * ad_cor_mat_chol;
40		}
41		// Heterogeneous Ante-dependence Cholesky factor.
42		template <class T>
43	238x	matrix<T> get_ante_dependence_heterogeneous(const vector<T>& theta, int n_visits) {
44	238x	vector<T> sd_values = exp(theta.head(n_visits));
45	238x	corr_fun_ante_dependence<T> fun(theta.tail(n_visits - 1));
46	476x	return get_heterogeneous_cov(sd_values, fun);
47		}
48
49		// Toeplitz:
50
51		// Correlation function.
52		template <class T>
53		struct corr_fun_toeplitz : generic_corr_fun<T> {
54		using generic_corr_fun<T>::generic_corr_fun;
55	2538x	const T operator() (int i, int j) const {
56	2538x	int index = (i - j) - 1; // Note: We need to start at 0.
57	2538x	return this->corr_values(index);
58		}
59		};
60		// Homogeneous Toeplitz Cholesky factor.
61		template <class T>
62	208x	matrix<T> get_toeplitz(const vector<T>& theta, int n_visits) {
63	208x	T const_sd = exp(theta(0));
64	208x	corr_fun_toeplitz<T> fun(theta.tail(n_visits - 1));
65	208x	matrix<T> toep_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
66	416x	return const_sd * toep_cor_mat_chol;
67		}
68		// Heterogeneous Toeplitz Cholesky factor.
69		template <class T>
70	214x	matrix<T> get_toeplitz_heterogeneous(const vector<T>& theta, int n_visits) {
71	214x	vector<T> sd_values = exp(theta.head(n_visits));
72	214x	corr_fun_toeplitz<T> fun(theta.tail(n_visits - 1));
73	428x	return get_heterogeneous_cov(sd_values, fun);
74		}
75
76		// Autoregressive:
77
78		// Correlation function.
79		template <class T>
80		struct corr_fun_autoregressive : generic_corr_fun<T> {
81		using generic_corr_fun<T>::generic_corr_fun;
82	11524x	const T operator() (int i, int j) const {
83	11524x	T diff = T((i - j) * 1.0);
84	14752x	return pow(this->corr_values(0), diff); // rho^{\|i-j\|}
85		}
86		};
87		// Homogeneous autoregressive Cholesky factor.
88		template <class T>
89	1702x	matrix<T> get_auto_regressive(const vector<T>& theta, int n_visits) {
90	1702x	T const_sd = exp(theta(0));
91	1702x	corr_fun_autoregressive<T> fun(theta.tail(1));
92	1702x	matrix<T> ar1_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
93	3404x	return const_sd * ar1_cor_mat_chol;
94		}
95		// Heterogeneous autoregressive Cholesky factor.
96		template <class T>
97	214x	matrix<T> get_auto_regressive_heterogeneous(const vector<T>& theta, int n_visits) {
98	214x	vector<T> sd_values = exp(theta.head(n_visits));
99	214x	corr_fun_autoregressive<T> fun(theta.tail(1));
100	428x	return get_heterogeneous_cov(sd_values, fun);
101		}
102
103		// Compound symmetry:
104
105		// Correlation function.
106		template <class T>
107		struct corr_fun_compound_symmetry {
108		const vector<T> corr_values;
109		// Custom constructor here because we need the special mapping from theta to rho.
110	616x	corr_fun_compound_symmetry(const vector<T>& theta, int n_visits) :
111	616x	corr_values(map_to_cs_cor(theta, n_visits)) {}
112	3636x	const T operator() (int i, int j) const {
113	3636x	return this->corr_values(0); // rho (constant)
114		}
115		};
116		// Homogeneous compound symmetry Cholesky factor.
117		template <class T>
118	350x	matrix<T> get_compound_symmetry(const vector<T>& theta, int n_visits) {
119	350x	T const_sd = exp(theta(0));
120	350x	corr_fun_compound_symmetry<T> fun(theta.tail(1), n_visits);
121	350x	matrix<T> cs_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
122	700x	return const_sd * cs_cor_mat_chol;
123		}
124		// Heterogeneous compound symmetry Cholesky factor.
125		template <class T>
126	262x	matrix<T> get_compound_symmetry_heterogeneous(const vector<T>& theta, int n_visits) {
127	262x	vector<T> sd_values = exp(theta.head(n_visits));
128	262x	corr_fun_compound_symmetry<T> fun(theta.tail(1), n_visits);
129	524x	return get_heterogeneous_cov(sd_values, fun);
130		}
131
132		// Spatial Exponential Cholesky factor.
133		template <class T>
134	23736x	matrix<T> get_spatial_exponential(const vector<T>& theta, const matrix<T>& distance) {
135	23736x	T const_sd = exp(theta(0));
136	23736x	T rho = invlogit(theta(1));
137	23736x	matrix<T> expdist = exp(distance.array() * log(rho));
138	23736x	matrix<T> result = expdist * const_sd;
139	23736x	Eigen::LLT<Eigen::Matrix<T,Eigen::Dynamic,Eigen::Dynamic> > cov_i_chol(result);
140	47472x	return cov_i_chol.matrixL();
141		}
142
143		// Creates a new correlation object dynamically.
144		template <class T>
145	12394x	matrix<T> get_covariance_lower_chol(const vector<T>& theta, int n_visits, std::string cov_type) {
146	12394x	matrix<T> result;
147
148	12394x	if (cov_type == "us") {
149	9060x	result = get_unstructured<T>(theta, n_visits);
150	3334x	} else if (cov_type == "toep") {
151	206x	result = get_toeplitz<T>(theta, n_visits);
152	3128x	} else if (cov_type == "toeph") {
153	212x	result = get_toeplitz_heterogeneous<T>(theta, n_visits);
154	2916x	} else if (cov_type == "ar1") {
155	1700x	result = get_auto_regressive<T>(theta, n_visits);
156	1216x	} else if (cov_type == "ar1h") {
157	212x	result = get_auto_regressive_heterogeneous<T>(theta, n_visits);
158	1004x	} else if (cov_type == "ad") {
159	156x	result = get_ante_dependence<T>(theta, n_visits);
160	848x	} else if (cov_type == "adh") {
161	236x	result = get_ante_dependence_heterogeneous<T>(theta, n_visits);
162	612x	} else if (cov_type == "cs") {
163	348x	result = get_compound_symmetry<T>(theta, n_visits);
164	264x	} else if (cov_type == "csh") {
165	260x	result = get_compound_symmetry_heterogeneous<T>(theta, n_visits);
166		} else {
167	2x	Rf_error("%s", ("Unknown covariance type '" + cov_type + "'.").c_str());
168		}
169
170	12390x	return result;
171		}
172
173		// Creates a new spatial covariance cholesky.
174		template <class T>
175	23736x	matrix<T> get_spatial_covariance_lower_chol(const vector<T>& theta, const matrix<T>& distance, std::string cov_type) {
176	23736x	matrix<T> result;
177	23736x	if (cov_type == "sp_exp") {
178	23736x	result = get_spatial_exponential<T>(theta, distance);
179		} else {
180	!	Rf_error("%s", ("Unknown spatial covariance type '" + cov_type + "'.").c_str());
181		}
182	23736x	return result;
183		}
184
185		#endif

1		#ifndef DERIVATIVE_INCLUDED_
2		#define DERIVATIVE_INCLUDED_
3
4		#include "chol_cache.h"
5
6		using namespace Rcpp;
7		using std::string;
8		// Struct chol to obtain the cholesky factor given theta.
9		// The reason to have it is that we need a functor that need only theta to
10		// obtain the derivatives from autodiff.
11		// Only non-spatial covariance structure here.
12		struct chol {
13		int dim_cov_mat;
14		string cov_type;
15	412x	chol(int dim, string cov): dim_cov_mat(dim), cov_type(cov) {};
16		template <class T>
17	1234x	vector<T> operator() (vector<T> &theta) {
18	1234x	return get_covariance_lower_chol(theta, this->dim_cov_mat, this->cov_type).vec();
19		}
20		};
21		// Struct chol_jacobian that has jacobian of the cholesky factor given theta.
22		// The reason to have it is that we need hessian so we use jacobian twice.
23		struct chol_jacobian {
24		int dim_cov_mat;
25		string cov_type;
26		chol mychol;
27	207x	chol_jacobian(int dim, string cov): dim_cov_mat(dim), cov_type(cov), mychol(dim, cov) {};
28		template<class T>
29	207x	vector<T> operator() (vector<T> &theta) {
30	207x	return autodiff::jacobian(this->mychol, theta).vec();
31		}
32		};
33
34		// Template function to obtain derivatives from visits, cov_type and theta.
35		// Basically this is calculating the derivatives for the sigma
36		// from the derivatives for the cholesky factor.
37		template <class Type>
38	205x	std::map<std::string, matrix<Type>> derivatives(int n_visits, std::string cov_type, vector<Type> theta) {
39	205x	std::map<std::string, matrix<Type>> ret;
40	205x	chol chol_obj(n_visits, cov_type);
41	205x	chol_jacobian chol_jac_obj(n_visits, cov_type);
42	205x	matrix<Type> l = chol_obj(theta).matrix();
43	205x	l.resize(n_visits, n_visits);
44	205x	vector<Type> chol_d1_vec = autodiff::jacobian(chol_obj, theta).vec(); // chol_d1_vec is (dim * dim * l_theta)
45	205x	vector<Type> chol_d2_vec = autodiff::jacobian(chol_jac_obj, theta).vec(); // chol_d2_vec is (dim * dim * l_theta * l_theta)
46	205x	matrix<Type> ret_d1 = matrix<Type>(n_visits * theta.size(), n_visits);
47	205x	matrix<Type> ret_d2 = matrix<Type>(n_visits * theta.size() * theta.size(), n_visits);
48	205x	int n_visits_sq = n_visits * n_visits;
49	1358x	for (int i = 0; i < theta.size(); i++) {
50	1153x	matrix<Type> ld1 = chol_d1_vec.segment(i * n_visits_sq, n_visits_sq).matrix();
51	1153x	ld1.resize(n_visits, n_visits);
52	1153x	matrix<Type> ld1_lt = ld1 * l.transpose();
53	1153x	auto sigma_d1_i = ld1_lt + ld1_lt.transpose();
54	1153x	ret_d1.block(i * n_visits, 0, n_visits, n_visits) = sigma_d1_i;
55	10334x	for (int j = 0; j < theta.size(); j++) {
56	9181x	matrix<Type> ld2 = chol_d2_vec.segment( (j * theta.size() + i) * n_visits_sq, n_visits_sq).matrix();
57	9181x	matrix<Type> ld1_j = chol_d1_vec.segment(j * n_visits_sq, n_visits_sq).matrix();
58	9181x	ld2.resize(n_visits, n_visits);
59	9181x	ld1_j.resize(n_visits, n_visits);
60	9181x	auto ld2_lt = ld2 * l.transpose();
61	9181x	auto ld1_ld1j = ld1 * ld1_j.transpose();
62	9181x	auto sigma_d2_ij = ld2_lt + ld2_lt.transpose() + ld1_ld1j + ld1_ld1j.transpose();
63	9181x	ret_d2.block((i * theta.size() + j) * n_visits, 0, n_visits, n_visits) = sigma_d2_ij;
64		}
65		}
66	410x	ret["derivative1"] = ret_d1;
67	205x	ret["derivative2"] = ret_d2;
68	410x	return ret;
69		}
70		// Base class of spatial and non-spatial derivatives.
71		template <class Type>
72		struct derivatives_base: virtual lower_chol_base<Type> {
73		virtual matrix<Type> get_inverse_chol(std::vector<int> visits, matrix<Type> dist) = 0;
74		virtual matrix<Type> get_sigma_derivative1(std::vector<int> visits, matrix<Type> dist) = 0;
75		virtual matrix<Type> get_sigma_derivative2(std::vector<int> visits, matrix<Type> dist) = 0;
76		virtual matrix<Type> get_inverse_derivative(std::vector<int> visits, matrix<Type> dist) = 0;
77		// Create virtual destructor to avoid the default desctructor being called.
78	217x	virtual ~derivatives_base() {};
79		};
80
81		// Struct derivatives_nonspatial is created to get the derivatives with cache.
82		// The main reason to have it is that we nearly always have duplicated visits
83		// and the inverse of a matrix is calculation expensive. In addition, we can save
84		// the resource needed for select matrix calculations.
85		template <class Type>
86		struct derivatives_nonspatial: public lower_chol_nonspatial<Type>, virtual derivatives_base<Type> {
87		std::map<std::vector<int>, matrix<Type>> inverse_chol_cache;
88		std::map<std::vector<int>, matrix<Type>> sigmad1_cache;
89		std::map<std::vector<int>, matrix<Type>> sigmad2_cache;
90		std::map<std::vector<int>, matrix<Type>> sigma_inverse_d1_cache;
91		derivatives_nonspatial() {
92		// This default constructor is needed because the use of `[]` in map.
93		}
94		// Constructor from theta, n_visits and cov_type, and cache full_visits values.
95	205x	derivatives_nonspatial(vector<Type> theta, int n_visits, std::string cov_type): lower_chol_nonspatial<Type>(theta, n_visits, cov_type) {
96	205x	std::map<std::string, tmbutils::matrix<Type>> allret = derivatives<Type>(this->n_visits, this->cov_type, this->theta);
97	410x	matrix<Type> sigma_d1 = allret["derivative1"];
98	205x	matrix<Type> sigma_d2 = allret["derivative2"];
99	205x	this->sigmad1_cache[this->full_visit] = sigma_d1;
100	205x	this->sigmad2_cache[this->full_visit] = sigma_d2;
101		}
102		// Cache and return the first order derivatives using select matrix.
103	10875x	matrix<Type> get_sigma_derivative1(std::vector<int> visits, matrix<Type> dist) override {
104	10875x	auto target = this->sigmad1_cache.find(visits);
105	10875x	if (target != this->sigmad1_cache.end()) {
106	8700x	return target->second;
107		} else {
108	2175x	int n_visits_i = visits.size();
109	2175x	matrix<Type> ret = matrix<Type>(this->n_theta * n_visits_i, n_visits_i);
110	14848x	for (int i = 0; i < this->n_theta; i++) {
111	12673x	ret.block(i * n_visits_i, 0, n_visits_i, n_visits_i) = subset_matrix<matrix<Type>, vector<int>>(this->sigmad1_cache[this->full_visit].block(i * this->n_visits, 0, this->n_visits, this->n_visits), visits, visits);
112		}
113	2175x	this->sigmad1_cache[visits] = ret;
114	2175x	return ret;
115		}
116		}
117		// Cache and return the second order derivatives using select matrix.
118	8534x	matrix<Type> get_sigma_derivative2(std::vector<int> visits, matrix<Type> dist) override {
119	8534x	auto target = this->sigmad2_cache.find(visits);
120	8534x	if (target != this->sigmad2_cache.end()) {
121	7787x	return target->second;
122		} else {
123	747x	int n_visits_i = visits.size();
124	747x	int theta_sq = this->n_theta * this->n_theta;
125	747x	matrix<Type> ret = matrix<Type>(theta_sq * n_visits_i, n_visits_i);
126	23747x	for (int i = 0; i < theta_sq; i++) {
127	23000x	ret.block(i * n_visits_i, 0, n_visits_i, n_visits_i) = subset_matrix<matrix<Type>, vector<int>>(this->sigmad2_cache[this->full_visit].block(i * this->n_visits, 0, this->n_visits, this->n_visits), visits, visits);
128		}
129	747x	this->sigmad2_cache[visits] = ret;
130	747x	return ret;
131		}
132		}
133		// Cache and return the lower cholesky factor of inverse of sigma using select matrix.
134	7289x	matrix<Type> get_inverse_chol(std::vector<int> visits, matrix<Type> dist) override {
135	7289x	auto target = this->inverse_chol_cache.find(visits);
136	7289x	if (target != this->inverse_chol_cache.end()) {
137	6734x	return target->second;
138		} else {
139	555x	matrix<Type> sigmainv = this->get_sigma_inverse(visits, dist);
140	555x	Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > sigma_inv_chol(sigmainv);
141	555x	matrix<Type> Li = sigma_inv_chol.matrixL();
142	555x	this->inverse_chol_cache[visits] = Li;
143	555x	return Li;
144		}
145		}
146		// Cache and return the first order derivatives of inverse of sigma using select matrix.
147	28630x	matrix<Type> get_inverse_derivative(std::vector<int> visits, matrix<Type> dist) override {
148	28630x	auto target = this->sigma_inverse_d1_cache.find(visits);
149	28630x	if (target != this->sigma_inverse_d1_cache.end()) {
150	26289x	return target->second;
151		} else {
152	2341x	auto sigma_d1 = this->get_sigma_derivative1(visits, dist);
153	2341x	matrix<Type> sigma_inv_d1(sigma_d1.rows(), sigma_d1.cols());
154	2341x	int n_visits_i = visits.size();
155	2341x	auto sigma_inv = this->get_sigma_inverse(visits, dist);
156	15997x	for (int r = 0; r < this->n_theta; r++) {
157	13656x	sigma_inv_d1.block(r * n_visits_i, 0, n_visits_i, n_visits_i) = - sigma_inv * sigma_d1.block(r * n_visits_i, 0, n_visits_i, n_visits_i) *sigma_inv;
158		}
159	2341x	this->sigma_inverse_d1_cache[visits] = sigma_inv_d1;
160	2341x	return sigma_inv_d1;
161		}
162		}
163		};
164
165		// derivatives_sp_exp struct is created to obtain the exact derivatives of spatial exponential
166		// covariance structure, and its inverse.
167		// No caching is used because the distance can be hardly the same for spatial covariance
168		// structures.
169		template <class Type>
170		struct derivatives_sp_exp: public lower_chol_spatial<Type>, virtual derivatives_base<Type> {
171		Type const_sd;
172		Type rho;
173		Type logrho;
174		derivatives_sp_exp() {
175		// This default constructor is needed because the use of `[]` in maps.
176		}
177		// Initialize the theta values; the reason to have theta is that for a fit, the theta
178		// is the same for all subjects, while the distance between each visits for each subject
179		// can be different.
180	12x	derivatives_sp_exp(vector<Type> theta, std::string cov_type): lower_chol_spatial<Type>(theta, cov_type) ,const_sd(exp(theta(0))), rho(invlogit(theta(1))) {
181	12x	this->logrho = log(this->rho);
182		}
183		// Obtain first order derivatives
184	2562x	matrix<Type> get_sigma_derivative1(std::vector<int> visits, matrix<Type> dist) override {
185	2562x	matrix<Type> ret(2 * dist.rows(), dist.cols());
186		// partial sigma / partial theta_1 = sigma.
187	2562x	auto sigma = this->get_sigma(visits, dist);
188	2562x	ret.block(0, 0, dist.rows(), dist.cols()) = sigma;
189	2562x	ret.block(dist.rows(), 0, dist.rows(), dist.cols()) = sigma.array() * dist.array() * (1 - this->rho);
190	5124x	return ret;
191		}
192		// Obtain second order derivatives.
193	986x	matrix<Type> get_sigma_derivative2(std::vector<int> visits, matrix<Type> dist) override {
194	986x	matrix<Type> ret(4 * dist.rows(), dist.cols());
195	986x	auto sigma = this->get_sigma(visits, dist);
196	986x	ret.block(0, 0, dist.rows(), dist.cols()) = sigma;
197	986x	Type rho_r = 1 - this->rho;
198	986x	auto dtheta1dtheta2 = sigma.array() * dist.array() * rho_r;
199	986x	ret.block(dist.rows(), 0, dist.rows(), dist.cols()) = dtheta1dtheta2;
200	986x	ret.block(dist.rows() * 2, 0, dist.rows(), dist.cols()) = dtheta1dtheta2;
201	986x	matrix<Type> dtheta2s = dtheta1dtheta2 * (dist.array() * rho_r - this->rho);
202	986x	ret.block(dist.rows() * 3, 0, dist.rows(), dist.cols()) = dtheta2s;
203	1972x	return ret;
204		}
205		// Obtain the lower cholesky factor of inverse of sigma using select matrix.
206	394x	matrix<Type> get_inverse_chol(std::vector<int> visits, matrix<Type> dist) override {
207	394x	auto sigmainv = this->get_sigma_inverse(visits, dist);
208	394x	Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > sigma_inv_chol(sigmainv);
209	394x	matrix<Type> Li = sigma_inv_chol.matrixL();
210	788x	return Li;
211		}
212		// Obtain first order derivatives for inverse of sigma.
213	1576x	matrix<Type> get_inverse_derivative(std::vector<int> visits, matrix<Type> dist) override {
214	1576x	matrix<Type> sigma_inv_d1 = matrix<Type>::Zero(2 * dist.rows(), dist.cols());
215	1576x	auto sigma_inv = this->get_sigma_inverse(visits, dist);
216	1576x	auto sigma_d1 = this->get_sigma_derivative1(visits, dist);
217	4728x	for (int r = 0; r < 2; r++) {
218	3152x	sigma_inv_d1.block(r * dist.rows(), 0, dist.rows(), dist.cols()) = - sigma_inv * sigma_d1.block(r * dist.rows(), 0, dist.rows(), dist.cols()) *sigma_inv;
219		}
220	3152x	return sigma_inv_d1;
221		}
222		};
223
224		#endif

1		#ifndef UTILS_INCLUDED_
2		#define UTILS_INCLUDED_
3		#include <Rcpp.h>
4		#define INCLUDE_RCPP
5		#include "tmb_includes.h"
6
7		#define as_num_matrix_tmb as_matrix<matrix<double>, NumericMatrix>
8		#define as_num_matrix_rcpp as_matrix<NumericMatrix, matrix<double>>
9		#define as_num_vector_tmb as_vector<vector<double>, NumericVector>
10		#define as_num_vector_rcpp as_vector<NumericVector, vector<double>>
11
12		// Obtain submatrix from index
13
14		template <typename T1, typename T2>
15	224104x	T1 subset_matrix(T1 input, T2 index1, T2 index2) {
16		#if EIGEN_VERSION_AT_LEAST(3,4,0)
17	224104x	T1 ret = input(index1, index2);
18		#else
19		T1 ret(index1.size(), index2.size());
20		for (decltype(index1.size()) i = 0; i < index1.size(); i++) {
21		for (decltype(index2.size()) j = 0; j < index2.size(); j++) {
22		ret(i, j) = input(index1[i], index2[j]);
23		}
24		}
25		#endif
26	224104x	return ret;
27		}
28
29		template <typename T1, typename T2>
30	242229x	T1 subset_matrix(T1 input, T2 index1) {
31		#if EIGEN_VERSION_AT_LEAST(3,4,0)
32	242229x	T1 ret = input(index1, Eigen::all);
33		#else
34		T1 ret(index1.size(), input.cols());
35		for (decltype(index1.size()) i = 0; i < index1.size(); i++) {
36		for (int j = 0; j < input.cols(); j++) {
37		ret(i, j) = input(index1[i], j);
38		}
39		}
40		#endif
41	242229x	return ret;
42		}
43
44
45		// Conversion from Rcpp vector/matrix to eigen vector/matrix
46		template <typename T1, typename T2>
47	488336x	T1 as_vector(T2 input) {
48	488336x	T1 ret(input.size());
49	1458558x	for (int i = 0; i < input.size(); i++) {
50	970222x	ret(i) = input(i);
51		}
52	488336x	return ret;
53		}
54
55		template <typename T1, typename T2>
56	210790x	T1 as_matrix(T2 input) {
57	210790x	T1 ret(input.rows(), input.cols());
58	3025392x	for (int i = 0; i < input.rows(); i++) {
59	31959010x	for (int j = 0; j < input.cols(); j++) {
60	29144408x	ret(i,j) = input(i,j);
61		}
62		}
63	210790x	return ret;
64		}
65
66		template <typename T>
67	726686x	T segment(T input, int start, int n) {
68	726686x	T ret(n);
69	3630604x	for (int i = 0, j = start; i < n; i++, j++) {
70	2903918x	ret(i) = input(j);
71		}
72	726686x	return ret;
73		}
74
75		// Calculate tcrossprod(lower_chol) = lower_chol * t(lower_chol).
76		// If complete, then adds the upper triangular part to the result as well.
77		// By default only the lower triangular part is populated, as this should be
78		// sufficient for downstream use of the result in most cases.
79		template <class Type>
80	659970x	matrix<Type> tcrossprod(const matrix<Type>& lower_chol, bool complete = false) {
81	659970x	int n = lower_chol.rows();
82	659970x	matrix<Type> result = matrix<Type>::Zero(n, n);
83	659970x	result.template selfadjointView<Eigen::Lower>().rankUpdate(lower_chol);
84	659970x	if (complete) {
85	13468x	result.template triangularView<Eigen::Upper>() = result.transpose();
86		}
87	659970x	return result;
88		}
89
90		// Calculate crossprod(x) = t(x) * x.
91		// Only the lower triangular part is populated, as this should be
92		// sufficient for downstream use of the result in most cases.
93		// Note that x does not need to be symmetric or square.
94		template <class Type>
95	1293004x	matrix<Type> crossprod(const matrix<Type>& x) {
96	1293004x	int n = x.cols();
97	1293004x	matrix<Type> result = matrix<Type>::Zero(n, n);
98	1293004x	result.template selfadjointView<Eigen::Lower>().rankUpdate(x.transpose());
99	1293004x	return result;
100		}
101
102		// Mapping from real values to correlation parameters in (-1, 1).
103		template <class T>
104	2746x	vector<T> map_to_cor(const vector<T>& theta) {
105	5492x	return theta / sqrt(T(1.0) + theta * theta);
106		}
107
108		// Mapping from real values to correlation parameters in (-1/(m-1), 1) for compound symmetry.
109		template <class T>
110	616x	vector<T> map_to_cs_cor(const vector<T>& theta, int n_visits) {
111	616x	T a = T(1.0) / (T(n_visits) - T(1.0));
112	616x	return invlogit(theta) * (T(1.0) + a) - a;
113		}
114
115		// Generic correlation function class containing and initializing correlation
116		// values from variance parameters theta.
117		template <class T>
118		struct generic_corr_fun {
119		const vector<T> corr_values;
120
121	2742x	generic_corr_fun(const vector<T>& theta) :
122	2742x	corr_values(map_to_cor(theta)) {}
123		};
124
125		// Correlation function based Cholesky factor of correlation matrix.
126		// This is used directly for homogeneous covariance matrices.
127		template <class T, template<class> class F>
128	3350x	matrix<T> get_corr_mat_chol(int n_visits, const F<T>& corr_fun) {
129	3350x	matrix<T> correlation(n_visits, n_visits);
130	3350x	correlation.setIdentity();
131	16666x	for(int i = 0; i < n_visits; i++) {
132	33180x	for(int j = 0; j < i; j++){
133	19864x	correlation(i, j) = corr_fun(i, j);
134		}
135		}
136	3350x	Eigen::LLT<Eigen::Matrix<T,Eigen::Dynamic,Eigen::Dynamic> > correlation_chol(correlation);
137	3350x	matrix<T> L = correlation_chol.matrixL();
138	6700x	return L;
139		}
140
141		// Heterogeneous covariance matrix calculation given vector of standard deviations (sd_values)
142		// and a correlation function (corr_fun).
143		template <class T, template<class> class F>
144	929x	matrix<T> get_heterogeneous_cov(const vector<T>& sd_values, const F<T>& corr_fun) {
145	929x	matrix<T> correlation_chol = get_corr_mat_chol(sd_values.size(), corr_fun);
146	929x	Eigen::DiagonalMatrix<T,Eigen::Dynamic,Eigen::Dynamic> D = sd_values.matrix().asDiagonal();
147	929x	matrix<T> result = D * correlation_chol;
148	1858x	return result;
149		}
150
151		// Obtain the Euclidean distance
152		template <class T>
153	18132x	matrix<T> euclidean(const matrix<T>& coordinates) {
154	18132x	matrix<T> result(coordinates.rows(), coordinates.rows());
155	68070x	for (int i = 0; i < coordinates.rows(); i++) {
156	49938x	result(i, i) = 0;
157	101220x	for (int j = 0; j < i; j ++) {
158	51282x	vector<T> diff = coordinates.row(i) - coordinates.row(j);
159	51282x	T d = sqrt((diff * diff).sum());
160	51282x	result(i, j) = d;
161	51282x	result(j, i) = d;
162		}
163		}
164	18132x	return result;
165		}
166
167		// Element wise power function of a matrix
168		template <class T>
169	792x	Eigen::Matrix<T, -1, -1> cpow(const Eigen::Matrix<T, -1, -1> & input, double p) {
170	792x	Eigen::Matrix<T, -1, -1> ret = Eigen::Matrix<T, -1, -1>(input.rows(), input.cols());
171	2954x	for (int i = 0; i < ret.rows(); i ++) {
172	4332x	for (int j = 0; j < ret.cols(); j++) {
173	2170x	ret(i, j) = std::pow(input(i, j), p);
174		}
175		}
176	792x	return ret;
177		}
178
179		// Calculate the square root of the pseudo inverse of a matrix
180		// adapted from the method for calculating the pseudo-Inverse as recommended by the Eigen developers
181		template<typename T>
182	790x	matrix<T> pseudoInverseSqrt(const matrix<T> &input, double epsilon = std::numeric_limits<double>::epsilon()) {
183	790x	Eigen::Matrix<T, -1, -1> eigen_mat = as_matrix<Eigen::Matrix<T, -1, -1>, matrix<T>>(input);
184	790x	Eigen::JacobiSVD< Eigen::Matrix<T, -1, -1> > svd(eigen_mat ,Eigen::ComputeFullU \| Eigen::ComputeFullV);
185	790x	double tolerance = epsilon * std::max(input.cols(), input.rows()) *svd.singularValues().array().abs()(0);
186	790x	auto singular_vals = Matrix<T,-1,-1>((svd.singularValues().array() > tolerance).select(svd.singularValues().array().inverse(), 0).matrix());
187	790x	Eigen::Matrix<T, -1, -1> ret_eigen = svd.matrixV() * cpow(singular_vals, 0.5).asDiagonal() * svd.matrixU().adjoint();
188	1580x	return as_matrix<matrix<T>, Eigen::Matrix<T, -1, -1>>(ret_eigen);
189		}
190
191		#endif

1		#include <RcppEigen.h>
2		#include "utils.h"
3
4		using namespace Rcpp;
5
6		#ifdef RCPP_USE_GLOBAL_ROSTREAM
7		Rcpp::Rostream<true>& Rcpp::Rcout = Rcpp::Rcpp_cout_get();
8		Rcpp::Rostream<false>& Rcpp::Rcerr = Rcpp::Rcpp_cerr_get();
9		#endif
10
11		List get_pqr(List mmrm_fit, NumericVector theta);
12	50x	RcppExport SEXP _mmrm_get_pqr(SEXP mmrm_fit_SEXP, SEXP theta_SEXP) {
13	50x	BEGIN_RCPP
14	50x	Rcpp::RObject rcpp_result_gen;
15	50x	Rcpp::RNGScope rcpp_rngScope_gen;
16	50x	Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
17	50x	Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
18	50x	rcpp_result_gen = Rcpp::wrap(get_pqr(mmrm_fit, theta));
19	50x	return rcpp_result_gen;
20	50x	END_RCPP
21		}
22
23		List get_jacobian(List mmrm_fit, NumericVector theta, NumericMatrix beta_vcov);
24	106x	RcppExport SEXP _mmrm_get_jacobian(SEXP mmrm_fit_SEXP, SEXP theta_SEXP, SEXP beta_vcov_SEXP) {
25	106x	BEGIN_RCPP
26	106x	Rcpp::RObject rcpp_result_gen;
27	106x	Rcpp::RNGScope rcpp_rngScope_gen;
28	106x	Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
29	106x	Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
30	106x	Rcpp::traits::input_parameter< NumericMatrix >::type beta_vcov(beta_vcov_SEXP);
31	106x	rcpp_result_gen = Rcpp::wrap(get_jacobian(mmrm_fit, theta, beta_vcov));
32	106x	return rcpp_result_gen;
33	106x	END_RCPP
34		}
35
36		List get_empirical(List mmrm_fit, NumericVector theta, NumericVector beta, NumericMatrix beta_vcov, std::string type);
37	39x	RcppExport SEXP _mmrm_get_empirical(SEXP mmrm_fit_SEXP, SEXP theta_SEXP, SEXP beta_SEXP, SEXP beta_vcov_SEXP, SEXP type_SEXP) {
38	39x	BEGIN_RCPP
39	39x	Rcpp::RObject rcpp_result_gen;
40	39x	Rcpp::RNGScope rcpp_rngScope_gen;
41	39x	Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
42	39x	Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
43	39x	Rcpp::traits::input_parameter< NumericVector >::type beta(beta_SEXP);
44	39x	Rcpp::traits::input_parameter< NumericMatrix >::type beta_vcov(beta_vcov_SEXP);
45	39x	Rcpp::traits::input_parameter< std::string >::type type(type_SEXP);
46	39x	rcpp_result_gen = Rcpp::wrap(get_empirical(mmrm_fit, theta, beta, beta_vcov, type));
47	39x	return rcpp_result_gen;
48	39x	END_RCPP
49		}
50
51		List predict(List mmrm_fit, NumericVector theta, NumericVector beta, NumericMatrix beta_vcov);
52	1704x	RcppExport SEXP _mmrm_predict(SEXP mmrm_fit_SEXP, SEXP theta_SEXP, SEXP beta_SEXP, SEXP beta_vcov_SEXP) {
53	1704x	BEGIN_RCPP
54	1704x	Rcpp::RObject rcpp_result_gen;
55	1704x	Rcpp::RNGScope rcpp_rngScope_gen;
56	1704x	Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
57	1704x	Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
58	1704x	Rcpp::traits::input_parameter< NumericVector >::type beta(beta_SEXP);
59	1704x	Rcpp::traits::input_parameter< NumericMatrix >::type beta_vcov(beta_vcov_SEXP);
60	1704x	rcpp_result_gen = Rcpp::wrap(predict(mmrm_fit, theta, beta, beta_vcov));
61	1704x	return rcpp_result_gen;
62	1704x	END_RCPP
63		}
64
65
66		RcppExport SEXP run_testthat_tests(SEXP);
67
68		static const R_CallMethodDef CallEntries[] = {
69		{"_mmrm_get_pqr", (DL_FUNC) &_mmrm_get_pqr, 2},
70		{"_mmrm_get_jacobian", (DL_FUNC) &_mmrm_get_jacobian, 3},
71		{"_mmrm_get_empirical", (DL_FUNC) &_mmrm_get_empirical, 5},
72		{"_mmrm_predict", (DL_FUNC) &_mmrm_predict, 4},
73		{"run_testthat_tests", (DL_FUNC) &run_testthat_tests, 1},
74		TMB_CALLDEFS,
75		{NULL, NULL, 0}
76		};
77
78	5x	RcppExport void R_init_mmrm(DllInfo *dll) {
79	5x	R_registerRoutines(dll, NULL, CallEntries, NULL, NULL);
80	5x	R_useDynamicSymbols(dll, FALSE);
81		#ifdef TMB_CCALLABLES
82	5x	TMB_CCALLABLES("mmrm");
83		#endif
84		}

1		#include "covariance.h"
2		#include "chol_cache.h"
3		// Definition:
4		//
5		// Y_i = X_i * beta + epsilon_i, i = 1, ..., n_subjects
6		// where Y_i = (Y_i1, ..., Y_im) are the observations of subject i over the m
7		// timepoints,
8		//
9		// and for the epsilon_i's :
10		// epsilon_i ~iid N(0, Sigma) where Sigma is a covariance matrix
11		// parameterized by a vector theta.
12		//
13		// Note: This is a special generalized least squares model
14		// Y = X * beta + epsilon,
15		// where we have a block structure for the covariance matrix of the epsilon
16		// vector.
17		//
18		// beta itself is not a parameter for TMB here:
19		// - For maximum likelihood estimation:
20		// Given theta and therefore Sigma, and writing W = Sigma^-1, we can determine
21		// the beta optimizing the likelihood via the weighted least squares equation
22		// (X^T W X) beta = X^T W Y.
23		// - For restricted maximum likelihood estimation:
24		// Given theta, beta is integrated out from the likelihood. Weighted least
25		// squares results are used to calculate integrated log likelihood.
26
27		template<class Type>
28	6620x	Type objective_function<Type>::operator() ()
29		{
30		// Read data from R.
31	6620x	DATA_MATRIX(x_matrix); // Model matrix (dimension n x p).
32		DATA_VECTOR(y_vector); // Response vector (length n).
33		DATA_VECTOR(weights_vector); // Weights vector (length n).
34	6620x	DATA_MATRIX(coordinates); // Coordinates matrix.
35	6620x	DATA_INTEGER(n_visits); // Number of visits, which is the dimension of the covariance matrix.
36	6620x	DATA_INTEGER(n_subjects); // Number of subjects.
37	6620x	DATA_IVECTOR(subject_zero_inds); // Starting indices for each subject (0-based) (length n_subjects).
38	6620x	DATA_IVECTOR(subject_n_visits); // Number of observed visits for each subject (length n_subjects).
39	6620x	DATA_STRING(cov_type); // Covariance type name.
40	6620x	DATA_INTEGER(is_spatial_int); // Spatial covariance (1)? Otherwise non-spatial covariance.
41	6620x	DATA_INTEGER(reml); // REML (1)? Otherwise ML (0).
42	6620x	DATA_FACTOR(subject_groups); // subject groups vector(0-based) (length n_subjects).
43	6620x	DATA_INTEGER(n_groups); // number of total groups.
44		// Read parameters from R.
45	6620x	PARAMETER_VECTOR(theta); // Covariance parameters (length k). Contents depend on covariance type.
46
47		// X^T W X will be calculated incrementally into here.
48	6620x	matrix<Type> XtWX = matrix<Type>::Zero(x_matrix.cols(), x_matrix.cols());
49		// X^T W Y will be calculated incrementally into here.
50	6620x	matrix<Type> XtWY = matrix<Type>::Zero(x_matrix.cols(), 1);
51		// W^T/2 X will be saved into here.
52	6620x	matrix<Type> x_mat_tilde = matrix<Type>::Zero(x_matrix.rows(), x_matrix.cols());
53		// W^T/2 Y will be saved into here.
54	6620x	vector<Type> y_vec_tilde = vector<Type>::Zero(y_vector.rows());
55		// Sum of the log determinant will be incrementally calculated here.
56	6620x	Type sum_log_det = 0.0;
57
58		// Convert is_spatial_int to bool.
59	6620x	bool is_spatial = (is_spatial_int == 1);
60		// Diagonal of weighted covariance
61	6620x	vector<Type> diag_cov_inv_sqrt(x_matrix.rows());
62		// Cholesky group object
63	6620x	auto chols_group = chol_cache_groups<Type>(theta, n_groups, is_spatial, cov_type, n_visits);
64		// Go through all subjects and calculate quantities initialized above.
65	1299618x	for (int i = 0; i < n_subjects; i++) {
66		// Start index and number of visits for this subject.
67	1293002x	int start_i = subject_zero_inds(i);
68	1293002x	int n_visits_i = subject_n_visits(i);
69	1293002x	std::vector<int> visit_i(n_visits_i);
70	1293002x	matrix<Type> dist_i(n_visits_i, n_visits_i);
71	1293002x	if (!is_spatial) {
72	4710126x	for (int j = 0; j < n_visits_i; j++) {
73	3448644x	visit_i[j] = int(asDouble(coordinates(start_i + j, 0)));
74		}
75		} else {
76	31520x	dist_i = euclidean(matrix<Type>(coordinates.block(start_i, 0, n_visits_i, coordinates.cols())));
77		}
78		// Obtain Cholesky factor Li.
79	1293002x	matrix<Type> Li = chols_group.cache[subject_groups[i]]->get_chol(visit_i, dist_i);
80		// Calculate weighted Cholesky factor for this subject.
81	1293002x	Eigen::DiagonalMatrix<Type,Eigen::Dynamic,Eigen::Dynamic> Gi_inv_sqrt = weights_vector.segment(start_i, n_visits_i).cwiseInverse().sqrt().matrix().asDiagonal();
82	1293002x	Li = Gi_inv_sqrt * Li;
83		// Calculate scaled design matrix and response vector for this subject.
84	1293002x	matrix<Type> Xi = x_matrix.block(start_i, 0, n_visits_i, x_matrix.cols());
85	1293002x	matrix<Type> XiTilde = Li.template triangularView<Eigen::Lower>().solve(Xi);
86	1293002x	matrix<Type> Yi = y_vector.segment(start_i, n_visits_i).matrix();
87	1293002x	matrix<Type> YiTilde = Li.template triangularView<Eigen::Lower>().solve(Yi);
88
89		// Increment quantities.
90	1293002x	matrix<Type> XiTildeCrossprod = crossprod(XiTilde);
91	1293002x	XtWX += XiTildeCrossprod.template triangularView<Eigen::Lower>();
92	1293002x	XtWY += XiTilde.transpose() * YiTilde;
93	1293002x	vector<Type> LiDiag = Li.diagonal();
94	1293002x	sum_log_det += sum(log(LiDiag));
95		// Cache the reciprocal of square root of diagonal of covariance
96	1293002x	diag_cov_inv_sqrt.segment(start_i, n_visits_i) = vector<Type>(tcrossprod(Li).diagonal()).rsqrt();
97		// Save stuff.
98	1293002x	x_mat_tilde.block(start_i, 0, n_visits_i, x_matrix.cols()) = XiTilde;
99	1293002x	y_vec_tilde.segment(start_i, n_visits_i) = YiTilde.col(0);
100		}
101
102		// Solve for beta.
103	6616x	Eigen::LDLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > XtWX_decomposition(XtWX);
104	6616x	matrix<Type> beta_mat = XtWX_decomposition.solve(XtWY);
105	6616x	vector<Type> beta = beta_mat.col(0);
106
107		// Define scaled residuals.
108	6616x	vector<Type> x_mat_tilde_beta = x_mat_tilde * beta;
109	6616x	vector<Type> epsilonTilde = y_vec_tilde - x_mat_tilde_beta;
110
111		// Calculate negative log-likelihood.
112	552x	Type neg_log_lik;
113
114		// Always extract the D vector since we want to report this below.
115	6616x	vector<Type> XtWX_D = XtWX_decomposition.vectorD();
116
117	6616x	if (reml == 1) {
118		// Use restricted maximum likelihood.
119	6282x	Type XtWX_log_det = XtWX_D.log().sum();
120	6282x	neg_log_lik = (x_matrix.rows() - x_matrix.cols()) / 2.0 * log(2.0 * M_PI) +
121	6282x	sum_log_det +
122	12564x	XtWX_log_det / 2.0 +
123	6778x	0.5 * (y_vec_tilde * y_vec_tilde).sum() - 0.5 * (x_mat_tilde_beta * x_mat_tilde_beta).sum();
124		} else {
125		// Use maximum likelihood.
126	334x	neg_log_lik = x_matrix.rows() / 2.0 * log(2.0 * M_PI) +
127	56x	sum_log_det +
128	390x	0.5 * (epsilonTilde * epsilonTilde).sum();
129		}
130
131		// Report quantities to R.
132	6064x	REPORT(beta);
133
134		// We already compute the inverse of XtWX here because we already did the
135		// matrix decomposition above.
136	6616x	matrix<Type> Identity(XtWX.rows(), XtWX.cols());
137	6616x	Identity.setIdentity();
138	6616x	matrix<Type> beta_vcov = XtWX_decomposition.solve(Identity);
139	6064x	REPORT(beta_vcov);
140
141		// Also return the decomposition components L and D.
142	6616x	matrix<Type> XtWX_L(XtWX.rows(), XtWX.cols());
143	6616x	XtWX_L = XtWX_decomposition.matrixL();
144	6064x	REPORT(XtWX_L);
145	6064x	REPORT(XtWX_D);
146
147		// normalized residual
148	6064x	REPORT(epsilonTilde);
149		// inverse square root of diagonal of covariance
150	6064x	REPORT(diag_cov_inv_sqrt);
151	6616x	matrix<Type> covariance_lower_chol = chols_group.get_default_chol();
152	6064x	REPORT(covariance_lower_chol);
153
154	6616x	return neg_log_lik;
155		}

1		#ifndef TESTTHAT_WRAP_H
2		#define TESTTHAT_WRAP_H
3		#include <testthat.h>
4		#include <limits>
5		#include "utils.h"
6
7		// Expect equal: Here use a default epsilon which gives around 1e-4 on
8		// my computer here.
9		#define expect_equal(TARGET, CURRENT) \
10		{ \
11		double const eps = \
12		std::pow(std::numeric_limits<double>::epsilon(), 0.25); \
13		\
14		if(std::abs((TARGET)) > eps) \
15		expect_true(std::abs((TARGET) - (CURRENT)) / \
16		std::abs((TARGET)) < eps); \
17		else \
18		expect_true(std::abs((TARGET) - (CURRENT)) < eps); \
19		}
20
21		#define expect_equal_eps(TARGET, CURRENT, EPS) \
22		{ \
23		if(std::abs((TARGET)) > (EPS)) \
24		expect_true(std::abs((TARGET) - (CURRENT)) / \
25		std::abs((TARGET)) < (EPS)); \
26		else \
27		expect_true(std::abs((TARGET) - (CURRENT)) < (EPS)); \
28		}
29
30		template <class T>
31	48x	void expect_equal_matrix(const T& target, const T& current)
32		{
33	48x	int nrow = target.rows();
34	48x	int ncol = target.cols();
35
36	!	expect_true(nrow == current.rows());
37	!	expect_true(ncol == current.cols());
38
39	183x	for (int i = 0; i < nrow; i++) {
40	510x	for (int j = 0; j < ncol; j++) {
41	!	expect_equal(target(i, j), current(i, j));
42		}
43		}
44		}
45
46		template <class T>
47	9x	void expect_equal_vector(const T& target, const T& current)
48		{
49	9x	int n = target.size();
50	!	expect_true(n == current.size());
51
52	54x	for (int i = 0; i < n; i++) {
53	!	expect_equal(target(i), current(i));
54		}
55		}
56
57		#endif