mmrm coverage - 97.44%

Files
Source

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

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

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

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

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

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

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

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

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

  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,
    contrasts = contrasts,
    emmeans_gcomp_vars = control$emmeans_gcomp_vars
  )
  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
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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


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

  colnames_basic <- colnames(data_basic)

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

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


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

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

  cols <- lapply(datasets, colnames)

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

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

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

  TRUE
}


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


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


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

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

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

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

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

  TRUE
}


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

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

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

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

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

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

  cov_struct_nesting <- h_check_cov_struct_nesting(model_basic, model_augmented)

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

  TRUE
}


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

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

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

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

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

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

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

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

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


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

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

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

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


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

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

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

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


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

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

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

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

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

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

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

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

  fit
}


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

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

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

  design_matrix <- component(object, "x_matrix")

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

  assign_coef_term <- factor(assign_coef_term, labels = term_labels)

  coef_identity_matrix <- diag(length(assign_coef_term))

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

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

  df_md_results_data_frame <- do.call(rbind, df_md_results_per_term)

  df_md_results_data_frame
}

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

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

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

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

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

  invisible(FALSE)
}

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

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

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

  invisible(TRUE)
}

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

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

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

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

  model_formula <- h_drop_covariance_terms(formula)

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

#' Data for `TMB` Fit
#'
#' @param formula_parts (`mmrm_tmb_formula_parts`)\cr list with formula parts
#'   from [h_mmrm_tmb_formula_parts()].
#' @param data (`data.frame`)\cr which contains variables used in `formula_parts`.
#' @param weights (`vector`)\cr weights to be used in the fitting process.
#' @param reml (`flag`)\cr whether restricted maximum likelihood (REML) estimation is used,
#'   otherwise maximum likelihood (ML) is used.
#' @param singular (`string`)\cr choices of method deal with rank-deficient matrices. "error" to
#'   stop the function return the error, "drop" to drop these columns, and "keep" to keep all the columns.
#' @param drop_visit_levels (`flag`)\cr whether to drop levels for visit variable, if visit variable is a factor.
#' @param allow_na_response (`flag`)\cr whether NA in response is allowed.
#' @param drop_levels (`flag`)\cr whether drop levels for covariates. If not dropped could lead to singular matrix.
#' @param xlev (`list` or `NULL`)\cr list of X levels produced by stats::.getXlevels
#' @param contrasts (`list` or `NULL`)\cr an optional named list of contrast
#'   matrices or contrast functions (like [stats::contr.sum] or
#'   [stats::contr.poly]) for specific factor variables, matching the
#'   `contrasts` argument in [stats::lm()]. The list names must correspond to
#'   factor variable names in the model formula. When `NULL` (the default),
#'   the contrasts set on the factor variables in `data` are used. If a
#'   contrast matrix has rownames that include levels not present in `data`,
#'   those levels are preserved and the corresponding model matrix columns
#'   are marked as aliased (not estimable), enabling prediction on new data
#'   containing those levels.
#' @param emmeans_gcomp_vars (`character` or `NULL`)
#'   treated as fixed for G-computation correction. Stored as a character
#'   vector in the returned list for downstream use in the emmeans hook.
#'
#' @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
#' - `emmeans_gcomp_vars`: `character` or `NULL` with the G-computation variable names.
#'
#' @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,
  emmeans_gcomp_vars = NULL
) {
  assert_class(formula_parts, "mmrm_tmb_formula_parts")
  assert_data_frame(data)
  assert_numeric(weights, len = nrow(data))
  assert_flag(reml)
  singular <- match.arg(singular)
  assert_flag(drop_visit_levels)
  assert_character(emmeans_gcomp_vars, min.len = 1L, null.ok = TRUE)
  if (!is.null(emmeans_gcomp_vars)) {
    assert_subset(emmeans_gcomp_vars, names(data))
  }

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

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

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

  # Contrasts following the same as what's done in lm(). The user gives contrast
  # matrices or functions which are passed to model.matrix(). With lm(), missing
  # factor levels are silently dropped. With contrasts, when a contrast matrix
  # includes levels not present in the data, those levels are kept and the model
  # matrix columns are marked aliased. This allows for prediction on new data
  # that have those levels.
  no_drop_lvls <- NULL
  if (!is.null(contrasts)) {
    formula_factors <- intersect(
      names(contrasts),
      names(which(vapply(
        full_frame,
        \(x) is.factor(x) || is.character(x) || is.logical(x),
        logical(1L)
      )))
    )
    # don't drop factorlike variables that have explicit contrasts
    no_drop_lvls <- names(which(vapply(
      formula_factors,
      \(x) !is.null(contrasts[[x]]),
      logical(1L)
    )))
  }

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

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

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

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

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

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

#' Low-Level Fitting Function for MMRM
#'
#' @description `r lifecycle::badge("stable")`
#'
#' This is the low-level function to fit an MMRM. Note that this does not
#' try different optimizers or adds Jacobian information etc. in contrast to
#' [mmrm()].
#'
#' @param formula (`formula`)\cr model formula with exactly one special term
#'   specifying the visits within subjects, see details.
#' @param data (`data.frame`)\cr input data containing the variables used in
#'   `formula`.
#' @param weights (`vector`)\cr input vector containing the weights.
#' @inheritParams h_mmrm_tmb_data
#' @param covariance (`cov_struct`)\cr A covariance structure type definition,
#'   or value that can be coerced to a covariance structure using
#'   [as.cov_struct()]. If no value is provided, a structure is derived from
#'   the provided formula.
#' @param contrasts (`list` or `NULL`)\cr an optional named list of contrast
#'   matrices or contrast functions for specific factor variables, see
#'   [mmrm()] for details.
#' @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
#' result <- fit_mmrm(formula, data, rep(1, nrow(fev_data)))
fit_mmrm <- function(
  formula,
  data,
  weights,
  reml = TRUE,
  covariance = NULL,
  tmb_data,
  formula_parts,
  contrasts = NULL,
  control = mmrm_control()
) {
  if (missing(formula_parts) || missing(tmb_data)) {
    covariance <- h_reconcile_cov_struct(formula, covariance)
    formula_parts <- h_mmrm_tmb_formula_parts(formula, covariance)

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

#' 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_along(pl)) {
    for (j in seq_along(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
}

#' G-computation Covariance Correction
#'
#' When the MMRM includes covariate-by-treatment interactions, the standard
#' emmeans output conditions on covariates rather than averaging over them.
#' These internal helpers compute subject-level potential outcomes under each
#' treatment and use their distribution to correct the L matrix and covariance.
#'
#' @keywords internal
#' @noRd
NULL

#' Get Subject-Level Covariate Data
#'
#' @description Returns one row per subject with covariate data. Pulls
#'   subject level data from the original dataset stored in `tmb_data$data`,
#'   and includes subjects with missing outcomes.
#'
#' @note If time varying covariates are present, only the first observed
#'   value per subject is retained. The intended use is for baseline
#'   (pre-randomization) covariates that are constant across visits.
#'
#' @param object (`mmrm`)\cr the fitted MMRM.
#'
#' @return A `data.frame` with one row per subject.
#'
#' @keywords internal
#' @noRd
h_get_subject_data <- function(object) {
  assert_class(object, "mmrm")
  original_data <- object$tmb_data$data
  subject_var <- object$formula_parts$subject_var
  model_vars <- all.vars(object$formula_parts$model_formula)
  keep_vars <- intersect(c(subject_var, model_vars), names(original_data))
  subj_data <- original_data[
    !duplicated(original_data[[subject_var]]),
    keep_vars,
    drop = FALSE
  ]
  rownames(subj_data) <- NULL
  subj_data
}

#' Compute Subject-Level Potential Outcomes
#'
#' @description For a given visit, computes each subject's fitted potential
#'   outcome under each combination of fixed variable levels by constructing
#'   the counterfactual design matrix. Each subject's own covariate values
#'   are used for all non-fixed variables.
#'
#' @param object (`mmrm`)\cr the fitted MMRM.
#' @param subj_data (`data.frame`)\cr one row per subject with covariates.
#' @param visit_value (`string`)\cr the visit level to evaluate at.
#' @param visit_var (`string`)\cr name of the visit variable.
#' @param counterfactual_grid (`data.frame`)\cr unique combinations of fixed
#'   variables (excluding visit) to loop over.
#'
#' @return A list with:
#'   - `vhat`: `matrix` of fitted potential outcomes (n_subjects x K),
#'     where n_subjects is the number of rows in `subj_data` and K is the
#'     number of rows in `counterfactual_grid`.
#'   - `L_global`: `matrix` of globally weighted contrast rows (K x p),
#'     where p is the number of estimable model coefficients.
#'
#' @keywords internal
#' @noRd
h_compute_potential_outcomes <- function(
  object,
  subj_data,
  visit_value,
  visit_var,
  counterfactual_grid
) {
  assert_class(object, "mmrm")
  assert_data_frame(subj_data, min.rows = 1L)
  assert_string(visit_value)
  assert_string(visit_var)
  assert_data_frame(counterfactual_grid, min.rows = 1L)
  beta_hat <- component(object, "beta_est")
  beta_clean <- beta_hat
  beta_clean[is.na(beta_clean)] <- 0
  contrasts_arg <- component(object, "contrasts")
  full_frame <- object$tmb_data$full_frame

  model_formula <- object$formula_parts$model_formula
  model_terms <- stats::delete.response(stats::terms(model_formula))

  K <- nrow(counterfactual_grid) # nolint
  n_subj <- nrow(subj_data)
  p <- length(beta_hat)

  # Stack all K counterfactual assignments into one data frame.
  # (for efficiency reasons, we can just do one model.frame and model.matrix call below instead of K).
  cf_list <- vector("list", K)
  for (k in seq_len(K)) {
    subj_cf <- subj_data
    subj_cf[[visit_var]] <- factor(
      visit_value,
      levels = levels(full_frame[[visit_var]])
    )
    for (var_name in names(counterfactual_grid)) {
      val <- counterfactual_grid[[var_name]][k]
      subj_cf[[var_name]] <- if (is.factor(full_frame[[var_name]])) {
        factor(val, levels = levels(full_frame[[var_name]]))
      } else {
        val
      }
    }
    cf_list[[k]] <- subj_cf
  }
  cf_stacked <- do.call(rbind, cf_list)

  mf <- stats::model.frame(
    model_terms,
    data = cf_stacked,
    na.action = stats::na.pass
  )
  # nolint next
  X_all <- stats::model.matrix(
    model_terms,
    data = mf,
    contrasts.arg = contrasts_arg
  )
  X_all <- X_all[, names(beta_hat), drop = FALSE] # nolint
  X_all[is.na(X_all)] <- 0 # nolint

  # Split back into K blocks of n_subj rows.
  vhat <- matrix(NA_real_, nrow = n_subj, ncol = K)
  L_global <- matrix(NA_real_, nrow = K, ncol = p) # nolint
  for (k in seq_len(K)) {
    rows <- ((k - 1L) * n_subj + 1L):(k * n_subj)
    X_k <- X_all[rows, , drop = FALSE] # nolint
    vhat[, k] <- as.vector(X_k %*% beta_clean)
    L_global[k, ] <- colMeans(X_k, na.rm = TRUE) # nolint
  }

  list(vhat = vhat, L_global = L_global)
}

#' Compute Per Visit Variance Contributions and Global L Matrix
#'
#' @description For each visit, computes subject level potential outcomes,
#'   the sample covariance of those outcomes (divided by n), and the
#'   globally weighted L matrix rows. Assembles these into the full
#'   block diagonal S matrix and global L matrix across all visits.
#'
#' @param object (`mmrm`)\cr the fitted MMRM.
#' @param model_mat (`matrix`)\cr the L matrix from the emmeans reference
#'   grid (n_grid x p).
#' @param grid (`data.frame`)\cr the emmeans reference grid.
#' @param fixed_vars (`character`)\cr names of the fixed (intervention)
#'   variables.
#' @param visit_var (`string`)\cr name of the visit variable.
#'
#' @return A list with:
#'   - `S_full`: `matrix` (n_grid x n_grid) block-diagonal sample covariance
#'     of potential outcomes divided by n, with blocks corresponding to visits.
#'   - `L_global`: `matrix` (n_grid x p) globally-weighted L matrix rows.
#'
#' @keywords internal
#' @noRd
h_gcomp_visit_contributions <- function(
  object,
  model_mat,
  grid,
  fixed_vars,
  visit_var
) {
  subj_data <- h_get_subject_data(object)
  full_frame <- object$tmb_data$full_frame
  n_subj <- nrow(subj_data)
  visit_levels <- levels(full_frame[[visit_var]])
  cell_vars <- setdiff(fixed_vars, visit_var)

  n_grid <- nrow(model_mat)
  grid_visit <- as.character(grid[[visit_var]])
  S_full <- matrix(0, nrow = n_grid, ncol = n_grid) # nolint
  L_global <- model_mat # nolint

  for (v in visit_levels) {
    v_mask <- grid_visit == v
    if (!any(v_mask)) {
      next
    }
    v_indices <- which(v_mask)

    grid_at_v <- grid[v_indices, cell_vars, drop = FALSE]
    cf_grid <- unique(grid_at_v)
    rownames(cf_grid) <- NULL

    po <- h_compute_potential_outcomes(
      object = object,
      subj_data = subj_data,
      visit_value = v,
      visit_var = visit_var,
      counterfactual_grid = cf_grid
    )

    sigma_v <- stats::cov(po$vhat, use = "pairwise.complete.obs")
    S_t <- sigma_v / n_subj # nolint

    cf_keys <- apply(cf_grid, 1, function(r) {
      paste(as.character(r), collapse = "|")
    })
    grid_keys <- apply(grid_at_v, 1, function(r) {
      paste(as.character(r), collapse = "|")
    })

    for (ii in seq_along(v_indices)) {
      ki <- match(grid_keys[ii], cf_keys)
      if (is.na(ki)) next # nocov

      L_global[v_indices[ii], ] <- po$L_global[ki, ] # nolint

      for (jj in seq_along(v_indices)) {
        kj <- match(grid_keys[jj], cf_keys)
        if (is.na(kj)) next # nocov
        S_full[v_indices[ii], v_indices[jj]] <- S_t[ki, kj] # nolint
      }
    }
  }

  list(S_full = S_full, L_global = L_global)
}

#' Compute G-computation Correction for emmeans Integration
#'
#' @description Computes the corrected L matrix (global covariate weighting)
#'   and the variance correction matrix delta. Together these produce the
#'   G-computation estimator with correct standard errors. Delta is a
#'   p x p positive semi-definite matrix (where p is the number of estimable
#'   model coefficients) that is added to the model-based covariance.
#'
#' @param object (`mmrm`)\cr the fitted MMRM with
#'   `emmeans_gcomp_vars` set.
#' @param model_mat (`matrix`)\cr the L matrix from the emmeans reference
#'   grid (n_grid x p).
#' @param grid (`data.frame`)\cr the reference grid from emmeans.
#'
#' @return A list with:
#'   - `delta`: `matrix` (p x p) positive semi-definite correction, where
#'     p is the number of estimable model coefficients.
#'   - `L_global`: `matrix` (n_grid x p) globally-weighted L matrix.
#'
#' @keywords internal
#' @noRd
h_gcomp_emm_correction <- function(object, model_mat, grid) {
  assert_class(object, "mmrm")
  assert_matrix(model_mat, mode = "numeric")
  assert_data_frame(grid)

  fixed_vars <- object$tmb_data$emmeans_gcomp_vars
  visit_var <- object$formula_parts$visit_var

  assert_character(fixed_vars, min.len = 1L)
  assert_string(visit_var)
  assert_subset(visit_var, names(object$tmb_data$full_frame))

  beta_names <- names(component(object, "beta_est"))
  # The emmeans grid model matrix may have extra columns (e.g. from
  # interaction reordering) that do not appear in beta_est. Trim to
  # match the coefficient vector, since estimability::nonest.basis
  # only handles rank deficiency within the fitted coefficient set.
  if (ncol(model_mat) > length(beta_names)) {
    keep_cols <- match(beta_names, colnames(model_mat))
    keep_cols <- keep_cols[!is.na(keep_cols)]
    model_mat <- model_mat[, keep_cols, drop = FALSE]
  }

  contributions <- h_gcomp_visit_contributions(
    object,
    model_mat,
    grid,
    fixed_vars,
    visit_var
  )
  S_full <- contributions$S_full # nolint
  L <- contributions$L_global # nolint

  llt_ginv_l <- MASS::ginv(tcrossprod(L)) %*% L
  delta <- t(llt_ginv_l) %*% S_full %*% llt_ginv_l
  V <- component(object, "beta_vcov") # nolint
  dimnames(delta) <- dimnames(V)

  list(delta = delta, L_global = L)
}

#' 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",
    sp_gau = "spatial Gaussian"
  )

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

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

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


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


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

  fits <- list(object, ...)

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

    assert_flag(test)
    assert_flag(refit)

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

      fit_data <- lapply(fits, h_get_minimal_fit_data)

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

      needs_refit <- !nested_in_common_data

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

    log_likelihood_vec <- lapply(fits, logLik)

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

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

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

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

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

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

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

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

  out
}

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

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

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

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

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

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

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

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

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

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

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


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


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

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

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

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

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

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

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

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

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

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

  ret_df
}

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

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

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

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

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

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

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

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

#' Match coefficient names against model matrix columns
#'
#' This helper is robust against re-ordering of interaction terms, e.g.
#' `A:B` vs `B:A` and `A:B:C` vs any permutation.
#'
#' @keywords internal
#' @noRd
h_match_coefs <- function(coef_names, model_colnames) {
  canon_interaction <- function(x) {
    split_terms <- strsplit(x, ":", fixed = TRUE)
    vapply(
      split_terms,
      FUN.VALUE = character(1),
      function(term_parts) {
        if (length(term_parts) <= 1L) {
          term_parts
        } else {
          paste(sort(term_parts), collapse = ":")
        }
      }
    )
  }

  kept <- match(coef_names, model_colnames)
  idx_na <- is.na(kept)

  if (any(idx_na)) {
    kept[idx_na] <- match(
      canon_interaction(coef_names[idx_na]),
      canon_interaction(model_colnames)
    )
  }

  kept
}

#' Returns a `data.frame` for `emmeans` Purposes
#'
#' @seealso See [emmeans::recover_data()] for background.
#' @keywords internal
#' @noRd
# nolint start
recover_data.mmrm <- function(object, ...) {
  # nolint end
  fun_call <- stats::getCall(object)
  # subject_var is excluded because it should not contain fixed effect.
  # visit_var is only excluded for spatial covariance structures -
  # for non-spatial covariance structures, emmeans can provide marginal mean
  # by each visit so we need visit_var. Also exclude visit_var when it has only
  # one level (emmeans cannot contrast it).
  visit_var <- object$formula_parts$visit_var
  exclude_visit <- object$formula_parts$is_spatial ||
    nlevels(object$tmb_data$full_frame[[visit_var]]) < 2L

  model_frame <- stats::model.frame(
    object,
    exclude = c(
      "subject_var",
      if (exclude_visit) "visit_var"
    )
  )
  model_terms <- stats::delete.response(stats::terms(model_frame))
  emmeans::recover_data(
    fun_call,
    trms = model_terms,
    na.action = "na.omit",
    frame = model_frame,
    ...
  )
}

#' Returns a List of Model Details for `emmeans` Purposes
#'
#' @seealso See [emmeans::emm_basis()] for background.
#' @keywords internal
#' @noRd
# nolint start
emm_basis.mmrm <- function(
  # nolint end
  object,
  trms,
  xlev,
  grid,
  ...
) {
  model_frame <- stats::model.frame(
    trms,
    grid,
    na.action = stats::na.pass,
    xlev = xlev
  )
  contrasts <- component(object, "contrasts")
  model_mat <- stats::model.matrix(trms, model_frame, contrasts.arg = contrasts)
  beta_hat <- component(object, "beta_est")
  nbasis <- if (length(beta_hat) < ncol(model_mat)) {
    kept <- h_match_coefs(names(beta_hat), colnames(model_mat))
    if (anyNA(kept)) {
      unmatched <- names(beta_hat)[is.na(kept)]
      stop(
        paste0(
          "Failed to match coefficient name(s) to model matrix columns: ",
          paste(unmatched, collapse = ", ")
        ),
        call. = FALSE
      )
    }
    beta_hat <- NA * model_mat[1L, ]
    beta_hat[kept] <- component(object, "beta_est")
    orig_model_mat <- stats::model.matrix(
      trms,
      stats::model.frame(
        object,
        exclude = c(
          "subject_var",
          if (object$formula_parts$is_spatial) "visit_var"
        )
      ),
      contrasts.arg = contrasts
    )
    estimability::nonest.basis(orig_model_mat)
  } else {
    estimability::all.estble
  }

  # Start with the model-based covariance.
  V <- component(object, "beta_vcov") # nolint

  if (!is.null(object$tmb_data$emmeans_gcomp_vars)) {
    gcomp_result <- h_emm_basis_gcomp(
      object,
      model_mat,
      beta_hat,
      nbasis,
      V,
      grid
    )
    model_mat <- gcomp_result$model_mat
    beta_hat <- gcomp_result$beta_hat
    nbasis <- gcomp_result$nbasis
    V <- gcomp_result$V # nolint
    dfargs <- gcomp_result$dfargs
    dffun <- gcomp_result$dffun
    message(
      "G-computation correction applied, so emmeans will produce the correct average treatment effect."
    )
  } else {
    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 = V,
    dffun = dffun,
    dfargs = dfargs
  )
}

#' Apply G-computation Correction in emm_basis
#'
#' @description Applies the G-computation correction to the emmeans basis
#'   components. Replaces the L matrix with globally-weighted rows, adds
#'   the variance correction to V, and handles aliased coefficient expansion.
#'
#' @param object (`mmrm`)\cr the fitted MMRM.
#' @param model_mat (`matrix`)\cr the L matrix from emmeans.
#' @param beta_hat (`numeric`)\cr coefficient estimates (may contain NA for
#'   aliased positions).
#' @param nbasis (`matrix` or `NULL`)\cr estimability basis from emmeans.
#' @param V (`matrix`)\cr model-based covariance of beta.
#' @param grid (`data.frame`)\cr the emmeans reference grid.
#'
#' @return A list with modified `model_mat`, `beta_hat`, `nbasis`, `V`,
#'   `dfargs`, and `dffun`.
#'
#' @keywords internal
#' @noRd
# nolint next
h_emm_basis_gcomp <- function(object, model_mat, beta_hat, nbasis, V, grid) {
  # Detect aliased coefficient positions before any modifications.
  aliased_pos <- which(is.na(beta_hat))

  visit_var <- object$formula_parts$visit_var
  V_corrected_est <- V # nolint

  if (visit_var %in% names(grid)) {
    # Save original model_mat dimensions. When aliased coefficients exist,
    # h_match_coefs expands model_mat beyond length(beta_est). The correction
    # works in estimable-coefficient space, then we re-expand to match.
    orig_ncol <- ncol(model_mat)
    orig_colnames <- colnames(model_mat)

    correction <- h_gcomp_emm_correction(object, model_mat, grid)
    V_corrected_est <- V + correction$delta # nolint

    # Re-expand L_global if h_match_coefs added columns for aliased coefs.
    # Use h_match_coefs (not match) to handle A:B vs B:A reordering.
    L_est <- correction$L_global # nolint
    if (ncol(L_est) < orig_ncol && length(aliased_pos) > 0) {
      est_pos <- setdiff(seq_len(orig_ncol), aliased_pos)

      # Expand L_global to full width with zeros for aliased columns.
      L_exp <- matrix(0, nrow = nrow(L_est), ncol = orig_ncol) # nolint
      colnames(L_exp) <- orig_colnames # nolint
      col_map <- h_match_coefs(colnames(L_est), orig_colnames)
      col_map <- col_map[!is.na(col_map)]
      L_exp[, col_map] <- L_est # nolint
      model_mat <- L_exp

      # Expand V to full width with zeros for aliased row/column.
      V_full <- matrix(0, nrow = orig_ncol, ncol = orig_ncol) # nolint
      V_full[est_pos, est_pos] <- V_corrected_est # nolint
      V <- V_full # nolint

      # Zero NA in bhat so 0*NA doesn't propagate.
      beta_hat[aliased_pos] <- 0

      # All grid rows are estimable through L_global (aliased col is zero).
      nbasis <- estimability::all.estble
    } else {
      model_mat <- L_est
      V <- V_corrected_est # nolint
    }
  }

  # dffun operates in estimable-coefficient space (p_est dimensions).
  # Strip aliased positions from k before passing to df_md.
  dfargs <- list(
    object = object,
    V_corrected = V_corrected_est,
    aliased_pos = aliased_pos
  )
  dffun <- function(k, dfargs) {
    obj <- dfargs$object
    if (is.null(obj$vcov) || identical(obj$vcov, "Asymptotic")) {
      obj$beta_vcov <- dfargs$V_corrected
    } else {
      obj$beta_vcov_adj <- dfargs$V_corrected
    }
    k_est <- k
    if (
      length(dfargs$aliased_pos) > 0 && length(k) > nrow(dfargs$V_corrected)
    ) {
      k_est <- k[-dfargs$aliased_pos]
    }
    mmrm::df_md(obj, contrast = k_est)$denom_df
  }

  list(
    model_mat = model_mat,
    beta_hat = beta_hat,
    nbasis = nbasis,
    V = V,
    dfargs = dfargs,
    dffun = dffun
  )
}

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

  as.data.frame(
    col.names = names(formals(type)),
    rbind(
      type("unstructured", "us", NA, FALSE, FALSE),
      type("Toeplitz", "toep", "toeph", TRUE, FALSE),
      type("auto-regressive order one", "ar1", "ar1h", TRUE, FALSE),
      type("ante-dependence", "ad", "adh", TRUE, FALSE),
      type("compound symmetry", "cs", "csh", TRUE, FALSE),
      type("spatial exponential", "sp_exp", NA, FALSE, TRUE),
      type("spatial Gaussian", "sp_gau", 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}}}
#' \cr
#'
#' sp_gau
#' \tab spatial Gaussian
#' \tab \eqn{2}
#' \tab \eqn{\sigma^{2}\rho^{d_{ij}^2}}
#'
#' }
#'
#' 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 exp", 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 spatial covariance structures may provide multiple
#' coordinates, which identify the Euclidean distance between the time points.
#'
#' @param x an object from which to derive a covariance structure. See object
#'   specific sections for details.
#' @param warn_partial (`flag`)\cr whether to emit a warning when parts of the
#'   formula are disregarded.
#' @param ... additional arguments unused.
#'
#' @return A [cov_struct()] object.
#'
#' @examples
#' # provide a covariance structure as a right-sided formula
#' as.cov_struct(~ csh(visit | group / subject))
#'
#' # when part of a full formula, suppress warnings using `warn_partial = FALSE`
#' as.cov_struct(y ~ x + csh(visit | group / subject), warn_partial = FALSE)
#'
#' @family covariance types
#' @export
# nolint start
as.cov_struct <- function(x, ...) {
  # nolint end
  UseMethod("as.cov_struct")
}

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

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

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

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

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

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

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

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

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

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

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


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

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

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

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

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

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

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

  num_df <- nrow(contrast)

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

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

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

  out
}

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

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

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

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

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

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

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


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

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

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

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

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

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

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

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

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


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

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

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

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

  n_subjects <- component(object, "n_subjects")

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

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

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

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

  component(x, name = "varcor")
}

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

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

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

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

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


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

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

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

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

  cat("\nModel Inference Optimization:")

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


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

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

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

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

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

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

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

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

  ret
}

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

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

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

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

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

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

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

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

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

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

