mmrm coverage - 97.36%

Files
Source

#' 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
}

#' 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.
#'
#' @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
  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(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.
#'   Must be `NULL` or one or more of `c("subject_var", "visit_var", "group_var", "response_var")`.
#' @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 returned model frame will include.
#' Possible options are "response_var", "subject_var", "visit_var" and "group_var", representing the
#' response variable, subject variable, visit variable or group variable.
#' `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"),
                                 full, na.action = "na.omit", ...) { # 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 to return to users.
    ret <-
      stats::model.frame(
        formula = lst_formula_and_data$formula,
        data = h_get_na_action(na.action)(lst_formula_and_data$data),
        na.action = na.action,
        xlev = stats::.getXlevels(terms(formula), formula$tmb_data$full_frame)
      )
  }
  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 mmrm fit object.
#' @param data optional data frame that will be
#'   passed to `model.frame()` or `model.matrix()`
#' @param include (`character`)\cr names of variable to include
#' @param full (`flag`)\cr indicator whether to return full model frame (deprecated).
#'
#' @return named list with four elements:
#' - `"formula"`: the formula including the columns requested in the `include=` argument.
#' - `"data"`: a data frame including all columns needed in the formula.
#'   full formula are identical
#' @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[include])
  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)
  assert_names(colnames(data), must.include = all_vars)
  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.
#'
#' @examples
#' # Model matrix:
#' model.matrix(object)
model.matrix.mmrm_tmb <- function(object, data, use_response = TRUE, ...) { # nolint
  # 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 = attr(object$tmb_data$x_matrix, "contrasts"),
    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.
#' @importFrom stats logLik
#' @exportS3Method
#' @examples
#' # Log likelihood given the estimated parameters:
#' logLik(object)
logLik.mmrm_tmb <- function(object, ...) {
  -component(object, "neg_log_lik")
}

#' @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
}

#' 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
    )
  )
}

#' 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)
  }
  return(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"))
    return(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 data (`data.frame`)\cr data used for model fitting.
#' @param model_formula (`formula`)\cr the formula in mmrm model without covariance structure part.
#' @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.
#'
#' @note `data` is used instead of `full_frame` because `full_frame` is already
#' transformed if model contains transformations, e.g. `log(FEV1) ~ exp(FEV1_BL)` will
#' drop `FEV1` and `FEV1_BL` but add `log(FEV1)` and `exp(FEV1_BL)` in `full_frame`.
#'
#' @return A numeric vector of starting values.
#'
#' @export
emp_start <- function(data, model_formula, visit_var, subject_var, subject_groups, ...) {
  assert_formula(model_formula)
  assert_data_frame(data)
  assert_subset(all.vars(model_formula), colnames(data))
  assert_string(visit_var)
  assert_string(subject_var)
  assert_factor(data[[visit_var]])
  n_visits <- length(levels(data[[visit_var]]))
  assert_factor(data[[subject_var]])
  subjects <- droplevels(data[[subject_var]])
  n_subjects <- length(levels(subjects))
  fit <- stats::lm(formula = model_formula, data = data)
  res <- rep(NA, n_subjects * n_visits)
  res[
    n_visits * as.integer(subjects) - n_visits + as.integer(data[[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)
  }
}

#' 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
}

#' 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`.
#'
#' @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)
    ),
    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)
  varname <- formula_parts[grepl("_var", names(formula_parts))]
  assert_names(
    names(data),
    must.include = unlist(varname, use.names = FALSE)
  )
  assert_true(is.factor(data[[formula_parts$subject_var]]) || is.character(data[[formula_parts$subject_var]]))
  assert_numeric(weights, len = nrow(data))
  assert_flag(reml)
  singular <- match.arg(singular)
  assert_flag(drop_visit_levels)

  if (is.character(data[[formula_parts$subject_var]])) {
    data[[formula_parts$subject_var]] <- factor(
      data[[formula_parts$subject_var]],
      levels = stringr::str_sort(unique(data[[formula_parts$subject_var]]), numeric = TRUE)
    )
  }
  data_order <- if (formula_parts$is_spatial) {
    order(data[[formula_parts$subject_var]])
  } else {
    subject_visit_data <- data[, 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(data[[formula_parts$subject_var]], data[[formula_parts$visit_var]])
  }
  if (identical(formula_parts$is_spatial, FALSE)) {
    h_confirm_large_levels(length(levels(data[[formula_parts$visit_var]])))
  }
  data <- data[data_order, ]
  weights <- weights[data_order]
  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()
  }
  full_frame <- eval(
    bquote(stats::model.frame(
      formula_parts$full_formula,
      data = data,
      weights = .(as.symbol(weights_name)),
      na.action = "na.pass",
      xlev = xlev
    ))
  )
  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)
  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")) {
      assign_attr <- attr(x_matrix, "assign")
      contrasts_attr <- attr(x_matrix, "contrasts")
      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(data[[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))
    assert_factor(coordinates[[1L]])
    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,
      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(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)
      return(ret)
    }
  )
  if (identical(tmb_data$n_groups, 1L)) {
    cov <- cov[[1]]
  } else {
    names(cov) <- levels(tmb_data$subject_groups)
  }
  return(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()].
#'
#' @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.
#'   - `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) {
  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 <- 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)

    if (!formula_parts$is_spatial && !is.factor(data[[formula_parts$visit_var]])) {
      stop("Time variable must be a factor for non-spatial covariance structures")
    }

    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)
  h_mmrm_tmb_check_conv(tmb_opt, fit)
  fit$call <- match.call()
  fit$call$formula <- formula_parts$formula
  fit$optimizer <- used_optimizer_name
  fit
}

#' 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()].
#' @param drop_visit_levels (`flag`)\cr whether to drop levels for visit variable,
#'   if visit variable is a factor, 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`.
#'
#' @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,
    ...,
    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_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!"
    ))
  }
  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
    ),
    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
}

#' 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")
  }
}

#' 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
}

#' 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
COV_TYPES <- local({ # nolint
  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
as.cov_struct <- function(x, ...) { # nolint
  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)
}

#' 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).
#' - `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", "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 {
        object$beta_vcov
      },
    "varcor" = object$cov,
    "score_per_subject" = object$score_per_subject,
    "x_matrix" = object$tmb_data$x_matrix,
    "xlev" = stats::.getXlevels(terms(object), object$tmb_data$full_frame),
    "contrasts" = attr(object$tmb_data$x_matrix, "contrasts"),
    "theta_vcov" = 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
}

#' 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
}

#' 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))
  }))
}

#' 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 Contrast for Specified Effect
#'
#' This is support function to obtain contrast matrix for type II/III testing.
#'
#' @param object (`mmrm`)\cr the fitted MMRM.
#' @param effect (`string`) the name of the effect.
#' @param type (`string`) type of test, "II", "III", '2', or '3'.
#' @param tol (`numeric`) threshold blow which values are treated as 0.
#'
#' @return A `matrix` of the contrast.
#'
#' @keywords internal
h_get_contrast <- function(object, effect, type = c("II", "III", "2", "3"), tol = sqrt(.Machine$double.eps)) {
  assert_class(object, "mmrm")
  assert_string(effect)
  assert_double(tol, finite = TRUE, len = 1L)
  type <- match.arg(type)
  mx <- component(object, "x_matrix")
  asg <- attr(mx, "assign")
  formula <- object$formula_parts$model_formula
  tms <- terms(formula)
  fcts <- attr(tms, "factors")[-1L, , drop = FALSE] # Discard the response.
  ods <- attr(tms, "order")
  assert_subset(effect, colnames(fcts))
  idx <- which(effect == colnames(fcts))
  cols <- which(asg == idx)
  xlev <- component(object, "xlev")
  contains_intercept <- (!0 %in% asg) && h_first_contain_categorical(effect, fcts, names(xlev))
  coef_rows <- length(cols) - as.integer(contains_intercept)
  l_mx <- matrix(0, nrow = coef_rows, ncol = length(asg))
  if (coef_rows == 0L) {
    return(l_mx)
  }
  if (contains_intercept) {
    l_mx[, cols] <- cbind(-1, diag(rep(1, coef_rows)))
  } else {
    l_mx[, cols] <- diag(rep(1, coef_rows))
  }
  for (i in setdiff(seq_len(ncol(fcts)), idx)) {
    additional_vars <- names(which(fcts[, i] > fcts[, idx]))
    additional_numeric <- any(!additional_vars %in% names(xlev))
    current_col <- which(asg == i)
    if (ods[i] >= ods[idx] && all(fcts[, i] >= fcts[, idx]) && !additional_numeric) {
      sub_mat <- switch(type,
        "2" = ,
        "II" = {
          x1 <- mx[, cols, drop = FALSE]
          x0 <- mx[, -c(cols, current_col), drop = FALSE]
          x2 <- mx[, current_col, drop = FALSE]
          m <- diag(rep(1, nrow(x0))) - x0 %*% solve(t(x0) %*% x0) %*% t(x0)
          ret <- solve(t(x1) %*% m %*% x1) %*% t(x1) %*% m %*% x2
          if (contains_intercept) {
            ret[-1, ] - ret[1, ]
          } else {
            ret
          }
        },
        "3" = ,
        "III" = {
          lvls <- h_obtain_lvls(effect, additional_vars, xlev)
          t_levels <- lvls$total
          nms_base <- colnames(mx)[cols]
          nms <- colnames(mx)[current_col]
          nms_base_split <- strsplit(nms_base, ":")
          nms_split <- strsplit(nms, ":")
          base_idx <- h_get_index(nms_split, nms_base_split)
          mt <- l_mx[, cols, drop = FALSE] / t_levels
          ret <- mt[, base_idx, drop = FALSE]
          # if there is extra levels, replace it with -1/t_levels
          ret[is.na(ret)] <- -1 / t_levels
          ret
        }
      )
      l_mx[, current_col] <- sub_mat
    }
  }
  l_mx[abs(l_mx) < tol] <- 0
  l_mx
}

#' Conduct type II/III hypothesis testing on the MMRM fit results.
#'
#' @param mod (`mmrm`)\cr the fitted MMRM.
#' @param ... not used.
#' @inheritParams h_get_contrast
#'
#' @details
#' `Anova` will return `anova` object with one row per variable and columns
#' `Num Df`(numerator degrees of freedom), `Denom Df`(denominator degrees of freedom),
#' `F Statistic` and `Pr(>=F)`.
#'
#' @keywords internal
# Please do not load `car` and then create the documentation. The Rd file will be different.
Anova.mmrm <- function(mod, type = c("II", "III", "2", "3"), tol = sqrt(.Machine$double.eps), ...) { # nolint
  assert_double(tol, finite = TRUE, len = 1L)
  type <- match.arg(type)
  vars <- colnames(attr(terms(mod$formula_parts$model_formula), "factors"))
  ret <- lapply(
    vars,
    function(x) df_md(mod, h_get_contrast(mod, x, type, tol))
  )
  ret_df <- do.call(rbind.data.frame, ret)
  row.names(ret_df) <- vars
  colnames(ret_df) <- c("Num Df", "Denom Df", "F Statistic", "Pr(>=F)")
  class(ret_df) <- c("anova", "data.frame")
  attr(ret_df, "heading") <- sprintf(
    "Analysis of Fixed Effect Table (Type %s F tests)",
    switch(type,
      "2" = ,
      "II" = "II",
      "3" = ,
      "III" = "III"
    )
  )
  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
  )
}

#' Check if the Effect is the First Categorical Effect
#' @param effect (`string`) name of the effect.
#' @param categorical (`character`) names of the categorical values.
#' @param factors (`matrix`) the factor matrix.
#' @keywords internal
h_first_contain_categorical <- function(effect, factors, categorical) {
  assert_string(effect)
  assert_matrix(factors)
  assert_character(categorical)
  mt <- match(effect, colnames(factors))
  varnms <- row.names(factors)
  # if the effect is not categorical in any value, return FALSE
  if (!any(varnms[factors[, mt] > 0] %in% categorical)) {
    return(FALSE)
  }
  # keep only categorical rows that is in front of the current factor
  factors <- factors[row.names(factors) %in% categorical, seq_len(mt - 1L), drop = FALSE]
  # if previous cols are all numerical, return TRUE
  if (ncol(factors) < 1L) {
    return(TRUE)
  }
  col_ind <- apply(factors, 2, prod)
  # if any of the previous cols are categorical, return FALSE
  return(!any(col_ind > 0))
}

#' 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
  )
}

#' 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)
  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 one-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) {
    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-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.
#'
#' @inheritParams df_md
#' @param contrast (`matrix`)\cr numeric contrast matrix.
#' @param df (`number`)\cr denominator degrees of freedom for the multi-dimensional contrast.
#' @param f_stat_factor (`number`)\cr optional scaling factor on top of the standard F-statistic.
#'
#' @return List with `num_df`, `denom_df`, `f_stat` and `p_val` (2-sided p-value).
#'
#' @keywords internal
h_test_md <- function(object,
                      contrast,
                      df,
                      f_stat_factor = 1) {
  assert_class(object, "mmrm")
  assert_matrix(contrast, ncols = length(component(object, "beta_est")))
  num_df <- nrow(contrast)
  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 <- component(object, "beta_est") %*% t(contrast)
  f_statistic <- as.numeric(f_stat_factor / num_df * h_quad_form_mat(contrast_est, prec_contrast))
  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
  )
}

