mmrm coverage - 97.70%

Files
Source

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

  model_formula <- h_drop_covariance_terms(formula)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


#' Calculation of Degrees of Freedom for Multi-Dimensional Contrast
#'
#' @description `r lifecycle::badge("stable")`
#' Calculates the estimate, standard error, degrees of freedom,
#' t statistic and p-value for 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) {
    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.
#'   Must be `NULL` or one or more of `c("subject_var", "visit_var", "group_var", "response_var")`.
#' @param full (`flag`)\cr indicator whether to return full model frame (deprecated).
#' @param na.action (`string`)\cr na action.
#' @importFrom stats model.frame
#' @exportS3Method
#'
#' @details
#' `include` argument controls the variables the returned model frame will include.
#' Possible options are "response_var", "subject_var", "visit_var" and "group_var", representing the
#' response variable, subject variable, visit variable or group variable.
#' `character` values in new data will always be factorized according to the data in the fit
#' to avoid mismatched in levels or issues in `model.matrix`.
#'
#' @examples
#' # Model frame:
#' model.frame(object)
#' model.frame(object, include = "subject_var")
model.frame.mmrm_tmb <- function(
  formula,
  data,
  include = c("subject_var", "visit_var", "group_var", "response_var"),
  full,
  na.action = "na.omit", # nolint
  ...
) {
  # nolint
  # Construct updated formula and data arguments.
  lst_formula_and_data <-
    h_construct_model_frame_inputs(
      formula = formula,
      data = data,
      include = include,
      full = full
    )
  # Only if include is default (full) and also data is missing, and also na.action is na.omit we will
  # use the model frame from the tmb_data.
  include_choice <- c("subject_var", "visit_var", "group_var", "response_var")
  if (
    missing(data) &&
      setequal(include, include_choice) &&
      identical(h_get_na_action(na.action), stats::na.omit)
  ) {
    ret <- formula$tmb_data$full_frame
    # Remove weights column.
    ret[, "(weights)"] <- NULL
    ret
  } else {
    # Construct data frame to return to users.
    ret <-
      stats::model.frame(
        formula = lst_formula_and_data$formula,
        data = h_get_na_action(na.action)(lst_formula_and_data$data),
        na.action = na.action,
        xlev = component(formula, "xlev")
      )
  }
  ret
}


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

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

  drop_response <- !"response_var" %in% include
  add_vars <- unlist(formula$formula_parts[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)
  assert_names(colnames(data), must.include = all_vars)
  data <- data[, all_vars, drop = FALSE]

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

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

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

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


#' @describeIn mmrm_tmb_methods obtains the attained log likelihood value.
#'   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
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  h_test_1d(object, contrast, df$m)
}

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

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

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

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


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


#' Conduct type II/III hypothesis testing on the MMRM fit results.
#'
#' @param mod (`mmrm`)\cr the fitted MMRM.
#' @param test.statistic (`string`)\cr either `"F` or `"Chisq"`, indicating the
#'   kind of test to perform.
#' @param ... arguments passed from other methods.
#' @inheritParams h_get_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.
Anova.mmrm <- function(
  mod, # nolint
  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"))
  contrasts <-
    lapply(vars, h_get_contrast, object = mod, type = type, tol = tol)

  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) <- vars
  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
  )
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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


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

  colnames_basic <- colnames(data_basic)

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

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


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

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

  cols <- lapply(datasets, colnames)

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

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

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

  TRUE
}


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


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


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

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

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

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

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

  TRUE
}


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

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

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

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

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

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

  cov_struct_nesting <- h_check_cov_struct_nesting(model_basic, model_augmented)

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

  TRUE
}


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

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

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

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

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

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

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

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

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



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

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

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

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


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

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

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

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


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

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

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

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

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

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

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

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

  fit
}




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

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

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

  design_matrix <- component(object, "x_matrix")

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

  assign_coef_term <- factor(assign_coef_term, labels = term_labels)

  coef_identity_matrix <- diag(length(assign_coef_term))

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

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

  df_md_results_data_frame <- do.call(rbind, df_md_results_per_term)

  df_md_results_data_frame
}

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

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

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

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

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

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

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

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

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

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

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

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

  validate_cov_struct(x)
}

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

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

  as.cov_struct(formula, warn_partial = FALSE)
}

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

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

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

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

  reportAssertions(checks)
  x
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

# Internal functions used for skipping tests or examples.

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

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

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

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

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

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

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

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

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

  coef_table
}

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

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

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

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

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

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


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



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

  assert_class(object, "mmrm")

  fits <- list(object, ...)

  if (length(fits) == 1L) {
    out <- h_anova_single_mmrm_model(object)

  } else {

    # Ensure all objects in ... are mmrm fits.
    lapply(fits[-1L], assert_class, classes = "mmrm", .var.name = "...")

    assert_flag(test)
    assert_flag(refit)

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

      fit_data <- lapply(fits, h_get_minimal_fit_data)

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

      needs_refit <- !nested_in_common_data

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

    log_likelihood_vec <- lapply(fits, logLik)

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

    if (test) {

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

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

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

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

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

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

  }

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

  out
}

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

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

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

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

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

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

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

  TRUE
}

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

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

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

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