#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)); // theta (-Inf, Inf), rho (0, 1), log(rho) (-Inf, 0)
  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();
}

// Spatial Gaussian Cholesky factor.
template <class T>
matrix<T> get_spatial_gaussian(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().square() * log(rho)); // only take square of the dist
  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 if (cov_type == "sp_gau") {
    result = get_spatial_gaussian<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;
  }
};


// derivatives_sp_gau struct is created to obtain the exact derivatives of spatial Gaussian
// covariance structure, and its inverse.
// Inherit from sp_exp because of similar parameterization; only the sigma
// derivatives differ (square the distance). Inverse / inverse-derivative logic
// is reused from the base via the virtual get_sigma_derivative1 dispatch.
template <class Type>
struct derivatives_sp_gau: public derivatives_sp_exp<Type> {
  derivatives_sp_gau() {
    // This default constructor is needed because the use of `[]` in maps.
  }
  derivatives_sp_gau(vector<Type> theta, std::string cov_type): derivatives_sp_exp<Type>(theta, cov_type) {
  }
  // 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().square() * (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().square() * 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().square() * rho_r - this->rho);
    ret.block(dist.rows() * 3, 0, dist.rows(), dist.cols()) = dtheta2s;
    return ret;
  }
};

// Cache of per-group derivatives objects, with cov_type-aware dispatch for
// spatial structures (sp_exp / sp_gau). Mirrors `cache_obj` but selects the
// concrete spatial subclass at runtime so analytic derivatives match cov_type.
template <class Type>
struct derivatives_cache {
  std::map<int, std::shared_ptr<derivatives_base<Type>>> cache;
  int n_groups;
  bool is_spatial;
  int n_visits;
  derivatives_cache(vector<Type> 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) {
    int theta_one_group_size = theta.size() / n_groups;
    for (int r = 0; r < n_groups; r++) {
      vector<Type> theta_r = theta.segment(r * theta_one_group_size, theta_one_group_size);
      if (is_spatial) {
        if (cov_type == "sp_exp") {
          this->cache[r] = std::make_shared<derivatives_sp_exp<Type>>(theta_r, cov_type);
        } else if (cov_type == "sp_gau") {
          this->cache[r] = std::make_shared<derivatives_sp_gau<Type>>(theta_r, cov_type);
        } else {
          Rf_error("%s", ("Unknown spatial covariance type '" + cov_type + "'.").c_str());
        }
      } else {
        this->cache[r] = std::make_shared<derivatives_nonspatial<Type>>(theta_r, n_visits, cov_type);
      }
    }
  }
};
#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 = derivatives_cache<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 = derivatives_cache<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 = derivatives_cache<double>(theta_v, n_groups, is_spatial, cov_type, n_visits);
  for (int i = 0; i < n_subjects; i++) {
    int start_i = subject_zero_inds[i];
    int n_visits_i = subject_n_visits[i];
    std::vector<int> visit_i(n_visits_i);
    matrix<double> dist_i(n_visits_i, n_visits_i);
    if (!is_spatial) {
      for (int i = 0; i < n_visits_i; i++) {
        visit_i[i] = int(coordinates(i + start_i, 0));
      }
    } else {
      dist_i = euclidean(matrix<double>(coords.block(start_i, 0, n_visits_i, coordinates.cols())));
    }
    int subject_group_i = subject_groups[i] - 1;
    matrix<double> sigma_inv, sigma_d1, sigma_d2, sigma, sigma_inv_d1;

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

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

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

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

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

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

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

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

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

  // Calculate negative log-likelihood.
  Type neg_log_lik;

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

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

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

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

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

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

  return neg_log_lik;
}

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

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

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

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

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

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

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

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

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

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

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

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

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