#' 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)
}

#' 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
recover_data.mmrm <- function(object, ...) { # nolint
  fun_call <- stats::getCall(object)
  # subject_var is excluded because it should not contain fixed effect.
  # visit_var is not excluded because emmeans can provide marginal mean
  # by each visit if visit_var is not spatial.
  model_frame <- stats::model.frame(
    object,
    include = c(
      if (!object$formula_parts$is_spatial) "visit_var" else NULL,
      "response_var", "group_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
emm_basis.mmrm <- function(object, # nolint
                           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,
        include = c(
          if (!object$formula_parts$is_spatial) "visit_var" else NULL,
          "response_var", "group_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
  )
}

#' 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
}

#' 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)
}

# 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()
}

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

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	47x	assert_class(tmb_data, "mmrm_tmb_data")
26	47x	assert_class(theta, "numeric")
27	47x	.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	6x	assert_class(object, "mmrm")
39	6x	assert_matrix(contrast, mode = "numeric", any.missing = FALSE, ncols = length(component(object, "beta_est")))
40	6x	if (component(object, "reml") != 1) {
41	!	stop("Kenward-Roger is only for REML")
42		}
43	6x	kr_comp <- object$kr_comp
44	6x	w <- component(object, "theta_vcov")
45	6x	v_adj <- object$beta_vcov_adj
46	6x	df <- h_kr_df(v0 = object$beta_vcov, l = contrast, w = w, p = kr_comp$P)
47
48	6x	h_test_md(object, contrast, df = df$m, f_stat_factor = df$lambda)
49		}
50
51		#' Calculation of Kenward-Roger Degrees of Freedom for One-Dimensional Contrast
52		#'
53		#' @description Used in [df_1d()] if method is
54		#' "Kenward-Roger" or "Kenward-Roger-Linear".
55		#'
56		#' @inheritParams h_df_1d_sat
57		#' @inherit h_df_1d_sat return
58		#' @keywords internal
59		h_df_1d_kr <- function(object, contrast) {
60	21x	assert_class(object, "mmrm")
61	21x	assert_numeric(contrast, len = length(component(object, "beta_est")))
62	21x	if (component(object, "reml") != 1) {
63	!	stop("Kenward-Roger is only for REML!")
64		}
65
66	21x	df <- h_kr_df(
67	21x	v0 = object$beta_vcov,
68	21x	l = matrix(contrast, nrow = 1),
69	21x	w = component(object, "theta_vcov"),
70	21x	p = object$kr_comp$P
71		)
72
73	21x	h_test_1d(object, contrast, df$m)
74		}
75
76		#' Obtain the Adjusted Kenward-Roger degrees of freedom
77		#'
78		#' @description Obtains the adjusted Kenward-Roger degrees of freedom and F statistic scale parameter.
79		#' Used in [h_df_md_kr()] or [h_df_1d_kr].
80		#'
81		#' @param v0 (`matrix`)\cr unadjusted covariance matrix.
82		#' @param l (`matrix`)\cr linear combination matrix.
83		#' @param w (`matrix`)\cr hessian matrix.
84		#' @param p (`matrix`)\cr P matrix from [h_get_kr_comp()].
85		#'
86		#' @return Named list with elements:
87		#' - `m`: `numeric` degrees of freedom.
88		#' - `lambda`: `numeric` F statistic scale parameter.
89		#'
90		#' @keywords internal
91		h_kr_df <- function(v0, l, w, p) {
92	28x	n_beta <- ncol(v0)
93	28x	assert_matrix(v0, ncols = n_beta, nrows = n_beta)
94	28x	assert_matrix(l, ncols = n_beta)
95	28x	n_theta <- ncol(w)
96	28x	assert_matrix(w, ncols = n_theta, nrows = n_theta)
97	28x	n_visits <- ncol(p)
98	28x	assert_matrix(p, nrows = n_visits * n_theta)
99		# see vignettes/kenward.Rmd#279
100	28x	slvol <- solve(h_quad_form_mat(l, v0))
101	28x	m <- h_quad_form_mat(t(l), slvol)
102	28x	nl <- nrow(l)
103	28x	mv0 <- m %*% v0
104	28x	pl <- lapply(seq_len(nrow(p) / ncol(p)), function(x) {
105	108x	ii <- (x - 1) * ncol(p) + 1
106	108x	jj <- x * ncol(p)
107	108x	p[ii:jj, ]
108		})
109	28x	mv0pv0 <- lapply(pl, function(x) {
110	108x	mv0 %% x %% v0
111		})
112	28x	a1 <- 0
113	28x	a2 <- 0
114		# see vignettes/kenward.Rmd#283
115	28x	for (i in seq_len(length(pl))) {
116	108x	for (j in seq_len(length(pl))) {
117	592x	a1 <- a1 + w[i, j] * h_tr(mv0pv0[[i]]) * h_tr(mv0pv0[[j]])
118	592x	a2 <- a2 + w[i, j] * h_tr(mv0pv0[[i]] %*% mv0pv0[[j]])
119		}
120		}
121	28x	b <- 1 / (2 * nl) * (a1 + 6 * a2)
122	28x	e <- 1 + a2 / nl
123	28x	e_star <- 1 / (1 - a2 / nl)
124	28x	g <- ((nl + 1) * a1 - (nl + 4) * a2) / ((nl + 2) * a2)
125	28x	denom <- (3 * nl + 2 - 2 * g)
126	28x	c1 <- g / denom
127	28x	c2 <- (nl - g) / denom
128	28x	c3 <- (nl + 2 - g) / denom
129	28x	v_star <- 2 / nl * (1 + c1 * b) / (1 - c2 * b)^2 / (1 - c3 * b)
130	28x	rho <- v_star / (2 * e_star^2)
131	28x	m <- 4 + (nl + 2) / (nl * rho - 1)
132	28x	lambda <- m / (e_star * (m - 2))
133	28x	list(m = m, lambda = lambda)
134		}
135
136		#' Obtain the Adjusted Covariance Matrix
137		#'
138		#' @description Obtains the Kenward-Roger adjusted covariance matrix for the
139		#' coefficient estimates.
140		#' Used in [mmrm()] fitting if method is "Kenward-Roger" or "Kenward-Roger-Linear".
141		#'
142		#' @param v (`matrix`)\cr unadjusted covariance matrix.
143		#' @param w (`matrix`)\cr hessian matrix.
144		#' @param p (`matrix`)\cr P matrix from [h_get_kr_comp()].
145		#' @param q (`matrix`)\cr Q matrix from [h_get_kr_comp()].
146		#' @param r (`matrix`)\cr R matrix from [h_get_kr_comp()].
147		#' @param linear (`flag`)\cr whether to use linear Kenward-Roger approximation.
148		#'
149		#' @return The matrix of adjusted covariance matrix.
150		#'
151		#' @keywords internal
152		h_var_adj <- function(v, w, p, q, r, linear = FALSE) {
153	49x	assert_flag(linear)
154	49x	n_beta <- ncol(v)
155	49x	assert_matrix(v, nrows = n_beta)
156	49x	n_theta <- ncol(w)
157	49x	assert_matrix(w, nrows = n_theta)
158	49x	n_visits <- ncol(p)
159	49x	theta_per_group <- nrow(q) / nrow(p)
160	49x	n_groups <- n_theta / theta_per_group
161	49x	assert_matrix(p, nrows = n_theta * n_visits)
162	49x	assert_matrix(q, nrows = theta_per_group^2 * n_groups * n_visits, ncols = n_visits)
163	49x	assert_matrix(r, nrows = theta_per_group^2 * n_groups * n_visits, ncols = n_visits)
164	49x	if (linear) {
165	13x	r <- matrix(0, nrow = nrow(r), ncol = ncol(r))
166		}
167
168		# see vignettes/kenward.Rmd#131
169	49x	ret <- v
170	49x	for (i in seq_len(n_theta)) {
171	264x	for (j in seq_len(n_theta)) {
172	2164x	gi <- ceiling(i / theta_per_group)
173	2164x	gj <- ceiling(j / theta_per_group)
174	2164x	iid <- (i - 1) * n_beta + 1
175	2164x	jid <- (j - 1) * n_beta + 1
176	2164x	ii <- i - (gi - 1) * theta_per_group
177	2164x	jj <- j - (gi - 1) * theta_per_group
178	2164x	ijid <- ((ii - 1) * theta_per_group + jj - 1) * n_beta + (gi - 1) * n_beta * theta_per_group^2 + 1
179	2164x	if (gi != gj) {
180	592x	ret <- ret + 2 * w[i, j] * v %% (-p[iid:(iid + n_beta - 1), ] %% v %% p[jid:(jid + n_beta - 1), ]) %% v
181		} else {
182	1572x	ret <- ret + 2 * w[i, j] * v %*% (
183	1572x	q[ijid:(ijid + n_beta - 1), ] -
184	1572x	p[iid:(iid + n_beta - 1), ] %% v %% p[jid:(jid + n_beta - 1), ] -
185	1572x	1 / 4 * r[ijid:(ijid + n_beta - 1), ]
186	1572x	) %*% v
187		}
188		}
189		}
190	49x	ret
191		}

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

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		#' - `xlev`: a named list of character vectors giving the full set of levels to be assumed for each factor.
45		#' - `contrasts`: a list of contrasts used for each factor.
46		#' - `y_vector`: response vector used.
47		#' - `jac_list`: Jacobian, see [h_jac_list()] for details.
48		#' - `full_frame`: `data.frame` with `n` rows containing all variables needed in the model.
49		#'
50		#' @seealso In the `lme4` package there is a similar function `getME()`.
51		#'
52		#' @examples
53		#' fit <- mmrm(
54		#' formula = FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT \| USUBJID),
55		#' data = fev_data
56		#' )
57		#' # Get all available components.
58		#' component(fit)
59		#' # Get convergence code and message.
60		#' component(fit, c("convergence", "conv_message"))
61		#' # Get modeled formula as a string.
62		#' component(fit, c("formula"))
63		#'
64		#' @export
65		component <- function(object,
66		name = c(
67		"cov_type", "subject_var", "n_theta", "n_subjects", "n_timepoints",
68		"n_obs", "beta_vcov", "beta_vcov_complete",
69		"varcor", "score_per_subject", "formula", "dataset", "n_groups",
70		"reml", "convergence", "evaluations", "method", "optimizer",
71		"conv_message", "call", "theta_est",
72		"beta_est", "beta_est_complete", "beta_aliased",
73		"x_matrix", "y_vector", "neg_log_lik", "jac_list", "theta_vcov",
74		"full_frame", "xlev", "contrasts"
75		)) {
76	5143x	assert_class(object, "mmrm_tmb")
77	5143x	name <- match.arg(name, several.ok = TRUE)
78
79	5143x	list_components <- sapply(
80	5143x	X = name,
81	5143x	FUN = switch,
82	5143x	"call" = object$call,
83		# Strings.
84	5143x	"cov_type" = object$formula_parts$cov_type,
85	5143x	"subject_var" = object$formula_parts$subject_var,
86	5143x	"formula" = deparse(object$call$formula),
87	5143x	"dataset" = object$call$data,
88	5143x	"reml" = object$reml,
89	5143x	"conv_message" = object$opt_details$message,
90		# Numeric of length 1.
91	5143x	"convergence" = object$opt_details$convergence,
92	5143x	"neg_log_lik" = object$neg_log_lik,
93	5143x	"n_theta" = length(object$theta_est),
94	5143x	"n_subjects" = object$tmb_data$n_subjects,
95	5143x	"n_timepoints" = object$tmb_data$n_visits,
96	5143x	"n_obs" = length(object$tmb_data$y_vector),
97	5143x	"n_groups" = ifelse(is.list(object$cov), length(object$cov), 1L),
98		# Numeric of length > 1.
99	5143x	"evaluations" = unlist(ifelse(is.null(object$opt_details$evaluations),
100	5143x	list(object$opt_details$counts),
101	5143x	list(object$opt_details$evaluations)
102		)),
103	5143x	"method" = object$method,
104	5143x	"optimizer" = object$optimizer,
105	5143x	"beta_est" = object$beta_est,
106	5143x	"beta_est_complete" =
107	5143x	if (any(object$tmb_data$x_cols_aliased)) {
108	8x	stats::setNames(
109	8x	object$beta_est[names(object$tmb_data$x_cols_aliased)],
110	8x	names(object$tmb_data$x_cols_aliased)
111		)
112		} else {
113	54x	object$beta_est
114		},
115	5143x	"beta_aliased" = object$tmb_data$x_cols_aliased,
116	5143x	"theta_est" = object$theta_est,
117	5143x	"y_vector" = object$tmb_data$y_vector,
118	5143x	"jac_list" = object$jac_list,
119		# Matrices.
120	5143x	"beta_vcov" =
121	5143x	if (is.null(object$vcov) \|\| identical(object$vcov, "Asymptotic")) {
122	985x	object$beta_vcov
123		} else {
124	68x	object$beta_vcov_adj
125		},
126	5143x	"beta_vcov_complete" =
127	5143x	if (any(object$tmb_data$x_cols_aliased)) {
128	2x	stats::.vcov.aliased(
129	2x	aliased = object$tmb_data$x_cols_aliased,
130	2x	vc = component(object, "beta_vcov"),
131	2x	complete = TRUE
132		)
133		} else {
134	4x	object$beta_vcov
135		},
136	5143x	"varcor" = object$cov,
137	5143x	"score_per_subject" = object$score_per_subject,
138	5143x	"x_matrix" = object$tmb_data$x_matrix,
139	5143x	"xlev" = stats::.getXlevels(terms(object), object$tmb_data$full_frame),
140	5143x	"contrasts" = attr(object$tmb_data$x_matrix, "contrasts"),
141	5143x	"theta_vcov" = object$theta_vcov,
142	5143x	"full_frame" = object$tmb_data$full_frame,
143		# If not found.
144	5143x	"..foo.." =
145	5143x	stop(sprintf(
146	5143x	"component '%s' is not available",
147	5143x	name, paste0(class(object), collapse = ", ")
148		)),
149	5143x	simplify = FALSE
150		)
151
152	23x	if (length(name) == 1) list_components[[1]] else list_components
153		}

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(contrast, len = length(component(object, "beta_est")), any.missing = FALSE)
25	338x	contrast <- as.vector(contrast)
26	338x	switch(object$method,
27	318x	"Satterthwaite" = h_df_1d_sat(object, contrast),
28	19x	"Kenward-Roger" = h_df_1d_kr(object, contrast),
29	!	"Residual" = h_df_1d_res(object, contrast),
30	1x	"Between-Within" = h_df_1d_bw(object, contrast),
31	!	stop("Unrecognized degrees of freedom method: ", object$method)
32		)
33		}
34
35
36		#' Calculation of Degrees of Freedom for Multi-Dimensional Contrast
37		#'
38		#' @description `r lifecycle::badge("stable")`
39		#' Calculates the estimate, standard error, degrees of freedom,
40		#' t statistic and p-value for one-dimensional contrast, depending on the method
41		#' used in [mmrm()].
42		#'
43		#' @param object (`mmrm`)\cr the MMRM fit.
44		#' @param contrast (`matrix`)\cr numeric contrast matrix, if given a `numeric`
45		#' then this is coerced to a row vector. Note that this should not include
46		#' elements for singular coefficient estimates, i.e. only refer to the
47		#' actually estimated coefficients.
48		#'
49		#' @return List with `num_df`, `denom_df`, `f_stat` and `p_val` (2-sided p-value).
50		#' @export
51		#'
52		#' @examples
53		#' object <- mmrm(
54		#' formula = FEV1 ~ RACE + SEX + ARMCD * AVISIT + us(AVISIT \| USUBJID),
55		#' data = fev_data
56		#' )
57		#' contrast <- matrix(data = 0, nrow = 2, ncol = length(object$beta_est))
58		#' contrast[1, 2] <- contrast[2, 3] <- 1
59		#' df_md(object, contrast)
60		df_md <- function(object, contrast) {
61	150x	assert_class(object, "mmrm")
62	150x	assert_numeric(contrast, any.missing = FALSE)
63	150x	if (!is.matrix(contrast)) {
64	113x	contrast <- matrix(contrast, ncol = length(contrast))
65		}
66	150x	assert_matrix(contrast, ncols = length(component(object, "beta_est")))
67	150x	if (nrow(contrast) == 0) {
68	1x	return(
69	1x	list(
70	1x	num_df = 0,
71	1x	denom_df = NA_real_,
72	1x	f_stat = NA_real_,
73	1x	p_val = NA_real_
74		)
75		)
76		}
77	149x	switch(object$method,
78	145x	"Satterthwaite" = h_df_md_sat(object, contrast),
79	3x	"Kenward-Roger" = h_df_md_kr(object, contrast),
80	!	"Residual" = h_df_md_res(object, contrast),
81	1x	"Between-Within" = h_df_md_bw(object, contrast),
82	!	stop("Unrecognized degrees of freedom method: ", object$method)
83		)
84		}
85
86		#' Creating T-Statistic Test Results For One-Dimensional Contrast
87		#'
88		#' @description Creates a list of results for one-dimensional contrasts using
89		#' a t-test statistic and the given degrees of freedom.
90		#'
91		#' @inheritParams df_1d
92		#' @param df (`number`)\cr degrees of freedom for the one-dimensional contrast.
93		#'
94		#' @return List with `est`, `se`, `df`, `t_stat` and `p_val` (2-sided p-value).
95		#'
96		#' @keywords internal
97		h_test_1d <- function(object,
98		contrast,
99		df) {
100	487x	assert_class(object, "mmrm")
101	487x	assert_numeric(contrast, len = length(component(object, "beta_est")))
102	487x	assert_number(df, lower = .Machine$double.xmin)
103
104	487x	est <- sum(contrast * component(object, "beta_est"))
105	487x	var <- h_quad_form_vec(contrast, component(object, "beta_vcov"))
106	487x	se <- sqrt(var)
107	487x	t_stat <- est / se
108	487x	p_val <- 2 * stats::pt(q = abs(t_stat), df = df, lower.tail = FALSE)
109
110	487x	list(
111	487x	est = est,
112	487x	se = se,
113	487x	df = df,
114	487x	t_stat = t_stat,
115	487x	p_val = p_val
116		)
117		}
118
119		#' Creating F-Statistic Test Results For Multi-Dimensional Contrast
120		#'
121		#' @description Creates a list of results for multi-dimensional contrasts using
122		#' an F-test statistic and the given degrees of freedom.
123		#'
124		#' @inheritParams df_md
125		#' @param contrast (`matrix`)\cr numeric contrast matrix.
126		#' @param df (`number`)\cr denominator degrees of freedom for the multi-dimensional contrast.
127		#' @param f_stat_factor (`number`)\cr optional scaling factor on top of the standard F-statistic.
128		#'
129		#' @return List with `num_df`, `denom_df`, `f_stat` and `p_val` (2-sided p-value).
130		#'
131		#' @keywords internal
132		h_test_md <- function(object,
133		contrast,
134		df,
135		f_stat_factor = 1) {
136	15x	assert_class(object, "mmrm")
137	15x	assert_matrix(contrast, ncols = length(component(object, "beta_est")))
138	15x	num_df <- nrow(contrast)
139	15x	assert_number(df, lower = .Machine$double.xmin)
140	15x	assert_number(f_stat_factor, lower = .Machine$double.xmin)
141
142	15x	prec_contrast <- solve(h_quad_form_mat(contrast, component(object, "beta_vcov")))
143	15x	contrast_est <- component(object, "beta_est") %*% t(contrast)
144	15x	f_statistic <- as.numeric(f_stat_factor / num_df * h_quad_form_mat(contrast_est, prec_contrast))
145	15x	p_val <- stats::pf(
146	15x	q = f_statistic,
147	15x	df1 = num_df,
148	15x	df2 = df,
149	15x	lower.tail = FALSE
150		)
151
152	15x	list(
153	15x	num_df = num_df,
154	15x	denom_df = df,
155	15x	f_stat = f_statistic,
156	15x	p_val = p_val
157		)
158		}

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	82x	assert_class(tmb_data, "mmrm_tmb_data")
16	82x	assert_numeric(theta_est)
17	82x	assert_matrix(beta_vcov)
18	82x	.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	5608x	vec <- as.vector(vec)
40	5608x	assert_numeric(vec, any.missing = FALSE)
41	5608x	assert_matrix(
42	5608x	center,
43	5608x	mode = "numeric",
44	5608x	any.missing = FALSE,
45	5608x	nrows = length(vec),
46	5608x	ncols = length(vec)
47		)
48
49	5608x	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	109x	assert_matrix(mat, mode = "numeric", any.missing = FALSE, min.cols = 1L)
63	109x	assert_matrix(
64	109x	center,
65	109x	mode = "numeric",
66	109x	any.missing = FALSE,
67	109x	nrows = ncol(center),
68	109x	ncols = ncol(center)
69		)
70
71	109x	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	491x	assert_list(jac_list)
90	491x	assert_numeric(contrast)
91
92	491x	vapply(
93	491x	jac_list,
94	491x	h_quad_form_vec,
95	491x	vec = contrast,
96	491x	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	457x	assert_class(object, "mmrm")
154	457x	contrast <- as.numeric(contrast)
155	457x	assert_numeric(contrast, len = length(component(object, "beta_est")))
156
157	457x	df <- if (identical(object$vcov, "Asymptotic")) {
158	444x	grad <- h_gradient(component(object, "jac_list"), contrast)
159	444x	v_num <- 2 * h_quad_form_vec(contrast, component(object, "beta_vcov"))^2
160	444x	v_denom <- h_quad_form_vec(grad, component(object, "theta_vcov"))
161	444x	v_num / v_denom
162	457x	} else if (
163	457x	object$vcov %in%
164	457x	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	457x	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	25x	assert_numeric(
193	25x	t_stat_df,
194	25x	min.len = 1L,
195	25x	lower = .Machine$double.xmin,
196	25x	any.missing = FALSE
197		)
198
199	25x	if (test_scalar(t_stat_df)) {
200	1x	t_stat_df
201	24x	} else if (all(abs(diff(t_stat_df)) < sqrt(.Machine$double.eps))) {
202	1x	mean(t_stat_df)
203	23x	} else if (any(t_stat_df <= 2)) {
204	2x	2
205		} else {
206	21x	e <- sum(t_stat_df / (t_stat_df - 2))
207	21x	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	134x	res_1d <- h_df_1d_sat(object, contrast)
224	134x	list(
225	134x	num_df = 1,
226	134x	denom_df = res_1d$df,
227	134x	f_stat = res_1d$t_stat^2,
228	134x	p_val = stats::pf(
229	134x	q = res_1d$t_stat^2,
230	134x	df1 = 1,
231	134x	df2 = res_1d$df,
232	134x	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	152x	assert_class(object, "mmrm")
251	152x	assert_matrix(
252	152x	contrast,
253	152x	mode = "numeric",
254	152x	any.missing = FALSE,
255	152x	ncols = length(component(object, "beta_est"))
256		)
257		# Early return if we are in the one-dimensional case.
258	152x	if (identical(nrow(contrast), 1L)) {
259	132x	return(h_df_md_from_1d(object, contrast))
260		}
261
262	20x	contrast_cov <- h_quad_form_mat(contrast, component(object, "beta_vcov"))
263	20x	eigen_cont_cov <- eigen(contrast_cov)
264	20x	eigen_cont_cov_vctrs <- eigen_cont_cov$vectors
265	20x	eigen_cont_cov_vals <- eigen_cont_cov$values
266
267	20x	eps <- sqrt(.Machine$double.eps)
268	20x	tol <- max(eps * eigen_cont_cov_vals[1], 0)
269	20x	rank_cont_cov <- sum(eigen_cont_cov_vals > tol)
270	20x	assert_number(rank_cont_cov, lower = .Machine$double.xmin)
271	20x	rank_seq <- seq_len(rank_cont_cov)
272	20x	vctrs_cont_prod <- crossprod(eigen_cont_cov_vctrs, contrast)[
273	20x	rank_seq, ,
274	20x	drop = FALSE
275		]
276
277		# Early return if rank 1.
278	20x	if (identical(rank_cont_cov, 1L)) {
279	1x	return(h_df_md_from_1d(object, vctrs_cont_prod))
280		}
281
282	19x	t_squared_nums <- drop(vctrs_cont_prod %*% object$beta_est)^2
283	19x	t_squared_denoms <- eigen_cont_cov_vals[rank_seq]
284	19x	t_squared <- t_squared_nums / t_squared_denoms
285	19x	f_stat <- sum(t_squared) / rank_cont_cov
286	19x	t_stat_df_nums <- 2 * eigen_cont_cov_vals^2
287	19x	t_stat_df <- if (identical(object$vcov, "Asymptotic")) {
288	18x	grads_vctrs_cont_prod <- lapply(
289	18x	rank_seq,
290	18x	function(m) {
291	46x	h_gradient(
292	46x	component(object, "jac_list"),
293	46x	contrast = vctrs_cont_prod[m, ]
294		)
295		}
296		)
297	18x	t_stat_df_denoms <- vapply(
298	18x	grads_vctrs_cont_prod,
299	18x	h_quad_form_vec,
300	18x	center = component(object, "theta_vcov"),
301	18x	numeric(1)
302		)
303	18x	t_stat_df_nums / t_stat_df_denoms
304		} else {
305	1x	vapply(
306	1x	rank_seq,
307	1x	function(m) {
308	2x	contrast_matrix <- Matrix::.bdiag(
309	2x	rep(
310	2x	list(vctrs_cont_prod[m, , drop = FALSE]),
311	2x	component(object, "n_subjects")
312		)
313		)
314	2x	contrast_matrix <- as.matrix(contrast_matrix)
315	2x	h_df_1d_sat_empirical(object, contrast_matrix)
316		},
317	1x	FUN.VALUE = 0
318		)
319		}
320	19x	denom_df <- h_md_denom_df(t_stat_df)
321
322	19x	list(
323	19x	num_df = rank_cont_cov,
324	19x	denom_df = denom_df,
325	19x	f_stat = f_stat,
326	19x	p_val = stats::pf(
327	19x	q = f_stat,
328	19x	df1 = rank_cont_cov,
329	19x	df2 = denom_df,
330	19x	lower.tail = FALSE
331		)
332		)
333		}

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("Estimate", "Std. Error", "df", "t value", "Pr(>\|t\|)")
62
63	40x	coef_aliased <- component(object, "beta_aliased")
64	40x	if (any(coef_aliased)) {
65	2x	names_coef_na <- names(which(coef_aliased))
66	2x	coef_na_table <- matrix(
67	2x	data = NA,
68	2x	nrow = length(names_coef_na),
69	2x	ncol = ncol(coef_table),
70	2x	dimnames = list(names_coef_na, colnames(coef_table))
71		)
72	2x	coef_table <- rbind(coef_table, coef_na_table)[names(coef_aliased), ]
73		}
74
75	40x	coef_table
76		}
77
78		#' @describeIn mmrm_methods summarizes the MMRM fit results.
79		#' @exportS3Method
80		#' @examples
81		#' # Summary:
82		#' summary(object)
83		summary.mmrm <- function(object, ...) {
84	20x	aic_list <- list(
85	20x	AIC = AIC(object),
86	20x	BIC = BIC(object),
87	20x	logLik = logLik(object),
88	20x	deviance = deviance(object)
89		)
90	20x	coefficients <- h_coef_table(object)
91	20x	call <- stats::getCall(object)
92	20x	components <- component(object, c(
93	20x	"cov_type", "reml", "n_groups", "n_theta",
94	20x	"n_subjects", "n_timepoints", "n_obs",
95	20x	"beta_vcov", "varcor"
96		))
97	20x	components$method <- object$method
98	20x	components$vcov <- object$vcov
99	20x	structure(
100	20x	c(
101	20x	components,
102	20x	list(
103	20x	coefficients = coefficients,
104	20x	n_singular_coefs = sum(component(object, "beta_aliased")),
105	20x	aic_list = aic_list,
106	20x	call = call
107		)
108		),
109	20x	class = "summary.mmrm"
110		)
111		}
112
113		#' Printing MMRM Function Call
114		#'
115		#' This is used in [print.summary.mmrm()].
116		#'
117		#' @param call (`call`)\cr original [mmrm()] function call.
118		#' @param n_obs (`int`)\cr number of observations.
119		#' @param n_subjects (`int`)\cr number of subjects.
120		#' @param n_timepoints (`int`)\cr number of timepoints.
121		#'
122		#' @keywords internal
123		h_print_call <- function(call, n_obs, n_subjects, n_timepoints) {
124	9x	pass <- 0
125	9x	if (!is.null(tmp <- call$formula)) {
126	9x	cat("Formula: ", deparse(tmp), fill = TRUE)
127	9x	rhs <- tmp[[2]]
128	9x	pass <- nchar(deparse(rhs))
129		}
130	9x	if (!is.null(call$data)) {
131	9x	cat(
132	9x	"Data: ", deparse(call$data), "(used", n_obs, "observations from",
133	9x	n_subjects, "subjects with maximum", n_timepoints, "timepoints)",
134	9x	fill = TRUE
135		)
136		}
137		# Display the expression of weights
138	9x	if (!is.null(call$weights)) {
139	4x	cat("Weights: ", deparse(call$weights), fill = TRUE)
140		}
141		}
142
143		#' Printing MMRM Covariance Type
144		#'
145		#' This is used in [print.summary.mmrm()].
146		#'
147		#' @param cov_type (`string`)\cr covariance structure abbreviation.
148		#' @param n_theta (`count`)\cr number of variance parameters.
149		#' @param n_groups (`count`)\cr number of groups.
150		#' @keywords internal
151		h_print_cov <- function(cov_type, n_theta, n_groups) {
152	9x	assert_string(cov_type)
153	9x	assert_count(n_theta, positive = TRUE)
154	9x	assert_count(n_groups, positive = TRUE)
155	9x	cov_definition <- switch(cov_type,
156	9x	us = "unstructured",
157	9x	toep = "Toeplitz",
158	9x	toeph = "heterogeneous Toeplitz",
159	9x	ar1 = "auto-regressive order one",
160	9x	ar1h = "heterogeneous auto-regressive order one",
161	9x	ad = "ante-dependence",
162	9x	adh = "heterogeneous ante-dependence",
163	9x	cs = "compound symmetry",
164	9x	csh = "heterogeneous compound symmetry",
165	9x	sp_exp = "spatial exponential"
166		)
167
168	9x	catstr <- sprintf(
169	9x	"Covariance: %s (%d variance parameters%s)\n",
170	9x	cov_definition,
171	9x	n_theta,
172	9x	ifelse(n_groups == 1L, "", sprintf(" of %d groups", n_groups))
173		)
174	9x	cat(catstr)
175		}
176
177		#' Printing AIC and other Model Fit Criteria
178		#'
179		#' This is used in [print.summary.mmrm()].
180		#'
181		#' @param aic_list (`list`)\cr list as part of from [summary.mmrm()].
182		#' @param digits (`number`)\cr number of decimal places used with [round()].
183		#'
184		#' @keywords internal
185		h_print_aic_list <- function(aic_list,
186		digits = 1) {
187	6x	diag_vals <- round(unlist(aic_list), digits)
188	6x	diag_vals <- format(diag_vals)
189	6x	print(diag_vals, quote = FALSE)
190		}
191
192		#' @describeIn mmrm_methods prints the MMRM fit summary.
193		#' @exportS3Method
194		#' @keywords internal
195		print.summary.mmrm <- function(x,
196		digits = max(3, getOption("digits") - 3),
197		signif.stars = getOption("show.signif.stars"), # nolint
198		...) {
199	5x	cat("mmrm fit\n\n")
200	5x	h_print_call(x$call, x$n_obs, x$n_subjects, x$n_timepoints)
201	5x	h_print_cov(x$cov_type, x$n_theta, x$n_groups)
202	5x	cat("Method: ", x$method, "\n", sep = "")
203	5x	cat("Vcov Method: ", x$vcov, "\n", sep = "")
204	5x	cat("Inference: ")
205	5x	cat(ifelse(x$reml, "REML", "ML"))
206	5x	cat("\n\n")
207	5x	cat("Model selection criteria:\n")
208	5x	h_print_aic_list(x$aic_list)
209	5x	cat("\n")
210	5x	cat("Coefficients: ")
211	5x	if (x$n_singular_coefs > 0) {
212	1x	cat("(", x$n_singular_coefs, " not defined because of singularities)", sep = "")
213		}
214	5x	cat("\n")
215	5x	stats::printCoefmat(
216	5x	x$coefficients,
217	5x	zap.ind = 3,
218	5x	digits = digits,
219	5x	signif.stars = signif.stars
220		)
221	5x	cat("\n")
222	5x	cat("Covariance estimate:\n")
223	5x	if (is.list(x$varcor)) {
224	1x	for (v in names(x$varcor)) {
225	2x	cat(sprintf("Group: %s\n", v))
226	2x	print(round(x$varcor[[v]], digits = digits))
227		}
228		} else {
229	4x	print(round(x$varcor, digits = digits))
230		}
231	5x	cat("\n")
232	5x	invisible(x)
233		}
234
235
236		#' @describeIn mmrm_methods obtain the confidence intervals for the coefficients.
237		#' @exportS3Method
238		#' @examples
239		#' # Confidence Interval:
240		#' confint(object)
241		confint.mmrm <- function(object, parm, level = 0.95, ...) {
242	20x	cf <- coef(object)
243	20x	pnames <- names(cf)
244	20x	if (missing(parm)) {
245	15x	parm <- pnames
246		}
247	20x	assert(
248	20x	check_subset(parm, pnames),
249	20x	check_integerish(parm, lower = 1L, upper = length(cf))
250		)
251	2x	if (is.numeric(parm)) parm <- pnames[parm]
252	18x	assert_number(level, lower = 0, upper = 1)
253	18x	a <- (1 - level) / 2
254	18x	pct <- paste(format(100 * c(a, 1 - a), trim = TRUE, scientific = FALSE, digits = 3), "%")
255	18x	coef_table <- h_coef_table(object)
256	18x	df <- coef_table[parm, "df"]
257	18x	ses <- coef_table[parm, "Std. Error"]
258	18x	fac <- stats::qt(a, df = df)
259	18x	ci <- array(NA, dim = c(length(parm), 2L), dimnames = list(parm, pct))
260	18x	sefac <- ses * fac
261	18x	ci[] <- cf[parm] + c(sefac, -sefac)
262	18x	ci
263		}

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

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	296x	specials <- cov_types(c("abbr", "habbr"))
12	296x	terms <- stats::terms(formula_rhs(f), specials = specials)
13	296x	covariance_terms <- Filter(length, attr(terms, "specials"))
14	296x	variables <- attr(terms, "variables")
15	296x	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	279x	specials <- cov_types(c("abbr", "habbr"))
31
32	279x	terms <- stats::terms(f, specials = specials)
33	279x	covariance_terms <- Filter(Negate(is.null), attr(terms, "specials"))
34
35		# if no covariance terms were found, return original formula
36	279x	if (length(covariance_terms) == 0) {
37	6x	return(f)
38		}
39	273x	if (length(f) != 3) {
40	1x	update_str <- "~ . -"
41		} else {
42	272x	update_str <- ". ~ . -"
43		}
44	273x	stats::update(
45	273x	f,
46	273x	stats::as.formula(paste(update_str, deparse(attr(terms, "variables")[[covariance_terms[[1]] + 1]])))
47		)
48		}
49
50		#' Add Individual Covariance Variables As Terms to Formula
51		#'
52		#' @param f (`formula`)\cr a formula to which covariance structure terms should
53		#' be added.
54		#' @param covariance (`cov_struct`)\cr a covariance structure object from which
55		#' additional variables should be sourced.
56		#'
57		#' @return A new formula with included covariance terms.
58		#'
59		#' @details [stats::update()] is used to append the covariance structure and the environment
60		#' attribute will not be changed. This ensures the returned formula and the input formula
61		#' have the same environment.
62		#'
63		#' @keywords internal
64		h_add_covariance_terms <- function(f, covariance) {
65	277x	cov_terms <- with(covariance, c(subject, visits, group))
66	271x	cov_terms <- paste(cov_terms, collapse = " + ")
67	271x	stats::update(f, stats::as.formula(paste(". ~ . + ", cov_terms)))
68		}
69
70		#' Add Formula Terms with Character
71		#'
72		#' Add formula terms from the original formula with character representation.
73		#'
74		#' @param f (`formula`)\cr a formula to be updated.
75		#' @param adds (`character`)\cr representation of elements to be added.
76		#' @param drop_response (`flag`)\cr whether response should be dropped.
77		#'
78		#' @details Elements in `adds` will be added from the formula, while the environment
79		#' of the formula is unchanged. If `adds` is `NULL` or `character(0)`, the formula is
80		#' unchanged.
81		#' @return A new formula with elements in `drops` removed.
82		#'
83		#' @keywords internal
84		h_add_terms <- function(f, adds, drop_response = FALSE) {
85	599x	assert_character(adds, null.ok = TRUE)
86	599x	if (length(adds) > 0L) {
87	321x	add_terms <- stats::as.formula(sprintf(". ~ . + %s", paste(adds, collapse = "+")))
88	321x	f <- stats::update(f, add_terms)
89		}
90	599x	if (drop_response && length(f) == 3L) {
91	35x	f[[2]] <- NULL
92		}
93	599x	f
94		}

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	560x	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	280x	flattened_args <- unlist(lapply(call[-1], flatten_expr))
56	280x	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	1255x	if (length(expr) > 1 && is_infix(expr[[1]])) {
69	337x	op <- list(expr[[1]])
70	337x	lhs <- flatten_expr(expr[[2]])
71	337x	rhs <- flatten_expr(expr[[3]])
72	337x	c(lhs, op, rhs)
73		} else {
74	918x	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	299x	if (length(f) == 2) {
93	9x	f
94		} else {
95	290x	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	344x	"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(collapse = ", ", lapply(x, function(i) {
130	16x	utils::capture.output(as.symbol(i))
131		}))
132		}

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	1x	if (!requireNamespace("car", quietly = quietly)) {
9	!	return(FALSE)
10		}
11	1x	envir <- asNamespace("mmrm")
12	1x	h_register_s3("car", "Anova", "mmrm", envir)
13	1x	TRUE
14		}
15
16
17		#' Obtain Contrast for Specified Effect
18		#'
19		#' This is support function to obtain contrast matrix for type II/III testing.
20		#'
21		#' @param object (`mmrm`)\cr the fitted MMRM.
22		#' @param effect (`string`) the name of the effect.
23		#' @param type (`string`) type of test, "II", "III", '2', or '3'.
24		#' @param tol (`numeric`) threshold blow which values are treated as 0.
25		#'
26		#' @return A `matrix` of the contrast.
27		#'
28		#' @keywords internal
29		h_get_contrast <- function(object, effect, type = c("II", "III", "2", "3"), tol = sqrt(.Machine$double.eps)) {
30	45x	assert_class(object, "mmrm")
31	45x	assert_string(effect)
32	45x	assert_double(tol, finite = TRUE, len = 1L)
33	45x	type <- match.arg(type)
34	45x	mx <- component(object, "x_matrix")
35	45x	asg <- attr(mx, "assign")
36	45x	formula <- object$formula_parts$model_formula
37	45x	tms <- terms(formula)
38	45x	fcts <- attr(tms, "factors")[-1L, , drop = FALSE] # Discard the response.
39	45x	ods <- attr(tms, "order")
40	45x	assert_subset(effect, colnames(fcts))
41	45x	idx <- which(effect == colnames(fcts))
42	45x	cols <- which(asg == idx)
43	45x	xlev <- component(object, "xlev")
44	45x	contains_intercept <- (!0 %in% asg) && h_first_contain_categorical(effect, fcts, names(xlev))
45	45x	coef_rows <- length(cols) - as.integer(contains_intercept)
46	45x	l_mx <- matrix(0, nrow = coef_rows, ncol = length(asg))
47	45x	if (coef_rows == 0L) {
48	1x	return(l_mx)
49		}
50	44x	if (contains_intercept) {
51	4x	l_mx[, cols] <- cbind(-1, diag(rep(1, coef_rows)))
52		} else {
53	40x	l_mx[, cols] <- diag(rep(1, coef_rows))
54		}
55	44x	for (i in setdiff(seq_len(ncol(fcts)), idx)) {
56	120x	additional_vars <- names(which(fcts[, i] > fcts[, idx]))
57	120x	additional_numeric <- any(!additional_vars %in% names(xlev))
58	120x	current_col <- which(asg == i)
59	120x	if (ods[i] >= ods[idx] && all(fcts[, i] >= fcts[, idx]) && !additional_numeric) {
60	24x	sub_mat <- switch(type,
61	24x	"2" = ,
62	24x	"II" = {
63	8x	x1 <- mx[, cols, drop = FALSE]
64	8x	x0 <- mx[, -c(cols, current_col), drop = FALSE]
65	8x	x2 <- mx[, current_col, drop = FALSE]
66	8x	m <- diag(rep(1, nrow(x0))) - x0 %% solve(t(x0) %% x0) %*% t(x0)
67	8x	ret <- solve(t(x1) %% m %% x1) %% t(x1) %% m %*% x2
68	8x	if (contains_intercept) {
69	1x	ret[-1, ] - ret[1, ]
70		} else {
71	7x	ret
72		}
73		},
74	24x	"3" = ,
75	24x	"III" = {
76	16x	lvls <- h_obtain_lvls(effect, additional_vars, xlev)
77	16x	t_levels <- lvls$total
78	16x	nms_base <- colnames(mx)[cols]
79	16x	nms <- colnames(mx)[current_col]
80	16x	nms_base_split <- strsplit(nms_base, ":")
81	16x	nms_split <- strsplit(nms, ":")
82	16x	base_idx <- h_get_index(nms_split, nms_base_split)
83	16x	mt <- l_mx[, cols, drop = FALSE] / t_levels
84	16x	ret <- mt[, base_idx, drop = FALSE]
85		# if there is extra levels, replace it with -1/t_levels
86	16x	ret[is.na(ret)] <- -1 / t_levels
87	16x	ret
88		}
89		)
90	24x	l_mx[, current_col] <- sub_mat
91		}
92		}
93	44x	l_mx[abs(l_mx) < tol] <- 0
94	44x	l_mx
95		}
96
97		#' Conduct type II/III hypothesis testing on the MMRM fit results.
98		#'
99		#' @param mod (`mmrm`)\cr the fitted MMRM.
100		#' @param ... not used.
101		#' @inheritParams h_get_contrast
102		#'
103		#' @details
104		#' `Anova` will return `anova` object with one row per variable and columns
105		#' `Num Df`(numerator degrees of freedom), `Denom Df`(denominator degrees of freedom),
106		#' `F Statistic` and `Pr(>=F)`.
107		#'
108		#' @keywords internal
109		# Please do not load `car` and then create the documentation. The Rd file will be different.
110		Anova.mmrm <- function(mod, type = c("II", "III", "2", "3"), tol = sqrt(.Machine$double.eps), ...) { # nolint
111	9x	assert_double(tol, finite = TRUE, len = 1L)
112	9x	type <- match.arg(type)
113	9x	vars <- colnames(attr(terms(mod$formula_parts$model_formula), "factors"))
114	9x	ret <- lapply(
115	9x	vars,
116	9x	function(x) df_md(mod, h_get_contrast(mod, x, type, tol))
117		)
118	9x	ret_df <- do.call(rbind.data.frame, ret)
119	9x	row.names(ret_df) <- vars
120	9x	colnames(ret_df) <- c("Num Df", "Denom Df", "F Statistic", "Pr(>=F)")
121	9x	class(ret_df) <- c("anova", "data.frame")
122	9x	attr(ret_df, "heading") <- sprintf(
123	9x	"Analysis of Fixed Effect Table (Type %s F tests)",
124	9x	switch(type,
125	9x	"2" = ,
126	9x	"II" = "II",
127	9x	"3" = ,
128	9x	"III" = "III"
129		)
130		)
131	9x	ret_df
132		}
133
134
135		#' Obtain Levels Prior and Posterior
136		#' @param var (`string`) name of the effect.
137		#' @param additional_vars (`character`) names of additional variables.
138		#' @param xlev (`list`) named list of character levels.
139		#' @param factors (`matrix`) the factor matrix.
140		#' @keywords internal
141		h_obtain_lvls <- function(var, additional_vars, xlev, factors) {
142	18x	assert_string(var)
143	18x	assert_character(additional_vars)
144	18x	assert_list(xlev, types = "character")
145	18x	nms <- names(xlev)
146	18x	assert_subset(additional_vars, nms)
147	18x	if (var %in% nms) {
148	14x	prior_vars <- intersect(nms[seq_len(match(var, nms) - 1)], additional_vars)
149	14x	prior_lvls <- vapply(xlev[prior_vars], length, FUN.VALUE = 1L)
150	14x	post_vars <- intersect(nms[seq(match(var, nms) + 1, length(nms))], additional_vars)
151	14x	post_lvls <- vapply(xlev[post_vars], length, FUN.VALUE = 1L)
152	14x	total_lvls <- prod(prior_lvls) * prod(post_lvls)
153		} else {
154	4x	prior_lvls <- vapply(xlev[additional_vars], length, FUN.VALUE = 1L)
155	4x	post_lvls <- 2L
156	4x	total_lvls <- prod(prior_lvls)
157		}
158	18x	list(
159	18x	prior = prior_lvls,
160	18x	post = post_lvls,
161	18x	total = total_lvls
162		)
163		}
164
165		#' Check if the Effect is the First Categorical Effect
166		#' @param effect (`string`) name of the effect.
167		#' @param categorical (`character`) names of the categorical values.
168		#' @param factors (`matrix`) the factor matrix.
169		#' @keywords internal
170		h_first_contain_categorical <- function(effect, factors, categorical) {
171	9x	assert_string(effect)
172	9x	assert_matrix(factors)
173	9x	assert_character(categorical)
174	9x	mt <- match(effect, colnames(factors))
175	9x	varnms <- row.names(factors)
176		# if the effect is not categorical in any value, return FALSE
177	9x	if (!any(varnms[factors[, mt] > 0] %in% categorical)) {
178	2x	return(FALSE)
179		}
180		# keep only categorical rows that is in front of the current factor
181	7x	factors <- factors[row.names(factors) %in% categorical, seq_len(mt - 1L), drop = FALSE]
182		# if previous cols are all numerical, return TRUE
183	7x	if (ncol(factors) < 1L) {
184	4x	return(TRUE)
185		}
186	3x	col_ind <- apply(factors, 2, prod)
187		# if any of the previous cols are categorical, return FALSE
188	3x	return(!any(col_ind > 0))
189		}
190
191		#' Test if the First Vector is Subset of the Second Vector
192		#' @param x (`vector`) the first list.
193		#' @param y (`vector`) the second list.
194		#' @keywords internal
195		h_get_index <- function(x, y) {
196	18x	assert_list(x)
197	18x	assert_list(y)
198	18x	vapply(
199	18x	x,
200	18x	\(i) {
201	68x	r <- vapply(y, \(j) test_subset(j, i), FUN.VALUE = TRUE)
202	68x	if (sum(r) == 1L) {
203	65x	which(r)
204		} else {
205	18x	NA_integer_
206		}
207		},
208	18x	FUN.VALUE = 1L
209		)
210		}

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(names(bw_calc), identical.to = c("coefs_between_within", "ddf_between", "ddf_within"))
84	17x	assert_logical(is_coef_involved, len = length(bw_calc$coefs_between_within))
85	17x	assert_true(sum(is_coef_involved) > 0)
86
87	17x	coef_categories <- bw_calc$coefs_between_within[is_coef_involved]
88	17x	coef_dfs <- vapply(
89	17x	X = coef_categories,
90	17x	FUN = switch,
91	17x	intercept = bw_calc$ddf_within,
92	17x	between = bw_calc$ddf_between,
93	17x	within = bw_calc$ddf_within,
94	17x	FUN.VALUE = integer(1)
95		)
96	17x	min(coef_dfs)
97		}
98
99		#' Calculation of Between-Within Degrees of Freedom for One-Dimensional Contrast
100		#'
101		#' @description Used in [df_1d()] if method is "Between-Within".
102		#'
103		#' @inheritParams h_df_1d_sat
104		#' @inherit h_df_1d_sat return
105		#' @keywords internal
106		h_df_1d_bw <- function(object, contrast) {
107	7x	assert_class(object, "mmrm")
108	7x	assert_numeric(contrast, len = length(component(object, "beta_est")))
109
110	7x	bw_calc <- h_df_bw_calc(object)
111	7x	is_coef_involved <- contrast != 0
112	7x	df <- h_df_min_bw(bw_calc, is_coef_involved)
113	7x	h_test_1d(object, contrast, df)
114		}
115
116		#' Calculation of Between-Within Degrees of Freedom for Multi-Dimensional Contrast
117		#'
118		#' @description Used in [df_md()] if method is "Between-Within".
119		#'
120		#' @inheritParams h_df_md_sat
121		#' @inherit h_df_md_sat return
122		#' @keywords internal
123		h_df_md_bw <- function(object, contrast) {
124	7x	assert_class(object, "mmrm")
125	7x	assert_matrix(contrast, mode = "numeric", any.missing = FALSE, ncols = length(component(object, "beta_est")))
126
127	7x	bw_calc <- h_df_bw_calc(object)
128	7x	is_coef_involved <- apply(X = contrast != 0, MARGIN = 2L, FUN = any)
129	7x	df <- h_df_min_bw(bw_calc, is_coef_involved)
130	7x	h_test_md(object, contrast, df)
131		}

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		recover_data.mmrm <- function(object, ...) { # nolint
30	13x	fun_call <- stats::getCall(object)
31		# subject_var is excluded because it should not contain fixed effect.
32		# visit_var is not excluded because emmeans can provide marginal mean
33		# by each visit if visit_var is not spatial.
34	13x	model_frame <- stats::model.frame(
35	13x	object,
36	13x	include = c(
37	13x	if (!object$formula_parts$is_spatial) "visit_var" else NULL,
38	13x	"response_var", "group_var"
39		)
40		)
41	13x	model_terms <- stats::delete.response(stats::terms(model_frame))
42	13x	emmeans::recover_data(
43	13x	fun_call,
44	13x	trms = model_terms,
45	13x	na.action = "na.omit",
46	13x	frame = model_frame,
47		...
48		)
49		}
50
51		#' Returns a List of Model Details for `emmeans` Purposes
52		#'
53		#' @seealso See [emmeans::emm_basis()] for background.
54		#' @keywords internal
55		#' @noRd
56		emm_basis.mmrm <- function(object, # nolint
57		trms,
58		xlev,
59		grid,
60		...) {
61	13x	model_frame <- stats::model.frame(trms, grid, na.action = stats::na.pass, xlev = xlev)
62	13x	contrasts <- component(object, "contrasts")
63	13x	model_mat <- stats::model.matrix(trms, model_frame, contrasts.arg = contrasts)
64	13x	beta_hat <- component(object, "beta_est")
65	13x	nbasis <- if (length(beta_hat) < ncol(model_mat)) {
66	6x	kept <- match(names(beta_hat), colnames(model_mat))
67	6x	beta_hat <- NA * model_mat[1L, ]
68	6x	beta_hat[kept] <- component(object, "beta_est")
69	6x	orig_model_mat <- stats::model.matrix(
70	6x	trms,
71	6x	stats::model.frame(
72	6x	object,
73	6x	include = c(
74	6x	if (!object$formula_parts$is_spatial) "visit_var" else NULL,
75	6x	"response_var", "group_var"
76		)
77		),
78	6x	contrasts.arg = contrasts
79		)
80	6x	estimability::nonest.basis(orig_model_mat)
81		} else {
82	7x	estimability::all.estble
83		}
84	13x	dfargs <- list(object = object)
85	13x	dffun <- function(k, dfargs) {
86	113x	mmrm::df_md(dfargs$object, contrast = k)$denom_df
87		}
88	13x	list(
89	13x	X = model_mat,
90	13x	bhat = beta_hat,
91	13x	nbasis = nbasis,
92	13x	V = component(object, "beta_vcov"),
93	13x	dffun = dffun,
94	13x	dfargs = dfargs
95		)
96		}

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

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

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

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

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(contrast, mode = "numeric", any.missing = FALSE, ncols = length(component(object, "beta_est")))
28
29	1x	df <- component(object, "n_obs") - length(component(object, "beta_est"))
30
31	1x	h_test_md(object, contrast, df)
32		}

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		#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	10694x	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	32370x	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	10790x	this->theta = theta;
34	32370x	std::iota(std::begin(this->full_visit), std::end(this->full_visit), 0);
35	10790x	this->n_theta = theta.size();
36	10790x	this->chol_full = get_covariance_lower_chol(this->theta, this->n_visits, this->cov_type);
37	10786x	this->chols[full_visit] = this->chol_full;
38	10786x	this->sigma_full = tcrossprod(this->chol_full, true);
39		}
40	1174433x	matrix<Type> get_chol(std::vector<int> visits, matrix<Type> dist) {
41	1174433x	auto target = this->chols.find(visits);
42	1174433x	if (target != this->chols.end()) {
43	1084762x	return target->second;
44		} else {
45	89671x	matrix<Type> cov_i = this->get_sigma(visits, dist);
46	89671x	Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > cov_i_chol(cov_i);
47	89671x	matrix<Type> Li = cov_i_chol.matrixL();
48	89671x	this->chols[visits] = Li;
49	89671x	return Li;
50		}
51		}
52	610681x	matrix<Type> get_sigma(std::vector<int> visits, matrix<Type> dist) {
53	610681x	auto target = this->sigmas.find(visits);
54	610681x	if (target != this->sigmas.end()) {
55	484366x	return target->second;
56		} else {
57	126315x	matrix<Type> ret = subset_matrix<matrix<Type>, vector<int>>(sigma_full, visits, visits);
58	126315x	this->sigmas[visits] = ret;
59	126315x	return ret;
60		}
61		}
62	208882x	matrix<Type> get_sigma_inverse(std::vector<int> visits, matrix<Type> dist) {
63	208882x	auto target = this->sigmas_inv.find(visits);
64	208882x	if (target != this->sigmas_inv.end()) {
65	182476x	return target->second;
66		} else {
67	26406x	matrix<Type> ret = this->get_sigma(visits, dist).inverse();
68	26406x	this->sigmas_inv[visits] = ret;
69	26406x	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	200x	lower_chol_spatial(vector<Type> theta, std::string cov_type): theta(theta), cov_type(cov_type) {
85		}
86	44897x	matrix<Type> get_chol(std::vector<int> visits, matrix<Type> dist) {
87	44897x	return get_spatial_covariance_lower_chol(this->theta, dist, this->cov_type);
88		}
89	15780x	matrix<Type> get_sigma(std::vector<int> visits, matrix<Type> dist) {
90	15780x	return tcrossprod(this->get_chol(visits, dist), true);
91		}
92	5912x	matrix<Type> get_sigma_inverse(std::vector<int> visits, matrix<Type> dist) {
93	5912x	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	10094x	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	10094x	int theta_one_group_size = theta.size() / n_groups;
106	21074x	for (int r = 0; r < n_groups; r++) {
107		// Use unique pointers here to better manage resource.
108	10984x	if (is_spatial) {
109	198x	this->cache[r] = std::make_shared<D1>(theta.segment(r * theta_one_group_size, theta_one_group_size), cov_type);
110		} else {
111	10786x	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	9758x	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	6360x	matrix<Type> get_default_chol() {
125	12720x	std::vector<int> visit(this->n_visits);
126	6360x	std::iota(std::begin(visit), std::end(visit), 0);
127	6360x	matrix<Type> dist(2, 2);
128	6360x	dist << 0, 1, 1, 0;
129	6360x	int dim = this->is_spatial?2:this->n_visits;
130	6360x	matrix<Type> covariance_lower_chol = matrix<Type>::Zero(dim * this->n_groups, dim);
131	13368x	for (int r = 0; r < this->n_groups; r++) {
132	7008x	covariance_lower_chol.block(r * dim, 0, dim, dim) = this->cache[r]->get_chol(visit, dist);
133		}
134	12720x	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	17520x	matrix<T> get_unstructured(const vector<T>& theta, int n_visits) {
10	17520x	vector<T> sd_values = exp(theta.head(n_visits));
11	17520x	vector<T> lower_tri_chol_values = theta.tail(theta.size() - n_visits);
12	17520x	matrix<T> covariance_lower_chol = matrix<T>::Zero(n_visits, n_visits);
13	17520x	int k = 0;
14	87568x	for(int i = 0; i < n_visits; i++) {
15	70048x	covariance_lower_chol(i, i) = sd_values(i);
16	175184x	for(int j = 0; j < i; j++){
17	105136x	covariance_lower_chol(i, j) = sd_values(i) * lower_tri_chol_values(k++);
18		}
19		}
20	35040x	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	4452x	const T operator() (int i, int j) const {
30	4452x	return this->corr_values.segment(j, i - j).prod();
31		}
32		};
33		// Homogeneous Ante-dependence Cholesky factor.
34		template <class T>
35	316x	matrix<T> get_ante_dependence(const vector<T>& theta, int n_visits) {
36	316x	T const_sd = exp(theta(0));
37	316x	corr_fun_ante_dependence<T> fun(theta.tail(n_visits - 1));
38	316x	matrix<T> ad_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
39	632x	return const_sd * ad_cor_mat_chol;
40		}
41		// Heterogeneous Ante-dependence Cholesky factor.
42		template <class T>
43	476x	matrix<T> get_ante_dependence_heterogeneous(const vector<T>& theta, int n_visits) {
44	476x	vector<T> sd_values = exp(theta.head(n_visits));
45	476x	corr_fun_ante_dependence<T> fun(theta.tail(n_visits - 1));
46	952x	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	5076x	const T operator() (int i, int j) const {
56	5076x	int index = (i - j) - 1; // Note: We need to start at 0.
57	5076x	return this->corr_values(index);
58		}
59		};
60		// Homogeneous Toeplitz Cholesky factor.
61		template <class T>
62	416x	matrix<T> get_toeplitz(const vector<T>& theta, int n_visits) {
63	416x	T const_sd = exp(theta(0));
64	416x	corr_fun_toeplitz<T> fun(theta.tail(n_visits - 1));
65	416x	matrix<T> toep_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
66	832x	return const_sd * toep_cor_mat_chol;
67		}
68		// Heterogeneous Toeplitz Cholesky factor.
69		template <class T>
70	428x	matrix<T> get_toeplitz_heterogeneous(const vector<T>& theta, int n_visits) {
71	428x	vector<T> sd_values = exp(theta.head(n_visits));
72	428x	corr_fun_toeplitz<T> fun(theta.tail(n_visits - 1));
73	856x	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	20600x	const T operator() (int i, int j) const {
83	20600x	T diff = T((i - j) * 1.0);
84	26336x	return pow(this->corr_values(0), diff); // rho^{\|i-j\|}
85		}
86		};
87		// Homogeneous autoregressive Cholesky factor.
88		template <class T>
89	2996x	matrix<T> get_auto_regressive(const vector<T>& theta, int n_visits) {
90	2996x	T const_sd = exp(theta(0));
91	2996x	corr_fun_autoregressive<T> fun(theta.tail(1));
92	2996x	matrix<T> ar1_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
93	5992x	return const_sd * ar1_cor_mat_chol;
94		}
95		// Heterogeneous autoregressive Cholesky factor.
96		template <class T>
97	428x	matrix<T> get_auto_regressive_heterogeneous(const vector<T>& theta, int n_visits) {
98	428x	vector<T> sd_values = exp(theta.head(n_visits));
99	428x	corr_fun_autoregressive<T> fun(theta.tail(1));
100	856x	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	1152x	corr_fun_compound_symmetry(const vector<T>& theta, int n_visits) :
111	1152x	corr_values(map_to_cs_cor(theta, n_visits)) {}
112	6912x	const T operator() (int i, int j) const {
113	6912x	return this->corr_values(0); // rho (constant)
114		}
115		};
116		// Homogeneous compound symmetry Cholesky factor.
117		template <class T>
118	620x	matrix<T> get_compound_symmetry(const vector<T>& theta, int n_visits) {
119	620x	T const_sd = exp(theta(0));
120	620x	corr_fun_compound_symmetry<T> fun(theta.tail(1), n_visits);
121	620x	matrix<T> cs_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
122	1240x	return const_sd * cs_cor_mat_chol;
123		}
124		// Heterogeneous compound symmetry Cholesky factor.
125		template <class T>
126	524x	matrix<T> get_compound_symmetry_heterogeneous(const vector<T>& theta, int n_visits) {
127	524x	vector<T> sd_values = exp(theta.head(n_visits));
128	524x	corr_fun_compound_symmetry<T> fun(theta.tail(1), n_visits);
129	1048x	return get_heterogeneous_cov(sd_values, fun);
130		}
131
132		// Spatial Exponential Cholesky factor.
133		template <class T>
134	44897x	matrix<T> get_spatial_exponential(const vector<T>& theta, const matrix<T>& distance) {
135	44897x	T const_sd = exp(theta(0));
136	44897x	T rho = invlogit(theta(1));
137	44897x	matrix<T> expdist = exp(distance.array() * log(rho));
138	44897x	matrix<T> result = expdist * const_sd;
139	44897x	Eigen::LLT<Eigen::Matrix<T,Eigen::Dynamic,Eigen::Dynamic> > cov_i_chol(result);
140	89794x	return cov_i_chol.matrixL();
141		}
142
143		// Creates a new correlation object dynamically.
144		template <class T>
145	23692x	matrix<T> get_covariance_lower_chol(const vector<T>& theta, int n_visits, std::string cov_type) {
146	23692x	matrix<T> result;
147
148	23692x	if (cov_type == "us") {
149	17512x	result = get_unstructured<T>(theta, n_visits);
150	6180x	} else if (cov_type == "toep") {
151	412x	result = get_toeplitz<T>(theta, n_visits);
152	5768x	} else if (cov_type == "toeph") {
153	424x	result = get_toeplitz_heterogeneous<T>(theta, n_visits);
154	5344x	} else if (cov_type == "ar1") {
155	2992x	result = get_auto_regressive<T>(theta, n_visits);
156	2352x	} else if (cov_type == "ar1h") {
157	424x	result = get_auto_regressive_heterogeneous<T>(theta, n_visits);
158	1928x	} else if (cov_type == "ad") {
159	312x	result = get_ante_dependence<T>(theta, n_visits);
160	1616x	} else if (cov_type == "adh") {
161	472x	result = get_ante_dependence_heterogeneous<T>(theta, n_visits);
162	1144x	} else if (cov_type == "cs") {
163	616x	result = get_compound_symmetry<T>(theta, n_visits);
164	528x	} else if (cov_type == "csh") {
165	520x	result = get_compound_symmetry_heterogeneous<T>(theta, n_visits);
166		} else {
167	4x	Rf_error("%s", ("Unknown covariance type '" + cov_type + "'.").c_str());
168		}
169
170	23684x	return result;
171		}
172
173		// Creates a new spatial covariance cholesky.
174		template <class T>
175	44897x	matrix<T> get_spatial_covariance_lower_chol(const vector<T>& theta, const matrix<T>& distance, std::string cov_type) {
176	44897x	matrix<T> result;
177	44897x	if (cov_type == "sp_exp") {
178	44897x	result = get_spatial_exponential<T>(theta, distance);
179		} else {
180	!	Rf_error("%s", ("Unknown spatial covariance type '" + cov_type + "'.").c_str());
181		}
182	44897x	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	704x	chol(int dim, string cov): dim_cov_mat(dim), cov_type(cov) {};
16		template <class T>
17	2108x	vector<T> operator() (vector<T> &theta) {
18	2108x	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	354x	chol_jacobian(int dim, string cov): dim_cov_mat(dim), cov_type(cov), mychol(dim, cov) {};
28		template<class T>
29	356x	vector<T> operator() (vector<T> &theta) {
30	356x	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	350x	std::map<std::string, matrix<Type>> derivatives(int n_visits, std::string cov_type, vector<Type> theta) {
39	350x	std::map<std::string, matrix<Type>> ret;
40	350x	chol chol_obj(n_visits, cov_type);
41	350x	chol_jacobian chol_jac_obj(n_visits, cov_type);
42	350x	matrix<Type> l = chol_obj(theta).matrix();
43	350x	l.resize(n_visits, n_visits);
44	350x	vector<Type> chol_d1_vec = autodiff::jacobian(chol_obj, theta).vec(); // chol_d1_vec is (dim * dim * l_theta)
45	350x	vector<Type> chol_d2_vec = autodiff::jacobian(chol_jac_obj, theta).vec(); // chol_d2_vec is (dim * dim * l_theta * l_theta)
46	350x	matrix<Type> ret_d1 = matrix<Type>(n_visits * theta.size(), n_visits);
47	350x	matrix<Type> ret_d2 = matrix<Type>(n_visits * theta.size() * theta.size(), n_visits);
48	350x	int n_visits_sq = n_visits * n_visits;
49	2284x	for (int i = 0; i < theta.size(); i++) {
50	1934x	matrix<Type> ld1 = chol_d1_vec.segment(i * n_visits_sq, n_visits_sq).matrix();
51	1934x	ld1.resize(n_visits, n_visits);
52	1934x	matrix<Type> ld1_lt = ld1 * l.transpose();
53	1934x	auto sigma_d1_i = ld1_lt + ld1_lt.transpose();
54	1934x	ret_d1.block(i * n_visits, 0, n_visits, n_visits) = sigma_d1_i;
55	17092x	for (int j = 0; j < theta.size(); j++) {
56	15158x	matrix<Type> ld2 = chol_d2_vec.segment( (j * theta.size() + i) * n_visits_sq, n_visits_sq).matrix();
57	15158x	matrix<Type> ld1_j = chol_d1_vec.segment(j * n_visits_sq, n_visits_sq).matrix();
58	15158x	ld2.resize(n_visits, n_visits);
59	15158x	ld1_j.resize(n_visits, n_visits);
60	15158x	auto ld2_lt = ld2 * l.transpose();
61	15158x	auto ld1_ld1j = ld1 * ld1_j.transpose();
62	15158x	auto sigma_d2_ij = ld2_lt + ld2_lt.transpose() + ld1_ld1j + ld1_ld1j.transpose();
63	15158x	ret_d2.block((i * theta.size() + j) * n_visits, 0, n_visits, n_visits) = sigma_d2_ij;
64		}
65		}
66	700x	ret["derivative1"] = ret_d1;
67	350x	ret["derivative2"] = ret_d2;
68	700x	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	374x	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	350x	derivatives_nonspatial(vector<Type> theta, int n_visits, std::string cov_type): lower_chol_nonspatial<Type>(theta, n_visits, cov_type) {
96	350x	std::map<std::string, tmbutils::matrix<Type>> allret = derivatives<Type>(this->n_visits, this->cov_type, this->theta);
97	700x	matrix<Type> sigma_d1 = allret["derivative1"];
98	350x	matrix<Type> sigma_d2 = allret["derivative2"];
99	350x	this->sigmad1_cache[this->full_visit] = sigma_d1;
100	350x	this->sigmad2_cache[this->full_visit] = sigma_d2;
101		}
102		// Cache and return the first order derivatives using select matrix.
103	20600x	matrix<Type> get_sigma_derivative1(std::vector<int> visits, matrix<Type> dist) override {
104	20600x	auto target = this->sigmad1_cache.find(visits);
105	20600x	if (target != this->sigmad1_cache.end()) {
106	16822x	return target->second;
107		} else {
108	3778x	int n_visits_i = visits.size();
109	3778x	matrix<Type> ret = matrix<Type>(this->n_theta * n_visits_i, n_visits_i);
110	26068x	for (int i = 0; i < this->n_theta; i++) {
111	22290x	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	3778x	this->sigmad1_cache[visits] = ret;
114	3778x	return ret;
115		}
116		}
117		// Cache and return the second order derivatives using select matrix.
118	16550x	matrix<Type> get_sigma_derivative2(std::vector<int> visits, matrix<Type> dist) override {
119	16550x	auto target = this->sigmad2_cache.find(visits);
120	16550x	if (target != this->sigmad2_cache.end()) {
121	15092x	return target->second;
122		} else {
123	1458x	int n_visits_i = visits.size();
124	1458x	int theta_sq = this->n_theta * this->n_theta;
125	1458x	matrix<Type> ret = matrix<Type>(theta_sq * n_visits_i, n_visits_i);
126	44586x	for (int i = 0; i < theta_sq; i++) {
127	43128x	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	1458x	this->sigmad2_cache[visits] = ret;
130	1458x	return ret;
131		}
132		}
133		// Cache and return the lower cholesky factor of inverse of sigma using select matrix.
134	14578x	matrix<Type> get_inverse_chol(std::vector<int> visits, matrix<Type> dist) override {
135	14578x	auto target = this->inverse_chol_cache.find(visits);
136	14578x	if (target != this->inverse_chol_cache.end()) {
137	13468x	return target->second;
138		} else {
139	1110x	matrix<Type> sigmainv = this->get_sigma_inverse(visits, dist);
140	1110x	Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > sigma_inv_chol(sigmainv);
141	1110x	matrix<Type> Li = sigma_inv_chol.matrixL();
142	1110x	this->inverse_chol_cache[visits] = Li;
143	1110x	return Li;
144		}
145		}
146		// Cache and return the first order derivatives of inverse of sigma using select matrix.
147	47282x	matrix<Type> get_inverse_derivative(std::vector<int> visits, matrix<Type> dist) override {
148	47282x	auto target = this->sigma_inverse_d1_cache.find(visits);
149	47282x	if (target != this->sigma_inverse_d1_cache.end()) {
150	43232x	return target->second;
151		} else {
152	4050x	auto sigma_d1 = this->get_sigma_derivative1(visits, dist);
153	4050x	matrix<Type> sigma_inv_d1(sigma_d1.rows(), sigma_d1.cols());
154	4050x	int n_visits_i = visits.size();
155	4050x	auto sigma_inv = this->get_sigma_inverse(visits, dist);
156	27934x	for (int r = 0; r < this->n_theta; r++) {
157	23884x	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	4050x	this->sigma_inverse_d1_cache[visits] = sigma_inv_d1;
160	4050x	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	24x	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	24x	this->logrho = log(this->rho);
182		}
183		// Obtain first order derivatives
184	5124x	matrix<Type> get_sigma_derivative1(std::vector<int> visits, matrix<Type> dist) override {
185	5124x	matrix<Type> ret(2 * dist.rows(), dist.cols());
186		// partial sigma / partial theta_1 = sigma.
187	5124x	auto sigma = this->get_sigma(visits, dist);
188	5124x	ret.block(0, 0, dist.rows(), dist.cols()) = sigma;
189	5124x	ret.block(dist.rows(), 0, dist.rows(), dist.cols()) = sigma.array() * dist.array() * (1 - this->rho);
190	10248x	return ret;
191		}
192		// Obtain second order derivatives.
193	1972x	matrix<Type> get_sigma_derivative2(std::vector<int> visits, matrix<Type> dist) override {
194	1972x	matrix<Type> ret(4 * dist.rows(), dist.cols());
195	1972x	auto sigma = this->get_sigma(visits, dist);
196	1972x	ret.block(0, 0, dist.rows(), dist.cols()) = sigma;
197	1972x	Type rho_r = 1 - this->rho;
198	1972x	auto dtheta1dtheta2 = sigma.array() * dist.array() * rho_r;
199	1972x	ret.block(dist.rows(), 0, dist.rows(), dist.cols()) = dtheta1dtheta2;
200	1972x	ret.block(dist.rows() * 2, 0, dist.rows(), dist.cols()) = dtheta1dtheta2;
201	1972x	matrix<Type> dtheta2s = dtheta1dtheta2 * (dist.array() * rho_r - this->rho);
202	1972x	ret.block(dist.rows() * 3, 0, dist.rows(), dist.cols()) = dtheta2s;
203	3944x	return ret;
204		}
205		// Obtain the lower cholesky factor of inverse of sigma using select matrix.
206	788x	matrix<Type> get_inverse_chol(std::vector<int> visits, matrix<Type> dist) override {
207	788x	auto sigmainv = this->get_sigma_inverse(visits, dist);
208	788x	Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > sigma_inv_chol(sigmainv);
209	788x	matrix<Type> Li = sigma_inv_chol.matrixL();
210	1576x	return Li;
211		}
212		// Obtain first order derivatives for inverse of sigma.
213	3152x	matrix<Type> get_inverse_derivative(std::vector<int> visits, matrix<Type> dist) override {
214	3152x	matrix<Type> sigma_inv_d1 = matrix<Type>::Zero(2 * dist.rows(), dist.cols());
215	3152x	auto sigma_inv = this->get_sigma_inverse(visits, dist);
216	3152x	auto sigma_d1 = this->get_sigma_derivative1(visits, dist);
217	9456x	for (int r = 0; r < 2; r++) {
218	6304x	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	6304x	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	312363x	T1 subset_matrix(T1 input, T2 index1, T2 index2) {
16		#if EIGEN_VERSION_AT_LEAST(3,4,0)
17	312363x	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	312363x	return ret;
27		}
28
29		template <typename T1, typename T2>
30	239826x	T1 subset_matrix(T1 input, T2 index1) {
31		#if EIGEN_VERSION_AT_LEAST(3,4,0)
32	239826x	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	239826x	return ret;
42		}
43
44
45		// Conversion from Rcpp vector/matrix to eigen vector/matrix
46		template <typename T1, typename T2>
47	607179x	T1 as_vector(T2 input) {
48	607179x	T1 ret(input.size());
49	2017590x	for (int i = 0; i < input.size(); i++) {
50	1410411x	ret(i) = input(i);
51		}
52	607179x	return ret;
53		}
54
55		template <typename T1, typename T2>
56	415552x	T1 as_matrix(T2 input) {
57	415552x	T1 ret(input.rows(), input.cols());
58	5852568x	for (int i = 0; i < input.rows(); i++) {
59	62737972x	for (int j = 0; j < input.cols(); j++) {
60	57300956x	ret(i,j) = input(i,j);
61		}
62		}
63	415552x	return ret;
64		}
65
66		template <typename T>
67	719474x	T segment(T input, int start, int n) {
68	719474x	T ret(n);
69	3594556x	for (int i = 0, j = start; i < n; i++, j++) {
70	2875082x	ret(i) = input(j);
71		}
72	719474x	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	1223118x	matrix<Type> tcrossprod(const matrix<Type>& lower_chol, bool complete = false) {
81	1223118x	int n = lower_chol.rows();
82	1223118x	matrix<Type> result = matrix<Type>::Zero(n, n);
83	1223118x	result.template selfadjointView<Eigen::Lower>().rankUpdate(lower_chol);
84	1223118x	if (complete) {
85	26290x	result.template triangularView<Eigen::Upper>() = result.transpose();
86		}
87	1223118x	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	1246532x	matrix<Type> crossprod(const matrix<Type>& x) {
96	1246532x	int n = x.cols();
97	1246532x	matrix<Type> result = matrix<Type>::Zero(n, n);
98	1246532x	result.template selfadjointView<Eigen::Lower>().rankUpdate(x.transpose());
99	1246532x	return result;
100		}
101
102		// Mapping from real values to correlation parameters in (-1, 1).
103		template <class T>
104	5084x	vector<T> map_to_cor(const vector<T>& theta) {
105	10168x	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	1152x	vector<T> map_to_cs_cor(const vector<T>& theta, int n_visits) {
111	1152x	T a = T(1.0) / (T(n_visits) - T(1.0));
112	1152x	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	5076x	generic_corr_fun(const vector<T>& theta) :
122	5076x	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	6212x	matrix<T> get_corr_mat_chol(int n_visits, const F<T>& corr_fun) {
129	6212x	matrix<T> correlation(n_visits, n_visits);
130	6212x	correlation.setIdentity();
131	30924x	for(int i = 0; i < n_visits; i++) {
132	61608x	for(int j = 0; j < i; j++){
133	36896x	correlation(i, j) = corr_fun(i, j);
134		}
135		}
136	6212x	Eigen::LLT<Eigen::Matrix<T,Eigen::Dynamic,Eigen::Dynamic> > correlation_chol(correlation);
137	6212x	matrix<T> L = correlation_chol.matrixL();
138	12424x	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	1858x	matrix<T> get_heterogeneous_cov(const vector<T>& sd_values, const F<T>& corr_fun) {
145	1858x	matrix<T> correlation_chol = get_corr_mat_chol(sd_values.size(), corr_fun);
146	1858x	Eigen::DiagonalMatrix<T,Eigen::Dynamic,Eigen::Dynamic> D = sd_values.matrix().asDiagonal();
147	1858x	matrix<T> result = D * correlation_chol;
148	3716x	return result;
149		}
150
151		// Obtain the Euclidean distance
152		template <class T>
153	33703x	matrix<T> euclidean(const matrix<T>& coordinates) {
154	33703x	matrix<T> result(coordinates.rows(), coordinates.rows());
155	126598x	for (int i = 0; i < coordinates.rows(); i++) {
156	92895x	result(i, i) = 0;
157	188400x	for (int j = 0; j < i; j ++) {
158	95505x	vector<T> diff = coordinates.row(i) - coordinates.row(j);
159	95505x	T d = sqrt((diff * diff).sum());
160	95505x	result(i, j) = d;
161	95505x	result(j, i) = d;
162		}
163		}
164	33703x	return result;
165		}
166
167		// Element wise power function of a matrix
168		template <class T>
169	1584x	Eigen::Matrix<T, -1, -1> cpow(const Eigen::Matrix<T, -1, -1> & input, double p) {
170	1584x	Eigen::Matrix<T, -1, -1> ret = Eigen::Matrix<T, -1, -1>(input.rows(), input.cols());
171	5908x	for (int i = 0; i < ret.rows(); i ++) {
172	8664x	for (int j = 0; j < ret.cols(); j++) {
173	4340x	ret(i, j) = std::pow(input(i, j), p);
174		}
175		}
176	1584x	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	1580x	matrix<T> pseudoInverseSqrt(const matrix<T> &input, double epsilon = std::numeric_limits<double>::epsilon()) {
183	1580x	Eigen::Matrix<T, -1, -1> eigen_mat = as_matrix<Eigen::Matrix<T, -1, -1>, matrix<T>>(input);
184	1580x	Eigen::JacobiSVD< Eigen::Matrix<T, -1, -1> > svd(eigen_mat ,Eigen::ComputeFullU \| Eigen::ComputeFullV);
185	1580x	double tolerance = epsilon * std::max(input.cols(), input.rows()) *svd.singularValues().array().abs()(0);
186	1580x	auto singular_vals = Matrix<T,-1,-1>((svd.singularValues().array() > tolerance).select(svd.singularValues().array().inverse(), 0).matrix());
187	1580x	Eigen::Matrix<T, -1, -1> ret_eigen = svd.matrixV() * cpow(singular_vals, 0.5).asDiagonal() * svd.matrixU().adjoint();
188	3160x	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	517x	RcppExport SEXP _mmrm_get_pqr(SEXP mmrm_fit_SEXP, SEXP theta_SEXP) {
13	517x	BEGIN_RCPP
14	517x	Rcpp::RObject rcpp_result_gen;
15	517x	Rcpp::RNGScope rcpp_rngScope_gen;
16	517x	Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
17	517x	Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
18	517x	rcpp_result_gen = Rcpp::wrap(get_pqr(mmrm_fit, theta));
19	517x	return rcpp_result_gen;
20	517x	END_RCPP
21		}
22
23		List get_jacobian(List mmrm_fit, NumericVector theta, NumericMatrix beta_vcov);
24	902x	RcppExport SEXP _mmrm_get_jacobian(SEXP mmrm_fit_SEXP, SEXP theta_SEXP, SEXP beta_vcov_SEXP) {
25	902x	BEGIN_RCPP
26	902x	Rcpp::RObject rcpp_result_gen;
27	902x	Rcpp::RNGScope rcpp_rngScope_gen;
28	902x	Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
29	902x	Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
30	902x	Rcpp::traits::input_parameter< NumericMatrix >::type beta_vcov(beta_vcov_SEXP);
31	902x	rcpp_result_gen = Rcpp::wrap(get_jacobian(mmrm_fit, theta, beta_vcov));
32	902x	return rcpp_result_gen;
33	902x	END_RCPP
34		}
35
36		List get_empirical(List mmrm_fit, NumericVector theta, NumericVector beta, NumericMatrix beta_vcov, std::string type);
37	429x	RcppExport SEXP _mmrm_get_empirical(SEXP mmrm_fit_SEXP, SEXP theta_SEXP, SEXP beta_SEXP, SEXP beta_vcov_SEXP, SEXP type_SEXP) {
38	429x	BEGIN_RCPP
39	429x	Rcpp::RObject rcpp_result_gen;
40	429x	Rcpp::RNGScope rcpp_rngScope_gen;
41	429x	Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
42	429x	Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
43	429x	Rcpp::traits::input_parameter< NumericVector >::type beta(beta_SEXP);
44	429x	Rcpp::traits::input_parameter< NumericMatrix >::type beta_vcov(beta_vcov_SEXP);
45	429x	Rcpp::traits::input_parameter< std::string >::type type(type_SEXP);
46	429x	rcpp_result_gen = Rcpp::wrap(get_empirical(mmrm_fit, theta, beta, beta_vcov, type));
47	429x	return rcpp_result_gen;
48	429x	END_RCPP
49		}
50
51		List predict(List mmrm_fit, NumericVector theta, NumericVector beta, NumericMatrix beta_vcov);
52	18656x	RcppExport SEXP _mmrm_predict(SEXP mmrm_fit_SEXP, SEXP theta_SEXP, SEXP beta_SEXP, SEXP beta_vcov_SEXP) {
53	18656x	BEGIN_RCPP
54	18656x	Rcpp::RObject rcpp_result_gen;
55	18656x	Rcpp::RNGScope rcpp_rngScope_gen;
56	18656x	Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
57	18656x	Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
58	18656x	Rcpp::traits::input_parameter< NumericVector >::type beta(beta_SEXP);
59	18656x	Rcpp::traits::input_parameter< NumericMatrix >::type beta_vcov(beta_vcov_SEXP);
60	18656x	rcpp_result_gen = Rcpp::wrap(predict(mmrm_fit, theta, beta, beta_vcov));
61	18656x	return rcpp_result_gen;
62	18656x	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	44x	RcppExport void R_init_mmrm(DllInfo *dll) {
79	44x	R_registerRoutines(dll, NULL, CallEntries, NULL, NULL);
80	44x	R_useDynamicSymbols(dll, FALSE);
81		#ifdef TMB_CCALLABLES
82	44x	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	50912x	Type objective_function<Type>::operator() ()
29		{
30		// Read data from R.
31	50912x	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	50912x	DATA_MATRIX(coordinates); // Coordinates matrix.
35	50912x	DATA_INTEGER(n_visits); // Number of visits, which is the dimension of the covariance matrix.
36	50912x	DATA_INTEGER(n_subjects); // Number of subjects.
37	50912x	DATA_IVECTOR(subject_zero_inds); // Starting indices for each subject (0-based) (length n_subjects).
38	50912x	DATA_IVECTOR(subject_n_visits); // Number of observed visits for each subject (length n_subjects).
39	50912x	DATA_STRING(cov_type); // Covariance type name.
40	50912x	DATA_INTEGER(is_spatial_int); // Spatial covariance (1)? Otherwise non-spatial covariance.
41	50912x	DATA_INTEGER(reml); // REML (1)? Otherwise ML (0).
42	50912x	DATA_FACTOR(subject_groups); // subject groups vector(0-based) (length n_subjects).
43	50912x	DATA_INTEGER(n_groups); // number of total groups.
44		// Read parameters from R.
45	50912x	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	50912x	matrix<Type> XtWX = matrix<Type>::Zero(x_matrix.cols(), x_matrix.cols());
49		// X^T W Y will be calculated incrementally into here.
50	50912x	matrix<Type> XtWY = matrix<Type>::Zero(x_matrix.cols(), 1);
51		// W^T/2 X will be saved into here.
52	50912x	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	50912x	vector<Type> y_vec_tilde = vector<Type>::Zero(y_vector.rows());
55		// Sum of the log determinant will be incrementally calculated here.
56	50912x	Type sum_log_det = 0.0;
57
58		// Convert is_spatial_int to bool.
59	50912x	bool is_spatial = (is_spatial_int == 1);
60		// Diagonal of weighted covariance
61	50912x	vector<Type> diag_cov_inv_sqrt(x_matrix.rows());
62		// Cholesky group object
63	50912x	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	10023120x	for (int i = 0; i < n_subjects; i++) {
66		// Start index and number of visits for this subject.
67	9972240x	int start_i = subject_zero_inds(i);
68	9972240x	int n_visits_i = subject_n_visits(i);
69	9972240x	std::vector<int> visit_i(n_visits_i);
70	9972240x	matrix<Type> dist_i(n_visits_i, n_visits_i);
71	9972240x	if (!is_spatial) {
72	36230864x	for (int j = 0; j < n_visits_i; j++) {
73	26510784x	visit_i[j] = int(asDouble(coordinates(start_i + j, 0)));
74		}
75		} else {
76	252160x	dist_i = euclidean(matrix<Type>(coordinates.block(start_i, 0, n_visits_i, coordinates.cols())));
77		}
78		// Obtain Cholesky factor Li.
79	9972240x	matrix<Type> Li = chols_group.cache[subject_groups[i]]->get_chol(visit_i, dist_i);
80		// Calculate weighted Cholesky factor for this subject.
81	9972240x	Eigen::DiagonalMatrix<Type,Eigen::Dynamic,Eigen::Dynamic> Gi_inv_sqrt = weights_vector.segment(start_i, n_visits_i).cwiseInverse().sqrt().matrix().asDiagonal();
82	9972240x	Li = Gi_inv_sqrt * Li;
83		// Calculate scaled design matrix and response vector for this subject.
84	9972240x	matrix<Type> Xi = x_matrix.block(start_i, 0, n_visits_i, x_matrix.cols());
85	9972240x	matrix<Type> XiTilde = Li.template triangularView<Eigen::Lower>().solve(Xi);
86	9972240x	matrix<Type> Yi = y_vector.segment(start_i, n_visits_i).matrix();
87	9972240x	matrix<Type> YiTilde = Li.template triangularView<Eigen::Lower>().solve(Yi);
88
89		// Increment quantities.
90	9972240x	matrix<Type> XiTildeCrossprod = crossprod(XiTilde);
91	9972240x	XtWX += XiTildeCrossprod.template triangularView<Eigen::Lower>();
92	9972240x	XtWY += XiTilde.transpose() * YiTilde;
93	9972240x	vector<Type> LiDiag = Li.diagonal();
94	9972240x	sum_log_det += sum(log(LiDiag));
95		// Cache the reciprocal of square root of diagonal of covariance
96	9972240x	diag_cov_inv_sqrt.segment(start_i, n_visits_i) = vector<Type>(tcrossprod(Li).diagonal()).rsqrt();
97		// Save stuff.
98	9972240x	x_mat_tilde.block(start_i, 0, n_visits_i, x_matrix.cols()) = XiTilde;
99	9972240x	y_vec_tilde.segment(start_i, n_visits_i) = YiTilde.col(0);
100		}
101
102		// Solve for beta.
103	50880x	Eigen::LDLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > XtWX_decomposition(XtWX);
104	50880x	matrix<Type> beta_mat = XtWX_decomposition.solve(XtWY);
105	50880x	vector<Type> beta = beta_mat.col(0);
106
107		// Define scaled residuals.
108	50880x	vector<Type> x_mat_tilde_beta = x_mat_tilde * beta;
109	50880x	vector<Type> epsilonTilde = y_vec_tilde - x_mat_tilde_beta;
110
111		// Calculate negative log-likelihood.
112	4080x	Type neg_log_lik;
113
114		// Always extract the D vector since we want to report this below.
115	50880x	vector<Type> XtWX_D = XtWX_decomposition.vectorD();
116
117	50880x	if (reml == 1) {
118		// Use restricted maximum likelihood.
119	48384x	Type XtWX_log_det = XtWX_D.log().sum();
120	48384x	neg_log_lik = (x_matrix.rows() - x_matrix.cols()) / 2.0 * log(2.0 * M_PI) +
121	48384x	sum_log_det +
122	96768x	XtWX_log_det / 2.0 +
123	52048x	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	2496x	neg_log_lik = x_matrix.rows() / 2.0 * log(2.0 * M_PI) +
127	416x	sum_log_det +
128	2912x	0.5 * (epsilonTilde * epsilonTilde).sum();
129		}
130
131		// Report quantities to R.
132	46800x	REPORT(beta);
133
134		// We already compute the inverse of XtWX here because we already did the
135		// matrix decomposition above.
136	50880x	matrix<Type> Identity(XtWX.rows(), XtWX.cols());
137	50880x	Identity.setIdentity();
138	50880x	matrix<Type> beta_vcov = XtWX_decomposition.solve(Identity);
139	46800x	REPORT(beta_vcov);
140
141		// Also return the decomposition components L and D.
142	50880x	matrix<Type> XtWX_L(XtWX.rows(), XtWX.cols());
143	50880x	XtWX_L = XtWX_decomposition.matrixL();
144	46800x	REPORT(XtWX_L);
145	46800x	REPORT(XtWX_D);
146
147		// normalized residual
148	46800x	REPORT(epsilonTilde);
149		// inverse square root of diagonal of covariance
150	46800x	REPORT(diag_cov_inv_sqrt);
151	50880x	matrix<Type> covariance_lower_chol = chols_group.get_default_chol();
152	46800x	REPORT(covariance_lower_chol);
153
154	50880x	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	49x	void expect_equal_matrix(const T& target, const T& current)
32		{
33	49x	int nrow = target.rows();
34	49x	int ncol = target.cols();
35
36	!	expect_true(nrow == current.rows());
37	!	expect_true(ncol == current.cols());
38
39	184x	for (int i = 0; i < nrow; i++) {
40	500x	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	18x	void expect_equal_vector(const T& target, const T& current)
48		{
49	18x	int n = target.size();
50	!	expect_true(n == current.size());
51
52	108x	for (int i = 0; i < n; i++) {
53	!	expect_equal(target(i), current(i));
54		}
55		}
56
57		#endif