#ifndef CHOL_CACHE_INCLUDED_
#define CHOL_CACHE_INCLUDED_

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

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


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

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

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

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

#endif

#ifndef COV_INCLUDED_
#define COV_INCLUDED_

#include "utils.h"

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

// Ante-dependence:

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

// Toeplitz:

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

// Autoregressive:

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

// Compound symmetry:

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

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

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

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

  return result;
}

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

#endif

#ifndef DERIVATIVE_INCLUDED_
#define DERIVATIVE_INCLUDED_

#include "chol_cache.h"

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

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

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

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

#endif

#include "derivatives.h"

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

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

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

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

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

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

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

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

// Obtain submatrix from index

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

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


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

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

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

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

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

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

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

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

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

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

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

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

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

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

#endif

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

using namespace Rcpp;

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

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

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

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

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


RcppExport SEXP run_testthat_tests(SEXP);

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

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

#include "derivatives.h"

using namespace Rcpp;
using std::string;

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

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

#include "derivatives.h"

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

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

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

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

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

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

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

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

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

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

  // Calculate negative log-likelihood.
  Type neg_log_lik;

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

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

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

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

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

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

  return neg_log_lik;
}

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

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

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

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

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

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

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

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

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

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

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

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

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

#endif

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1		#include "testthat-helpers.h"
2		#include "covariance.h"
3
4	5x	context("unstructured") {
5	15x	test_that("get_unstructured produces expected result") {
6	10x	vector<double> theta {{log(1.0), log(2.0), 3.0}};
7	5x	matrix<double> result = get_unstructured(theta, 2);
8	5x	matrix<double> expected(2, 2);
9	!	expected <<
10	5x	1.0, 0.0,
11	5x	6.0, 2.0;
12	5x	expect_equal_matrix(result, expected);
13		}
14		}
15
16	15x	context("ante_dependence") {
17	45x	test_that("corr_fun_ante_dependence works as expected") {
18	10x	vector<double> theta {{log(2.0), 1.0, 2.0}};
19	5x	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	45x	test_that("get_ante_dependence produces expected result") {
29	10x	vector<double> theta {{log(2.0), log(2.0), 1.0, 2.0}};
30	5x	matrix<double> result = get_ante_dependence(theta, 4);
31	5x	matrix<double> expected(4, 4);
32	!	expected <<
33	5x	2.0, 0.0, 0.0, 0.0,
34	5x	1.139353, 1.643738, 0.0, 0.0,
35	5x	0.805644, 1.162299, 1.414214, 0.0,
36	5x	0.720590, 1.039591, 1.264911, 0.894427;
37	5x	expect_equal_matrix(result, expected);
38		}
39
40	45x	test_that("get_ante_dependence_heterogeneous produces expected result") {

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

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

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

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

1		#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	10852x	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	32874x	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	10958x	this->theta = theta;
34	32874x	std::iota(std::begin(this->full_visit), std::end(this->full_visit), 0);
35	10958x	this->n_theta = theta.size();
36	10958x	this->chol_full = get_covariance_lower_chol(this->theta, this->n_visits, this->cov_type);
37	10954x	this->chols[full_visit] = this->chol_full;
38	10954x	this->sigma_full = tcrossprod(this->chol_full, true);
39		}
40	1190526x	matrix<Type> get_chol(std::vector<int> visits, matrix<Type> dist) {
41	1190526x	auto target = this->chols.find(visits);
42	1190526x	if (target != this->chols.end()) {
43	1100469x	return target->second;
44		} else {
45	90057x	matrix<Type> cov_i = this->get_sigma(visits, dist);
46	90057x	Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > cov_i_chol(cov_i);
47	90057x	matrix<Type> Li = cov_i_chol.matrixL();
48	90057x	this->chols[visits] = Li;
49	90057x	return Li;
50		}
51		}
52	616127x	matrix<Type> get_sigma(std::vector<int> visits, matrix<Type> dist) {
53	616127x	auto target = this->sigmas.find(visits);
54	616127x	if (target != this->sigmas.end()) {
55	489154x	return target->second;
56		} else {
57	126973x	matrix<Type> ret = subset_matrix<matrix<Type>, vector<int>>(sigma_full, visits, visits);
58	126973x	this->sigmas[visits] = ret;
59	126973x	return ret;
60		}
61		}
62	209142x	matrix<Type> get_sigma_inverse(std::vector<int> visits, matrix<Type> dist) {
63	209142x	auto target = this->sigmas_inv.find(visits);
64	209142x	if (target != this->sigmas_inv.end()) {
65	182476x	return target->second;
66		} else {
67	26666x	matrix<Type> ret = this->get_sigma(visits, dist).inverse();
68	26666x	this->sigmas_inv[visits] = ret;
69	26666x	return ret;
70		}
71		}
72		};
73
74
75		// Struct to obtain Cholesky for spatial exponential.
76		template <class Type>
77		struct lower_chol_spatial: virtual lower_chol_base<Type> {
78		vector<Type> theta;
79		std::string cov_type;
80		lower_chol_spatial() {
81		// This default constructor is needed because the use of `[]` in map.
82		}
83		// Constructor from theta. For now the cholesky does not need to be cached.
84	200x	lower_chol_spatial(vector<Type> theta, std::string cov_type): theta(theta), cov_type(cov_type) {
85		}
86	44897x	matrix<Type> get_chol(std::vector<int> visits, matrix<Type> dist) {
87	44897x	return get_spatial_covariance_lower_chol(this->theta, dist, this->cov_type);
88		}
89	15780x	matrix<Type> get_sigma(std::vector<int> visits, matrix<Type> dist) {
90	15780x	return tcrossprod(this->get_chol(visits, dist), true);
91		}
92	5912x	matrix<Type> get_sigma_inverse(std::vector<int> visits, matrix<Type> dist) {
93	5912x	return this->get_sigma(visits, dist).inverse();
94		}
95		};
96
97		template <class T, class Base, class D1, class D2>
98		struct cache_obj {
99		std::map<int, std::shared_ptr<Base>> cache;
100		int n_groups;
101		bool is_spatial;
102		int n_visits;
103	10220x	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	10220x	int theta_one_group_size = theta.size() / n_groups;
106	21368x	for (int r = 0; r < n_groups; r++) {
107		// Use unique pointers here to better manage resource.
108	11152x	if (is_spatial) {
109	198x	this->cache[r] = std::make_shared<D1>(theta.segment(r * theta_one_group_size, theta_one_group_size), cov_type);
110		} else {
111	10954x	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	9858x	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	6448x	matrix<Type> get_default_chol() {
125	12896x	std::vector<int> visit(this->n_visits);
126	6448x	std::iota(std::begin(visit), std::end(visit), 0);
127	6448x	matrix<Type> dist(2, 2);
128	6448x	dist << 0, 1, 1, 0;
129	6448x	int dim = this->is_spatial?2:this->n_visits;
130	6448x	matrix<Type> covariance_lower_chol = matrix<Type>::Zero(dim * this->n_groups, dim);
131	13580x	for (int r = 0; r < this->n_groups; r++) {
132	7132x	covariance_lower_chol.block(r * dim, 0, dim, dim) = this->cache[r]->get_chol(visit, dist);
133		}
134	12896x	return covariance_lower_chol;
135		}
136		};
137
138		#endif

1		#ifndef COV_INCLUDED_
2		#define COV_INCLUDED_
3
4		#include "utils.h"
5
6		// Unstructured covariance:
7		// Cholesky factor.
8		template <class T>
9	17928x	matrix<T> get_unstructured(const vector<T>& theta, int n_visits) {
10	17928x	vector<T> sd_values = exp(theta.head(n_visits));
11	17928x	vector<T> lower_tri_chol_values = theta.tail(theta.size() - n_visits);
12	17928x	matrix<T> covariance_lower_chol = matrix<T>::Zero(n_visits, n_visits);
13	17928x	int k = 0;
14	89568x	for(int i = 0; i < n_visits; i++) {
15	71640x	covariance_lower_chol(i, i) = sd_values(i);
16	179104x	for(int j = 0; j < i; j++){
17	107464x	covariance_lower_chol(i, j) = sd_values(i) * lower_tri_chol_values(k++);
18		}
19		}
20	35856x	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	4512x	const T operator() (int i, int j) const {
30	4512x	return this->corr_values.segment(j, i - j).prod();
31		}
32		};
33		// Homogeneous Ante-dependence Cholesky factor.
34		template <class T>
35	316x	matrix<T> get_ante_dependence(const vector<T>& theta, int n_visits) {
36	316x	T const_sd = exp(theta(0));
37	316x	corr_fun_ante_dependence<T> fun(theta.tail(n_visits - 1));
38	316x	matrix<T> ad_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
39	632x	return const_sd * ad_cor_mat_chol;
40		}
41		// Heterogeneous Ante-dependence Cholesky factor.
42		template <class T>
43	476x	matrix<T> get_ante_dependence_heterogeneous(const vector<T>& theta, int n_visits) {
44	476x	vector<T> sd_values = exp(theta.head(n_visits));
45	476x	corr_fun_ante_dependence<T> fun(theta.tail(n_visits - 1));
46	952x	return get_heterogeneous_cov(sd_values, fun);
47		}
48
49		// Toeplitz:
50
51		// Correlation function.
52		template <class T>
53		struct corr_fun_toeplitz : generic_corr_fun<T> {
54		using generic_corr_fun<T>::generic_corr_fun;
55	5076x	const T operator() (int i, int j) const {
56	5076x	int index = (i - j) - 1; // Note: We need to start at 0.
57	5076x	return this->corr_values(index);
58		}
59		};
60		// Homogeneous Toeplitz Cholesky factor.
61		template <class T>
62	416x	matrix<T> get_toeplitz(const vector<T>& theta, int n_visits) {
63	416x	T const_sd = exp(theta(0));
64	416x	corr_fun_toeplitz<T> fun(theta.tail(n_visits - 1));
65	416x	matrix<T> toep_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
66	832x	return const_sd * toep_cor_mat_chol;
67		}
68		// Heterogeneous Toeplitz Cholesky factor.
69		template <class T>
70	428x	matrix<T> get_toeplitz_heterogeneous(const vector<T>& theta, int n_visits) {
71	428x	vector<T> sd_values = exp(theta.head(n_visits));
72	428x	corr_fun_toeplitz<T> fun(theta.tail(n_visits - 1));
73	856x	return get_heterogeneous_cov(sd_values, fun);
74		}
75
76		// Autoregressive:
77
78		// Correlation function.
79		template <class T>
80		struct corr_fun_autoregressive : generic_corr_fun<T> {
81		using generic_corr_fun<T>::generic_corr_fun;
82	20840x	const T operator() (int i, int j) const {
83	20840x	T diff = T((i - j) * 1.0);
84	26648x	return pow(this->corr_values(0), diff); // rho^{\|i-j\|}
85		}
86		};
87		// Homogeneous autoregressive Cholesky factor.
88		template <class T>
89	3036x	matrix<T> get_auto_regressive(const vector<T>& theta, int n_visits) {
90	3036x	T const_sd = exp(theta(0));
91	3036x	corr_fun_autoregressive<T> fun(theta.tail(1));
92	3036x	matrix<T> ar1_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
93	6072x	return const_sd * ar1_cor_mat_chol;
94		}
95		// Heterogeneous autoregressive Cholesky factor.
96		template <class T>
97	428x	matrix<T> get_auto_regressive_heterogeneous(const vector<T>& theta, int n_visits) {
98	428x	vector<T> sd_values = exp(theta.head(n_visits));
99	428x	corr_fun_autoregressive<T> fun(theta.tail(1));
100	856x	return get_heterogeneous_cov(sd_values, fun);
101		}
102
103		// Compound symmetry:
104
105		// Correlation function.
106		template <class T>
107		struct corr_fun_compound_symmetry {
108		const vector<T> corr_values;
109		// Custom constructor here because we need the special mapping from theta to rho.
110	1232x	corr_fun_compound_symmetry(const vector<T>& theta, int n_visits) :
111	1232x	corr_values(map_to_cs_cor(theta, n_visits)) {}
112	7272x	const T operator() (int i, int j) const {
113	7272x	return this->corr_values(0); // rho (constant)
114		}
115		};
116		// Homogeneous compound symmetry Cholesky factor.
117		template <class T>
118	700x	matrix<T> get_compound_symmetry(const vector<T>& theta, int n_visits) {
119	700x	T const_sd = exp(theta(0));
120	700x	corr_fun_compound_symmetry<T> fun(theta.tail(1), n_visits);
121	700x	matrix<T> cs_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
122	1400x	return const_sd * cs_cor_mat_chol;
123		}
124		// Heterogeneous compound symmetry Cholesky factor.
125		template <class T>
126	524x	matrix<T> get_compound_symmetry_heterogeneous(const vector<T>& theta, int n_visits) {
127	524x	vector<T> sd_values = exp(theta.head(n_visits));
128	524x	corr_fun_compound_symmetry<T> fun(theta.tail(1), n_visits);
129	1048x	return get_heterogeneous_cov(sd_values, fun);
130		}
131
132		// Spatial Exponential Cholesky factor.
133		template <class T>
134	44897x	matrix<T> get_spatial_exponential(const vector<T>& theta, const matrix<T>& distance) {
135	44897x	T const_sd = exp(theta(0));
136	44897x	T rho = invlogit(theta(1));
137	44897x	matrix<T> expdist = exp(distance.array() * log(rho));
138	44897x	matrix<T> result = expdist * const_sd;
139	44897x	Eigen::LLT<Eigen::Matrix<T,Eigen::Dynamic,Eigen::Dynamic> > cov_i_chol(result);
140	89794x	return cov_i_chol.matrixL();
141		}
142
143		// Creates a new correlation object dynamically.
144		template <class T>
145	24220x	matrix<T> get_covariance_lower_chol(const vector<T>& theta, int n_visits, std::string cov_type) {
146	24220x	matrix<T> result;
147
148	24220x	if (cov_type == "us") {
149	17920x	result = get_unstructured<T>(theta, n_visits);
150	6300x	} else if (cov_type == "toep") {
151	412x	result = get_toeplitz<T>(theta, n_visits);
152	5888x	} else if (cov_type == "toeph") {
153	424x	result = get_toeplitz_heterogeneous<T>(theta, n_visits);
154	5464x	} else if (cov_type == "ar1") {
155	3032x	result = get_auto_regressive<T>(theta, n_visits);
156	2432x	} else if (cov_type == "ar1h") {
157	424x	result = get_auto_regressive_heterogeneous<T>(theta, n_visits);
158	2008x	} else if (cov_type == "ad") {
159	312x	result = get_ante_dependence<T>(theta, n_visits);
160	1696x	} else if (cov_type == "adh") {
161	472x	result = get_ante_dependence_heterogeneous<T>(theta, n_visits);
162	1224x	} else if (cov_type == "cs") {
163	696x	result = get_compound_symmetry<T>(theta, n_visits);
164	528x	} else if (cov_type == "csh") {
165	520x	result = get_compound_symmetry_heterogeneous<T>(theta, n_visits);
166		} else {
167	4x	Rf_error("%s", ("Unknown covariance type '" + cov_type + "'.").c_str());
168		}
169
170	24212x	return result;
171		}
172
173		// Creates a new spatial covariance cholesky.
174		template <class T>
175	44897x	matrix<T> get_spatial_covariance_lower_chol(const vector<T>& theta, const matrix<T>& distance, std::string cov_type) {
176	44897x	matrix<T> result;
177	44897x	if (cov_type == "sp_exp") {
178	44897x	result = get_spatial_exponential<T>(theta, distance);
179		} else {
180	!	Rf_error("%s", ("Unknown spatial covariance type '" + cov_type + "'.").c_str());
181		}
182	44897x	return result;
183		}
184
185		#endif

1		#ifndef DERIVATIVE_INCLUDED_
2		#define DERIVATIVE_INCLUDED_
3
4		#include "chol_cache.h"
5
6		using namespace Rcpp;
7		using std::string;
8		// Struct chol to obtain the cholesky factor given theta.
9		// The reason to have it is that we need a functor that need only theta to
10		// obtain the derivatives from autodiff.
11		// Only non-spatial covariance structure here.
12		struct chol {
13		int dim_cov_mat;
14		string cov_type;
15	768x	chol(int dim, string cov): dim_cov_mat(dim), cov_type(cov) {};
16		template <class T>
17	2300x	vector<T> operator() (vector<T> &theta) {
18	2300x	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	386x	chol_jacobian(int dim, string cov): dim_cov_mat(dim), cov_type(cov), mychol(dim, cov) {};
28		template<class T>
29	388x	vector<T> operator() (vector<T> &theta) {
30	388x	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	382x	std::map<std::string, matrix<Type>> derivatives(int n_visits, std::string cov_type, vector<Type> theta) {
39	382x	std::map<std::string, matrix<Type>> ret;
40	382x	chol chol_obj(n_visits, cov_type);
41	382x	chol_jacobian chol_jac_obj(n_visits, cov_type);
42	382x	matrix<Type> l = chol_obj(theta).matrix();
43	382x	l.resize(n_visits, n_visits);
44	382x	vector<Type> chol_d1_vec = autodiff::jacobian(chol_obj, theta).vec(); // chol_d1_vec is (dim * dim * l_theta)
45	382x	vector<Type> chol_d2_vec = autodiff::jacobian(chol_jac_obj, theta).vec(); // chol_d2_vec is (dim * dim * l_theta * l_theta)
46	382x	matrix<Type> ret_d1 = matrix<Type>(n_visits * theta.size(), n_visits);
47	382x	matrix<Type> ret_d2 = matrix<Type>(n_visits * theta.size() * theta.size(), n_visits);
48	382x	int n_visits_sq = n_visits * n_visits;
49	2580x	for (int i = 0; i < theta.size(); i++) {
50	2198x	matrix<Type> ld1 = chol_d1_vec.segment(i * n_visits_sq, n_visits_sq).matrix();
51	2198x	ld1.resize(n_visits, n_visits);
52	2198x	matrix<Type> ld1_lt = ld1 * l.transpose();
53	2198x	auto sigma_d1_i = ld1_lt + ld1_lt.transpose();
54	2198x	ret_d1.block(i * n_visits, 0, n_visits, n_visits) = sigma_d1_i;
55	19852x	for (int j = 0; j < theta.size(); j++) {
56	17654x	matrix<Type> ld2 = chol_d2_vec.segment( (j * theta.size() + i) * n_visits_sq, n_visits_sq).matrix();
57	17654x	matrix<Type> ld1_j = chol_d1_vec.segment(j * n_visits_sq, n_visits_sq).matrix();
58	17654x	ld2.resize(n_visits, n_visits);
59	17654x	ld1_j.resize(n_visits, n_visits);
60	17654x	auto ld2_lt = ld2 * l.transpose();
61	17654x	auto ld1_ld1j = ld1 * ld1_j.transpose();
62	17654x	auto sigma_d2_ij = ld2_lt + ld2_lt.transpose() + ld1_ld1j + ld1_ld1j.transpose();
63	17654x	ret_d2.block((i * theta.size() + j) * n_visits, 0, n_visits, n_visits) = sigma_d2_ij;
64		}
65		}
66	764x	ret["derivative1"] = ret_d1;
67	382x	ret["derivative2"] = ret_d2;
68	764x	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	406x	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	382x	derivatives_nonspatial(vector<Type> theta, int n_visits, std::string cov_type): lower_chol_nonspatial<Type>(theta, n_visits, cov_type) {
96	382x	std::map<std::string, tmbutils::matrix<Type>> allret = derivatives<Type>(this->n_visits, this->cov_type, this->theta);
97	764x	matrix<Type> sigma_d1 = allret["derivative1"];
98	382x	matrix<Type> sigma_d2 = allret["derivative2"];
99	382x	this->sigmad1_cache[this->full_visit] = sigma_d1;
100	382x	this->sigmad2_cache[this->full_visit] = sigma_d2;
101		}
102		// Cache and return the first order derivatives using select matrix.
103	20860x	matrix<Type> get_sigma_derivative1(std::vector<int> visits, matrix<Type> dist) override {
104	20860x	auto target = this->sigmad1_cache.find(visits);
105	20860x	if (target != this->sigmad1_cache.end()) {
106	16854x	return target->second;
107		} else {
108	4006x	int n_visits_i = visits.size();
109	4006x	matrix<Type> ret = matrix<Type>(this->n_theta * n_visits_i, n_visits_i);
110	27984x	for (int i = 0; i < this->n_theta; i++) {
111	23978x	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	4006x	this->sigmad1_cache[visits] = ret;
114	4006x	return ret;
115		}
116		}
117		// Cache and return the second order derivatives using select matrix.
118	16550x	matrix<Type> get_sigma_derivative2(std::vector<int> visits, matrix<Type> dist) override {
119	16550x	auto target = this->sigmad2_cache.find(visits);
120	16550x	if (target != this->sigmad2_cache.end()) {
121	15092x	return target->second;
122		} else {
123	1458x	int n_visits_i = visits.size();
124	1458x	int theta_sq = this->n_theta * this->n_theta;
125	1458x	matrix<Type> ret = matrix<Type>(theta_sq * n_visits_i, n_visits_i);
126	44586x	for (int i = 0; i < theta_sq; i++) {
127	43128x	ret.block(i * n_visits_i, 0, n_visits_i, n_visits_i) = subset_matrix<matrix<Type>, vector<int>>(this->sigmad2_cache[this->full_visit].block(i * this->n_visits, 0, this->n_visits, this->n_visits), visits, visits);
128		}
129	1458x	this->sigmad2_cache[visits] = ret;
130	1458x	return ret;
131		}
132		}
133		// Cache and return the lower cholesky factor of inverse of sigma using select matrix.
134	14578x	matrix<Type> get_inverse_chol(std::vector<int> visits, matrix<Type> dist) override {
135	14578x	auto target = this->inverse_chol_cache.find(visits);
136	14578x	if (target != this->inverse_chol_cache.end()) {
137	13468x	return target->second;
138		} else {
139	1110x	matrix<Type> sigmainv = this->get_sigma_inverse(visits, dist);
140	1110x	Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > sigma_inv_chol(sigmainv);
141	1110x	matrix<Type> Li = sigma_inv_chol.matrixL();
142	1110x	this->inverse_chol_cache[visits] = Li;
143	1110x	return Li;
144		}
145		}
146		// Cache and return the first order derivatives of inverse of sigma using select matrix.
147	52408x	matrix<Type> get_inverse_derivative(std::vector<int> visits, matrix<Type> dist) override {
148	52408x	auto target = this->sigma_inverse_d1_cache.find(visits);
149	52408x	if (target != this->sigma_inverse_d1_cache.end()) {
150	48098x	return target->second;
151		} else {
152	4310x	auto sigma_d1 = this->get_sigma_derivative1(visits, dist);
153	4310x	matrix<Type> sigma_inv_d1(sigma_d1.rows(), sigma_d1.cols());
154	4310x	int n_visits_i = visits.size();
155	4310x	auto sigma_inv = this->get_sigma_inverse(visits, dist);
156	30146x	for (int r = 0; r < this->n_theta; r++) {
157	25836x	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	4310x	this->sigma_inverse_d1_cache[visits] = sigma_inv_d1;
160	4310x	return sigma_inv_d1;
161		}
162		}
163		};
164
165		// derivatives_sp_exp struct is created to obtain the exact derivatives of spatial exponential
166		// covariance structure, and its inverse.
167		// No caching is used because the distance can be hardly the same for spatial covariance
168		// structures.
169		template <class Type>
170		struct derivatives_sp_exp: public lower_chol_spatial<Type>, virtual derivatives_base<Type> {
171		Type const_sd;
172		Type rho;
173		Type logrho;
174		derivatives_sp_exp() {
175		// This default constructor is needed because the use of `[]` in maps.
176		}
177		// Initialize the theta values; the reason to have theta is that for a fit, the theta
178		// is the same for all subjects, while the distance between each visits for each subject
179		// can be different.
180	24x	derivatives_sp_exp(vector<Type> theta, std::string cov_type): lower_chol_spatial<Type>(theta, cov_type) ,const_sd(exp(theta(0))), rho(invlogit(theta(1))) {
181	24x	this->logrho = log(this->rho);
182		}
183		// Obtain first order derivatives
184	5124x	matrix<Type> get_sigma_derivative1(std::vector<int> visits, matrix<Type> dist) override {
185	5124x	matrix<Type> ret(2 * dist.rows(), dist.cols());
186		// partial sigma / partial theta_1 = sigma.
187	5124x	auto sigma = this->get_sigma(visits, dist);
188	5124x	ret.block(0, 0, dist.rows(), dist.cols()) = sigma;
189	5124x	ret.block(dist.rows(), 0, dist.rows(), dist.cols()) = sigma.array() * dist.array() * (1 - this->rho);
190	10248x	return ret;
191		}
192		// Obtain second order derivatives.
193	1972x	matrix<Type> get_sigma_derivative2(std::vector<int> visits, matrix<Type> dist) override {
194	1972x	matrix<Type> ret(4 * dist.rows(), dist.cols());
195	1972x	auto sigma = this->get_sigma(visits, dist);
196	1972x	ret.block(0, 0, dist.rows(), dist.cols()) = sigma;
197	1972x	Type rho_r = 1 - this->rho;
198	1972x	auto dtheta1dtheta2 = sigma.array() * dist.array() * rho_r;
199	1972x	ret.block(dist.rows(), 0, dist.rows(), dist.cols()) = dtheta1dtheta2;
200	1972x	ret.block(dist.rows() * 2, 0, dist.rows(), dist.cols()) = dtheta1dtheta2;
201	1972x	matrix<Type> dtheta2s = dtheta1dtheta2 * (dist.array() * rho_r - this->rho);
202	1972x	ret.block(dist.rows() * 3, 0, dist.rows(), dist.cols()) = dtheta2s;
203	3944x	return ret;
204		}
205		// Obtain the lower cholesky factor of inverse of sigma using select matrix.
206	788x	matrix<Type> get_inverse_chol(std::vector<int> visits, matrix<Type> dist) override {
207	788x	auto sigmainv = this->get_sigma_inverse(visits, dist);
208	788x	Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > sigma_inv_chol(sigmainv);
209	788x	matrix<Type> Li = sigma_inv_chol.matrixL();
210	1576x	return Li;
211		}
212		// Obtain first order derivatives for inverse of sigma.
213	3152x	matrix<Type> get_inverse_derivative(std::vector<int> visits, matrix<Type> dist) override {
214	3152x	matrix<Type> sigma_inv_d1 = matrix<Type>::Zero(2 * dist.rows(), dist.cols());
215	3152x	auto sigma_inv = this->get_sigma_inverse(visits, dist);
216	3152x	auto sigma_d1 = this->get_sigma_derivative1(visits, dist);
217	9456x	for (int r = 0; r < 2; r++) {
218	6304x	sigma_inv_d1.block(r * dist.rows(), 0, dist.rows(), dist.cols()) = - sigma_inv * sigma_d1.block(r * dist.rows(), 0, dist.rows(), dist.cols()) *sigma_inv;
219		}
220	6304x	return sigma_inv_d1;
221		}
222		};
223
224		#endif

1		#ifndef UTILS_INCLUDED_
2		#define UTILS_INCLUDED_
3		#include <Rcpp.h>
4		#define INCLUDE_RCPP
5		#include "tmb_includes.h"
6
7		#define as_num_matrix_tmb as_matrix<matrix<double>, NumericMatrix>
8		#define as_num_matrix_rcpp as_matrix<NumericMatrix, matrix<double>>
9		#define as_num_vector_tmb as_vector<vector<double>, NumericVector>
10		#define as_num_vector_rcpp as_vector<NumericVector, vector<double>>
11
12		// Obtain submatrix from index
13
14		template <typename T1, typename T2>
15	315909x	T1 subset_matrix(T1 input, T2 index1, T2 index2) {
16		#if EIGEN_VERSION_AT_LEAST(3,4,0)
17	315909x	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	315909x	return ret;
27		}
28
29		template <typename T1, typename T2>
30	242226x	T1 subset_matrix(T1 input, T2 index1) {
31		#if EIGEN_VERSION_AT_LEAST(3,4,0)
32	242226x	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	242226x	return ret;
42		}
43
44
45		// Conversion from Rcpp vector/matrix to eigen vector/matrix
46		template <typename T1, typename T2>
47	613255x	T1 as_vector(T2 input) {
48	613255x	T1 ret(input.size());
49	2054086x	for (int i = 0; i < input.size(); i++) {
50	1440831x	ret(i) = input(i);
51		}
52	613255x	return ret;
53		}
54
55		template <typename T1, typename T2>
56	421160x	T1 as_matrix(T2 input) {
57	421160x	T1 ret(input.rows(), input.cols());
58	5984516x	for (int i = 0; i < input.rows(); i++) {
59	63454156x	for (int j = 0; j < input.cols(); j++) {
60	57890800x	ret(i,j) = input(i,j);
61		}
62		}
63	421160x	return ret;
64		}
65
66		template <typename T>
67	726674x	T segment(T input, int start, int n) {
68	726674x	T ret(n);
69	3630556x	for (int i = 0, j = start; i < n; i++, j++) {
70	2903882x	ret(i) = input(j);
71		}
72	726674x	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	1239255x	matrix<Type> tcrossprod(const matrix<Type>& lower_chol, bool complete = false) {
81	1239255x	int n = lower_chol.rows();
82	1239255x	matrix<Type> result = matrix<Type>::Zero(n, n);
83	1239255x	result.template selfadjointView<Eigen::Lower>().rankUpdate(lower_chol);
84	1239255x	if (complete) {
85	26448x	result.template triangularView<Eigen::Upper>() = result.transpose();
86		}
87	1239255x	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	1263892x	matrix<Type> crossprod(const matrix<Type>& x) {
96	1263892x	int n = x.cols();
97	1263892x	matrix<Type> result = matrix<Type>::Zero(n, n);
98	1263892x	result.template selfadjointView<Eigen::Lower>().rankUpdate(x.transpose());
99	1263892x	return result;
100		}
101
102		// Mapping from real values to correlation parameters in (-1, 1).
103		template <class T>
104	5124x	vector<T> map_to_cor(const vector<T>& theta) {
105	10248x	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	1232x	vector<T> map_to_cs_cor(const vector<T>& theta, int n_visits) {
111	1232x	T a = T(1.0) / (T(n_visits) - T(1.0));
112	1232x	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	5116x	generic_corr_fun(const vector<T>& theta) :
122	5116x	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	6332x	matrix<T> get_corr_mat_chol(int n_visits, const F<T>& corr_fun) {
129	6332x	matrix<T> correlation(n_visits, n_visits);
130	6332x	correlation.setIdentity();
131	31492x	for(int i = 0; i < n_visits; i++) {
132	62680x	for(int j = 0; j < i; j++){
133	37520x	correlation(i, j) = corr_fun(i, j);
134		}
135		}
136	6332x	Eigen::LLT<Eigen::Matrix<T,Eigen::Dynamic,Eigen::Dynamic> > correlation_chol(correlation);
137	6332x	matrix<T> L = correlation_chol.matrixL();
138	12664x	return L;
139		}
140
141		// Heterogeneous covariance matrix calculation given vector of standard deviations (sd_values)
142		// and a correlation function (corr_fun).
143		template <class T, template<class> class F>
144	1858x	matrix<T> get_heterogeneous_cov(const vector<T>& sd_values, const F<T>& corr_fun) {
145	1858x	matrix<T> correlation_chol = get_corr_mat_chol(sd_values.size(), corr_fun);
146	1858x	Eigen::DiagonalMatrix<T,Eigen::Dynamic,Eigen::Dynamic> D = sd_values.matrix().asDiagonal();
147	1858x	matrix<T> result = D * correlation_chol;
148	3716x	return result;
149		}
150
151		// Obtain the Euclidean distance
152		template <class T>
153	33703x	matrix<T> euclidean(const matrix<T>& coordinates) {
154	33703x	matrix<T> result(coordinates.rows(), coordinates.rows());
155	126598x	for (int i = 0; i < coordinates.rows(); i++) {
156	92895x	result(i, i) = 0;
157	188400x	for (int j = 0; j < i; j ++) {
158	95505x	vector<T> diff = coordinates.row(i) - coordinates.row(j);
159	95505x	T d = sqrt((diff * diff).sum());
160	95505x	result(i, j) = d;
161	95505x	result(j, i) = d;
162		}
163		}
164	33703x	return result;
165		}
166
167		// Element wise power function of a matrix
168		template <class T>
169	1584x	Eigen::Matrix<T, -1, -1> cpow(const Eigen::Matrix<T, -1, -1> & input, double p) {
170	1584x	Eigen::Matrix<T, -1, -1> ret = Eigen::Matrix<T, -1, -1>(input.rows(), input.cols());
171	5908x	for (int i = 0; i < ret.rows(); i ++) {
172	8664x	for (int j = 0; j < ret.cols(); j++) {
173	4340x	ret(i, j) = std::pow(input(i, j), p);
174		}
175		}
176	1584x	return ret;
177		}
178
179		// Calculate the square root of the pseudo inverse of a matrix
180		// adapted from the method for calculating the pseudo-Inverse as recommended by the Eigen developers
181		template<typename T>
182	1580x	matrix<T> pseudoInverseSqrt(const matrix<T> &input, double epsilon = std::numeric_limits<double>::epsilon()) {
183	1580x	Eigen::Matrix<T, -1, -1> eigen_mat = as_matrix<Eigen::Matrix<T, -1, -1>, matrix<T>>(input);
184	1580x	Eigen::JacobiSVD< Eigen::Matrix<T, -1, -1> > svd(eigen_mat ,Eigen::ComputeFullU \| Eigen::ComputeFullV);
185	1580x	double tolerance = epsilon * std::max(input.cols(), input.rows()) *svd.singularValues().array().abs()(0);
186	1580x	auto singular_vals = Matrix<T,-1,-1>((svd.singularValues().array() > tolerance).select(svd.singularValues().array().inverse(), 0).matrix());
187	1580x	Eigen::Matrix<T, -1, -1> ret_eigen = svd.matrixV() * cpow(singular_vals, 0.5).asDiagonal() * svd.matrixU().adjoint();
188	3160x	return as_matrix<matrix<T>, Eigen::Matrix<T, -1, -1>>(ret_eigen);
189		}
190
191		#endif

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

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

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