#endif

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

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

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

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

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	5869x	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	33717x	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	11239x	this->theta = theta;
34	33717x	std::iota(std::begin(this->full_visit), std::end(this->full_visit), 0);
35	11239x	this->n_theta = theta.size();
36	11239x	this->chol_full = get_covariance_lower_chol(this->theta, this->n_visits, this->cov_type);
37	11235x	this->chols[full_visit] = this->chol_full;
38	11235x	this->sigma_full = tcrossprod(this->chol_full, true);
39		}
40	658891x	matrix<Type> get_chol(std::vector<int> visits, matrix<Type> dist) {
41	658891x	auto target = this->chols.find(visits);
42	658891x	if (target != this->chols.end()) {
43	609345x	return target->second;
44		} else {
45	49546x	matrix<Type> cov_i = this->get_sigma(visits, dist);
46	49546x	Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > cov_i_chol(cov_i);
47	49546x	matrix<Type> Li = cov_i_chol.matrixL();
48	49546x	this->chols[visits] = Li;
49	49546x	return Li;
50		}
51		}
52	313517x	matrix<Type> get_sigma(std::vector<int> visits, matrix<Type> dist) {
53	313517x	auto target = this->sigmas.find(visits);
54	313517x	if (target != this->sigmas.end()) {
55	245033x	return target->second;
56		} else {
57	68484x	matrix<Type> ret = subset_matrix<matrix<Type>, vector<int>>(sigma_full, visits, visits);
58	68484x	this->sigmas[visits] = ret;
59	68484x	return ret;
60		}
61		}
62	105503x	matrix<Type> get_sigma_inverse(std::vector<int> visits, matrix<Type> dist) {
63	105503x	auto target = this->sigmas_inv.find(visits);
64	105503x	if (target != this->sigmas_inv.end()) {
65	91694x	return target->second;
66		} else {
67	13809x	matrix<Type> ret = this->get_sigma(visits, dist).inverse();
68	13809x	this->sigmas_inv[visits] = ret;
69	13809x	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	262x	lower_chol_spatial(vector<Type> theta, std::string cov_type): theta(theta), cov_type(cov_type) {
85		}
86	33228x	matrix<Type> get_chol(std::vector<int> visits, matrix<Type> dist) {
87	33228x	return get_spatial_covariance_lower_chol(this->theta, dist, this->cov_type);
88		}
89	10254x	matrix<Type> get_sigma(std::vector<int> visits, matrix<Type> dist) {
90	10254x	return tcrossprod(this->get_chol(visits, dist), true);
91		}
92	3744x	matrix<Type> get_sigma_inverse(std::vector<int> visits, matrix<Type> dist) {
93	3744x	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	10366x	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	10366x	int theta_one_group_size = theta.size() / n_groups;
106	21620x	for (int r = 0; r < n_groups; r++) {
107		// Use unique pointers here to better manage resource.
108	11258x	if (is_spatial) {
109	248x	this->cache[r] = std::make_shared<D1>(theta.segment(r * theta_one_group_size, theta_one_group_size), cov_type);
110		} else {
111	11010x	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	10366x	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	6952x	matrix<Type> get_default_chol() {
125	13904x	std::vector<int> visit(this->n_visits);
126	6952x	std::iota(std::begin(visit), std::end(visit), 0);
127	6952x	matrix<Type> dist(2, 2);
128	6952x	dist << 0, 1, 1, 0;
129	6952x	int dim = this->is_spatial?2:this->n_visits;
130	6952x	matrix<Type> covariance_lower_chol = matrix<Type>::Zero(dim * this->n_groups, dim);
131	14588x	for (int r = 0; r < this->n_groups; r++) {
132	7636x	covariance_lower_chol.block(r * dim, 0, dim, dim) = this->cache[r]->get_chol(visit, dist);
133		}
134	13904x	return covariance_lower_chol;
135		}
136		};
137
138		#endif

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

1		#' G-computation Covariance Correction
2		#'
3		#' When the MMRM includes covariate-by-treatment interactions, the standard
4		#' emmeans output conditions on covariates rather than averaging over them.
5		#' These internal helpers compute subject-level potential outcomes under each
6		#' treatment and use their distribution to correct the L matrix and covariance.
7		#'
8		#' @keywords internal
9		#' @noRd
10		NULL
11
12		#' Get Subject-Level Covariate Data
13		#'
14		#' @description Returns one row per subject with covariate data. Pulls
15		#' subject level data from the original dataset stored in `tmb_data$data`,
16		#' and includes subjects with missing outcomes.
17		#'
18		#' @note If time varying covariates are present, only the first observed
19		#' value per subject is retained. The intended use is for baseline
20		#' (pre-randomization) covariates that are constant across visits.
21		#'
22		#' @param object (`mmrm`)\cr the fitted MMRM.
23		#'
24		#' @return A `data.frame` with one row per subject.
25		#'
26		#' @keywords internal
27		#' @noRd
28		h_get_subject_data <- function(object) {
29	21x	assert_class(object, "mmrm")
30	21x	original_data <- object$tmb_data$data
31	21x	subject_var <- object$formula_parts$subject_var
32	21x	model_vars <- all.vars(object$formula_parts$model_formula)
33	21x	keep_vars <- intersect(c(subject_var, model_vars), names(original_data))
34	21x	subj_data <- original_data[
35	21x	!duplicated(original_data[[subject_var]]),
36	21x	keep_vars,
37	21x	drop = FALSE
38		]
39	21x	rownames(subj_data) <- NULL
40	21x	subj_data
41		}
42
43		#' Compute Subject-Level Potential Outcomes
44		#'
45		#' @description For a given visit, computes each subject's fitted potential
46		#' outcome under each combination of fixed variable levels by constructing
47		#' the counterfactual design matrix. Each subject's own covariate values
48		#' are used for all non-fixed variables.
49		#'
50		#' @param object (`mmrm`)\cr the fitted MMRM.
51		#' @param subj_data (`data.frame`)\cr one row per subject with covariates.
52		#' @param visit_value (`string`)\cr the visit level to evaluate at.
53		#' @param visit_var (`string`)\cr name of the visit variable.
54		#' @param counterfactual_grid (`data.frame`)\cr unique combinations of fixed
55		#' variables (excluding visit) to loop over.
56		#'
57		#' @return A list with:
58		#' - `vhat`: `matrix` of fitted potential outcomes (n_subjects x K),
59		#' where n_subjects is the number of rows in `subj_data` and K is the
60		#' number of rows in `counterfactual_grid`.
61		#' - `L_global`: `matrix` of globally weighted contrast rows (K x p),
62		#' where p is the number of estimable model coefficients.
63		#'
64		#' @keywords internal
65		#' @noRd
66		h_compute_potential_outcomes <- function(
67		object,
68		subj_data,
69		visit_value,
70		visit_var,
71		counterfactual_grid
72		) {
73	58x	assert_class(object, "mmrm")
74	58x	assert_data_frame(subj_data, min.rows = 1L)
75	58x	assert_string(visit_value)
76	58x	assert_string(visit_var)
77	58x	assert_data_frame(counterfactual_grid, min.rows = 1L)
78	58x	beta_hat <- component(object, "beta_est")
79	58x	beta_clean <- beta_hat
80	58x	beta_clean[is.na(beta_clean)] <- 0
81	58x	contrasts_arg <- component(object, "contrasts")
82	58x	full_frame <- object$tmb_data$full_frame
83
84	58x	model_formula <- object$formula_parts$model_formula
85	58x	model_terms <- stats::delete.response(stats::terms(model_formula))
86
87	58x	K <- nrow(counterfactual_grid) # nolint
88	58x	n_subj <- nrow(subj_data)
89	58x	p <- length(beta_hat)
90
91		# Stack all K counterfactual assignments into one data frame.
92		# (for efficiency reasons, we can just do one model.frame and model.matrix call below instead of K).
93	58x	cf_list <- vector("list", K)
94	58x	for (k in seq_len(K)) {
95	116x	subj_cf <- subj_data
96	116x	subj_cf[[visit_var]] <- factor(
97	116x	visit_value,
98	116x	levels = levels(full_frame[[visit_var]])
99		)
100	116x	for (var_name in names(counterfactual_grid)) {
101	116x	val <- counterfactual_grid[[var_name]][k]
102	116x	subj_cf[[var_name]] <- if (is.factor(full_frame[[var_name]])) {
103	106x	factor(val, levels = levels(full_frame[[var_name]]))
104		} else {
105	10x	val
106		}
107		}
108	116x	cf_list[[k]] <- subj_cf
109		}
110	58x	cf_stacked <- do.call(rbind, cf_list)
111
112	58x	mf <- stats::model.frame(
113	58x	model_terms,
114	58x	data = cf_stacked,
115	58x	na.action = stats::na.pass
116		)
117		# nolint next
118	58x	X_all <- stats::model.matrix(
119	58x	model_terms,
120	58x	data = mf,
121	58x	contrasts.arg = contrasts_arg
122		)
123	58x	X_all <- X_all[, names(beta_hat), drop = FALSE] # nolint
124	58x	X_all[is.na(X_all)] <- 0 # nolint
125
126		# Split back into K blocks of n_subj rows.
127	58x	vhat <- matrix(NA_real_, nrow = n_subj, ncol = K)
128	58x	L_global <- matrix(NA_real_, nrow = K, ncol = p) # nolint
129	58x	for (k in seq_len(K)) {
130	116x	rows <- ((k - 1L) * n_subj + 1L):(k * n_subj)
131	116x	X_k <- X_all[rows, , drop = FALSE] # nolint
132	116x	vhat[, k] <- as.vector(X_k %*% beta_clean)
133	116x	L_global[k, ] <- colMeans(X_k, na.rm = TRUE) # nolint
134		}
135
136	58x	list(vhat = vhat, L_global = L_global)
137		}
138
139		#' Compute Per Visit Variance Contributions and Global L Matrix
140		#'
141		#' @description For each visit, computes subject level potential outcomes,
142		#' the sample covariance of those outcomes (divided by n), and the
143		#' globally weighted L matrix rows. Assembles these into the full
144		#' block diagonal S matrix and global L matrix across all visits.
145		#'
146		#' @param object (`mmrm`)\cr the fitted MMRM.
147		#' @param model_mat (`matrix`)\cr the L matrix from the emmeans reference
148		#' grid (n_grid x p).
149		#' @param grid (`data.frame`)\cr the emmeans reference grid.
150		#' @param fixed_vars (`character`)\cr names of the fixed (intervention)
151		#' variables.
152		#' @param visit_var (`string`)\cr name of the visit variable.
153		#'
154		#' @return A list with:
155		#' - `S_full`: `matrix` (n_grid x n_grid) block-diagonal sample covariance
156		#' of potential outcomes divided by n, with blocks corresponding to visits.
157		#' - `L_global`: `matrix` (n_grid x p) globally-weighted L matrix rows.
158		#'
159		#' @keywords internal
160		#' @noRd
161		h_gcomp_visit_contributions <- function(
162		object,
163		model_mat,
164		grid,
165		fixed_vars,
166		visit_var
167		) {
168	16x	subj_data <- h_get_subject_data(object)
169	16x	full_frame <- object$tmb_data$full_frame
170	16x	n_subj <- nrow(subj_data)
171	16x	visit_levels <- levels(full_frame[[visit_var]])
172	16x	cell_vars <- setdiff(fixed_vars, visit_var)
173
174	16x	n_grid <- nrow(model_mat)
175	16x	grid_visit <- as.character(grid[[visit_var]])
176	16x	S_full <- matrix(0, nrow = n_grid, ncol = n_grid) # nolint
177	16x	L_global <- model_mat # nolint
178
179	16x	for (v in visit_levels) {
180	64x	v_mask <- grid_visit == v
181	64x	if (!any(v_mask)) {
182	9x	next
183		}
184	55x	v_indices <- which(v_mask)
185
186	55x	grid_at_v <- grid[v_indices, cell_vars, drop = FALSE]
187	55x	cf_grid <- unique(grid_at_v)
188	55x	rownames(cf_grid) <- NULL
189
190	55x	po <- h_compute_potential_outcomes(
191	55x	object = object,
192	55x	subj_data = subj_data,
193	55x	visit_value = v,
194	55x	visit_var = visit_var,
195	55x	counterfactual_grid = cf_grid
196		)
197
198	55x	sigma_v <- stats::cov(po$vhat, use = "pairwise.complete.obs")
199	55x	S_t <- sigma_v / n_subj # nolint
200
201	55x	cf_keys <- apply(cf_grid, 1, function(r) {
202	110x	paste(as.character(r), collapse = "\|")
203		})
204	55x	grid_keys <- apply(grid_at_v, 1, function(r) {
205	118x	paste(as.character(r), collapse = "\|")
206		})
207
208	55x	for (ii in seq_along(v_indices)) {
209	118x	ki <- match(grid_keys[ii], cf_keys)
210		if (is.na(ki)) next # nocov
211
212	118x	L_global[v_indices[ii], ] <- po$L_global[ki, ] # nolint
213
214	118x	for (jj in seq_along(v_indices)) {
215	268x	kj <- match(grid_keys[jj], cf_keys)
216		if (is.na(kj)) next # nocov
217	268x	S_full[v_indices[ii], v_indices[jj]] <- S_t[ki, kj] # nolint
218		}
219		}
220		}
221
222	16x	list(S_full = S_full, L_global = L_global)
223		}
224
225		#' Compute G-computation Correction for emmeans Integration
226		#'
227		#' @description Computes the corrected L matrix (global covariate weighting)
228		#' and the variance correction matrix delta. Together these produce the
229		#' G-computation estimator with correct standard errors. Delta is a
230		#' p x p positive semi-definite matrix (where p is the number of estimable
231		#' model coefficients) that is added to the model-based covariance.
232		#'
233		#' @param object (`mmrm`)\cr the fitted MMRM with
234		#' `emmeans_gcomp_vars` set.
235		#' @param model_mat (`matrix`)\cr the L matrix from the emmeans reference
236		#' grid (n_grid x p).
237		#' @param grid (`data.frame`)\cr the reference grid from emmeans.
238		#'
239		#' @return A list with:
240		#' - `delta`: `matrix` (p x p) positive semi-definite correction, where
241		#' p is the number of estimable model coefficients.
242		#' - `L_global`: `matrix` (n_grid x p) globally-weighted L matrix.
243		#'
244		#' @keywords internal
245		#' @noRd
246		h_gcomp_emm_correction <- function(object, model_mat, grid) {
247	15x	assert_class(object, "mmrm")
248	15x	assert_matrix(model_mat, mode = "numeric")
249	15x	assert_data_frame(grid)
250
251	15x	fixed_vars <- object$tmb_data$emmeans_gcomp_vars
252	15x	visit_var <- object$formula_parts$visit_var
253
254	15x	assert_character(fixed_vars, min.len = 1L)
255	15x	assert_string(visit_var)
256	15x	assert_subset(visit_var, names(object$tmb_data$full_frame))
257
258	15x	beta_names <- names(component(object, "beta_est"))
259		# The emmeans grid model matrix may have extra columns (e.g. from
260		# interaction reordering) that do not appear in beta_est. Trim to
261		# match the coefficient vector, since estimability::nonest.basis
262		# only handles rank deficiency within the fitted coefficient set.
263	15x	if (ncol(model_mat) > length(beta_names)) {
264	4x	keep_cols <- match(beta_names, colnames(model_mat))
265	4x	keep_cols <- keep_cols[!is.na(keep_cols)]
266	4x	model_mat <- model_mat[, keep_cols, drop = FALSE]
267		}
268
269	15x	contributions <- h_gcomp_visit_contributions(
270	15x	object,
271	15x	model_mat,
272	15x	grid,
273	15x	fixed_vars,
274	15x	visit_var
275		)
276	15x	S_full <- contributions$S_full # nolint
277	15x	L <- contributions$L_global # nolint
278
279	15x	llt_ginv_l <- MASS::ginv(tcrossprod(L)) %*% L
280	15x	delta <- t(llt_ginv_l) %% S_full %% llt_ginv_l
281	15x	V <- component(object, "beta_vcov") # nolint
282	15x	dimnames(delta) <- dimnames(V)
283
284	15x	list(delta = delta, L_global = L)
285		}

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

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	126x	assert_class(tmb_data, "mmrm_tmb_data")
16	126x	assert_numeric(theta_est)
17	126x	assert_matrix(beta_vcov)
18	126x	.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	7315x	vec <- as.vector(vec)
40	7315x	assert_numeric(vec, any.missing = FALSE)
41	7315x	assert_matrix(
42	7315x	center,
43	7315x	mode = "numeric",
44	7315x	any.missing = FALSE,
45	7315x	nrows = length(vec),
46	7315x	ncols = length(vec)
47		)
48
49	7315x	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	458x	assert_matrix(mat, mode = "numeric", any.missing = FALSE, min.cols = 1L)
63	458x	assert_matrix(
64	458x	center,
65	458x	mode = "numeric",
66	458x	any.missing = FALSE,
67	458x	nrows = ncol(center),
68	458x	ncols = ncol(center)
69		)
70
71	458x	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	662x	assert_list(jac_list)
90	662x	assert_numeric(contrast)
91
92	662x	vapply(
93	662x	jac_list,
94	662x	h_quad_form_vec,
95	662x	vec = contrast,
96	662x	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	556x	assert_class(object, "mmrm")
154	556x	contrast <- as.numeric(contrast)
155	556x	assert_numeric(contrast, len = length(component(object, "beta_est")))
156
157	556x	df <- if (identical(object$vcov, "Asymptotic")) {
158	543x	grad <- h_gradient(component(object, "jac_list"), contrast)
159	543x	v_num <- 2 * h_quad_form_vec(contrast, component(object, "beta_vcov"))^2
160	543x	v_denom <- h_quad_form_vec(grad, component(object, "theta_vcov"))
161	543x	v_num / v_denom
162	556x	} else if (
163	556x	object$vcov %in%
164	556x	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	556x	h_test_1d(object, contrast, df)
175		}
176
177		#' Calculating Denominator Degrees of Freedom for the Multi-Dimensional Case
178		#'
179		#' @description Calculates the degrees of freedom for multi-dimensional contrast.
180		#'
181		#' @param t_stat_df (`numeric`)\cr `n` t-statistic derived degrees of freedom.
182		#'
183		#' @return Usually the calculation is returning `2 * E / (E - n)` where
184		#' `E` is the sum of `t / (t - 2)` over all `t_stat_df` values `t`.
185		#'
186		#' @note If the input values are two similar to each other then just the average
187		#' of them is returned. If any of the inputs is not larger than 2 then 2 is
188		#' returned.
189		#'
190		#' @keywords internal
191		h_md_denom_df <- function(t_stat_df) {
192	52x	assert_numeric(
193	52x	t_stat_df,
194	52x	min.len = 1L,
195	52x	lower = .Machine$double.xmin,
196	52x	any.missing = FALSE
197		)
198
199	52x	if (test_scalar(t_stat_df)) {
200	1x	t_stat_df
201	51x	} else if (all(abs(diff(t_stat_df)) < sqrt(.Machine$double.eps))) {
202	1x	mean(t_stat_df)
203	50x	} else if (any(t_stat_df <= 2)) {
204	2x	2
205		} else {
206	48x	e <- sum(t_stat_df / (t_stat_df - 2))
207	48x	2 * e / (e - (length(t_stat_df)))
208		}
209		}
210
211		#' Creating F-Statistic Results from One-Dimensional Contrast
212		#'
213		#' @description Creates multi-dimensional result from one-dimensional contrast from [df_1d()].
214		#'
215		#' @param object (`mmrm`)\cr model fit.
216		#' @param contrast (`numeric`)\cr one-dimensional contrast.
217		#'
218		#' @return The one-dimensional degrees of freedom are calculated and then
219		#' based on that the p-value is calculated.
220		#'
221		#' @keywords internal
222		h_df_md_from_1d <- function(object, contrast) {
223	233x	res_1d <- h_df_1d_sat(object, contrast)
224	233x	list(
225	233x	num_df = 1,
226	233x	denom_df = res_1d$df,
227	233x	f_stat = res_1d$t_stat^2,
228	233x	p_val = stats::pf(
229	233x	q = res_1d$t_stat^2,
230	233x	df1 = 1,
231	233x	df2 = res_1d$df,
232	233x	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	278x	assert_class(object, "mmrm")
251	278x	assert_matrix(
252	278x	contrast,
253	278x	mode = "numeric",
254	278x	any.missing = FALSE,
255	278x	ncols = length(component(object, "beta_est"))
256		)
257		# Early return if we are in the one-dimensional case.
258	278x	if (identical(nrow(contrast), 1L)) {
259	231x	return(h_df_md_from_1d(object, contrast))
260		}
261
262	47x	contrast_cov <- h_quad_form_mat(contrast, component(object, "beta_vcov"))
263	47x	eigen_cont_cov <- eigen(contrast_cov)
264	47x	eigen_cont_cov_vctrs <- eigen_cont_cov$vectors
265	47x	eigen_cont_cov_vals <- eigen_cont_cov$values
266
267	47x	eps <- sqrt(.Machine$double.eps)
268	47x	tol <- max(eps * eigen_cont_cov_vals[1], 0)
269	47x	rank_cont_cov <- sum(eigen_cont_cov_vals > tol)
270	47x	assert_number(rank_cont_cov, lower = .Machine$double.xmin)
271	47x	rank_seq <- seq_len(rank_cont_cov)
272	47x	vctrs_cont_prod <- crossprod(eigen_cont_cov_vctrs, contrast)[
273	47x	rank_seq,
274		,
275	47x	drop = FALSE
276		]
277
278		# Early return if rank 1.
279	47x	if (identical(rank_cont_cov, 1L)) {
280	1x	return(h_df_md_from_1d(object, vctrs_cont_prod))
281		}
282
283	46x	t_squared_nums <- drop(vctrs_cont_prod %*% object$beta_est)^2
284	46x	t_squared_denoms <- eigen_cont_cov_vals[rank_seq]
285	46x	t_squared <- t_squared_nums / t_squared_denoms
286	46x	f_stat <- sum(t_squared) / rank_cont_cov
287	46x	t_stat_df_nums <- 2 * eigen_cont_cov_vals^2
288	46x	t_stat_df <- if (identical(object$vcov, "Asymptotic")) {
289	45x	grads_vctrs_cont_prod <- lapply(
290	45x	rank_seq,
291	45x	function(m) {
292	118x	h_gradient(
293	118x	component(object, "jac_list"),
294	118x	contrast = vctrs_cont_prod[m, ]
295		)
296		}
297		)
298	45x	t_stat_df_denoms <- vapply(
299	45x	grads_vctrs_cont_prod,
300	45x	h_quad_form_vec,
301	45x	center = component(object, "theta_vcov"),
302	45x	numeric(1)
303		)
304	45x	t_stat_df_nums / t_stat_df_denoms
305		} else {
306	1x	vapply(
307	1x	rank_seq,
308	1x	function(m) {
309	2x	contrast_matrix <- Matrix::.bdiag(
310	2x	rep(
311	2x	list(vctrs_cont_prod[m, , drop = FALSE]),
312	2x	component(object, "n_subjects")
313		)
314		)
315	2x	contrast_matrix <- as.matrix(contrast_matrix)
316	2x	h_df_1d_sat_empirical(object, contrast_matrix)
317		},
318	1x	FUN.VALUE = 0
319		)
320		}
321	46x	denom_df <- h_md_denom_df(t_stat_df)
322
323	46x	list(
324	46x	num_df = rank_cont_cov,
325	46x	denom_df = denom_df,
326	46x	f_stat = f_stat,
327	46x	p_val = stats::pf(
328	46x	q = f_stat,
329	46x	df1 = rank_cont_cov,
330	46x	df2 = denom_df,
331	46x	lower.tail = FALSE
332		)
333		)
334		}

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

1		#' 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		#' Match coefficient names against model matrix columns
25		#'
26		#' This helper is robust against re-ordering of interaction terms, e.g.
27		#' `A:B` vs `B:A` and `A:B:C` vs any permutation.
28		#'
29		#' @keywords internal
30		#' @noRd
31		h_match_coefs <- function(coef_names, model_colnames) {
32	18x	canon_interaction <- function(x) {
33	6x	split_terms <- strsplit(x, ":", fixed = TRUE)
34	6x	vapply(
35	6x	split_terms,
36	6x	FUN.VALUE = character(1),
37	6x	function(term_parts) {
38	26x	if (length(term_parts) <= 1L) {
39	8x	term_parts
40		} else {
41	18x	paste(sort(term_parts), collapse = ":")
42		}
43		}
44		)
45		}
46
47	18x	kept <- match(coef_names, model_colnames)
48	18x	idx_na <- is.na(kept)
49
50	18x	if (any(idx_na)) {
51	3x	kept[idx_na] <- match(
52	3x	canon_interaction(coef_names[idx_na]),
53	3x	canon_interaction(model_colnames)
54		)
55		}
56
57	18x	kept
58		}
59
60		#' Returns a `data.frame` for `emmeans` Purposes
61		#'
62		#' @seealso See [emmeans::recover_data()] for background.
63		#' @keywords internal
64		#' @noRd
65		# nolint start
66		recover_data.mmrm <- function(object, ...) {
67		# nolint end
68	36x	fun_call <- stats::getCall(object)
69		# subject_var is excluded because it should not contain fixed effect.
70		# visit_var is only excluded for spatial covariance structures -
71		# for non-spatial covariance structures, emmeans can provide marginal mean
72		# by each visit so we need visit_var. Also exclude visit_var when it has only
73		# one level (emmeans cannot contrast it).
74	36x	visit_var <- object$formula_parts$visit_var
75	36x	exclude_visit <- object$formula_parts$is_spatial \|\|
76	36x	nlevels(object$tmb_data$full_frame[[visit_var]]) < 2L
77
78	36x	model_frame <- stats::model.frame(
79	36x	object,
80	36x	exclude = c(
81	36x	"subject_var",
82	36x	if (exclude_visit) "visit_var"
83		)
84		)
85	36x	model_terms <- stats::delete.response(stats::terms(model_frame))
86	36x	emmeans::recover_data(
87	36x	fun_call,
88	36x	trms = model_terms,
89	36x	na.action = "na.omit",
90	36x	frame = model_frame,
91		...
92		)
93		}
94
95		#' Returns a List of Model Details for `emmeans` Purposes
96		#'
97		#' @seealso See [emmeans::emm_basis()] for background.
98		#' @keywords internal
99		#' @noRd
100		# nolint start
101		emm_basis.mmrm <- function(
102		# nolint end
103		object,
104		trms,
105		xlev,
106		grid,
107		...
108		) {
109	35x	model_frame <- stats::model.frame(
110	35x	trms,
111	35x	grid,
112	35x	na.action = stats::na.pass,
113	35x	xlev = xlev
114		)
115	35x	contrasts <- component(object, "contrasts")
116	35x	model_mat <- stats::model.matrix(trms, model_frame, contrasts.arg = contrasts)
117	35x	beta_hat <- component(object, "beta_est")
118	35x	nbasis <- if (length(beta_hat) < ncol(model_mat)) {
119	12x	kept <- h_match_coefs(names(beta_hat), colnames(model_mat))
120	12x	if (anyNA(kept)) {
121	!	unmatched <- names(beta_hat)[is.na(kept)]
122	!	stop(
123	!	paste0(
124	!	"Failed to match coefficient name(s) to model matrix columns: ",
125	!	paste(unmatched, collapse = ", ")
126		),
127	!	call. = FALSE
128		)
129		}
130	12x	beta_hat <- NA * model_mat[1L, ]
131	12x	beta_hat[kept] <- component(object, "beta_est")
132	12x	orig_model_mat <- stats::model.matrix(
133	12x	trms,
134	12x	stats::model.frame(
135	12x	object,
136	12x	exclude = c(
137	12x	"subject_var",
138	12x	if (object$formula_parts$is_spatial) "visit_var"
139		)
140		),
141	12x	contrasts.arg = contrasts
142		)
143	12x	estimability::nonest.basis(orig_model_mat)
144		} else {
145	23x	estimability::all.estble
146		}
147
148		# Start with the model-based covariance.
149	35x	V <- component(object, "beta_vcov") # nolint
150
151	35x	if (!is.null(object$tmb_data$emmeans_gcomp_vars)) {
152	12x	gcomp_result <- h_emm_basis_gcomp(
153	12x	object,
154	12x	model_mat,
155	12x	beta_hat,
156	12x	nbasis,
157	12x	V,
158	12x	grid
159		)
160	12x	model_mat <- gcomp_result$model_mat
161	12x	beta_hat <- gcomp_result$beta_hat
162	12x	nbasis <- gcomp_result$nbasis
163	12x	V <- gcomp_result$V # nolint
164	12x	dfargs <- gcomp_result$dfargs
165	12x	dffun <- gcomp_result$dffun
166	12x	message(
167	12x	"G-computation correction applied, so emmeans will produce the correct average treatment effect."
168		)
169		} else {
170	23x	dfargs <- list(object = object)
171	23x	dffun <- function(k, dfargs) {
172	177x	mmrm::df_md(dfargs$object, contrast = k)$denom_df
173		}
174		}
175
176	35x	list(
177	35x	X = model_mat,
178	35x	bhat = beta_hat,
179	35x	nbasis = nbasis,
180	35x	V = V,
181	35x	dffun = dffun,
182	35x	dfargs = dfargs
183		)
184		}
185
186		#' Apply G-computation Correction in emm_basis
187		#'
188		#' @description Applies the G-computation correction to the emmeans basis
189		#' components. Replaces the L matrix with globally-weighted rows, adds
190		#' the variance correction to V, and handles aliased coefficient expansion.
191		#'
192		#' @param object (`mmrm`)\cr the fitted MMRM.
193		#' @param model_mat (`matrix`)\cr the L matrix from emmeans.
194		#' @param beta_hat (`numeric`)\cr coefficient estimates (may contain NA for
195		#' aliased positions).
196		#' @param nbasis (`matrix` or `NULL`)\cr estimability basis from emmeans.
197		#' @param V (`matrix`)\cr model-based covariance of beta.
198		#' @param grid (`data.frame`)\cr the emmeans reference grid.
199		#'
200		#' @return A list with modified `model_mat`, `beta_hat`, `nbasis`, `V`,
201		#' `dfargs`, and `dffun`.
202		#'
203		#' @keywords internal
204		#' @noRd
205		# nolint next
206		h_emm_basis_gcomp <- function(object, model_mat, beta_hat, nbasis, V, grid) {
207		# Detect aliased coefficient positions before any modifications.
208	12x	aliased_pos <- which(is.na(beta_hat))
209
210	12x	visit_var <- object$formula_parts$visit_var
211	12x	V_corrected_est <- V # nolint
212
213	12x	if (visit_var %in% names(grid)) {
214		# Save original model_mat dimensions. When aliased coefficients exist,
215		# h_match_coefs expands model_mat beyond length(beta_est). The correction
216		# works in estimable-coefficient space, then we re-expand to match.
217	12x	orig_ncol <- ncol(model_mat)
218	12x	orig_colnames <- colnames(model_mat)
219
220	12x	correction <- h_gcomp_emm_correction(object, model_mat, grid)
221	12x	V_corrected_est <- V + correction$delta # nolint
222
223		# Re-expand L_global if h_match_coefs added columns for aliased coefs.
224		# Use h_match_coefs (not match) to handle A:B vs B:A reordering.
225	12x	L_est <- correction$L_global # nolint
226	12x	if (ncol(L_est) < orig_ncol && length(aliased_pos) > 0) {
227	3x	est_pos <- setdiff(seq_len(orig_ncol), aliased_pos)
228
229		# Expand L_global to full width with zeros for aliased columns.
230	3x	L_exp <- matrix(0, nrow = nrow(L_est), ncol = orig_ncol) # nolint
231	3x	colnames(L_exp) <- orig_colnames # nolint
232	3x	col_map <- h_match_coefs(colnames(L_est), orig_colnames)
233	3x	col_map <- col_map[!is.na(col_map)]
234	3x	L_exp[, col_map] <- L_est # nolint
235	3x	model_mat <- L_exp
236
237		# Expand V to full width with zeros for aliased row/column.
238	3x	V_full <- matrix(0, nrow = orig_ncol, ncol = orig_ncol) # nolint
239	3x	V_full[est_pos, est_pos] <- V_corrected_est # nolint
240	3x	V <- V_full # nolint
241
242		# Zero NA in bhat so 0*NA doesn't propagate.
243	3x	beta_hat[aliased_pos] <- 0
244
245		# All grid rows are estimable through L_global (aliased col is zero).
246	3x	nbasis <- estimability::all.estble
247		} else {
248	9x	model_mat <- L_est
249	9x	V <- V_corrected_est # nolint
250		}
251		}
252
253		# dffun operates in estimable-coefficient space (p_est dimensions).
254		# Strip aliased positions from k before passing to df_md.
255	12x	dfargs <- list(
256	12x	object = object,
257	12x	V_corrected = V_corrected_est,
258	12x	aliased_pos = aliased_pos
259		)
260	12x	dffun <- function(k, dfargs) {
261	54x	obj <- dfargs$object
262	54x	if (is.null(obj$vcov) \|\| identical(obj$vcov, "Asymptotic")) {
263	38x	obj$beta_vcov <- dfargs$V_corrected
264		} else {
265	16x	obj$beta_vcov_adj <- dfargs$V_corrected
266		}
267	54x	k_est <- k
268	12x	if (
269	54x	length(dfargs$aliased_pos) > 0 && length(k) > nrow(dfargs$V_corrected)
270		) {
271	24x	k_est <- k[-dfargs$aliased_pos]
272		}
273	54x	mmrm::df_md(obj, contrast = k_est)$denom_df
274		}
275
276	12x	list(
277	12x	model_mat = model_mat,
278	12x	beta_hat = beta_hat,
279	12x	nbasis = nbasis,
280	12x	V = V,
281	12x	dfargs = dfargs,
282	12x	dffun = dffun
283		)
284		}

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

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

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

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

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

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

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

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

1		# 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		#' 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	9x	assert_class(object, "mmrm")
27	9x	assert_matrix(
28	9x	contrast,
29	9x	mode = "numeric",
30	9x	any.missing = FALSE,
31	9x	ncols = length(component(object, "beta_est"))
32		)
33
34	9x	df <- component(object, "n_obs") - length(component(object, "beta_est"))
35
36	9x	h_test_md(object, contrast, df)
37		}

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	40x	assert_class(tmb_data, "mmrm_tmb_data")
26	40x	assert_numeric(theta)
27	40x	n_beta <- ncol(tmb_data$x_matrix)
28	40x	assert_numeric(beta, finite = TRUE, any.missing = FALSE, len = n_beta)
29	40x	assert_matrix(
30	40x	beta_vcov,
31	40x	mode = "numeric",
32	40x	any.missing = FALSE,
33	40x	nrows = n_beta,
34	40x	ncols = n_beta
35		)
36	40x	assert_subset(
37	40x	type,
38	40x	c("Empirical", "Empirical-Jackknife", "Empirical-Bias-Reduced")
39		)
40	40x	.Call(
41	40x	`_mmrm_get_empirical`,
42	40x	PACKAGE = "mmrm",
43	40x	tmb_data,
44	40x	theta,
45	40x	beta,
46	40x	beta_vcov,
47	40x	type
48		)
49		}

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	680x	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	340x	flattened_args <- unlist(lapply(call[-1], flatten_expr))
56	340x	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	1505x	if (length(expr) > 1 && is_infix(expr[[1]])) {
69	400x	op <- list(expr[[1]])
70	400x	lhs <- flatten_expr(expr[[2]])
71	400x	rhs <- flatten_expr(expr[[3]])
72	400x	c(lhs, op, rhs)
73		} else {
74	1105x	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	359x	if (length(f) == 2) {
93	9x	f
94		} else {
95	350x	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	412x	"Ops" %in% methods::getGroup(as.character(name), recursive = TRUE)
116		}
117
118		#' Format Symbol Objects
119		#'
120		#' For printing, variable names are converted to symbols and deparsed to use R's
121		#' built-in formatting of variables that may contain spaces or quote characters.
122		#'
123		#' @param x (`character`) A vector of variable names.
124		#'
125		#' @return A formatted string of comma-separated variables.
126		#'
127		#' @keywords internal
128		format_symbols <- function(x) {
129	12x	paste0(
130	12x	collapse = ", ",
131	12x	lapply(x, function(i) {
132	16x	utils::capture.output(as.symbol(i))
133		})
134		)
135		}

1		#ifndef COV_INCLUDED_
2		#define COV_INCLUDED_
3
4		#include "utils.h"
5
6		// Unstructured covariance:
7		// Cholesky factor.
8		template <class T>
9	9404x	matrix<T> get_unstructured(const vector<T>& theta, int n_visits) {
10	9404x	vector<T> sd_values = exp(theta.head(n_visits));
11	9404x	vector<T> lower_tri_chol_values = theta.tail(theta.size() - n_visits);
12	9404x	matrix<T> covariance_lower_chol = matrix<T>::Zero(n_visits, n_visits);
13	9404x	int k = 0;
14	46724x	for(int i = 0; i < n_visits; i++) {
15	37320x	covariance_lower_chol(i, i) = sd_values(i);
16	93132x	for(int j = 0; j < i; j++){
17	55812x	covariance_lower_chol(i, j) = sd_values(i) * lower_tri_chol_values(k++);
18		}
19		}
20	18808x	return covariance_lower_chol;
21		}
22
23		// Ante-dependence:
24
25		// Correlation function.
26		template <class T>
27		struct corr_fun_ante_dependence : generic_corr_fun<T> {
28		using generic_corr_fun<T>::generic_corr_fun;
29	2256x	const T operator() (int i, int j) const {
30	2256x	return this->corr_values.segment(j, i - j).prod();
31		}
32		};
33		// Homogeneous Ante-dependence Cholesky factor.
34		template <class T>
35	158x	matrix<T> get_ante_dependence(const vector<T>& theta, int n_visits) {
36	158x	T const_sd = exp(theta(0));
37	158x	corr_fun_ante_dependence<T> fun(theta.tail(n_visits - 1));
38	158x	matrix<T> ad_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
39	316x	return const_sd * ad_cor_mat_chol;
40		}
41		// Heterogeneous Ante-dependence Cholesky factor.
42		template <class T>
43	238x	matrix<T> get_ante_dependence_heterogeneous(const vector<T>& theta, int n_visits) {
44	238x	vector<T> sd_values = exp(theta.head(n_visits));
45	238x	corr_fun_ante_dependence<T> fun(theta.tail(n_visits - 1));
46	476x	return get_heterogeneous_cov(sd_values, fun);
47		}
48
49		// Toeplitz:
50
51		// Correlation function.
52		template <class T>
53		struct corr_fun_toeplitz : generic_corr_fun<T> {
54		using generic_corr_fun<T>::generic_corr_fun;
55	2538x	const T operator() (int i, int j) const {
56	2538x	int index = (i - j) - 1; // Note: We need to start at 0.
57	2538x	return this->corr_values(index);
58		}
59		};
60		// Homogeneous Toeplitz Cholesky factor.
61		template <class T>
62	208x	matrix<T> get_toeplitz(const vector<T>& theta, int n_visits) {
63	208x	T const_sd = exp(theta(0));
64	208x	corr_fun_toeplitz<T> fun(theta.tail(n_visits - 1));
65	208x	matrix<T> toep_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
66	416x	return const_sd * toep_cor_mat_chol;
67		}
68		// Heterogeneous Toeplitz Cholesky factor.
69		template <class T>
70	214x	matrix<T> get_toeplitz_heterogeneous(const vector<T>& theta, int n_visits) {
71	214x	vector<T> sd_values = exp(theta.head(n_visits));
72	214x	corr_fun_toeplitz<T> fun(theta.tail(n_visits - 1));
73	428x	return get_heterogeneous_cov(sd_values, fun);
74		}
75
76		// Autoregressive:
77
78		// Correlation function.
79		template <class T>
80		struct corr_fun_autoregressive : generic_corr_fun<T> {
81		using generic_corr_fun<T>::generic_corr_fun;
82	12124x	const T operator() (int i, int j) const {
83	12124x	T diff = T((i - j) * 1.0);
84	15532x	return pow(this->corr_values(0), diff); // rho^{\|i-j\|}
85		}
86		};
87		// Homogeneous autoregressive Cholesky factor.
88		template <class T>
89	1802x	matrix<T> get_auto_regressive(const vector<T>& theta, int n_visits) {
90	1802x	T const_sd = exp(theta(0));
91	1802x	corr_fun_autoregressive<T> fun(theta.tail(1));
92	1802x	matrix<T> ar1_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
93	3604x	return const_sd * ar1_cor_mat_chol;
94		}
95		// Heterogeneous autoregressive Cholesky factor.
96		template <class T>
97	214x	matrix<T> get_auto_regressive_heterogeneous(const vector<T>& theta, int n_visits) {
98	214x	vector<T> sd_values = exp(theta.head(n_visits));
99	214x	corr_fun_autoregressive<T> fun(theta.tail(1));
100	428x	return get_heterogeneous_cov(sd_values, fun);
101		}
102
103		// Compound symmetry:
104
105		// Correlation function.
106		template <class T>
107		struct corr_fun_compound_symmetry {
108		const vector<T> corr_values;
109		// Custom constructor here because we need the special mapping from theta to rho.
110	616x	corr_fun_compound_symmetry(const vector<T>& theta, int n_visits) :
111	616x	corr_values(map_to_cs_cor(theta, n_visits)) {}
112	3636x	const T operator() (int i, int j) const {
113	3636x	return this->corr_values(0); // rho (constant)
114		}
115		};
116		// Homogeneous compound symmetry Cholesky factor.
117		template <class T>
118	350x	matrix<T> get_compound_symmetry(const vector<T>& theta, int n_visits) {
119	350x	T const_sd = exp(theta(0));
120	350x	corr_fun_compound_symmetry<T> fun(theta.tail(1), n_visits);
121	350x	matrix<T> cs_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
122	700x	return const_sd * cs_cor_mat_chol;
123		}
124		// Heterogeneous compound symmetry Cholesky factor.
125		template <class T>
126	262x	matrix<T> get_compound_symmetry_heterogeneous(const vector<T>& theta, int n_visits) {
127	262x	vector<T> sd_values = exp(theta.head(n_visits));
128	262x	corr_fun_compound_symmetry<T> fun(theta.tail(1), n_visits);
129	524x	return get_heterogeneous_cov(sd_values, fun);
130		}
131
132		// Spatial Exponential Cholesky factor.
133		template <class T>
134	23736x	matrix<T> get_spatial_exponential(const vector<T>& theta, const matrix<T>& distance) {
135	23736x	T const_sd = exp(theta(0));
136	23736x	T rho = invlogit(theta(1)); // theta (-Inf, Inf), rho (0, 1), log(rho) (-Inf, 0)
137	23736x	matrix<T> expdist = exp(distance.array() * log(rho));
138	23736x	matrix<T> result = expdist * const_sd;
139	23736x	Eigen::LLT<Eigen::Matrix<T,Eigen::Dynamic,Eigen::Dynamic> > cov_i_chol(result);
140	47472x	return cov_i_chol.matrixL();
141		}
142
143		// Spatial Gaussian Cholesky factor.
144		template <class T>
145	9492x	matrix<T> get_spatial_gaussian(const vector<T>& theta, const matrix<T>& distance) {
146	9492x	T const_sd = exp(theta(0));
147	9492x	T rho = invlogit(theta(1));
148	9492x	matrix<T> expdist = exp(distance.array().square() * log(rho)); // only take square of the dist
149	9492x	matrix<T> result = expdist * const_sd;
150	9492x	Eigen::LLT<Eigen::Matrix<T,Eigen::Dynamic,Eigen::Dynamic> > cov_i_chol(result);
151	18984x	return cov_i_chol.matrixL();
152		}
153
154
155		// Creates a new correlation object dynamically.
156		template <class T>
157	12834x	matrix<T> get_covariance_lower_chol(const vector<T>& theta, int n_visits, std::string cov_type) {
158	12834x	matrix<T> result;
159
160	12834x	if (cov_type == "us") {
161	9400x	result = get_unstructured<T>(theta, n_visits);
162	3434x	} else if (cov_type == "toep") {
163	206x	result = get_toeplitz<T>(theta, n_visits);
164	3228x	} else if (cov_type == "toeph") {
165	212x	result = get_toeplitz_heterogeneous<T>(theta, n_visits);
166	3016x	} else if (cov_type == "ar1") {
167	1800x	result = get_auto_regressive<T>(theta, n_visits);
168	1216x	} else if (cov_type == "ar1h") {
169	212x	result = get_auto_regressive_heterogeneous<T>(theta, n_visits);
170	1004x	} else if (cov_type == "ad") {
171	156x	result = get_ante_dependence<T>(theta, n_visits);
172	848x	} else if (cov_type == "adh") {
173	236x	result = get_ante_dependence_heterogeneous<T>(theta, n_visits);
174	612x	} else if (cov_type == "cs") {
175	348x	result = get_compound_symmetry<T>(theta, n_visits);
176	264x	} else if (cov_type == "csh") {
177	260x	result = get_compound_symmetry_heterogeneous<T>(theta, n_visits);
178		} else {
179	2x	Rf_error("%s", ("Unknown covariance type '" + cov_type + "'.").c_str());
180		}
181
182	12830x	return result;
183		}
184
185		// Creates a new spatial covariance cholesky.
186		template <class T>
187	33228x	matrix<T> get_spatial_covariance_lower_chol(const vector<T>& theta, const matrix<T>& distance, std::string cov_type) {
188	33228x	matrix<T> result;
189	33228x	if (cov_type == "sp_exp") {
190	23736x	result = get_spatial_exponential<T>(theta, distance);
191	9492x	} else if (cov_type == "sp_gau") {
192	9492x	result = get_spatial_gaussian<T>(theta, distance);
193		} else {
194	!	Rf_error("%s", ("Unknown spatial covariance type '" + cov_type + "'.").c_str());
195		}
196	33228x	return result;
197		}
198
199		#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	456x	chol(int dim, string cov): dim_cov_mat(dim), cov_type(cov) {};
16		template <class T>
17	1366x	vector<T> operator() (vector<T> &theta) {
18	1366x	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	229x	chol_jacobian(int dim, string cov): dim_cov_mat(dim), cov_type(cov), mychol(dim, cov) {};
28		template<class T>
29	229x	vector<T> operator() (vector<T> &theta) {
30	229x	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	227x	std::map<std::string, matrix<Type>> derivatives(int n_visits, std::string cov_type, vector<Type> theta) {
39	227x	std::map<std::string, matrix<Type>> ret;
40	227x	chol chol_obj(n_visits, cov_type);
41	227x	chol_jacobian chol_jac_obj(n_visits, cov_type);
42	227x	matrix<Type> l = chol_obj(theta).matrix();
43	227x	l.resize(n_visits, n_visits);
44	227x	vector<Type> chol_d1_vec = autodiff::jacobian(chol_obj, theta).vec(); // chol_d1_vec is (dim * dim * l_theta)
45	227x	vector<Type> chol_d2_vec = autodiff::jacobian(chol_jac_obj, theta).vec(); // chol_d2_vec is (dim * dim * l_theta * l_theta)
46	227x	matrix<Type> ret_d1 = matrix<Type>(n_visits * theta.size(), n_visits);
47	227x	matrix<Type> ret_d2 = matrix<Type>(n_visits * theta.size() * theta.size(), n_visits);
48	227x	int n_visits_sq = n_visits * n_visits;
49	1533x	for (int i = 0; i < theta.size(); i++) {
50	1306x	matrix<Type> ld1 = chol_d1_vec.segment(i * n_visits_sq, n_visits_sq).matrix();
51	1306x	ld1.resize(n_visits, n_visits);
52	1306x	matrix<Type> ld1_lt = ld1 * l.transpose();
53	1306x	auto sigma_d1_i = ld1_lt + ld1_lt.transpose();
54	1306x	ret_d1.block(i * n_visits, 0, n_visits, n_visits) = sigma_d1_i;
55	11910x	for (int j = 0; j < theta.size(); j++) {
56	10604x	matrix<Type> ld2 = chol_d2_vec.segment( (j * theta.size() + i) * n_visits_sq, n_visits_sq).matrix();
57	10604x	matrix<Type> ld1_j = chol_d1_vec.segment(j * n_visits_sq, n_visits_sq).matrix();
58	10604x	ld2.resize(n_visits, n_visits);
59	10604x	ld1_j.resize(n_visits, n_visits);
60	10604x	auto ld2_lt = ld2 * l.transpose();
61	10604x	auto ld1_ld1j = ld1 * ld1_j.transpose();
62	10604x	auto sigma_d2_ij = ld2_lt + ld2_lt.transpose() + ld1_ld1j + ld1_ld1j.transpose();
63	10604x	ret_d2.block((i * theta.size() + j) * n_visits, 0, n_visits, n_visits) = sigma_d2_ij;
64		}
65		}
66	454x	ret["derivative1"] = ret_d1;
67	227x	ret["derivative2"] = ret_d2;
68	454x	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	241x	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	227x	derivatives_nonspatial(vector<Type> theta, int n_visits, std::string cov_type): lower_chol_nonspatial<Type>(theta, n_visits, cov_type) {
96	227x	std::map<std::string, tmbutils::matrix<Type>> allret = derivatives<Type>(this->n_visits, this->cov_type, this->theta);
97	454x	matrix<Type> sigma_d1 = allret["derivative1"];
98	227x	matrix<Type> sigma_d2 = allret["derivative2"];
99	227x	this->sigmad1_cache[this->full_visit] = sigma_d1;
100	227x	this->sigmad2_cache[this->full_visit] = sigma_d2;
101		}
102		// Cache and return the first order derivatives using select matrix.
103	11345x	matrix<Type> get_sigma_derivative1(std::vector<int> visits, matrix<Type> dist) override {
104	11345x	auto target = this->sigmad1_cache.find(visits);
105	11345x	if (target != this->sigmad1_cache.end()) {
106	8918x	return target->second;
107		} else {
108	2427x	int n_visits_i = visits.size();
109	2427x	matrix<Type> ret = matrix<Type>(this->n_theta * n_visits_i, n_visits_i);
110	17172x	for (int i = 0; i < this->n_theta; i++) {
111	14745x	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	2427x	this->sigmad1_cache[visits] = ret;
114	2427x	return ret;
115		}
116		}
117		// Cache and return the second order derivatives using select matrix.
118	8731x	matrix<Type> get_sigma_derivative2(std::vector<int> visits, matrix<Type> dist) override {
119	8731x	auto target = this->sigmad2_cache.find(visits);
120	8731x	if (target != this->sigmad2_cache.end()) {
121	7970x	return target->second;
122		} else {
123	761x	int n_visits_i = visits.size();
124	761x	int theta_sq = this->n_theta * this->n_theta;
125	761x	matrix<Type> ret = matrix<Type>(theta_sq * n_visits_i, n_visits_i);
126	23817x	for (int i = 0; i < theta_sq; i++) {
127	23056x	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	761x	this->sigmad2_cache[visits] = ret;
130	761x	return ret;
131		}
132		}
133		// Cache and return the lower cholesky factor of inverse of sigma using select matrix.
134	7486x	matrix<Type> get_inverse_chol(std::vector<int> visits, matrix<Type> dist) override {
135	7486x	auto target = this->inverse_chol_cache.find(visits);
136	7486x	if (target != this->inverse_chol_cache.end()) {
137	6916x	return target->second;
138		} else {
139	570x	matrix<Type> sigmainv = this->get_sigma_inverse(visits, dist);
140	570x	Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > sigma_inv_chol(sigmainv);
141	570x	matrix<Type> Li = sigma_inv_chol.matrixL();
142	570x	this->inverse_chol_cache[visits] = Li;
143	570x	return Li;
144		}
145		}
146		// Cache and return the first order derivatives of inverse of sigma using select matrix.
147	32508x	matrix<Type> get_inverse_derivative(std::vector<int> visits, matrix<Type> dist) override {
148	32508x	auto target = this->sigma_inverse_d1_cache.find(visits);
149	32508x	if (target != this->sigma_inverse_d1_cache.end()) {
150	29894x	return target->second;
151		} else {
152	2614x	auto sigma_d1 = this->get_sigma_derivative1(visits, dist);
153	2614x	matrix<Type> sigma_inv_d1(sigma_d1.rows(), sigma_d1.cols());
154	2614x	int n_visits_i = visits.size();
155	2614x	auto sigma_inv = this->get_sigma_inverse(visits, dist);
156	18493x	for (int r = 0; r < this->n_theta; r++) {
157	15879x	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	2614x	this->sigma_inverse_d1_cache[visits] = sigma_inv_d1;
160	2614x	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	28x	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	28x	this->logrho = log(this->rho);
182		}
183		// Obtain first order derivatives
184	2562x	matrix<Type> get_sigma_derivative1(std::vector<int> visits, matrix<Type> dist) override {
185	2562x	matrix<Type> ret(2 * dist.rows(), dist.cols());
186		// partial sigma / partial theta_1 = sigma.
187	2562x	auto sigma = this->get_sigma(visits, dist);
188	2562x	ret.block(0, 0, dist.rows(), dist.cols()) = sigma;
189	2562x	ret.block(dist.rows(), 0, dist.rows(), dist.cols()) = sigma.array() * dist.array() * (1 - this->rho);
190	5124x	return ret;
191		}
192		// Obtain second order derivatives.
193	986x	matrix<Type> get_sigma_derivative2(std::vector<int> visits, matrix<Type> dist) override {
194	986x	matrix<Type> ret(4 * dist.rows(), dist.cols());
195	986x	auto sigma = this->get_sigma(visits, dist);
196	986x	ret.block(0, 0, dist.rows(), dist.cols()) = sigma;
197	986x	Type rho_r = 1 - this->rho;
198	986x	auto dtheta1dtheta2 = sigma.array() * dist.array() * rho_r;
199	986x	ret.block(dist.rows(), 0, dist.rows(), dist.cols()) = dtheta1dtheta2;
200	986x	ret.block(dist.rows() * 2, 0, dist.rows(), dist.cols()) = dtheta1dtheta2;
201	986x	matrix<Type> dtheta2s = dtheta1dtheta2 * (dist.array() * rho_r - this->rho);
202	986x	ret.block(dist.rows() * 3, 0, dist.rows(), dist.cols()) = dtheta2s;
203	1972x	return ret;
204		}
205		// Obtain the lower cholesky factor of inverse of sigma using select matrix.
206	394x	matrix<Type> get_inverse_chol(std::vector<int> visits, matrix<Type> dist) override {
207	394x	auto sigmainv = this->get_sigma_inverse(visits, dist);
208	394x	Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > sigma_inv_chol(sigmainv);
209	394x	matrix<Type> Li = sigma_inv_chol.matrixL();
210	788x	return Li;
211		}
212		// Obtain first order derivatives for inverse of sigma.
213	1970x	matrix<Type> get_inverse_derivative(std::vector<int> visits, matrix<Type> dist) override {
214	1970x	matrix<Type> sigma_inv_d1 = matrix<Type>::Zero(2 * dist.rows(), dist.cols());
215	1970x	auto sigma_inv = this->get_sigma_inverse(visits, dist);
216	1970x	auto sigma_d1 = this->get_sigma_derivative1(visits, dist);
217	5910x	for (int r = 0; r < 2; r++) {
218	3940x	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	3940x	return sigma_inv_d1;
221		}
222		};
223
224
225		// derivatives_sp_gau struct is created to obtain the exact derivatives of spatial Gaussian
226		// covariance structure, and its inverse.
227		// Inherit from sp_exp because of similar parameterization; only the sigma
228		// derivatives differ (square the distance). Inverse / inverse-derivative logic
229		// is reused from the base via the virtual get_sigma_derivative1 dispatch.
230		template <class Type>
231		struct derivatives_sp_gau: public derivatives_sp_exp<Type> {
232		derivatives_sp_gau() {
233		// This default constructor is needed because the use of `[]` in maps.
234		}
235	2x	derivatives_sp_gau(vector<Type> theta, std::string cov_type): derivatives_sp_exp<Type>(theta, cov_type) {
236		}
237		// Obtain first order derivatives
238	788x	matrix<Type> get_sigma_derivative1(std::vector<int> visits, matrix<Type> dist) override {
239	788x	matrix<Type> ret(2 * dist.rows(), dist.cols());
240		// partial sigma / partial theta_1 = sigma.
241	788x	auto sigma = this->get_sigma(visits, dist);
242	788x	ret.block(0, 0, dist.rows(), dist.cols()) = sigma;
243	788x	ret.block(dist.rows(), 0, dist.rows(), dist.cols()) = sigma.array() * dist.array().square() * (1 - this->rho);
244	1576x	return ret;
245		}
246		// Obtain second order derivatives.
247	394x	matrix<Type> get_sigma_derivative2(std::vector<int> visits, matrix<Type> dist) override {
248	394x	matrix<Type> ret(4 * dist.rows(), dist.cols());
249	394x	auto sigma = this->get_sigma(visits, dist);
250	394x	ret.block(0, 0, dist.rows(), dist.cols()) = sigma;
251	394x	Type rho_r = 1 - this->rho;
252	394x	auto dtheta1dtheta2 = sigma.array() * dist.array().square() * rho_r;
253	394x	ret.block(dist.rows(), 0, dist.rows(), dist.cols()) = dtheta1dtheta2;
254	394x	ret.block(dist.rows() * 2, 0, dist.rows(), dist.cols()) = dtheta1dtheta2;
255	394x	matrix<Type> dtheta2s = dtheta1dtheta2 * (dist.array().square() * rho_r - this->rho);
256	394x	ret.block(dist.rows() * 3, 0, dist.rows(), dist.cols()) = dtheta2s;
257	788x	return ret;
258		}
259		};
260
261		// Cache of per-group derivatives objects, with cov_type-aware dispatch for
262		// spatial structures (sp_exp / sp_gau). Mirrors `cache_obj` but selects the
263		// concrete spatial subclass at runtime so analytic derivatives match cov_type.
264		template <class Type>
265		struct derivatives_cache {
266		std::map<int, std::shared_ptr<derivatives_base<Type>>> cache;
267		int n_groups;
268		bool is_spatial;
269		int n_visits;
270	219x	derivatives_cache(vector<Type> 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) {
271	219x	int theta_one_group_size = theta.size() / n_groups;
272	458x	for (int r = 0; r < n_groups; r++) {
273	239x	vector<Type> theta_r = theta.segment(r * theta_one_group_size, theta_one_group_size);
274	239x	if (is_spatial) {
275	13x	if (cov_type == "sp_exp") {
276	11x	this->cache[r] = std::make_shared<derivatives_sp_exp<Type>>(theta_r, cov_type);
277	2x	} else if (cov_type == "sp_gau") {
278	2x	this->cache[r] = std::make_shared<derivatives_sp_gau<Type>>(theta_r, cov_type);
279		} else {
280	!	Rf_error("%s", ("Unknown spatial covariance type '" + cov_type + "'.").c_str());
281		}
282		} else {
283	226x	this->cache[r] = std::make_shared<derivatives_nonspatial<Type>>(theta_r, n_visits, cov_type);
284		}
285		}
286		}
287		};
288		#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	228116x	T1 subset_matrix(T1 input, T2 index1, T2 index2) {
16		#if EIGEN_VERSION_AT_LEAST(3,4,0)
17	228116x	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	228116x	return ret;
27		}
28
29		template <typename T1, typename T2>
30	242229x	T1 subset_matrix(T1 input, T2 index1) {
31		#if EIGEN_VERSION_AT_LEAST(3,4,0)
32	242229x	T1 ret = input(index1, Eigen::all);
33		#else
34		T1 ret(index1.size(), input.cols());
35		for (decltype(index1.size()) i = 0; i < index1.size(); i++) {
36		for (int j = 0; j < input.cols(); j++) {
37		ret(i, j) = input(index1[i], j);
38		}
39		}
40		#endif
41	242229x	return ret;
42		}
43
44
45		// Conversion from Rcpp vector/matrix to eigen vector/matrix
46		template <typename T1, typename T2>
47	488386x	T1 as_vector(T2 input) {
48	488386x	T1 ret(input.size());
49	1470784x	for (int i = 0; i < input.size(); i++) {
50	982398x	ret(i) = input(i);
51		}
52	488386x	return ret;
53		}
54
55		template <typename T1, typename T2>
56	211290x	T1 as_matrix(T2 input) {
57	211290x	T1 ret(input.rows(), input.cols());
58	3078372x	for (int i = 0; i < input.rows(); i++) {
59	32869906x	for (int j = 0; j < input.cols(); j++) {
60	30002824x	ret(i,j) = input(i,j);
61		}
62		}
63	211290x	return ret;
64		}
65
66		template <typename T>
67	726686x	T segment(T input, int start, int n) {
68	726686x	T ret(n);
69	3630604x	for (int i = 0, j = start; i < n; i++, j++) {
70	2903918x	ret(i) = input(j);
71		}
72	726686x	return ret;
73		}
74
75		// Calculate tcrossprod(lower_chol) = lower_chol * t(lower_chol).
76		// If complete, then adds the upper triangular part to the result as well.
77		// By default only the lower triangular part is populated, as this should be
78		// sufficient for downstream use of the result in most cases.
79		template <class Type>
80	694030x	matrix<Type> tcrossprod(const matrix<Type>& lower_chol, bool complete = false) {
81	694030x	int n = lower_chol.rows();
82	694030x	matrix<Type> result = matrix<Type>::Zero(n, n);
83	694030x	result.template selfadjointView<Eigen::Lower>().rankUpdate(lower_chol);
84	694030x	if (complete) {
85	15986x	result.template triangularView<Eigen::Upper>() = result.transpose();
86		}
87	694030x	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	1356088x	matrix<Type> crossprod(const matrix<Type>& x) {
96	1356088x	int n = x.cols();
97	1356088x	matrix<Type> result = matrix<Type>::Zero(n, n);
98	1356088x	result.template selfadjointView<Eigen::Lower>().rankUpdate(x.transpose());
99	1356088x	return result;
100		}
101
102		// Mapping from real values to correlation parameters in (-1, 1).
103		template <class T>
104	2846x	vector<T> map_to_cor(const vector<T>& theta) {
105	5692x	return theta / sqrt(T(1.0) + theta * theta);
106		}
107
108		// Mapping from real values to correlation parameters in (-1/(m-1), 1) for compound symmetry.
109		template <class T>
110	616x	vector<T> map_to_cs_cor(const vector<T>& theta, int n_visits) {
111	616x	T a = T(1.0) / (T(n_visits) - T(1.0));
112	616x	return invlogit(theta) * (T(1.0) + a) - a;
113		}
114
115		// Generic correlation function class containing and initializing correlation
116		// values from variance parameters theta.
117		template <class T>
118		struct generic_corr_fun {
119		const vector<T> corr_values;
120
121	2842x	generic_corr_fun(const vector<T>& theta) :
122	2842x	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	3450x	matrix<T> get_corr_mat_chol(int n_visits, const F<T>& corr_fun) {
129	3450x	matrix<T> correlation(n_visits, n_visits);
130	3450x	correlation.setIdentity();
131	17166x	for(int i = 0; i < n_visits; i++) {
132	34180x	for(int j = 0; j < i; j++){
133	20464x	correlation(i, j) = corr_fun(i, j);
134		}
135		}
136	3450x	Eigen::LLT<Eigen::Matrix<T,Eigen::Dynamic,Eigen::Dynamic> > correlation_chol(correlation);
137	3450x	matrix<T> L = correlation_chol.matrixL();
138	6900x	return L;
139		}
140
141		// Heterogeneous covariance matrix calculation given vector of standard deviations (sd_values)
142		// and a correlation function (corr_fun).
143		template <class T, template<class> class F>
144	929x	matrix<T> get_heterogeneous_cov(const vector<T>& sd_values, const F<T>& corr_fun) {
145	929x	matrix<T> correlation_chol = get_corr_mat_chol(sd_values.size(), corr_fun);
146	929x	Eigen::DiagonalMatrix<T,Eigen::Dynamic,Eigen::Dynamic> D = sd_values.matrix().asDiagonal();
147	929x	matrix<T> result = D * correlation_chol;
148	1858x	return result;
149		}
150
151		// Obtain the Euclidean distance
152		template <class T>
153	25618x	matrix<T> euclidean(const matrix<T>& coordinates) {
154	25618x	matrix<T> result(coordinates.rows(), coordinates.rows());
155	95962x	for (int i = 0; i < coordinates.rows(); i++) {
156	70344x	result(i, i) = 0;
157	142260x	for (int j = 0; j < i; j ++) {
158	71916x	vector<T> diff = coordinates.row(i) - coordinates.row(j);
159	71916x	T d = sqrt((diff * diff).sum());
160	71916x	result(i, j) = d;
161	71916x	result(j, i) = d;
162		}
163		}
164	25618x	return result;
165		}
166
167		// Element wise power function of a matrix
168		template <class T>
169	792x	Eigen::Matrix<T, -1, -1> cpow(const Eigen::Matrix<T, -1, -1> & input, double p) {
170	792x	Eigen::Matrix<T, -1, -1> ret = Eigen::Matrix<T, -1, -1>(input.rows(), input.cols());
171	2954x	for (int i = 0; i < ret.rows(); i ++) {
172	4332x	for (int j = 0; j < ret.cols(); j++) {
173	2170x	ret(i, j) = std::pow(input(i, j), p);
174		}
175		}
176	792x	return ret;
177		}
178
179		// Calculate the square root of the pseudo inverse of a matrix
180		// adapted from the method for calculating the pseudo-Inverse as recommended by the Eigen developers
181		template<typename T>
182	790x	matrix<T> pseudoInverseSqrt(const matrix<T> &input, double epsilon = std::numeric_limits<double>::epsilon()) {
183	790x	Eigen::Matrix<T, -1, -1> eigen_mat = as_matrix<Eigen::Matrix<T, -1, -1>, matrix<T>>(input);
184	790x	Eigen::JacobiSVD< Eigen::Matrix<T, -1, -1> > svd(eigen_mat ,Eigen::ComputeFullU \| Eigen::ComputeFullV);
185	790x	double tolerance = epsilon * std::max(input.cols(), input.rows()) *svd.singularValues().array().abs()(0);
186	790x	auto singular_vals = Matrix<T,-1,-1>((svd.singularValues().array() > tolerance).select(svd.singularValues().array().inverse(), 0).matrix());
187	790x	Eigen::Matrix<T, -1, -1> ret_eigen = svd.matrixV() * cpow(singular_vals, 0.5).asDiagonal() * svd.matrixU().adjoint();
188	1580x	return as_matrix<matrix<T>, Eigen::Matrix<T, -1, -1>>(ret_eigen);
189		}
190
191		#endif

1		#include <RcppEigen.h>
2		#include "utils.h"
3
4		using namespace Rcpp;
5
6		#ifdef RCPP_USE_GLOBAL_ROSTREAM
7		Rcpp::Rostream<true>& Rcpp::Rcout = Rcpp::Rcpp_cout_get();
8		Rcpp::Rostream<false>& Rcpp::Rcerr = Rcpp::Rcpp_cerr_get();
9		#endif
10
11		List get_pqr(List mmrm_fit, NumericVector theta);
12	53x	RcppExport SEXP _mmrm_get_pqr(SEXP mmrm_fit_SEXP, SEXP theta_SEXP) {
13	53x	BEGIN_RCPP
14	53x	Rcpp::RObject rcpp_result_gen;
15	53x	Rcpp::RNGScope rcpp_rngScope_gen;
16	53x	Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
17	53x	Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
18	53x	rcpp_result_gen = Rcpp::wrap(get_pqr(mmrm_fit, theta));
19	53x	return rcpp_result_gen;
20	53x	END_RCPP
21		}
22
23		List get_jacobian(List mmrm_fit, NumericVector theta, NumericMatrix beta_vcov);
24	126x	RcppExport SEXP _mmrm_get_jacobian(SEXP mmrm_fit_SEXP, SEXP theta_SEXP, SEXP beta_vcov_SEXP) {
25	126x	BEGIN_RCPP
26	126x	Rcpp::RObject rcpp_result_gen;
27	126x	Rcpp::RNGScope rcpp_rngScope_gen;
28	126x	Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
29	126x	Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
30	126x	Rcpp::traits::input_parameter< NumericMatrix >::type beta_vcov(beta_vcov_SEXP);
31	126x	rcpp_result_gen = Rcpp::wrap(get_jacobian(mmrm_fit, theta, beta_vcov));
32	126x	return rcpp_result_gen;
33	126x	END_RCPP
34		}
35
36		List get_empirical(List mmrm_fit, NumericVector theta, NumericVector beta, NumericMatrix beta_vcov, std::string type);
37	40x	RcppExport SEXP _mmrm_get_empirical(SEXP mmrm_fit_SEXP, SEXP theta_SEXP, SEXP beta_SEXP, SEXP beta_vcov_SEXP, SEXP type_SEXP) {
38	40x	BEGIN_RCPP
39	40x	Rcpp::RObject rcpp_result_gen;
40	40x	Rcpp::RNGScope rcpp_rngScope_gen;
41	40x	Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
42	40x	Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
43	40x	Rcpp::traits::input_parameter< NumericVector >::type beta(beta_SEXP);
44	40x	Rcpp::traits::input_parameter< NumericMatrix >::type beta_vcov(beta_vcov_SEXP);
45	40x	Rcpp::traits::input_parameter< std::string >::type type(type_SEXP);
46	40x	rcpp_result_gen = Rcpp::wrap(get_empirical(mmrm_fit, theta, beta, beta_vcov, type));
47	40x	return rcpp_result_gen;
48	40x	END_RCPP
49		}
50
51		List predict(List mmrm_fit, NumericVector theta, NumericVector beta, NumericMatrix beta_vcov);
52	1704x	RcppExport SEXP _mmrm_predict(SEXP mmrm_fit_SEXP, SEXP theta_SEXP, SEXP beta_SEXP, SEXP beta_vcov_SEXP) {
53	1704x	BEGIN_RCPP
54	1704x	Rcpp::RObject rcpp_result_gen;
55	1704x	Rcpp::RNGScope rcpp_rngScope_gen;
56	1704x	Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
57	1704x	Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
58	1704x	Rcpp::traits::input_parameter< NumericVector >::type beta(beta_SEXP);
59	1704x	Rcpp::traits::input_parameter< NumericMatrix >::type beta_vcov(beta_vcov_SEXP);
60	1704x	rcpp_result_gen = Rcpp::wrap(predict(mmrm_fit, theta, beta, beta_vcov));
61	1704x	return rcpp_result_gen;
62	1704x	END_RCPP
63		}
64
65
66		RcppExport SEXP run_testthat_tests(SEXP);
67
68		static const R_CallMethodDef CallEntries[] = {
69		{"_mmrm_get_pqr", (DL_FUNC) &_mmrm_get_pqr, 2},
70		{"_mmrm_get_jacobian", (DL_FUNC) &_mmrm_get_jacobian, 3},
71		{"_mmrm_get_empirical", (DL_FUNC) &_mmrm_get_empirical, 5},
72		{"_mmrm_predict", (DL_FUNC) &_mmrm_predict, 4},
73		{"run_testthat_tests", (DL_FUNC) &run_testthat_tests, 1},
74		TMB_CALLDEFS,
75		{NULL, NULL, 0}
76		};
77
78	4x	RcppExport void R_init_mmrm(DllInfo *dll) {
79	4x	R_registerRoutines(dll, NULL, CallEntries, NULL, NULL);
80	4x	R_useDynamicSymbols(dll, FALSE);
81		#ifdef TMB_CCALLABLES
82	4x	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	6956x	Type objective_function<Type>::operator() ()
29		{
30		// Read data from R.
31	6956x	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	6956x	DATA_MATRIX(coordinates); // Coordinates matrix.
35	6956x	DATA_INTEGER(n_visits); // Number of visits, which is the dimension of the covariance matrix.
36	6956x	DATA_INTEGER(n_subjects); // Number of subjects.
37	6956x	DATA_IVECTOR(subject_zero_inds); // Starting indices for each subject (0-based) (length n_subjects).
38	6956x	DATA_IVECTOR(subject_n_visits); // Number of observed visits for each subject (length n_subjects).
39	6956x	DATA_STRING(cov_type); // Covariance type name.
40	6956x	DATA_INTEGER(is_spatial_int); // Spatial covariance (1)? Otherwise non-spatial covariance.
41	6956x	DATA_INTEGER(reml); // REML (1)? Otherwise ML (0).
42	6956x	DATA_FACTOR(subject_groups); // subject groups vector(0-based) (length n_subjects).
43	6956x	DATA_INTEGER(n_groups); // number of total groups.
44		// Read parameters from R.
45	6956x	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	6956x	matrix<Type> XtWX = matrix<Type>::Zero(x_matrix.cols(), x_matrix.cols());
49		// X^T W Y will be calculated incrementally into here.
50	6956x	matrix<Type> XtWY = matrix<Type>::Zero(x_matrix.cols(), 1);
51		// W^T/2 X will be saved into here.
52	6956x	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	6956x	vector<Type> y_vec_tilde = vector<Type>::Zero(y_vector.rows());
55		// Sum of the log determinant will be incrementally calculated here.
56	6956x	Type sum_log_det = 0.0;
57
58		// Convert is_spatial_int to bool.
59	6956x	bool is_spatial = (is_spatial_int == 1);
60		// Diagonal of weighted covariance
61	6956x	vector<Type> diag_cov_inv_sqrt(x_matrix.rows());
62		// Cholesky group object
63	6956x	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	1363038x	for (int i = 0; i < n_subjects; i++) {
66		// Start index and number of visits for this subject.
67	1356086x	int start_i = subject_zero_inds(i);
68	1356086x	int n_visits_i = subject_n_visits(i);
69	1356086x	std::vector<int> visit_i(n_visits_i);
70	1356086x	matrix<Type> dist_i(n_visits_i, n_visits_i);
71	1356086x	if (!is_spatial) {
72	4883886x	for (int j = 0; j < n_visits_i; j++) {
73	3573504x	visit_i[j] = int(asDouble(coordinates(start_i + j, 0)));
74		}
75		} else {
76	45704x	dist_i = euclidean(matrix<Type>(coordinates.block(start_i, 0, n_visits_i, coordinates.cols())));
77		}
78		// Obtain Cholesky factor Li.
79	1356086x	matrix<Type> Li = chols_group.cache[subject_groups[i]]->get_chol(visit_i, dist_i);
80		// Calculate weighted Cholesky factor for this subject.
81	1356086x	Eigen::DiagonalMatrix<Type,Eigen::Dynamic,Eigen::Dynamic> Gi_inv_sqrt = weights_vector.segment(start_i, n_visits_i).cwiseInverse().sqrt().matrix().asDiagonal();
82	1356086x	Li = Gi_inv_sqrt * Li;
83		// Calculate scaled design matrix and response vector for this subject.
84	1356086x	matrix<Type> Xi = x_matrix.block(start_i, 0, n_visits_i, x_matrix.cols());
85	1356086x	matrix<Type> XiTilde = Li.template triangularView<Eigen::Lower>().solve(Xi);
86	1356086x	matrix<Type> Yi = y_vector.segment(start_i, n_visits_i).matrix();
87	1356086x	matrix<Type> YiTilde = Li.template triangularView<Eigen::Lower>().solve(Yi);
88
89		// Increment quantities.
90	1356086x	matrix<Type> XiTildeCrossprod = crossprod(XiTilde);
91	1356086x	XtWX += XiTildeCrossprod.template triangularView<Eigen::Lower>();
92	1356086x	XtWY += XiTilde.transpose() * YiTilde;
93	1356086x	vector<Type> LiDiag = Li.diagonal();
94	1356086x	sum_log_det += sum(log(LiDiag));
95		// Cache the reciprocal of square root of diagonal of covariance
96	1356086x	diag_cov_inv_sqrt.segment(start_i, n_visits_i) = vector<Type>(tcrossprod(Li).diagonal()).rsqrt();
97		// Save stuff.
98	1356086x	x_mat_tilde.block(start_i, 0, n_visits_i, x_matrix.cols()) = XiTilde;
99	1356086x	y_vec_tilde.segment(start_i, n_visits_i) = YiTilde.col(0);
100		}
101
102		// Solve for beta.
103	6952x	Eigen::LDLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > XtWX_decomposition(XtWX);
104	6952x	matrix<Type> beta_mat = XtWX_decomposition.solve(XtWY);
105	6952x	vector<Type> beta = beta_mat.col(0);
106
107		// Define scaled residuals.
108	6952x	vector<Type> x_mat_tilde_beta = x_mat_tilde * beta;
109	6952x	vector<Type> epsilonTilde = y_vec_tilde - x_mat_tilde_beta;
110
111		// Calculate negative log-likelihood.
112	608x	Type neg_log_lik;
113
114		// Always extract the D vector since we want to report this below.
115	6952x	vector<Type> XtWX_D = XtWX_decomposition.vectorD();
116
117	6952x	if (reml == 1) {
118		// Use restricted maximum likelihood.
119	6594x	Type XtWX_log_det = XtWX_D.log().sum();
120	6594x	neg_log_lik = (x_matrix.rows() - x_matrix.cols()) / 2.0 * log(2.0 * M_PI) +
121	6594x	sum_log_det +
122	13188x	XtWX_log_det / 2.0 +
123	7142x	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	358x	neg_log_lik = x_matrix.rows() / 2.0 * log(2.0 * M_PI) +
127	60x	sum_log_det +
128	418x	0.5 * (epsilonTilde * epsilonTilde).sum();
129		}
130
131		// Report quantities to R.
132	6344x	REPORT(beta);
133
134		// We already compute the inverse of XtWX here because we already did the
135		// matrix decomposition above.
136	6952x	matrix<Type> Identity(XtWX.rows(), XtWX.cols());
137	6952x	Identity.setIdentity();
138	6952x	matrix<Type> beta_vcov = XtWX_decomposition.solve(Identity);
139	6344x	REPORT(beta_vcov);
140
141		// Also return the decomposition components L and D.
142	6952x	matrix<Type> XtWX_L(XtWX.rows(), XtWX.cols());
143	6952x	XtWX_L = XtWX_decomposition.matrixL();
144	6344x	REPORT(XtWX_L);
145	6344x	REPORT(XtWX_D);
146
147		// normalized residual
148	6344x	REPORT(epsilonTilde);
149		// inverse square root of diagonal of covariance
150	6344x	REPORT(diag_cov_inv_sqrt);
151	6952x	matrix<Type> covariance_lower_chol = chols_group.get_default_chol();
152	6344x	REPORT(covariance_lower_chol);
153
154	6952x	return neg_log_lik;
155		}

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