mmrm coverage - 97.76%

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 = grepl("^sp_", covariance$type),
      visit_var = covariance$visits,
      subject_var = covariance$subject,
      group_var = if (length(covariance$group) < 1) NULL else covariance$group,
      model_var = setdiff(all.vars(formula[[3]]), covariance$subject),
      response_var = all.vars(formula[[2]])
    ),
    class = "mmrm_tmb_formula_parts"
  )
}

#' Data for `TMB` Fit
#'
#' @param formula_parts (`mmrm_tmb_formula_parts`)\cr list with formula parts
#'   from [h_mmrm_tmb_formula_parts()].
#' @param data (`data.frame`)\cr which contains variables used in `formula_parts`.
#' @param weights (`vector`)\cr weights to be used in the fitting process.
#' @param reml (`flag`)\cr whether restricted maximum likelihood (REML) estimation is used,
#'   otherwise maximum likelihood (ML) is used.
#' @param singular (`string`)\cr choices of method deal with rank-deficient matrices. "error" to
#'   stop the function return the error, "drop" to drop these columns, and "keep" to keep all the columns.
#' @param drop_visit_levels (`flag`)\cr whether to drop levels for visit variable, if visit variable is a factor.
#' @param allow_na_response (`flag`)\cr whether NA in response is allowed.
#' @param drop_levels (`flag`)\cr whether drop levels for covariates. If not dropped could lead to singular matrix.
#' @param xlev (`list` or `NULL`)\cr list of X levels produced by stats::.getXlevels
#' @param contrasts (`list` or `NULL`)\cr an optional named list of contrast
#'   matrices or contrast functions (like [stats::contr.sum] or
#'   [stats::contr.poly]) for specific factor variables, matching the
#'   `contrasts` argument in [stats::lm()]. The list names must correspond to
#'   factor variable names in the model formula. When `NULL` (the default),
#'   the contrasts set on the factor variables in `data` are used. If a
#'   contrast matrix has rownames that include levels not present in `data`,
#'   those levels are preserved and the corresponding model matrix columns
#'   are marked as aliased (not estimable), enabling prediction on new data
#'   containing those levels.
#' @param emmeans_gcomp_vars (`character` or `NULL`)
#'   treated as fixed for G-computation correction. Stored as a character
#'   vector in the returned list for downstream use in the emmeans hook.
#'
#' @return List of class `mmrm_tmb_data` with elements:
#' - `full_frame`: `data.frame` with `n` rows containing all variables needed in the model.
#' - `data`: `data.frame` of input dataset.
#' - `x_matrix`: `matrix` with `n` rows and `p` columns specifying the overall design matrix.
#' - `x_cols_aliased`: `logical` with potentially more than `p` elements indicating which
#'      columns in the original design matrix have been left out to obtain a full rank
#'      `x_matrix`.
#' - `y_vector`: length `n` `numeric` specifying the overall response vector.
#' - `weights_vector`: length `n` `numeric` specifying the weights vector.
#' - `n_visits`: `int` with the number of visits, which is the dimension of the
#'      covariance matrix.
#' - `n_subjects`: `int` with the number of subjects.
#' - `subject_zero_inds`: length `n_subjects` `integer` containing the zero-based start
#'     indices for each subject.
#' - `subject_n_visits`: length `n_subjects` `integer` containing the number of
#'     observed visits for each subjects. So the sum of this vector equals `n`.
#' - `cov_type`: `string` value specifying the covariance type.
#' - `is_spatial_int`: `int` specifying whether the covariance structure is spatial(1) or not(0).
#' - `reml`: `int` specifying whether REML estimation is used (1), otherwise ML (0).
#' - `subject_groups`: `factor` specifying the grouping for each subject.
#' - `n_groups`: `int` with the number of total groups
#' - `emmeans_gcomp_vars`: `character` or `NULL` with the G-computation variable names.
#'
#' @details Note that the `subject_var` must not be factor but can also be character.
#'   If it is character, then it will be converted to factor internally. Here
#'   the levels will be the unique values, sorted alphabetically and numerically if there
#'   is a common string prefix of numbers in the character elements. For full control
#'   on the order please use a factor.
#'
#' @keywords internal
h_mmrm_tmb_data <- function(
  formula_parts,
  data,
  weights,
  reml,
  singular = c("drop", "error", "keep"),
  drop_visit_levels,
  allow_na_response = FALSE,
  drop_levels = TRUE,
  xlev = NULL,
  contrasts = NULL,
  emmeans_gcomp_vars = NULL
) {
  assert_class(formula_parts, "mmrm_tmb_formula_parts")
  assert_data_frame(data)
  assert_numeric(weights, len = nrow(data))
  assert_flag(reml)
  singular <- match.arg(singular)
  assert_flag(drop_visit_levels)
  assert_character(emmeans_gcomp_vars, min.len = 1L, null.ok = TRUE)
  if (!is.null(emmeans_gcomp_vars)) {
    assert_subset(emmeans_gcomp_vars, names(data))
  }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  kept
}

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

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

#' Returns a List of Model Details for `emmeans` Purposes
#'
#' @seealso See [emmeans::emm_basis()] for background.
#' @keywords internal
#' @noRd
# nolint start
emm_basis.mmrm <- function(
  # nolint end
  object,
  trms,
  xlev,
  grid,
  ...
) {
  # Recovered terms include covariance variables for the reference grid.
  # Use the fitted fixed-effect terms to preserve the coefficient basis.
  trms <- stats::delete.response(stats::terms(object))
  # Be careful with the factor levels:
  # We only need those that are in the fitted terms.
  xlev_subset <- intersect(names(xlev), rownames(attr(trms, "factors")))
  xlev <- xlev[xlev_subset]
  model_frame <- stats::model.frame(
    trms,
    grid,
    na.action = stats::na.pass,
    xlev = xlev
  )
  contrasts <- component(object, "contrasts")
  model_mat <- stats::model.matrix(trms, model_frame, contrasts.arg = contrasts)
  beta_hat <- component(object, "beta_est")
  nbasis <- if (length(beta_hat) < ncol(model_mat)) {
    kept <- h_match_coefs(names(beta_hat), colnames(model_mat))
    if (anyNA(kept)) {
      unmatched <- names(beta_hat)[is.na(kept)]
      stop(
        paste0(
          "Failed to match coefficient name(s) to model matrix columns: ",
          paste(unmatched, collapse = ", ")
        ),
        call. = FALSE
      )
    }
    beta_hat <- NA * model_mat[1L, ]
    beta_hat[kept] <- component(object, "beta_est")
    orig_model_mat <- stats::model.matrix(
      trms,
      stats::model.frame(
        object,
        exclude = c(
          "subject_var",
          if (object$formula_parts$is_spatial) "visit_var"
        )
      ),
      contrasts.arg = contrasts
    )
    estimability::nonest.basis(orig_model_mat)
  } else {
    estimability::all.estble
  }

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

  if (!is.null(object$tmb_data$emmeans_gcomp_vars)) {
    gcomp_result <- h_emm_basis_gcomp(
      object,
      model_mat,
      beta_hat,
      nbasis,
      V,
      grid
    )
    model_mat <- gcomp_result$model_mat
    beta_hat <- gcomp_result$beta_hat
    nbasis <- gcomp_result$nbasis
    V <- gcomp_result$V # nolint
    dfargs <- gcomp_result$dfargs
    dffun <- gcomp_result$dffun
    message(
      "G-computation correction applied, so emmeans will produce the correct average treatment effect."
    )
  } else {
    dfargs <- list(object = object)
    dffun <- function(k, dfargs) {
      mmrm::df_md(dfargs$object, contrast = k)$denom_df
    }
  }

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

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

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

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

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

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

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

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

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

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

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

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

#' 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.
#' @param linear (`flag`)\cr whether to omit second derivatives and the R component.
#' @param w (`matrix` or `NULL`)\cr covariance of the covariance parameters.
#'   Supply with `linear = TRUE` to contract Q without constructing its blocks.
#'
#' @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.
#'
#' Supplying `w` for linear KR instead retains `P` and contracts Q into the single
#' matrix `S_Q`, without storing Q blocks, see the section "Contracted linear covariance
#' adjustment" in `vignette("kenward", package = "mmrm")`. The symmetric part of `w` is
#' used, and non-finite entries in `w` give a non-finite `S_Q`.
#' [mmrm()] always supplies `w` for linear KR. Calling with `linear = TRUE` but without `w`
#' gives the pairwise `P` and `Q` blocks, which tests use as an independent reference.
#'
#' @return Named list with elements:
#' - `P`: `matrix` of \eqn{P} component.
#' - `Q`: `matrix` of \eqn{Q} component.
#' - `R`: `matrix` of \eqn{R} component, or `NULL` when `linear = TRUE`.
#' - `S_Q`: contracted Q sum when `w` is supplied; `Q` and `R` are then `NULL`.
#'
#' @keywords internal
h_get_kr_comp <- function(tmb_data, theta, linear = FALSE, w = NULL) {
  assert_class(tmb_data, "mmrm_tmb_data")
  assert_class(theta, "numeric")
  assert_flag(linear)
  if (!is.null(w)) {
    if (!linear) {
      stop(
        "`w` can only be supplied with `linear = TRUE`, ",
        "because full Kenward-Roger needs the pairwise Q and R components"
      )
    }
    assert_matrix(w, mode = "numeric", nrows = length(theta), ncols = length(theta))
  }
  .Call(`_mmrm_get_pqr`, PACKAGE = "mmrm", tmb_data, theta, linear, w)
}

#' Calculation of Kenward-Roger Degrees of Freedom for Multi-Dimensional Contrast
#'
#' @description Used in [df_md()] if method is "Kenward-Roger", i.e. for both
#'   "Kenward-Roger" and "Kenward-Roger-Linear" vcov.
#'
#' @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")
  }
  df <- h_kr_df(
    v0 = object$beta_vcov,
    l = contrast,
    w = component(object, "theta_vcov"),
    p = object$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", i.e. for both
#'   "Kenward-Roger" and "Kenward-Roger-Linear" vcov.
#'
#' @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 covariance matrix of the estimated covariance parameters.
#' @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)
  assert_matrix(p, ncols = n_beta, nrows = n_beta * n_theta)
  nl <- nrow(l)
  lv0 <- l %*% v0
  if (nl == 1L) {
    # A1 = A2 = a' W a: scalar Satterthwaite based on unadjusted variance.
    variance <- sum(lv0 * l)
    if (!isTRUE(variance > 0)) {
      stop("contrast variance must be positive")
    }
    a <- vapply(seq_len(n_theta), function(h) {
      rows <- (h - 1L) * n_beta + seq_len(n_beta)
      sum(lv0 * (lv0 %*% p[rows, , drop = FALSE])) / variance
    }, numeric(1L))
    return(list(m = 2 / h_quad_form_vec(a, w), lambda = 1))
  }

  # Normalize in contrast space: H = T T', Z = T^{-1} L Phi.
  # See vignettes/kenward.Rmd, "Degrees of freedom in contrast space".
  contrast_vcov <- tcrossprod(lv0, l)
  # Cholesky alone can accept a singular hypothesis after roundoff, so check
  # the reciprocal condition number with the same threshold as solve().
  if (rcond(contrast_vcov) < .Machine$double.eps) {
    stop("contrast covariance is numerically singular")
  }
  contrast_chol <- tryCatch(
    chol(contrast_vcov),
    error = function(e) stop("contrast covariance is not positive definite", call. = FALSE)
  )
  z <- forwardsolve(t(contrast_chol), lv0)
  f <- vapply(seq_len(n_theta), function(h) {
    rows <- (h - 1L) * n_beta + seq_len(n_beta)
    as.vector(tcrossprod(z %*% p[rows, , drop = FALSE], z))
  }, numeric(nl * nl))
  traces <- colSums(f[seq.int(1L, nl * nl, by = nl + 1L), , drop = FALSE])
  a1 <- h_quad_form_vec(traces, w)
  a2 <- sum((f %*% w) * f)
  b <- 1 / (2 * nl) * (a1 + 6 * a2)
  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 vcov is "Kenward-Roger".
#'
#' @param v (`matrix`)\cr unadjusted covariance matrix.
#' @param w (`matrix`)\cr covariance matrix of the estimated covariance parameters.
#' @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` or `NULL`)\cr R matrix from [h_get_kr_comp()].
#'   May be `NULL` for the linear approximation.
#' @param linear (`flag`)\cr whether to use linear Kenward-Roger approximation.
#'
#' @details [mmrm()] uses this function for full Kenward-Roger only, and
#'   [h_var_adj_contracted()] for the linear approximation. The pairwise linear
#'   path is kept as an independent reference for tests.
#'
#' @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
  )
  if (!linear) {
    assert_matrix(
      r,
      nrows = theta_per_group^2 * n_groups * n_visits,
      ncols = n_visits
    )
  }

  # See the adjusted covariance formula in vignettes/kenward.Rmd.
  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 {
        adjustment <- q[ijid:(ijid + n_beta - 1), ] -
          p[iid:(iid + n_beta - 1), ] %*% v %*% p[jid:(jid + n_beta - 1), ]
        if (!linear) {
          adjustment <- adjustment - 1 / 4 * r[ijid:(ijid + n_beta - 1), ]
        }
        ret <- ret + 2 * w[i, j] * v %*% adjustment %*% v
      }
    }
  }
  ret
}

#' Obtain the Contracted Linear Adjusted Covariance Matrix
#'
#' @description Obtains the linear Kenward-Roger adjusted covariance matrix for the
#'   coefficient estimates from the `P` matrices and the contracted `S_Q`.
#' Used in [mmrm()] fitting if vcov is "Kenward-Roger-Linear".
#'
#' @inheritParams h_var_adj
#' @param s_q (`matrix`)\cr contracted Q sum `S_Q` from [h_get_kr_comp()] called with `w`.
#'
#' @details See the section "Contracted linear covariance adjustment" in
#'   `vignette("kenward", package = "mmrm")`. The full `w` is used, so that
#'   cross-group covariance parameter covariances contribute.
#'
#' @return The matrix of adjusted covariance matrix.
#'
#' @keywords internal
h_var_adj_contracted <- function(v, w, p, s_q) {
  n_beta <- ncol(v)
  assert_matrix(v, nrows = n_beta)
  n_theta <- ncol(w)
  assert_matrix(w, nrows = n_theta)
  assert_matrix(p, nrows = n_beta * n_theta, ncols = n_beta)
  assert_matrix(s_q, nrows = n_beta, ncols = n_beta)
  # Column h of p_flat is vec(P_h), so column h of p_bar is vec(sum_j W_hj P_j).
  p_flat <- matrix(
    aperm(array(p, c(n_beta, n_theta, n_beta)), c(1L, 3L, 2L)),
    ncol = n_theta
  )
  p_bar <- p_flat %*% t(w)
  s_p <- matrix(0, n_beta, n_beta)
  for (h in seq_len(n_theta)) {
    s_p <- s_p + matrix(p_flat[, h], n_beta) %*% v %*% matrix(p_bar[, h], n_beta)
  }
  v + 2 * v %*% (s_q - s_p) %*% v
}

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

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

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

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

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

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

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

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

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

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

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

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

  if (is.null(weights)) {
    weights <- rep(1, nrow(data))
  } else {
    attr(weights, which = "dataname") <- deparse(match.call()$weights)
  }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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


#' Predicate Indicating Whether Two Datasets Contain the Same Observations
#'
#' Checks whether two datasets contain the same observations and that the first
#' dataset's columns are a subset of the second dataset's columns.
#'
#' For efficiency, the inspection takes place in this order:
#'
#' 1. `FALSE` is returned early if the datasets do not have the same number of
#' rows.
#'
#' 1. `FALSE` is returned early if the first dataset has a column not in the
#' second dataset.
#'
#' 1. The columns in common are sorted and compared using [base::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 [base::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
}

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

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

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

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

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

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

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

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

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

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


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

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

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

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

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

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

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

  num_df <- nrow(contrast)

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

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

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

  out
}

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

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

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

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

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

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

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


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

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

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

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

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

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

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

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

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


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

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

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

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

  n_subjects <- component(object, "n_subjects")

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

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

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

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

  component(x, name = "varcor")
}

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

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

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

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

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


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

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

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

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

  cat("\nModel Inference Optimization:")

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


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

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

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

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

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

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

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

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

  ret
}

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

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

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

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

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

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

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

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

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

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


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


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

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

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

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

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

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

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

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

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

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

  ret_df
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  list(vhat = vhat, L_global = L_global)
}

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

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

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

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

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

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

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

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

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

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

  list(S_full = S_full, L_global = L_global)
}

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

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

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

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

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

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

  list(delta = delta, L_global = L)
}

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

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

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

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

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

  coef_table
}

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

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

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

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

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

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


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


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

  fits <- list(object, ...)

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

    assert_flag(test)
    assert_flag(refit)

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

      fit_data <- lapply(fits, h_get_minimal_fit_data)

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

      needs_refit <- !nested_in_common_data

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

    log_likelihood_vec <- lapply(fits, logLik)

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

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

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

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

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

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

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

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

  out
}

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

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

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

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

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

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

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

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

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

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

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

  validate_cov_struct(x)
}

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

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

  as.cov_struct(formula, warn_partial = FALSE)
}

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

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

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

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

  reportAssertions(checks)
  x
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

#ifndef CHOL_CACHE_INCLUDED_
#define CHOL_CACHE_INCLUDED_

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

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


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

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

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

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

#endif

#ifndef COV_INCLUDED_
#define COV_INCLUDED_

#include "utils.h"

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

// Ante-dependence:

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

// Toeplitz:

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

// Autoregressive:

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

// Compound symmetry:

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

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

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


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

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

  return result;
}

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

#endif

#ifndef DERIVATIVE_INCLUDED_
#define DERIVATIVE_INCLUDED_

#include "chol_cache.h"

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

// Template function to obtain derivatives from visits, cov_type and theta.
// Basically this is calculating the derivatives for the sigma
// from the derivatives for the cholesky factor.
template <class Type>
std::map<std::string, matrix<Type>> derivatives(int n_visits, std::string cov_type, vector<Type> theta, bool second_order) {
  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;
  if (second_order) {
    chol_d2_vec = autodiff::jacobian(chol_jac_obj, theta).vec();
  }
  matrix<Type> ret_d1 = matrix<Type>(n_visits * theta.size(), n_visits);
  matrix<Type> ret_d2;
  if (second_order) {
    ret_d2.resize(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;
    if (!second_order) {
      continue;
    }
    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;
  if (second_order) {
    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;
  // Whether second order derivatives are calculated and cached.
  bool second_order = false;
  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, bool second_order): lower_chol_nonspatial<Type>(theta, n_visits, cov_type), second_order(second_order) {
    std::map<std::string, tmbutils::matrix<Type>> allret = derivatives<Type>(this->n_visits, this->cov_type, this->theta, second_order);
    this->sigmad1_cache[this->full_visit] = allret.at("derivative1");
    if (second_order) {
      this->sigmad2_cache[this->full_visit] = allret.at("derivative2");
    }
  }
  // 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 {
    if (!this->second_order) {
      Rcpp::stop("Second derivatives were not initialized in this cache.");
    }
    auto target = this->sigmad2_cache.find(visits);
     if (target != this->sigmad2_cache.end()) {
      return target->second;
    } else {
      int n_visits_i = visits.size();
      int theta_sq = this->n_theta * this->n_theta;
      matrix<Type> ret = matrix<Type>(theta_sq * n_visits_i, n_visits_i);
      for (int i = 0; i < theta_sq; i++) {
        ret.block(i  * n_visits_i, 0, n_visits_i, n_visits_i) = subset_matrix<matrix<Type>, vector<int>>(this->sigmad2_cache[this->full_visit].block(i  * this->n_visits, 0, this->n_visits, this->n_visits), visits, visits);
      }
      this->sigmad2_cache[visits] = ret;
      return ret;
    }
  }
  // Cache and return the lower cholesky factor of inverse of sigma using select matrix.
  matrix<Type> get_inverse_chol(std::vector<int> visits, matrix<Type> dist) override {
    auto target = this->inverse_chol_cache.find(visits);
     if (target != this->inverse_chol_cache.end()) {
      return target->second;
    } else {
      matrix<Type> sigmainv = this->get_sigma_inverse(visits, dist);
      Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > sigma_inv_chol(sigmainv);
      matrix<Type> Li = sigma_inv_chol.matrixL();
      this->inverse_chol_cache[visits] = Li;
      return Li;
    }
  }
  // Cache and return the first order derivatives of inverse of sigma using select matrix.
  matrix<Type> get_inverse_derivative(std::vector<int> visits, matrix<Type> dist) override {
    auto target = this->sigma_inverse_d1_cache.find(visits);
     if (target != this->sigma_inverse_d1_cache.end()) {
      return target->second;
    } else {
      auto sigma_d1 = this->get_sigma_derivative1(visits, dist);
      matrix<Type> sigma_inv_d1(sigma_d1.rows(), sigma_d1.cols());
      int n_visits_i = visits.size();
      auto sigma_inv = this->get_sigma_inverse(visits, dist);
      for (int r = 0; r < this->n_theta; r++) {
        sigma_inv_d1.block(r * n_visits_i, 0, n_visits_i, n_visits_i) = - sigma_inv * sigma_d1.block(r * n_visits_i, 0, n_visits_i, n_visits_i) *sigma_inv;
      }
      this->sigma_inverse_d1_cache[visits] = sigma_inv_d1;
      return sigma_inv_d1;
    }
  }
};

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


// derivatives_sp_gau struct is created to obtain the exact derivatives of spatial Gaussian
// covariance structure, and its inverse.
// Inherit from sp_exp because of similar parameterization; only the sigma
// derivatives differ (square the distance). Inverse / inverse-derivative logic
// is reused from the base via the virtual get_sigma_derivative1 dispatch.
template <class Type>
struct derivatives_sp_gau: public derivatives_sp_exp<Type> {
  derivatives_sp_gau() {
    // This default constructor is needed because the use of `[]` in maps.
  }
  derivatives_sp_gau(vector<Type> theta, std::string cov_type): derivatives_sp_exp<Type>(theta, cov_type) {
  }
  // Obtain first order derivatives
  matrix<Type> get_sigma_derivative1(std::vector<int> visits, matrix<Type> dist) override {
    matrix<Type> ret(2 * dist.rows(), dist.cols());
    // partial sigma / partial theta_1 = sigma.
    auto sigma = this->get_sigma(visits, dist);
    ret.block(0, 0, dist.rows(), dist.cols()) = sigma;
    ret.block(dist.rows(), 0, dist.rows(), dist.cols()) = sigma.array() * dist.array().square() * (1 - this->rho);
    return ret;
  }
  // Obtain second order derivatives.
  matrix<Type> get_sigma_derivative2(std::vector<int> visits, matrix<Type> dist) override {
    matrix<Type> ret(4 * dist.rows(), dist.cols());
    auto sigma = this->get_sigma(visits, dist);
    ret.block(0, 0, dist.rows(), dist.cols()) = sigma;
    Type rho_r = 1 - this->rho;
    auto dtheta1dtheta2 = sigma.array() * dist.array().square() * rho_r;
    ret.block(dist.rows(), 0, dist.rows(), dist.cols()) =  dtheta1dtheta2;
    ret.block(dist.rows() * 2, 0, dist.rows(), dist.cols()) = dtheta1dtheta2;
    matrix<Type> dtheta2s = dtheta1dtheta2 * (dist.array().square() * rho_r - this->rho);
    ret.block(dist.rows() * 3, 0, dist.rows(), dist.cols()) = dtheta2s;
    return ret;
  }
};

// Cache of per-group derivatives objects, with cov_type-aware dispatch for
// spatial structures (sp_exp / sp_gau). Mirrors `cache_obj` but selects the
// concrete spatial subclass at runtime so analytic derivatives match cov_type.
template <class Type>
struct derivatives_cache {
  std::map<int, std::shared_ptr<derivatives_base<Type>>> cache;
  int n_groups;
  bool is_spatial;
  int n_visits;
  derivatives_cache(vector<Type> theta, int n_groups, bool is_spatial, std::string cov_type, int n_visits, bool second_order): n_groups(n_groups), is_spatial(is_spatial), n_visits(n_visits) {
    int theta_one_group_size = theta.size() / n_groups;
    for (int r = 0; r < n_groups; r++) {
      vector<Type> theta_r = theta.segment(r * theta_one_group_size, theta_one_group_size);
      if (is_spatial) {
        if (cov_type == "sp_exp") {
          this->cache[r] = std::make_shared<derivatives_sp_exp<Type>>(theta_r, cov_type);
        } else if (cov_type == "sp_gau") {
          this->cache[r] = std::make_shared<derivatives_sp_gau<Type>>(theta_r, cov_type);
        } else {
          Rcpp::stop("Unknown spatial covariance type '" + cov_type + "'.");
        }
      } else {
        this->cache[r] = std::make_shared<derivatives_nonspatial<Type>>(theta_r, n_visits, cov_type, second_order);
      }
    }
  }
};
#endif

#include "derivatives.h"

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

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

  // Use map to hold these base class pointers (can also work for child class objects).
  auto derivatives_by_group = derivatives_cache<double>(theta_v, n_groups, is_spatial, cov_type, n_visits, false);
  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(5,0,0)
    T1 ret = input(index1, Eigen::indexing::all);
  #elif 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, bool linear, Nullable<NumericMatrix> w);
RcppExport SEXP _mmrm_get_pqr(SEXP mmrm_fit_SEXP, SEXP theta_SEXP, SEXP linear_SEXP, SEXP w_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< bool >::type linear(linear_SEXP);
    rcpp_result_gen = Rcpp::wrap(get_pqr(mmrm_fit, theta, linear, Nullable<NumericMatrix>(w_SEXP)));
    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, 4},
    {"_mmrm_get_jacobian", (DL_FUNC) &_mmrm_get_jacobian, 3},
    {"_mmrm_get_empirical", (DL_FUNC) &_mmrm_get_empirical, 5},
    {"_mmrm_predict", (DL_FUNC) &_mmrm_predict, 4},
    {"run_testthat_tests", (DL_FUNC) &run_testthat_tests, 1},
    TMB_CALLDEFS,
    {NULL, NULL, 0}
};

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

#include "derivatives.h"

using namespace Rcpp;
using std::string;

// Obtain Jacobian from a mmrm fit, given theta.
List get_jacobian(List mmrm_fit, NumericVector theta, NumericMatrix beta_vcov) {
  NumericMatrix x = mmrm_fit["x_matrix"];
  matrix<double> x_matrix = as_num_matrix_tmb(x);
  IntegerVector subject_zero_inds = mmrm_fit["subject_zero_inds"];
  int n_subjects = mmrm_fit["n_subjects"];
  IntegerVector subject_n_visits = mmrm_fit["subject_n_visits"];
  int n_visits = mmrm_fit["n_visits"];
  String cov_type = mmrm_fit["cov_type"];
  int is_spatial_int = mmrm_fit["is_spatial_int"];
  bool is_spatial = is_spatial_int == 1;
  int n_groups = mmrm_fit["n_groups"];
  IntegerVector subject_groups = mmrm_fit["subject_groups"];
  NumericVector weights_vector = mmrm_fit["weights_vector"];
  NumericMatrix coordinates = mmrm_fit["coordinates"];
  matrix<double> coords = as_num_matrix_tmb(coordinates);
  matrix<double> beta_vcov_m = as_num_matrix_tmb(beta_vcov);
  vector<double> theta_v = as_num_vector_tmb(theta);
  vector<double> G_sqrt = as_num_vector_tmb(sqrt(weights_vector));
  int n_theta = theta.size();
  int theta_size_per_group = n_theta / n_groups;
  int p = x.cols();
  matrix<double> P = matrix<double>::Zero(p * n_theta, p);
  // Use map to hold these base class pointers (can also work for child class objects).
  auto derivatives_by_group = derivatives_cache<double>(theta_v, n_groups, is_spatial, cov_type, n_visits, false);
  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 "kr_comp.h"

using namespace Rcpp;
using std::string;
// Linear KR contracts Q before storing it, retaining P. See "Contracted linear
// covariance adjustment" in vignettes/kenward.Rmd.
List get_pqr_contracted(List mmrm_fit, NumericVector theta, NumericMatrix w) {
  kr_data data(mmrm_fit, theta);
  int k = data.theta_per_group;
  int p = data.p;
  matrix<double> W = as_num_matrix_tmb(w);
  // Non-finite W (a convergence problem that mmrm() only warns about) yields
  // a non-finite S_Q, as the pairwise contraction would.
  bool finite_w = W.allFinite();
  matrix<double> P = matrix<double>::Zero(p * data.n_theta, p);
  matrix<double> Q_sum = matrix<double>::Zero(p, p);
  auto derivatives = data.derivatives(false);
  std::vector<kr_directions> directions;
  std::vector<std::map<std::vector<int>, kr_pattern>> patterns(data.n_groups);
  for (int g = 0; g < data.n_groups; ++g) {
    // Only same-group Q blocks are nonzero. Cross-group W enters S_P in
    // h_var_adj_contracted(), using the full W.
    matrix<double> w_g = finite_w ? matrix<double>(W.block(g * k, g * k, k, k)) : matrix<double>(matrix<double>::Zero(k, k));
    directions.emplace_back(w_g);
  }
  for (int i = 0; i < data.n_subjects; ++i) {
    kr_subject subject = data.subject(i);
    int start = subject.start, m = subject.n_visits, g = subject.group;
    auto accumulate = [&](const kr_pattern& pattern) {
      matrix<double> X_tilde = data.weights_sqrt.segment(start, m).matrix().asDiagonal() * data.x.block(start, 0, m, p);
      for (int h = 0; h < k; ++h) {
        P.block((g * k + h) * p, 0, p, p) +=
          X_tilde.transpose() * pattern.inverse_d1.block(h * m, 0, m, m) * X_tilde;
      }
      for (int ell = 0; ell < k; ++ell) {
        matrix<double> Z = pattern.inverse_directions.block(ell * m, 0, m, m) * X_tilde;
        Q_sum += directions[g].signs(ell) * Z.transpose() * pattern.sigma * Z;
      }
    };
    if (data.spatial) {
      // Spatial distances vary by subject; do not key them by visit indices.
      accumulate(kr_pattern(derivatives.cache[g].get(), subject.visits, subject.dist, directions[g].B));
    } else {
      auto found = patterns[g].find(subject.visits);
      if (found == patterns[g].end()) {
        found = patterns[g].emplace(subject.visits,
          kr_pattern(derivatives.cache[g].get(), subject.visits, subject.dist, directions[g].B)).first;
      }
      accumulate(found->second);
    }
  }
  if (!finite_w) {
    Q_sum.setConstant(R_NaN);
  }
  return List::create(Named("P") = as_num_matrix_rcpp(P),
                      Named("Q") = R_NilValue, Named("R") = R_NilValue,
                      Named("S_Q") = as_num_matrix_rcpp(Q_sum));
}

// Obtain P,Q,R elements, or contracted linear components when W is supplied.
// mmrm() always supplies W for linear KR; the pairwise linear blocks are kept
// as an independent reference for tests.
List get_pqr(List mmrm_fit, NumericVector theta, bool linear, Nullable<NumericMatrix> w) {
  if (w.isNotNull()) {
    return get_pqr_contracted(mmrm_fit, theta, NumericMatrix(w));
  }
  kr_data data(mmrm_fit, theta);
  int n_theta = data.n_theta;
  int theta_size_per_group = data.theta_per_group;
  int p = data.p;
  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;
  if (!linear) {
    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 = data.derivatives(!linear);
  for (int i = 0; i < data.n_subjects; i++) {
    kr_subject subject = data.subject(i);
    int start_i = subject.start;
    int n_visits_i = subject.n_visits;
    int subject_group_i = subject.group;
    matrix<double> sigma_inv, sigma_d2, sigma, sigma_inv_d1;

    sigma_inv = derivatives_by_group.cache[subject_group_i]->get_sigma_inverse(subject.visits, subject.dist);
    if (!linear) {
      sigma_d2 = derivatives_by_group.cache[subject_group_i]->get_sigma_derivative2(subject.visits, subject.dist);
    }
    sigma = derivatives_by_group.cache[subject_group_i]->get_sigma(subject.visits, subject.dist);
    sigma_inv_d1 = derivatives_by_group.cache[subject_group_i]->get_inverse_derivative(subject.visits, subject.dist);

    matrix<double> Xi = data.x.block(start_i, 0, n_visits_i, p);
    auto gi_sqrt_root = data.weights_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;
        if (!linear) {
          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;
        }
      }
    }
  }
  List ret = List::create(
    Named("P") = as_num_matrix_rcpp(P),
    Named("Q") = as_num_matrix_rcpp(Q),
    Named("R") = R_NilValue
  );
  if (!linear) {
    ret["R"] = as_num_matrix_rcpp(R);
  }
  return ret;
}

#ifndef KR_COMP_INCLUDED_
#define KR_COMP_INCLUDED_

#include "derivatives.h"

// Observed visits of one subject, see kr_data::subject().
struct kr_subject {
  int start, n_visits, group;
  std::vector<int> visits;
  matrix<double> dist;
};

// Fitted model inputs shared by the pairwise and contracted backends.
struct kr_data {
  matrix<double> x, coords;
  vector<double> weights_sqrt, theta;
  IntegerVector starts, sizes, groups;
  int n_subjects, n_groups, n_visits, n_theta, theta_per_group, p;
  bool spatial;
  std::string cov_type;
  kr_data(List mmrm_fit, NumericVector theta_r) {
    x = as_num_matrix_tmb(as<NumericMatrix>(mmrm_fit["x_matrix"]));
    coords = as_num_matrix_tmb(as<NumericMatrix>(mmrm_fit["coordinates"]));
    weights_sqrt = as_num_vector_tmb(sqrt(as<NumericVector>(mmrm_fit["weights_vector"])));
    theta = as_num_vector_tmb(theta_r);
    starts = mmrm_fit["subject_zero_inds"];
    sizes = mmrm_fit["subject_n_visits"];
    groups = mmrm_fit["subject_groups"];
    n_subjects = mmrm_fit["n_subjects"];
    n_groups = mmrm_fit["n_groups"];
    n_visits = mmrm_fit["n_visits"];
    n_theta = theta_r.size();
    theta_per_group = n_theta / n_groups;
    p = x.cols();
    spatial = as<int>(mmrm_fit["is_spatial_int"]) == 1;
    cov_type = as<std::string>(mmrm_fit["cov_type"]);
  }
  derivatives_cache<double> derivatives(bool second_order) const {
    return derivatives_cache<double>(theta, n_groups, spatial, cov_type, n_visits, second_order);
  }
  // Zero-based group with visit indices (non-spatial) or distances (spatial).
  kr_subject subject(int i) const {
    kr_subject s{starts[i], sizes[i], groups[i] - 1, std::vector<int>(sizes[i]), matrix<double>(0, 0)};
    if (spatial) {
      s.dist = euclidean(matrix<double>(coords.block(s.start, 0, s.n_visits, coords.cols())));
    } else {
      for (int j = 0; j < s.n_visits; ++j) s.visits[j] = int(coords(s.start + j, 0));
    }
    return s;
  }
};

// W = B diag(signs) B'. Keep every direction, including negative eigenvalues
// when W is indefinite: clipping them would change the supplied contraction.
// W is an inverse Hessian, symmetric up to rounding, so use its symmetric part.
struct kr_directions {
  matrix<double> B;
  vector<double> signs;
  explicit kr_directions(const matrix<double>& w_r) {
    matrix<double> w = (w_r + w_r.transpose()) / 2;
    int k = w.rows();
    signs = vector<double>::Ones(k);
    Eigen::LLT<Eigen::MatrixXd> llt(w);
    if (llt.info() == Eigen::Success) {
      B = llt.matrixL();
    } else {
      Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> eig(w);
      if (eig.info() != Eigen::Success) {
        Rcpp::stop("Could not factor covariance parameter covariance matrix.");
      }
      B = eig.eigenvectors();
      for (int ell = 0; ell < k; ++ell) {
        double value = eig.eigenvalues()(ell);
        signs(ell) = value < 0 ? -1 : 1;
        B.col(ell) *= std::sqrt(std::abs(value));
      }
    }
  }
};

// Pattern-level products, reused across subjects in the same nonspatial group.
// Block ell of inverse_directions is sum_h B(h, ell) * d(Sigma^-1)/d(theta_h).
struct kr_pattern {
  matrix<double> sigma, inverse_d1, inverse_directions;
  kr_pattern(derivatives_base<double>* cache, const std::vector<int>& visits,
             const matrix<double>& dist, const matrix<double>& B) {
    sigma = cache->get_sigma(visits, dist);
    inverse_d1 = cache->get_inverse_derivative(visits, dist);
    int m = sigma.rows();
    int k = B.rows();
    inverse_directions = matrix<double>::Zero(k * m, m);
    for (int ell = 0; ell < k; ++ell) {
      for (int h = 0; h < k; ++h) {
        inverse_directions.block(ell * m, 0, m, m) +=
          B(h, ell) * inverse_d1.block(h * m, 0, m, m);
      }
    }
  }
};

#endif

#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

#include "testthat-helpers.h"
#include "kr_comp.h"

context("KR covariance directions") {
  test_that("directions reconstruct positive, semidefinite and indefinite W without truncation") {
    matrix<double> rotation(3, 3);
    rotation << 1, 2, -1, 0, 1, 2, 2, -1, 1;
    for (double last : {0.5, 0.0, -0.5, 1e-12}) {
      vector<double> values(3);
      values << 2, 1, last;
      matrix<double> w = rotation * values.matrix().asDiagonal() * rotation.transpose();
      kr_directions directions(w);
      matrix<double> reconstructed = directions.B * directions.signs.matrix().asDiagonal() * directions.B.transpose();
      expect_true(directions.B.cols() == 3);
      expect_true((reconstructed - w).norm() / w.norm() < 1e-12);
      if (last < 0) expect_true(directions.signs.minCoeff() == -1);
    }
  }
  test_that("pattern directions equal first-derivative combinations with no second derivatives") {
    vector<double> theta(3);
    theta << 0.1, 0.2, 0.4;
    derivatives_nonspatial<double> cache(theta, 2, "us", false);
    matrix<double> w = matrix<double>::Identity(3, 3);
    w(0, 1) = w(1, 0) = 0.3;
    kr_directions directions(w);
    matrix<double> dist(0, 0);
    for (std::vector<int> visits : {std::vector<int>{0, 1}, std::vector<int>{1}}) {
      kr_pattern pattern(&cache, visits, dist, directions.B);
      int m = visits.size();
      for (int ell = 0; ell < 3; ++ell) {
        matrix<double> expected = matrix<double>::Zero(m, m);
        for (int h = 0; h < 3; ++h) {
          expected += directions.B(h, ell) * pattern.inverse_d1.block(h * m, 0, m, m);
        }
        expect_true((expected - pattern.inverse_directions.block(ell * m, 0, m, m)).norm() < 1e-12);
      }
    }
    expect_true(cache.sigmad2_cache.empty());
  }
}

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

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

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

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

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

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

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		#' @param linear (`flag`)\cr whether to omit second derivatives and the R component.
9		#' @param w (`matrix` or `NULL`)\cr covariance of the covariance parameters.
10		#' Supply with `linear = TRUE` to contract Q without constructing its blocks.
11		#'
12		#' @details the function returns a named list, \eqn{P}, \eqn{Q} and \eqn{R}, which corresponds to the
13		#' paper in 1997. The matrices are stacked in columns so that \eqn{P}, \eqn{Q} and \eqn{R} has the same
14		#' column number(number of beta parameters). The number of rows, is dependent on
15		#' the total number of theta and number of groups, if the fit is a grouped mmrm.
16		#' For \eqn{P} matrix, it is stacked sequentially. For \eqn{Q} and \eqn{R} matrix, it is stacked so
17		#' 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.
18		#' \eqn{Q} and \eqn{R} only contains intra-group results and inter-group results should be all zero matrices
19		#' so they are not stacked in the result.
20		#'
21		#' Supplying `w` for linear KR instead retains `P` and contracts Q into the single
22		#' matrix `S_Q`, without storing Q blocks, see the section "Contracted linear covariance
23		#' adjustment" in `vignette("kenward", package = "mmrm")`. The symmetric part of `w` is
24		#' used, and non-finite entries in `w` give a non-finite `S_Q`.
25		#' [mmrm()] always supplies `w` for linear KR. Calling with `linear = TRUE` but without `w`
26		#' gives the pairwise `P` and `Q` blocks, which tests use as an independent reference.
27		#'
28		#' @return Named list with elements:
29		#' - `P`: `matrix` of \eqn{P} component.
30		#' - `Q`: `matrix` of \eqn{Q} component.
31		#' - `R`: `matrix` of \eqn{R} component, or `NULL` when `linear = TRUE`.
32		#' - `S_Q`: contracted Q sum when `w` is supplied; `Q` and `R` are then `NULL`.
33		#'
34		#' @keywords internal
35		h_get_kr_comp <- function(tmb_data, theta, linear = FALSE, w = NULL) {
36	105x	assert_class(tmb_data, "mmrm_tmb_data")
37	105x	assert_class(theta, "numeric")
38	105x	assert_flag(linear)
39	105x	if (!is.null(w)) {
40	37x	if (!linear) {
41	1x	stop(
42	1x	"`w` can only be supplied with `linear = TRUE`, ",
43	1x	"because full Kenward-Roger needs the pairwise Q and R components"
44		)
45		}
46	36x	assert_matrix(w, mode = "numeric", nrows = length(theta), ncols = length(theta))
47		}
48	103x	.Call(`_mmrm_get_pqr`, PACKAGE = "mmrm", tmb_data, theta, linear, w)
49		}
50
51		#' Calculation of Kenward-Roger Degrees of Freedom for Multi-Dimensional Contrast
52		#'
53		#' @description Used in [df_md()] if method is "Kenward-Roger", i.e. for both
54		#' "Kenward-Roger" and "Kenward-Roger-Linear" vcov.
55		#'
56		#' @inheritParams h_df_md_sat
57		#' @inherit h_df_md_sat return
58		#' @keywords internal
59		h_df_md_kr <- function(object, contrast) {
60	85x	assert_class(object, "mmrm")
61	85x	assert_matrix(
62	85x	contrast,
63	85x	mode = "numeric",
64	85x	any.missing = FALSE,
65	85x	ncols = length(component(object, "beta_est"))
66		)
67	85x	if (component(object, "reml") != 1) {
68	1x	stop("Kenward-Roger is only for REML")
69		}
70	84x	df <- h_kr_df(
71	84x	v0 = object$beta_vcov,
72	84x	l = contrast,
73	84x	w = component(object, "theta_vcov"),
74	84x	p = object$kr_comp$P
75		)
76
77	83x	h_test_md(object, contrast, df = df$m, f_stat_factor = df$lambda)
78		}
79
80		#' Calculation of Kenward-Roger Degrees of Freedom for One-Dimensional Contrast
81		#'
82		#' @description Used in [df_1d()] if method is "Kenward-Roger", i.e. for both
83		#' "Kenward-Roger" and "Kenward-Roger-Linear" vcov.
84		#'
85		#' @inheritParams h_df_1d_sat
86		#' @inherit h_df_1d_sat return
87		#' @keywords internal
88		h_df_1d_kr <- function(object, contrast) {
89	47x	assert_class(object, "mmrm")
90	47x	assert_numeric(contrast, len = length(component(object, "beta_est")))
91	47x	if (component(object, "reml") != 1) {
92	1x	stop("Kenward-Roger is only for REML")
93		}
94
95	46x	df <- h_kr_df(
96	46x	v0 = object$beta_vcov,
97	46x	l = matrix(contrast, nrow = 1),
98	46x	w = component(object, "theta_vcov"),
99	46x	p = object$kr_comp$P
100		)
101
102	46x	h_test_1d(object, contrast, df$m)
103		}
104
105		#' Obtain the Adjusted Kenward-Roger degrees of freedom
106		#'
107		#' @description Obtains the adjusted Kenward-Roger degrees of freedom and F statistic scale parameter.
108		#' Used in [h_df_md_kr()] or [h_df_1d_kr].
109		#'
110		#' @param v0 (`matrix`)\cr unadjusted covariance matrix.
111		#' @param l (`matrix`)\cr linear combination matrix.
112		#' @param w (`matrix`)\cr covariance matrix of the estimated covariance parameters.
113		#' @param p (`matrix`)\cr P matrix from [h_get_kr_comp()].
114		#'
115		#' @return Named list with elements:
116		#' - `m`: `numeric` degrees of freedom.
117		#' - `lambda`: `numeric` F statistic scale parameter.
118		#'
119		#' @keywords internal
120		h_kr_df <- function(v0, l, w, p) {
121	159x	n_beta <- ncol(v0)
122	159x	assert_matrix(v0, ncols = n_beta, nrows = n_beta)
123	159x	assert_matrix(l, ncols = n_beta)
124	159x	n_theta <- ncol(w)
125	159x	assert_matrix(w, ncols = n_theta, nrows = n_theta)
126	159x	assert_matrix(p, ncols = n_beta, nrows = n_beta * n_theta)
127	157x	nl <- nrow(l)
128	157x	lv0 <- l %*% v0
129	157x	if (nl == 1L) {
130		# A1 = A2 = a' W a: scalar Satterthwaite based on unadjusted variance.
131	119x	variance <- sum(lv0 * l)
132	119x	if (!isTRUE(variance > 0)) {
133	2x	stop("contrast variance must be positive")
134		}
135	117x	a <- vapply(seq_len(n_theta), function(h) {
136	458x	rows <- (h - 1L) * n_beta + seq_len(n_beta)
137	458x	sum(lv0 * (lv0 %*% p[rows, , drop = FALSE])) / variance
138	117x	}, numeric(1L))
139	117x	return(list(m = 2 / h_quad_form_vec(a, w), lambda = 1))
140		}
141
142		# Normalize in contrast space: H = T T', Z = T^{-1} L Phi.
143		# See vignettes/kenward.Rmd, "Degrees of freedom in contrast space".
144	38x	contrast_vcov <- tcrossprod(lv0, l)
145		# Cholesky alone can accept a singular hypothesis after roundoff, so check
146		# the reciprocal condition number with the same threshold as solve().
147	38x	if (rcond(contrast_vcov) < .Machine$double.eps) {
148	2x	stop("contrast covariance is numerically singular")
149		}
150	36x	contrast_chol <- tryCatch(
151	36x	chol(contrast_vcov),
152	36x	error = function(e) stop("contrast covariance is not positive definite", call. = FALSE)
153		)
154	35x	z <- forwardsolve(t(contrast_chol), lv0)
155	35x	f <- vapply(seq_len(n_theta), function(h) {
156	219x	rows <- (h - 1L) * n_beta + seq_len(n_beta)
157	219x	as.vector(tcrossprod(z %*% p[rows, , drop = FALSE], z))
158	35x	}, numeric(nl * nl))
159	35x	traces <- colSums(f[seq.int(1L, nl * nl, by = nl + 1L), , drop = FALSE])
160	35x	a1 <- h_quad_form_vec(traces, w)
161	35x	a2 <- sum((f %% w) f)
162	35x	b <- 1 / (2 * nl) * (a1 + 6 * a2)
163	35x	e_star <- 1 / (1 - a2 / nl)
164	35x	g <- ((nl + 1) * a1 - (nl + 4) * a2) / ((nl + 2) * a2)
165	35x	denom <- (3 * nl + 2 - 2 * g)
166	35x	c1 <- g / denom
167	35x	c2 <- (nl - g) / denom
168	35x	c3 <- (nl + 2 - g) / denom
169	35x	v_star <- 2 / nl * (1 + c1 * b) / (1 - c2 * b)^2 / (1 - c3 * b)
170	35x	rho <- v_star / (2 * e_star^2)
171	35x	m <- 4 + (nl + 2) / (nl * rho - 1)
172	35x	lambda <- m / (e_star * (m - 2))
173	35x	list(m = m, lambda = lambda)
174		}
175
176		#' Obtain the Adjusted Covariance Matrix
177		#'
178		#' @description Obtains the Kenward-Roger adjusted covariance matrix for the
179		#' coefficient estimates.
180		#' Used in [mmrm()] fitting if vcov is "Kenward-Roger".
181		#'
182		#' @param v (`matrix`)\cr unadjusted covariance matrix.
183		#' @param w (`matrix`)\cr covariance matrix of the estimated covariance parameters.
184		#' @param p (`matrix`)\cr P matrix from [h_get_kr_comp()].
185		#' @param q (`matrix`)\cr Q matrix from [h_get_kr_comp()].
186		#' @param r (`matrix` or `NULL`)\cr R matrix from [h_get_kr_comp()].
187		#' May be `NULL` for the linear approximation.
188		#' @param linear (`flag`)\cr whether to use linear Kenward-Roger approximation.
189		#'
190		#' @details [mmrm()] uses this function for full Kenward-Roger only, and
191		#' [h_var_adj_contracted()] for the linear approximation. The pairwise linear
192		#' path is kept as an independent reference for tests.
193		#'
194		#' @return The matrix of adjusted covariance matrix.
195		#'
196		#' @keywords internal
197		h_var_adj <- function(v, w, p, q, r, linear = FALSE) {
198	63x	assert_flag(linear)
199	63x	n_beta <- ncol(v)
200	63x	assert_matrix(v, nrows = n_beta)
201	63x	n_theta <- ncol(w)
202	63x	assert_matrix(w, nrows = n_theta)
203	63x	n_visits <- ncol(p)
204	63x	theta_per_group <- nrow(q) / nrow(p)
205	63x	n_groups <- n_theta / theta_per_group
206	63x	assert_matrix(p, nrows = n_theta * n_visits)
207	63x	assert_matrix(
208	63x	q,
209	63x	nrows = theta_per_group^2 * n_groups * n_visits,
210	63x	ncols = n_visits
211		)
212	63x	if (!linear) {
213	52x	assert_matrix(
214	52x	r,
215	52x	nrows = theta_per_group^2 * n_groups * n_visits,
216	52x	ncols = n_visits
217		)
218		}
219
220		# See the adjusted covariance formula in vignettes/kenward.Rmd.
221	62x	ret <- v
222	62x	for (i in seq_len(n_theta)) {
223	392x	for (j in seq_len(n_theta)) {
224	3852x	gi <- ceiling(i / theta_per_group)
225	3852x	gj <- ceiling(j / theta_per_group)
226	3852x	iid <- (i - 1) * n_beta + 1
227	3852x	jid <- (j - 1) * n_beta + 1
228	3852x	ii <- i - (gi - 1) * theta_per_group
229	3852x	jj <- j - (gi - 1) * theta_per_group
230	3852x	ijid <- ((ii - 1) * theta_per_group + jj - 1) *
231	3852x	n_beta +
232	3852x	(gi - 1) * n_beta * theta_per_group^2 +
233	3852x	1
234	3852x	if (gi != gj) {
235	1008x	ret <- ret +
236	1008x	2 *
237	1008x	w[i, j] *
238	1008x	v %*%
239	1008x	(-p[iid:(iid + n_beta - 1), ] %*%
240	1008x	v %*%
241	1008x	p[jid:(jid + n_beta - 1), ]) %*%
242	1008x	v
243		} else {
244	2844x	adjustment <- q[ijid:(ijid + n_beta - 1), ] -
245	2844x	p[iid:(iid + n_beta - 1), ] %% v %% p[jid:(jid + n_beta - 1), ]
246	2844x	if (!linear) {
247	2038x	adjustment <- adjustment - 1 / 4 * r[ijid:(ijid + n_beta - 1), ]
248		}
249	2844x	ret <- ret + 2 * w[i, j] * v %% adjustment %% v
250		}
251		}
252		}
253	62x	ret
254		}
255
256		#' Obtain the Contracted Linear Adjusted Covariance Matrix
257		#'
258		#' @description Obtains the linear Kenward-Roger adjusted covariance matrix for the
259		#' coefficient estimates from the `P` matrices and the contracted `S_Q`.
260		#' Used in [mmrm()] fitting if vcov is "Kenward-Roger-Linear".
261		#'
262		#' @inheritParams h_var_adj
263		#' @param s_q (`matrix`)\cr contracted Q sum `S_Q` from [h_get_kr_comp()] called with `w`.
264		#'
265		#' @details See the section "Contracted linear covariance adjustment" in
266		#' `vignette("kenward", package = "mmrm")`. The full `w` is used, so that
267		#' cross-group covariance parameter covariances contribute.
268		#'
269		#' @return The matrix of adjusted covariance matrix.
270		#'
271		#' @keywords internal
272		h_var_adj_contracted <- function(v, w, p, s_q) {
273	35x	n_beta <- ncol(v)
274	35x	assert_matrix(v, nrows = n_beta)
275	35x	n_theta <- ncol(w)
276	35x	assert_matrix(w, nrows = n_theta)
277	35x	assert_matrix(p, nrows = n_beta * n_theta, ncols = n_beta)
278	35x	assert_matrix(s_q, nrows = n_beta, ncols = n_beta)
279		# Column h of p_flat is vec(P_h), so column h of p_bar is vec(sum_j W_hj P_j).
280	34x	p_flat <- matrix(
281	34x	aperm(array(p, c(n_beta, n_theta, n_beta)), c(1L, 3L, 2L)),
282	34x	ncol = n_theta
283		)
284	34x	p_bar <- p_flat %*% t(w)
285	34x	s_p <- matrix(0, n_beta, n_beta)
286	34x	for (h in seq_len(n_theta)) {
287	220x	s_p <- s_p + matrix(p_flat[, h], n_beta) %% v %% matrix(p_bar[, h], n_beta)
288		}
289	34x	v + 2 * v %% (s_q - s_p) %% v
290		}

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

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

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

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

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

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

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

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

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

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

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

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

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

1		#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	6122x	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	34620x	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	11540x	this->theta = theta;
34	34620x	std::iota(std::begin(this->full_visit), std::end(this->full_visit), 0);
35	11540x	this->n_theta = theta.size();
36	11540x	this->chol_full = get_covariance_lower_chol(this->theta, this->n_visits, this->cov_type);
37	11536x	this->chols[full_visit] = this->chol_full;
38	11536x	this->sigma_full = tcrossprod(this->chol_full, true);
39		}
40	678349x	matrix<Type> get_chol(std::vector<int> visits, matrix<Type> dist) {
41	678349x	auto target = this->chols.find(visits);
42	678349x	if (target != this->chols.end()) {
43	627291x	return target->second;
44		} else {
45	51058x	matrix<Type> cov_i = this->get_sigma(visits, dist);
46	51058x	Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > cov_i_chol(cov_i);
47	51058x	matrix<Type> Li = cov_i_chol.matrixL();
48	51058x	this->chols[visits] = Li;
49	51058x	return Li;
50		}
51		}
52	317894x	matrix<Type> get_sigma(std::vector<int> visits, matrix<Type> dist) {
53	317894x	auto target = this->sigmas.find(visits);
54	317894x	if (target != this->sigmas.end()) {
55	247104x	return target->second;
56		} else {
57	70790x	matrix<Type> ret = subset_matrix<matrix<Type>, vector<int>>(sigma_full, visits, visits);
58	70790x	this->sigmas[visits] = ret;
59	70790x	return ret;
60		}
61		}
62	107944x	matrix<Type> get_sigma_inverse(std::vector<int> visits, matrix<Type> dist) {
63	107944x	auto target = this->sigmas_inv.find(visits);
64	107944x	if (target != this->sigmas_inv.end()) {
65	93341x	return target->second;
66		} else {
67	14603x	matrix<Type> ret = this->get_sigma(visits, dist).inverse();
68	14603x	this->sigmas_inv[visits] = ret;
69	14603x	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	382x	lower_chol_spatial(vector<Type> theta, std::string cov_type): theta(theta), cov_type(cov_type) {
85		}
86	47657x	matrix<Type> get_chol(std::vector<int> visits, matrix<Type> dist) {
87	47657x	return get_spatial_covariance_lower_chol(this->theta, dist, this->cov_type);
88		}
89	18725x	matrix<Type> get_sigma(std::vector<int> visits, matrix<Type> dist) {
90	18725x	return tcrossprod(this->get_chol(visits, dist), true);
91		}
92	7881x	matrix<Type> get_sigma_inverse(std::vector<int> visits, matrix<Type> dist) {
93	7881x	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	10630x	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	10630x	int theta_one_group_size = theta.size() / n_groups;
106	22220x	for (int r = 0; r < n_groups; r++) {
107		// Use unique pointers here to better manage resource.
108	11594x	if (is_spatial) {
109	344x	this->cache[r] = std::make_shared<D1>(theta.segment(r * theta_one_group_size, theta_one_group_size), cov_type);
110		} else {
111	11250x	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	10630x	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	7216x	matrix<Type> get_default_chol() {
125	14432x	std::vector<int> visit(this->n_visits);
126	7216x	std::iota(std::begin(visit), std::end(visit), 0);
127	7216x	matrix<Type> dist(2, 2);
128	7216x	dist << 0, 1, 1, 0;
129	7216x	int dim = this->is_spatial?2:this->n_visits;
130	7216x	matrix<Type> covariance_lower_chol = matrix<Type>::Zero(dim * this->n_groups, dim);
131	15188x	for (int r = 0; r < this->n_groups; r++) {
132	7972x	covariance_lower_chol.block(r * dim, 0, dim, dim) = this->cache[r]->get_chol(visit, dist);
133		}
134	14432x	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	9786x	matrix<T> get_unstructured(const vector<T>& theta, int n_visits) {
10	9786x	vector<T> sd_values = exp(theta.head(n_visits));
11	9786x	vector<T> lower_tri_chol_values = theta.tail(theta.size() - n_visits);
12	9786x	matrix<T> covariance_lower_chol = matrix<T>::Zero(n_visits, n_visits);
13	9786x	int k = 0;
14	48472x	for(int i = 0; i < n_visits; i++) {
15	38686x	covariance_lower_chol(i, i) = sd_values(i);
16	96440x	for(int j = 0; j < i; j++){
17	57754x	covariance_lower_chol(i, j) = sd_values(i) * lower_tri_chol_values(k++);
18		}
19		}
20	19572x	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	2196x	const T operator() (int i, int j) const {
30	2196x	return this->corr_values.segment(j, i - j).prod();
31		}
32		};
33		// Homogeneous Ante-dependence Cholesky factor.
34		template <class T>
35	152x	matrix<T> get_ante_dependence(const vector<T>& theta, int n_visits) {
36	152x	T const_sd = exp(theta(0));
37	152x	corr_fun_ante_dependence<T> fun(theta.tail(n_visits - 1));
38	152x	matrix<T> ad_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
39	304x	return const_sd * ad_cor_mat_chol;
40		}
41		// Heterogeneous Ante-dependence Cholesky factor.
42		template <class T>
43	232x	matrix<T> get_ante_dependence_heterogeneous(const vector<T>& theta, int n_visits) {
44	232x	vector<T> sd_values = exp(theta.head(n_visits));
45	232x	corr_fun_ante_dependence<T> fun(theta.tail(n_visits - 1));
46	464x	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	2442x	const T operator() (int i, int j) const {
56	2442x	int index = (i - j) - 1; // Note: We need to start at 0.
57	2442x	return this->corr_values(index);
58		}
59		};
60		// Homogeneous Toeplitz Cholesky factor.
61		template <class T>
62	198x	matrix<T> get_toeplitz(const vector<T>& theta, int n_visits) {
63	198x	T const_sd = exp(theta(0));
64	198x	corr_fun_toeplitz<T> fun(theta.tail(n_visits - 1));
65	198x	matrix<T> toep_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
66	396x	return const_sd * toep_cor_mat_chol;
67		}
68		// Heterogeneous Toeplitz Cholesky factor.
69		template <class T>
70	208x	matrix<T> get_toeplitz_heterogeneous(const vector<T>& theta, int n_visits) {
71	208x	vector<T> sd_values = exp(theta.head(n_visits));
72	208x	corr_fun_toeplitz<T> fun(theta.tail(n_visits - 1));
73	416x	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	11656x	const T operator() (int i, int j) const {
83	11656x	T diff = T((i - j) * 1.0);
84	14324x	return pow(this->corr_values(0), diff); // rho^{\|i-j\|}
85		}
86		};
87		// Homogeneous autoregressive Cholesky factor.
88		template <class T>
89	1728x	matrix<T> get_auto_regressive(const vector<T>& theta, int n_visits) {
90	1728x	T const_sd = exp(theta(0));
91	1728x	corr_fun_autoregressive<T> fun(theta.tail(1));
92	1728x	matrix<T> ar1_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
93	3456x	return const_sd * ar1_cor_mat_chol;
94		}
95		// Heterogeneous autoregressive Cholesky factor.
96		template <class T>
97	220x	matrix<T> get_auto_regressive_heterogeneous(const vector<T>& theta, int n_visits) {
98	220x	vector<T> sd_values = exp(theta.head(n_visits));
99	220x	corr_fun_autoregressive<T> fun(theta.tail(1));
100	440x	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	600x	corr_fun_compound_symmetry(const vector<T>& theta, int n_visits) :
111	600x	corr_values(map_to_cs_cor(theta, n_visits)) {}
112	3546x	const T operator() (int i, int j) const {
113	3546x	return this->corr_values(0); // rho (constant)
114		}
115		};
116		// Homogeneous compound symmetry Cholesky factor.
117		template <class T>
118	340x	matrix<T> get_compound_symmetry(const vector<T>& theta, int n_visits) {
119	340x	T const_sd = exp(theta(0));
120	340x	corr_fun_compound_symmetry<T> fun(theta.tail(1), n_visits);
121	340x	matrix<T> cs_cor_mat_chol = get_corr_mat_chol(n_visits, fun);
122	680x	return const_sd * cs_cor_mat_chol;
123		}
124		// Heterogeneous compound symmetry Cholesky factor.
125		template <class T>
126	256x	matrix<T> get_compound_symmetry_heterogeneous(const vector<T>& theta, int n_visits) {
127	256x	vector<T> sd_values = exp(theta.head(n_visits));
128	256x	corr_fun_compound_symmetry<T> fun(theta.tail(1), n_visits);
129	512x	return get_heterogeneous_cov(sd_values, fun);
130		}
131
132		// Spatial Exponential Cholesky factor.
133		template <class T>
134	29067x	matrix<T> get_spatial_exponential(const vector<T>& theta, const matrix<T>& distance) {
135	29067x	T const_sd = exp(theta(0));
136	29067x	T rho = invlogit(theta(1)); // theta (-Inf, Inf), rho (0, 1), log(rho) (-Inf, 0)
137	29067x	matrix<T> expdist = exp(distance.array() * log(rho));
138	29067x	matrix<T> result = expdist * const_sd;
139	29067x	Eigen::LLT<Eigen::Matrix<T,Eigen::Dynamic,Eigen::Dynamic> > cov_i_chol(result);
140	58134x	return cov_i_chol.matrixL();
141		}
142
143		// Spatial Gaussian Cholesky factor.
144		template <class T>
145	18590x	matrix<T> get_spatial_gaussian(const vector<T>& theta, const matrix<T>& distance) {
146	18590x	T const_sd = exp(theta(0));
147	18590x	T rho = invlogit(theta(1));
148	18590x	matrix<T> expdist = exp(distance.array().square() * log(rho)); // only take square of the dist
149	18590x	matrix<T> result = expdist * const_sd;
150	18590x	Eigen::LLT<Eigen::Matrix<T,Eigen::Dynamic,Eigen::Dynamic> > cov_i_chol(result);
151	37180x	return cov_i_chol.matrixL();
152		}
153
154
155		// Creates a new correlation object dynamically.
156		template <class T>
157	13104x	matrix<T> get_covariance_lower_chol(const vector<T>& theta, int n_visits, std::string cov_type) {
158	13104x	matrix<T> result;
159
160	13104x	if (cov_type == "us") {
161	9782x	result = get_unstructured<T>(theta, n_visits);
162	3322x	} else if (cov_type == "toep") {
163	196x	result = get_toeplitz<T>(theta, n_visits);
164	3126x	} else if (cov_type == "toeph") {
165	206x	result = get_toeplitz_heterogeneous<T>(theta, n_visits);
166	2920x	} else if (cov_type == "ar1") {
167	1726x	result = get_auto_regressive<T>(theta, n_visits);
168	1194x	} else if (cov_type == "ar1h") {
169	218x	result = get_auto_regressive_heterogeneous<T>(theta, n_visits);
170	976x	} else if (cov_type == "ad") {
171	150x	result = get_ante_dependence<T>(theta, n_visits);
172	826x	} else if (cov_type == "adh") {
173	230x	result = get_ante_dependence_heterogeneous<T>(theta, n_visits);
174	596x	} else if (cov_type == "cs") {
175	338x	result = get_compound_symmetry<T>(theta, n_visits);
176	258x	} else if (cov_type == "csh") {
177	254x	result = get_compound_symmetry_heterogeneous<T>(theta, n_visits);
178		} else {
179	2x	Rf_error("%s", ("Unknown covariance type '" + cov_type + "'.").c_str());
180		}
181
182	13100x	return result;
183		}
184
185		// Creates a new spatial covariance cholesky.
186		template <class T>
187	47657x	matrix<T> get_spatial_covariance_lower_chol(const vector<T>& theta, const matrix<T>& distance, std::string cov_type) {
188	47657x	matrix<T> result;
189	47657x	if (cov_type == "sp_exp") {
190	29067x	result = get_spatial_exponential<T>(theta, distance);
191	18590x	} else if (cov_type == "sp_gau") {
192	18590x	result = get_spatial_gaussian<T>(theta, distance);
193		} else {
194	!	Rf_error("%s", ("Unknown spatial covariance type '" + cov_type + "'.").c_str());
195		}
196	47657x	return result;
197		}
198
199		#endif

1		#ifndef DERIVATIVE_INCLUDED_
2		#define DERIVATIVE_INCLUDED_
3
4		#include "chol_cache.h"
5
6		using namespace Rcpp;
7		using std::string;
8		// Struct chol to obtain the cholesky factor given theta.
9		// The reason to have it is that we need a functor that need only theta to
10		// obtain the derivatives from autodiff.
11		// Only non-spatial covariance structure here.
12		struct chol {
13		int dim_cov_mat;
14		string cov_type;
15	580x	chol(int dim, string cov): dim_cov_mat(dim), cov_type(cov) {};
16		template <class T>
17	1274x	vector<T> operator() (vector<T> &theta) {
18	1274x	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	291x	chol_jacobian(int dim, string cov): dim_cov_mat(dim), cov_type(cov), mychol(dim, cov) {};
28		template<class T>
29	59x	vector<T> operator() (vector<T> &theta) {
30	59x	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	289x	std::map<std::string, matrix<Type>> derivatives(int n_visits, std::string cov_type, vector<Type> theta, bool second_order) {
39	289x	std::map<std::string, matrix<Type>> ret;
40	289x	chol chol_obj(n_visits, cov_type);
41	289x	chol_jacobian chol_jac_obj(n_visits, cov_type);
42	289x	matrix<Type> l = chol_obj(theta).matrix();
43	289x	l.resize(n_visits, n_visits);
44	289x	vector<Type> chol_d1_vec = autodiff::jacobian(chol_obj, theta).vec(); // chol_d1_vec is (dim * dim * l_theta)
45	289x	vector<Type> chol_d2_vec;
46	289x	if (second_order) {
47	57x	chol_d2_vec = autodiff::jacobian(chol_jac_obj, theta).vec();
48		}
49	289x	matrix<Type> ret_d1 = matrix<Type>(n_visits * theta.size(), n_visits);
50	289x	matrix<Type> ret_d2;
51	289x	if (second_order) {
52	57x	ret_d2.resize(n_visits * theta.size() * theta.size(), n_visits);
53		}
54	289x	int n_visits_sq = n_visits * n_visits;
55	5423x	for (int i = 0; i < theta.size(); i++) {
56	1809x	matrix<Type> ld1 = chol_d1_vec.segment(i * n_visits_sq, n_visits_sq).matrix();
57	1809x	ld1.resize(n_visits, n_visits);
58	1809x	matrix<Type> ld1_lt = ld1 * l.transpose();
59	1809x	auto sigma_d1_i = ld1_lt + ld1_lt.transpose();
60	1809x	ret_d1.block(i * n_visits, 0, n_visits, n_visits) = sigma_d1_i;
61	1809x	if (!second_order) {
62	1516x	continue;
63		}
64	2324x	for (int j = 0; j < theta.size(); j++) {
65	2031x	matrix<Type> ld2 = chol_d2_vec.segment( (j * theta.size() + i) * n_visits_sq, n_visits_sq).matrix();
66	2031x	matrix<Type> ld1_j = chol_d1_vec.segment(j * n_visits_sq, n_visits_sq).matrix();
67	2031x	ld2.resize(n_visits, n_visits);
68	2031x	ld1_j.resize(n_visits, n_visits);
69	2031x	auto ld2_lt = ld2 * l.transpose();
70	2031x	auto ld1_ld1j = ld1 * ld1_j.transpose();
71	2031x	auto sigma_d2_ij = ld2_lt + ld2_lt.transpose() + ld1_ld1j + ld1_ld1j.transpose();
72	2031x	ret_d2.block((i * theta.size() + j) * n_visits, 0, n_visits, n_visits) = sigma_d2_ij;
73		}
74		}
75	289x	ret["derivative1"] = ret_d1;
76	289x	if (second_order) {
77	114x	ret["derivative2"] = ret_d2;
78		}
79	578x	return ret;
80		}
81		// Base class of spatial and non-spatial derivatives.
82		template <class Type>
83		struct derivatives_base: virtual lower_chol_base<Type> {
84		virtual matrix<Type> get_inverse_chol(std::vector<int> visits, matrix<Type> dist) = 0;
85		virtual matrix<Type> get_sigma_derivative1(std::vector<int> visits, matrix<Type> dist) = 0;
86		virtual matrix<Type> get_sigma_derivative2(std::vector<int> visits, matrix<Type> dist) = 0;
87		virtual matrix<Type> get_inverse_derivative(std::vector<int> visits, matrix<Type> dist) = 0;
88		// Create virtual destructor to avoid the default desctructor being called.
89	326x	virtual ~derivatives_base() {};
90		};
91
92		// Struct derivatives_nonspatial is created to get the derivatives with cache.
93		// The main reason to have it is that we nearly always have duplicated visits
94		// and the inverse of a matrix is calculation expensive. In addition, we can save
95		// the resource needed for select matrix calculations.
96		template <class Type>
97		struct derivatives_nonspatial: public lower_chol_nonspatial<Type>, virtual derivatives_base<Type> {
98		std::map<std::vector<int>, matrix<Type>> inverse_chol_cache;
99		std::map<std::vector<int>, matrix<Type>> sigmad1_cache;
100		std::map<std::vector<int>, matrix<Type>> sigmad2_cache;
101		std::map<std::vector<int>, matrix<Type>> sigma_inverse_d1_cache;
102		// Whether second order derivatives are calculated and cached.
103		bool second_order = false;
104		derivatives_nonspatial() {
105		// This default constructor is needed because the use of `[]` in map.
106		}
107		// Constructor from theta, n_visits and cov_type, and cache full_visits values.
108	288x	derivatives_nonspatial(vector<Type> theta, int n_visits, std::string cov_type, bool second_order): lower_chol_nonspatial<Type>(theta, n_visits, cov_type), second_order(second_order) {
109	288x	std::map<std::string, tmbutils::matrix<Type>> allret = derivatives<Type>(this->n_visits, this->cov_type, this->theta, second_order);
110	288x	this->sigmad1_cache[this->full_visit] = allret.at("derivative1");
111	288x	if (second_order) {
112	114x	this->sigmad2_cache[this->full_visit] = allret.at("derivative2");
113		}
114		}
115		// Cache and return the first order derivatives using select matrix.
116	3426x	matrix<Type> get_sigma_derivative1(std::vector<int> visits, matrix<Type> dist) override {
117	3426x	auto target = this->sigmad1_cache.find(visits);
118	3426x	if (target != this->sigmad1_cache.end()) {
119	264x	return target->second;
120		} else {
121	3162x	int n_visits_i = visits.size();
122	3162x	matrix<Type> ret = matrix<Type>(this->n_theta * n_visits_i, n_visits_i);
123	24554x	for (int i = 0; i < this->n_theta; i++) {
124	21392x	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);
125		}
126	3162x	this->sigmad1_cache[visits] = ret;
127	3162x	return ret;
128		}
129		}
130		// Cache and return the second order derivatives using select matrix.
131	8077x	matrix<Type> get_sigma_derivative2(std::vector<int> visits, matrix<Type> dist) override {
132	8077x	if (!this->second_order) {
133	!	Rcpp::stop("Second derivatives were not initialized in this cache.");
134		}
135	8077x	auto target = this->sigmad2_cache.find(visits);
136	8077x	if (target != this->sigmad2_cache.end()) {
137	7344x	return target->second;
138		} else {
139	733x	int n_visits_i = visits.size();
140	733x	int theta_sq = this->n_theta * this->n_theta;
141	733x	matrix<Type> ret = matrix<Type>(theta_sq * n_visits_i, n_visits_i);
142	28381x	for (int i = 0; i < theta_sq; i++) {
143	27648x	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);
144		}
145	733x	this->sigmad2_cache[visits] = ret;
146	733x	return ret;
147		}
148		}
149		// Cache and return the lower cholesky factor of inverse of sigma using select matrix.
150	7486x	matrix<Type> get_inverse_chol(std::vector<int> visits, matrix<Type> dist) override {
151	7486x	auto target = this->inverse_chol_cache.find(visits);
152	7486x	if (target != this->inverse_chol_cache.end()) {
153	6916x	return target->second;
154		} else {
155	570x	matrix<Type> sigmainv = this->get_sigma_inverse(visits, dist);
156	570x	Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > sigma_inv_chol(sigmainv);
157	570x	matrix<Type> Li = sigma_inv_chol.matrixL();
158	570x	this->inverse_chol_cache[visits] = Li;
159	570x	return Li;
160		}
161		}
162		// Cache and return the first order derivatives of inverse of sigma using select matrix.
163	36565x	matrix<Type> get_inverse_derivative(std::vector<int> visits, matrix<Type> dist) override {
164	36565x	auto target = this->sigma_inverse_d1_cache.find(visits);
165	36565x	if (target != this->sigma_inverse_d1_cache.end()) {
166	33157x	return target->second;
167		} else {
168	3408x	auto sigma_d1 = this->get_sigma_derivative1(visits, dist);
169	3408x	matrix<Type> sigma_inv_d1(sigma_d1.rows(), sigma_d1.cols());
170	3408x	int n_visits_i = visits.size();
171	3408x	auto sigma_inv = this->get_sigma_inverse(visits, dist);
172	26425x	for (int r = 0; r < this->n_theta; r++) {
173	23017x	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;
174		}
175	3408x	this->sigma_inverse_d1_cache[visits] = sigma_inv_d1;
176	3408x	return sigma_inv_d1;
177		}
178		}
179		};
180
181		// derivatives_sp_exp struct is created to obtain the exact derivatives of spatial exponential
182		// covariance structure, and its inverse.
183		// No caching is used because the distance can be hardly the same for spatial covariance
184		// structures.
185		template <class Type>
186		struct derivatives_sp_exp: public lower_chol_spatial<Type>, virtual derivatives_base<Type> {
187		Type const_sd;
188		Type rho;
189		Type logrho;
190		derivatives_sp_exp() {
191		// This default constructor is needed because the use of `[]` in maps.
192		}
193		// Initialize the theta values; the reason to have theta is that for a fit, the theta
194		// is the same for all subjects, while the distance between each visits for each subject
195		// can be different.
196	76x	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))) {
197	76x	this->logrho = log(this->rho);
198		}
199		// Obtain first order derivatives
200	2759x	matrix<Type> get_sigma_derivative1(std::vector<int> visits, matrix<Type> dist) override {
201	2759x	matrix<Type> ret(2 * dist.rows(), dist.cols());
202		// partial sigma / partial theta_1 = sigma.
203	2759x	auto sigma = this->get_sigma(visits, dist);
204	2759x	ret.block(0, 0, dist.rows(), dist.cols()) = sigma;
205	2759x	ret.block(dist.rows(), 0, dist.rows(), dist.cols()) = sigma.array() * dist.array() * (1 - this->rho);
206	5518x	return ret;
207		}
208		// Obtain second order derivatives.
209	986x	matrix<Type> get_sigma_derivative2(std::vector<int> visits, matrix<Type> dist) override {
210	986x	matrix<Type> ret(4 * dist.rows(), dist.cols());
211	986x	auto sigma = this->get_sigma(visits, dist);
212	986x	ret.block(0, 0, dist.rows(), dist.cols()) = sigma;
213	986x	Type rho_r = 1 - this->rho;
214	986x	auto dtheta1dtheta2 = sigma.array() * dist.array() * rho_r;
215	986x	ret.block(dist.rows(), 0, dist.rows(), dist.cols()) = dtheta1dtheta2;
216	986x	ret.block(dist.rows() * 2, 0, dist.rows(), dist.cols()) = dtheta1dtheta2;
217	986x	matrix<Type> dtheta2s = dtheta1dtheta2 * (dist.array() * rho_r - this->rho);
218	986x	ret.block(dist.rows() * 3, 0, dist.rows(), dist.cols()) = dtheta2s;
219	1972x	return ret;
220		}
221		// Obtain the lower cholesky factor of inverse of sigma using select matrix.
222	394x	matrix<Type> get_inverse_chol(std::vector<int> visits, matrix<Type> dist) override {
223	394x	auto sigmainv = this->get_sigma_inverse(visits, dist);
224	394x	Eigen::LLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > sigma_inv_chol(sigmainv);
225	394x	matrix<Type> Li = sigma_inv_chol.matrixL();
226	788x	return Li;
227		}
228		// Obtain first order derivatives for inverse of sigma.
229	4925x	matrix<Type> get_inverse_derivative(std::vector<int> visits, matrix<Type> dist) override {
230	4925x	matrix<Type> sigma_inv_d1 = matrix<Type>::Zero(2 * dist.rows(), dist.cols());
231	4925x	auto sigma_inv = this->get_sigma_inverse(visits, dist);
232	4925x	auto sigma_d1 = this->get_sigma_derivative1(visits, dist);
233	14775x	for (int r = 0; r < 2; r++) {
234	9850x	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;
235		}
236	9850x	return sigma_inv_d1;
237		}
238		};
239
240
241		// derivatives_sp_gau struct is created to obtain the exact derivatives of spatial Gaussian
242		// covariance structure, and its inverse.
243		// Inherit from sp_exp because of similar parameterization; only the sigma
244		// derivatives differ (square the distance). Inverse / inverse-derivative logic
245		// is reused from the base via the virtual get_sigma_derivative1 dispatch.
246		template <class Type>
247		struct derivatives_sp_gau: public derivatives_sp_exp<Type> {
248		derivatives_sp_gau() {
249		// This default constructor is needed because the use of `[]` in maps.
250		}
251	20x	derivatives_sp_gau(vector<Type> theta, std::string cov_type): derivatives_sp_exp<Type>(theta, cov_type) {
252		}
253		// Obtain first order derivatives
254	2167x	matrix<Type> get_sigma_derivative1(std::vector<int> visits, matrix<Type> dist) override {
255	2167x	matrix<Type> ret(2 * dist.rows(), dist.cols());
256		// partial sigma / partial theta_1 = sigma.
257	2167x	auto sigma = this->get_sigma(visits, dist);
258	2167x	ret.block(0, 0, dist.rows(), dist.cols()) = sigma;
259	2167x	ret.block(dist.rows(), 0, dist.rows(), dist.cols()) = sigma.array() * dist.array().square() * (1 - this->rho);
260	4334x	return ret;
261		}
262		// Obtain second order derivatives.
263	591x	matrix<Type> get_sigma_derivative2(std::vector<int> visits, matrix<Type> dist) override {
264	591x	matrix<Type> ret(4 * dist.rows(), dist.cols());
265	591x	auto sigma = this->get_sigma(visits, dist);
266	591x	ret.block(0, 0, dist.rows(), dist.cols()) = sigma;
267	591x	Type rho_r = 1 - this->rho;
268	591x	auto dtheta1dtheta2 = sigma.array() * dist.array().square() * rho_r;
269	591x	ret.block(dist.rows(), 0, dist.rows(), dist.cols()) = dtheta1dtheta2;
270	591x	ret.block(dist.rows() * 2, 0, dist.rows(), dist.cols()) = dtheta1dtheta2;
271	591x	matrix<Type> dtheta2s = dtheta1dtheta2 * (dist.array().square() * rho_r - this->rho);
272	591x	ret.block(dist.rows() * 3, 0, dist.rows(), dist.cols()) = dtheta2s;
273	1182x	return ret;
274		}
275		};
276
277		// Cache of per-group derivatives objects, with cov_type-aware dispatch for
278		// spatial structures (sp_exp / sp_gau). Mirrors `cache_obj` but selects the
279		// concrete spatial subclass at runtime so analytic derivatives match cov_type.
280		template <class Type>
281		struct derivatives_cache {
282		std::map<int, std::shared_ptr<derivatives_base<Type>>> cache;
283		int n_groups;
284		bool is_spatial;
285		int n_visits;
286	282x	derivatives_cache(vector<Type> theta, int n_groups, bool is_spatial, std::string cov_type, int n_visits, bool second_order): n_groups(n_groups), is_spatial(is_spatial), n_visits(n_visits) {
287	282x	int theta_one_group_size = theta.size() / n_groups;
288	603x	for (int r = 0; r < n_groups; r++) {
289	321x	vector<Type> theta_r = theta.segment(r * theta_one_group_size, theta_one_group_size);
290	321x	if (is_spatial) {
291	37x	if (cov_type == "sp_exp") {
292	17x	this->cache[r] = std::make_shared<derivatives_sp_exp<Type>>(theta_r, cov_type);
293	20x	} else if (cov_type == "sp_gau") {
294	20x	this->cache[r] = std::make_shared<derivatives_sp_gau<Type>>(theta_r, cov_type);
295		} else {
296	!	Rcpp::stop("Unknown spatial covariance type '" + cov_type + "'.");
297		}
298		} else {
299	284x	this->cache[r] = std::make_shared<derivatives_nonspatial<Type>>(theta_r, n_visits, cov_type, second_order);
300		}
301		}
302		}
303		};
304		#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	241661x	T1 subset_matrix(T1 input, T2 index1, T2 index2) {
16		#if EIGEN_VERSION_AT_LEAST(3,4,0)
17	241661x	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	241661x	return ret;
27		}
28
29		template <typename T1, typename T2>
30	242229x	T1 subset_matrix(T1 input, T2 index1) {
31		#if EIGEN_VERSION_AT_LEAST(5,0,0)
32		T1 ret = input(index1, Eigen::indexing::all);
33		#elif EIGEN_VERSION_AT_LEAST(3,4,0)
34	242229x	T1 ret = input(index1, Eigen::all);
35		#else
36		T1 ret(index1.size(), input.cols());
37		for (decltype(index1.size()) i = 0; i < index1.size(); i++) {
38		for (int j = 0; j < input.cols(); j++) {
39		ret(i, j) = input(index1[i], j);
40		}
41		}
42		#endif
43	242229x	return ret;
44		}
45
46
47		// Conversion from Rcpp vector/matrix to eigen vector/matrix
48		template <typename T1, typename T2>
49	488510x	T1 as_vector(T2 input) {
50	488510x	T1 ret(input.size());
51	1503114x	for (int i = 0; i < input.size(); i++) {
52	1014604x	ret(i) = input(i);
53		}
54	488510x	return ret;
55		}
56
57		template <typename T1, typename T2>
58	212108x	T1 as_matrix(T2 input) {
59	212108x	T1 ret(input.rows(), input.cols());
60	3249888x	for (int i = 0; i < input.rows(); i++) {
61	33907096x	for (int j = 0; j < input.cols(); j++) {
62	30869316x	ret(i,j) = input(i,j);
63		}
64		}
65	212108x	return ret;
66		}
67
68		template <typename T>
69	726686x	T segment(T input, int start, int n) {
70	726686x	T ret(n);
71	3630604x	for (int i = 0, j = start; i < n; i++, j++) {
72	2903918x	ret(i) = input(j);
73		}
74	726686x	return ret;
75		}
76
77		// Calculate tcrossprod(lower_chol) = lower_chol * t(lower_chol).
78		// If complete, then adds the upper triangular part to the result as well.
79		// By default only the lower triangular part is populated, as this should be
80		// sufficient for downstream use of the result in most cases.
81		template <class Type>
82	727930x	matrix<Type> tcrossprod(const matrix<Type>& lower_chol, bool complete = false) {
83	727930x	int n = lower_chol.rows();
84	727930x	matrix<Type> result = matrix<Type>::Zero(n, n);
85	727930x	result.template selfadjointView<Eigen::Lower>().rankUpdate(lower_chol);
86	727930x	if (complete) {
87	24638x	result.template triangularView<Eigen::Upper>() = result.transpose();
88		}
89	727930x	return result;
90		}
91
92		// Calculate crossprod(x) = t(x) * x.
93		// Only the lower triangular part is populated, as this should be
94		// sufficient for downstream use of the result in most cases.
95		// Note that x does not need to be symmetric or square.
96		template <class Type>
97	1406584x	matrix<Type> crossprod(const matrix<Type>& x) {
98	1406584x	int n = x.cols();
99	1406584x	matrix<Type> result = matrix<Type>::Zero(n, n);
100	1406584x	result.template selfadjointView<Eigen::Lower>().rankUpdate(x.transpose());
101	1406584x	return result;
102		}
103
104		// Mapping from real values to correlation parameters in (-1, 1).
105		template <class T>
106	2750x	vector<T> map_to_cor(const vector<T>& theta) {
107	4866x	return theta / sqrt(T(1.0) + theta * theta);
108		}
109
110		// Mapping from real values to correlation parameters in (-1/(m-1), 1) for compound symmetry.
111		template <class T>
112	600x	vector<T> map_to_cs_cor(const vector<T>& theta, int n_visits) {
113	600x	T a = T(1.0) / (T(n_visits) - T(1.0));
114	730x	return invlogit(theta) * (T(1.0) + a) - a;
115		}
116
117		// Generic correlation function class containing and initializing correlation
118		// values from variance parameters theta.
119		template <class T>
120		struct generic_corr_fun {
121		const vector<T> corr_values;
122
123	2746x	generic_corr_fun(const vector<T>& theta) :
124	2746x	corr_values(map_to_cor(theta)) {}
125		};
126
127		// Correlation function based Cholesky factor of correlation matrix.
128		// This is used directly for homogeneous covariance matrices.
129		template <class T, template<class> class F>
130	3338x	matrix<T> get_corr_mat_chol(int n_visits, const F<T>& corr_fun) {
131	3338x	matrix<T> correlation(n_visits, n_visits);
132	3338x	correlation.setIdentity();
133	16588x	for(int i = 0; i < n_visits; i++) {
134	33000x	for(int j = 0; j < i; j++){
135	19750x	correlation(i, j) = corr_fun(i, j);
136		}
137		}
138	3338x	Eigen::LLT<Eigen::Matrix<T,Eigen::Dynamic,Eigen::Dynamic> > correlation_chol(correlation);
139	3338x	matrix<T> L = correlation_chol.matrixL();
140	6676x	return L;
141		}
142
143		// Heterogeneous covariance matrix calculation given vector of standard deviations (sd_values)
144		// and a correlation function (corr_fun).
145		template <class T, template<class> class F>
146	917x	matrix<T> get_heterogeneous_cov(const vector<T>& sd_values, const F<T>& corr_fun) {
147	917x	matrix<T> correlation_chol = get_corr_mat_chol(sd_values.size(), corr_fun);
148	917x	Eigen::DiagonalMatrix<T,Eigen::Dynamic,Eigen::Dynamic> D = sd_values.matrix().asDiagonal();
149	917x	matrix<T> result = D * correlation_chol;
150	1834x	return result;
151		}
152
153		// Obtain the Euclidean distance
154		template <class T>
155	34483x	matrix<T> euclidean(const matrix<T>& coordinates) {
156	34483x	matrix<T> result(coordinates.rows(), coordinates.rows());
157	128992x	for (int i = 0; i < coordinates.rows(); i++) {
158	94509x	result(i, i) = 0;
159	190860x	for (int j = 0; j < i; j ++) {
160	96351x	vector<T> diff = coordinates.row(i) - coordinates.row(j);
161	96351x	T d = sqrt((diff * diff).sum());
162	96351x	result(i, j) = d;
163	96351x	result(j, i) = d;
164		}
165		}
166	34483x	return result;
167		}
168
169		// Element wise power function of a matrix
170		template <class T>
171	792x	Eigen::Matrix<T, -1, -1> cpow(const Eigen::Matrix<T, -1, -1> & input, double p) {
172	792x	Eigen::Matrix<T, -1, -1> ret = Eigen::Matrix<T, -1, -1>(input.rows(), input.cols());
173	2954x	for (int i = 0; i < ret.rows(); i ++) {
174	4332x	for (int j = 0; j < ret.cols(); j++) {
175	2170x	ret(i, j) = std::pow(input(i, j), p);
176		}
177		}
178	792x	return ret;
179		}
180
181		// Calculate the square root of the pseudo inverse of a matrix
182		// adapted from the method for calculating the pseudo-Inverse as recommended by the Eigen developers
183		template<typename T>
184	790x	matrix<T> pseudoInverseSqrt(const matrix<T> &input, double epsilon = std::numeric_limits<double>::epsilon()) {
185	790x	Eigen::Matrix<T, -1, -1> eigen_mat = as_matrix<Eigen::Matrix<T, -1, -1>, matrix<T>>(input);
186	790x	Eigen::JacobiSVD< Eigen::Matrix<T, -1, -1> > svd(eigen_mat ,Eigen::ComputeFullU \| Eigen::ComputeFullV);
187	790x	double tolerance = epsilon * std::max(input.cols(), input.rows()) *svd.singularValues().array().abs()(0);
188	790x	auto singular_vals = Matrix<T,-1,-1>((svd.singularValues().array() > tolerance).select(svd.singularValues().array().inverse(), 0).matrix());
189	790x	Eigen::Matrix<T, -1, -1> ret_eigen = svd.matrixV() * cpow(singular_vals, 0.5).asDiagonal() * svd.matrixU().adjoint();
190	1580x	return as_matrix<matrix<T>, Eigen::Matrix<T, -1, -1>>(ret_eigen);
191		}
192
193		#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, bool linear, Nullable<NumericMatrix> w);
12	103x	RcppExport SEXP _mmrm_get_pqr(SEXP mmrm_fit_SEXP, SEXP theta_SEXP, SEXP linear_SEXP, SEXP w_SEXP) {
13	103x	BEGIN_RCPP
14	103x	Rcpp::RObject rcpp_result_gen;
15	103x	Rcpp::RNGScope rcpp_rngScope_gen;
16	103x	Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
17	103x	Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
18	103x	Rcpp::traits::input_parameter< bool >::type linear(linear_SEXP);
19	103x	rcpp_result_gen = Rcpp::wrap(get_pqr(mmrm_fit, theta, linear, Nullable<NumericMatrix>(w_SEXP)));
20	103x	return rcpp_result_gen;
21	103x	END_RCPP
22		}
23
24		List get_jacobian(List mmrm_fit, NumericVector theta, NumericMatrix beta_vcov);
25	138x	RcppExport SEXP _mmrm_get_jacobian(SEXP mmrm_fit_SEXP, SEXP theta_SEXP, SEXP beta_vcov_SEXP) {
26	138x	BEGIN_RCPP
27	138x	Rcpp::RObject rcpp_result_gen;
28	138x	Rcpp::RNGScope rcpp_rngScope_gen;
29	138x	Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
30	138x	Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
31	138x	Rcpp::traits::input_parameter< NumericMatrix >::type beta_vcov(beta_vcov_SEXP);
32	138x	rcpp_result_gen = Rcpp::wrap(get_jacobian(mmrm_fit, theta, beta_vcov));
33	138x	return rcpp_result_gen;
34	138x	END_RCPP
35		}
36
37		List get_empirical(List mmrm_fit, NumericVector theta, NumericVector beta, NumericMatrix beta_vcov, std::string type);
38	40x	RcppExport SEXP _mmrm_get_empirical(SEXP mmrm_fit_SEXP, SEXP theta_SEXP, SEXP beta_SEXP, SEXP beta_vcov_SEXP, SEXP type_SEXP) {
39	40x	BEGIN_RCPP
40	40x	Rcpp::RObject rcpp_result_gen;
41	40x	Rcpp::RNGScope rcpp_rngScope_gen;
42	40x	Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
43	40x	Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
44	40x	Rcpp::traits::input_parameter< NumericVector >::type beta(beta_SEXP);
45	40x	Rcpp::traits::input_parameter< NumericMatrix >::type beta_vcov(beta_vcov_SEXP);
46	40x	Rcpp::traits::input_parameter< std::string >::type type(type_SEXP);
47	40x	rcpp_result_gen = Rcpp::wrap(get_empirical(mmrm_fit, theta, beta, beta_vcov, type));
48	40x	return rcpp_result_gen;
49	40x	END_RCPP
50		}
51
52		List predict(List mmrm_fit, NumericVector theta, NumericVector beta, NumericMatrix beta_vcov);
53	1704x	RcppExport SEXP _mmrm_predict(SEXP mmrm_fit_SEXP, SEXP theta_SEXP, SEXP beta_SEXP, SEXP beta_vcov_SEXP) {
54	1704x	BEGIN_RCPP
55	1704x	Rcpp::RObject rcpp_result_gen;
56	1704x	Rcpp::RNGScope rcpp_rngScope_gen;
57	1704x	Rcpp::traits::input_parameter< List >::type mmrm_fit(mmrm_fit_SEXP);
58	1704x	Rcpp::traits::input_parameter< NumericVector >::type theta(theta_SEXP);
59	1704x	Rcpp::traits::input_parameter< NumericVector >::type beta(beta_SEXP);
60	1704x	Rcpp::traits::input_parameter< NumericMatrix >::type beta_vcov(beta_vcov_SEXP);
61	1704x	rcpp_result_gen = Rcpp::wrap(predict(mmrm_fit, theta, beta, beta_vcov));
62	1704x	return rcpp_result_gen;
63	1704x	END_RCPP
64		}
65
66
67		RcppExport SEXP run_testthat_tests(SEXP);
68
69		static const R_CallMethodDef CallEntries[] = {
70		{"_mmrm_get_pqr", (DL_FUNC) &_mmrm_get_pqr, 4},
71		{"_mmrm_get_jacobian", (DL_FUNC) &_mmrm_get_jacobian, 3},
72		{"_mmrm_get_empirical", (DL_FUNC) &_mmrm_get_empirical, 5},
73		{"_mmrm_predict", (DL_FUNC) &_mmrm_predict, 4},
74		{"run_testthat_tests", (DL_FUNC) &run_testthat_tests, 1},
75		TMB_CALLDEFS,
76		{NULL, NULL, 0}
77		};
78
79	4x	RcppExport void R_init_mmrm(DllInfo *dll) {
80	4x	R_registerRoutines(dll, NULL, CallEntries, NULL, NULL);
81	4x	R_useDynamicSymbols(dll, FALSE);
82		#ifdef TMB_CCALLABLES
83	4x	TMB_CCALLABLES("mmrm");
84		#endif
85		}

1		#include "kr_comp.h"
2
3		using namespace Rcpp;
4		using std::string;
5		// Linear KR contracts Q before storing it, retaining P. See "Contracted linear
6		// covariance adjustment" in vignettes/kenward.Rmd.
7	35x	List get_pqr_contracted(List mmrm_fit, NumericVector theta, NumericMatrix w) {
8	35x	kr_data data(mmrm_fit, theta);
9	35x	int k = data.theta_per_group;
10	35x	int p = data.p;
11	35x	matrix<double> W = as_num_matrix_tmb(w);
12		// Non-finite W (a convergence problem that mmrm() only warns about) yields
13		// a non-finite S_Q, as the pairwise contraction would.
14	35x	bool finite_w = W.allFinite();
15	35x	matrix<double> P = matrix<double>::Zero(p * data.n_theta, p);
16	35x	matrix<double> Q_sum = matrix<double>::Zero(p, p);
17	35x	auto derivatives = data.derivatives(false);
18	35x	std::vector<kr_directions> directions;
19	35x	std::vector<std::map<std::vector<int>, kr_pattern>> patterns(data.n_groups);
20	74x	for (int g = 0; g < data.n_groups; ++g) {
21		// Only same-group Q blocks are nonzero. Cross-group W enters S_P in
22		// h_var_adj_contracted(), using the full W.
23	39x	matrix<double> w_g = finite_w ? matrix<double>(W.block(g * k, g * k, k, k)) : matrix<double>(matrix<double>::Zero(k, k));
24	39x	directions.emplace_back(w_g);
25		}
26	6804x	for (int i = 0; i < data.n_subjects; ++i) {
27	6769x	kr_subject subject = data.subject(i);
28	6769x	int start = subject.start, m = subject.n_visits, g = subject.group;
29	6769x	auto accumulate = [&](const kr_pattern& pattern) {
30	6769x	matrix<double> X_tilde = data.weights_sqrt.segment(start, m).matrix().asDiagonal() * data.x.block(start, 0, m, p);
31	47225x	for (int h = 0; h < k; ++h) {
32	40456x	P.block((g * k + h) * p, 0, p, p) +=
33	80912x	X_tilde.transpose() * pattern.inverse_d1.block(h * m, 0, m, m) * X_tilde;
34		}
35	47225x	for (int ell = 0; ell < k; ++ell) {
36	40456x	matrix<double> Z = pattern.inverse_directions.block(ell * m, 0, m, m) * X_tilde;
37	40456x	Q_sum += directions[g].signs(ell) * Z.transpose() * pattern.sigma * Z;
38		}
39		};
40	6769x	if (data.spatial) {
41		// Spatial distances vary by subject; do not key them by visit indices.
42	1379x	accumulate(kr_pattern(derivatives.cache[g].get(), subject.visits, subject.dist, directions[g].B));
43		} else {
44	5390x	auto found = patterns[g].find(subject.visits);
45	5390x	if (found == patterns[g].end()) {
46	844x	found = patterns[g].emplace(subject.visits,
47	844x	kr_pattern(derivatives.cache[g].get(), subject.visits, subject.dist, directions[g].B)).first;
48		}
49	5390x	accumulate(found->second);
50		}
51		}
52	35x	if (!finite_w) {
53	3x	Q_sum.setConstant(R_NaN);
54		}
55	140x	return List::create(Named("P") = as_num_matrix_rcpp(P),
56	140x	Named("Q") = R_NilValue, Named("R") = R_NilValue,
57	210x	Named("S_Q") = as_num_matrix_rcpp(Q_sum));
58		}
59
60		// Obtain P,Q,R elements, or contracted linear components when W is supplied.
61		// mmrm() always supplies W for linear KR; the pairwise linear blocks are kept
62		// as an independent reference for tests.
63	103x	List get_pqr(List mmrm_fit, NumericVector theta, bool linear, Nullable<NumericMatrix> w) {
64	103x	if (w.isNotNull()) {
65	35x	return get_pqr_contracted(mmrm_fit, theta, NumericMatrix(w));
66		}
67	68x	kr_data data(mmrm_fit, theta);
68	68x	int n_theta = data.n_theta;
69	68x	int theta_size_per_group = data.theta_per_group;
70	68x	int p = data.p;
71	68x	matrix<double> P = matrix<double>::Zero(p * n_theta, p);
72	68x	matrix<double> Q = matrix<double>::Zero(p * theta_size_per_group * n_theta, p);
73	68x	matrix<double> R;
74	68x	if (!linear) {
75	51x	R = matrix<double>::Zero(p * theta_size_per_group * n_theta, p);
76		}
77		// Use map to hold these base class pointers (can also work for child class objects).
78	68x	auto derivatives_by_group = data.derivatives(!linear);
79	13006x	for (int i = 0; i < data.n_subjects; i++) {
80	12938x	kr_subject subject = data.subject(i);
81	12938x	int start_i = subject.start;
82	12938x	int n_visits_i = subject.n_visits;
83	12938x	int subject_group_i = subject.group;
84	12938x	matrix<double> sigma_inv, sigma_d2, sigma, sigma_inv_d1;
85
86	12938x	sigma_inv = derivatives_by_group.cache[subject_group_i]->get_sigma_inverse(subject.visits, subject.dist);
87	12938x	if (!linear) {
88	9652x	sigma_d2 = derivatives_by_group.cache[subject_group_i]->get_sigma_derivative2(subject.visits, subject.dist);
89		}
90	12938x	sigma = derivatives_by_group.cache[subject_group_i]->get_sigma(subject.visits, subject.dist);
91	12938x	sigma_inv_d1 = derivatives_by_group.cache[subject_group_i]->get_inverse_derivative(subject.visits, subject.dist);
92
93	12938x	matrix<double> Xi = data.x.block(start_i, 0, n_visits_i, p);
94	12938x	auto gi_sqrt_root = data.weights_sqrt.segment(start_i, n_visits_i).matrix().asDiagonal();
95	78402x	for (int r = 0; r < theta_size_per_group; r ++) {
96	65464x	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;
97	65464x	P.block(r * p + theta_size_per_group * subject_group_i * p, 0, p, p) += Pi;
98	546970x	for (int j = 0; j < theta_size_per_group; j++) {
99	481506x	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;
100		// switch the order so that in the matrix partial(i) and partial(j) increase j first
101	481506x	Q.block((r * theta_size_per_group + j + theta_size_per_group * theta_size_per_group * subject_group_i) * p, 0, p, p) += Qij;
102	481506x	if (!linear) {
103	298556x	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;
104	298556x	R.block((r * theta_size_per_group + j + theta_size_per_group * theta_size_per_group * subject_group_i) * p, 0, p, p) += Rij;
105		}
106		}
107		}
108		}
109		List ret = List::create(
110	272x	Named("P") = as_num_matrix_rcpp(P),
111	272x	Named("Q") = as_num_matrix_rcpp(Q),
112	136x	Named("R") = R_NilValue
113		);
114	68x	if (!linear) {
115	153x	ret["R"] = as_num_matrix_rcpp(R);
116		}
117	68x	return ret;
118		}

1		#ifndef KR_COMP_INCLUDED_
2		#define KR_COMP_INCLUDED_
3
4		#include "derivatives.h"
5
6		// Observed visits of one subject, see kr_data::subject().
7		struct kr_subject {
8		int start, n_visits, group;
9		std::vector<int> visits;
10		matrix<double> dist;
11		};
12
13		// Fitted model inputs shared by the pairwise and contracted backends.
14		struct kr_data {
15		matrix<double> x, coords;
16		vector<double> weights_sqrt, theta;
17		IntegerVector starts, sizes, groups;
18		int n_subjects, n_groups, n_visits, n_theta, theta_per_group, p;
19		bool spatial;
20		std::string cov_type;
21	103x	kr_data(List mmrm_fit, NumericVector theta_r) {
22	206x	x = as_num_matrix_tmb(as<NumericMatrix>(mmrm_fit["x_matrix"]));
23	206x	coords = as_num_matrix_tmb(as<NumericMatrix>(mmrm_fit["coordinates"]));
24	103x	weights_sqrt = as_num_vector_tmb(sqrt(as<NumericVector>(mmrm_fit["weights_vector"])));
25	103x	theta = as_num_vector_tmb(theta_r);
26	206x	starts = mmrm_fit["subject_zero_inds"];
27	206x	sizes = mmrm_fit["subject_n_visits"];
28	206x	groups = mmrm_fit["subject_groups"];
29	206x	n_subjects = mmrm_fit["n_subjects"];
30	206x	n_groups = mmrm_fit["n_groups"];
31	103x	n_visits = mmrm_fit["n_visits"];
32	103x	n_theta = theta_r.size();
33	103x	theta_per_group = n_theta / n_groups;
34	103x	p = x.cols();
35	206x	spatial = as<int>(mmrm_fit["is_spatial_int"]) == 1;
36	103x	cov_type = as<std::string>(mmrm_fit["cov_type"]);
37		}
38	103x	derivatives_cache<double> derivatives(bool second_order) const {
39	103x	return derivatives_cache<double>(theta, n_groups, spatial, cov_type, n_visits, second_order);
40		}
41		// Zero-based group with visit indices (non-spatial) or distances (spatial).
42	19707x	kr_subject subject(int i) const {
43	39414x	kr_subject s{starts[i], sizes[i], groups[i] - 1, std::vector<int>(sizes[i]), matrix<double>(0, 0)};
44	19707x	if (spatial) {
45	3940x	s.dist = euclidean(matrix<double>(coords.block(s.start, 0, s.n_visits, coords.cols())));
46		} else {
47	57724x	for (int j = 0; j < s.n_visits; ++j) s.visits[j] = int(coords(s.start + j, 0));
48		}
49	19707x	return s;
50		}
51		};
52
53		// W = B diag(signs) B'. Keep every direction, including negative eigenvalues
54		// when W is indefinite: clipping them would change the supplied contraction.
55		// W is an inverse Hessian, symmetric up to rounding, so use its symmetric part.
56		struct kr_directions {
57		matrix<double> B;
58		vector<double> signs;
59	44x	explicit kr_directions(const matrix<double>& w_r) {
60	44x	matrix<double> w = (w_r + w_r.transpose()) / 2;
61	44x	int k = w.rows();
62	44x	signs = vector<double>::Ones(k);
63	44x	Eigen::LLT<Eigen::MatrixXd> llt(w);
64	44x	if (llt.info() == Eigen::Success) {
65	36x	B = llt.matrixL();
66		} else {
67	8x	Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> eig(w);
68	8x	if (eig.info() != Eigen::Success) {
69	!	Rcpp::stop("Could not factor covariance parameter covariance matrix.");
70		}
71	8x	B = eig.eigenvectors();
72	74x	for (int ell = 0; ell < k; ++ell) {
73	66x	double value = eig.eigenvalues()(ell);
74	66x	signs(ell) = value < 0 ? -1 : 1;
75	66x	B.col(ell) *= std::sqrt(std::abs(value));
76		}
77		}
78		}
79		};
80
81		// Pattern-level products, reused across subjects in the same nonspatial group.
82		// Block ell of inverse_directions is sum_h B(h, ell) * d(Sigma^-1)/d(theta_h).
83		struct kr_pattern {
84		matrix<double> sigma, inverse_d1, inverse_directions;
85	1803x	kr_pattern(derivatives_base<double>* cache, const std::vector<int>& visits,
86	1803x	const matrix<double>& dist, const matrix<double>& B) {
87	1803x	sigma = cache->get_sigma(visits, dist);
88	1803x	inverse_d1 = cache->get_inverse_derivative(visits, dist);
89	1803x	int m = sigma.rows();
90	1803x	int k = B.rows();
91	1803x	inverse_directions = matrix<double>::Zero(k * m, m);
92	7719x	for (int ell = 0; ell < k; ++ell) {
93	39772x	for (int h = 0; h < k; ++h) {
94	33856x	inverse_directions.block(ell * m, 0, m, m) +=
95	67712x	B(h, ell) * inverse_d1.block(h * m, 0, m, m);
96		}
97		}
98		}
99		};
100
101		#endif

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	7220x	Type objective_function<Type>::operator() ()
29		{
30		// Read data from R.
31	7220x	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	7220x	DATA_MATRIX(coordinates); // Coordinates matrix.
35	7220x	DATA_INTEGER(n_visits); // Number of visits, which is the dimension of the covariance matrix.
36	7220x	DATA_INTEGER(n_subjects); // Number of subjects.
37	7220x	DATA_IVECTOR(subject_zero_inds); // Starting indices for each subject (0-based) (length n_subjects).
38	7220x	DATA_IVECTOR(subject_n_visits); // Number of observed visits for each subject (length n_subjects).
39	7220x	DATA_STRING(cov_type); // Covariance type name.
40	7220x	DATA_INTEGER(is_spatial_int); // Spatial covariance (1)? Otherwise non-spatial covariance.
41	7220x	DATA_INTEGER(reml); // REML (1)? Otherwise ML (0).
42	7220x	DATA_FACTOR(subject_groups); // subject groups vector(0-based) (length n_subjects).
43	7220x	DATA_INTEGER(n_groups); // number of total groups.
44		// Read parameters from R.
45	7220x	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	7220x	matrix<Type> XtWX = matrix<Type>::Zero(x_matrix.cols(), x_matrix.cols());
49		// X^T W Y will be calculated incrementally into here.
50	7220x	matrix<Type> XtWY = matrix<Type>::Zero(x_matrix.cols(), 1);
51		// W^T/2 X will be saved into here.
52	7220x	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	7220x	vector<Type> y_vec_tilde = vector<Type>::Zero(y_vector.rows());
55		// Sum of the log determinant will be incrementally calculated here.
56	7220x	Type sum_log_det = 0.0;
57
58		// Convert is_spatial_int to bool.
59	7220x	bool is_spatial = (is_spatial_int == 1);
60		// Diagonal of weighted covariance
61	7220x	vector<Type> diag_cov_inv_sqrt(x_matrix.rows());
62		// Cholesky group object
63	7220x	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	1413798x	for (int i = 0; i < n_subjects; i++) {
66		// Start index and number of visits for this subject.
67	1406582x	int start_i = subject_zero_inds(i);
68	1406582x	int n_visits_i = subject_n_visits(i);
69	1406582x	std::vector<int> visit_i(n_visits_i);
70	1406582x	matrix<Type> dist_i(n_visits_i, n_visits_i);
71	1406582x	if (!is_spatial) {
72	5022438x	for (int j = 0; j < n_visits_i; j++) {
73	3673380x	visit_i[j] = int(asDouble(coordinates(start_i + j, 0)));
74		}
75		} else {
76	57524x	dist_i = euclidean(matrix<Type>(coordinates.block(start_i, 0, n_visits_i, coordinates.cols())));
77		}
78		// Obtain Cholesky factor Li.
79	1406582x	matrix<Type> Li = chols_group.cache[subject_groups[i]]->get_chol(visit_i, dist_i);
80		// Calculate weighted Cholesky factor for this subject.
81	1406582x	Eigen::DiagonalMatrix<Type,Eigen::Dynamic,Eigen::Dynamic> Gi_inv_sqrt = weights_vector.segment(start_i, n_visits_i).cwiseInverse().sqrt().matrix().asDiagonal();
82	1406582x	Li = Gi_inv_sqrt * Li;
83		// Calculate scaled design matrix and response vector for this subject.
84	1406582x	matrix<Type> Xi = x_matrix.block(start_i, 0, n_visits_i, x_matrix.cols());
85	1406582x	matrix<Type> XiTilde = Li.template triangularView<Eigen::Lower>().solve(Xi);
86	1406582x	matrix<Type> Yi = y_vector.segment(start_i, n_visits_i).matrix();
87	1406582x	matrix<Type> YiTilde = Li.template triangularView<Eigen::Lower>().solve(Yi);
88
89		// Increment quantities.
90	1406582x	matrix<Type> XiTildeCrossprod = crossprod(XiTilde);
91	1406582x	XtWX += XiTildeCrossprod.template triangularView<Eigen::Lower>();
92	1406582x	XtWY += XiTilde.transpose() * YiTilde;
93	1406582x	vector<Type> LiDiag = Li.diagonal();
94	1406582x	sum_log_det += sum(log(LiDiag));
95		// Cache the reciprocal of square root of diagonal of covariance
96	1406582x	diag_cov_inv_sqrt.segment(start_i, n_visits_i) = vector<Type>(tcrossprod(Li).diagonal()).rsqrt();
97		// Save stuff.
98	1406582x	x_mat_tilde.block(start_i, 0, n_visits_i, x_matrix.cols()) = XiTilde;
99	1406582x	y_vec_tilde.segment(start_i, n_visits_i) = YiTilde.col(0);
100		}
101
102		// Solve for beta.
103	7216x	Eigen::LDLT<Eigen::Matrix<Type,Eigen::Dynamic,Eigen::Dynamic> > XtWX_decomposition(XtWX);
104	7216x	matrix<Type> beta_mat = XtWX_decomposition.solve(XtWY);
105	7216x	vector<Type> beta = beta_mat.col(0);
106
107		// Define scaled residuals.
108	7216x	vector<Type> x_mat_tilde_beta = x_mat_tilde * beta;
109	7216x	vector<Type> epsilonTilde = y_vec_tilde - x_mat_tilde_beta;
110
111		// Calculate negative log-likelihood.
112	652x	Type neg_log_lik;
113
114		// Always extract the D vector since we want to report this below.
115	7216x	vector<Type> XtWX_D = XtWX_decomposition.vectorD();
116
117	7216x	if (reml == 1) {
118		// Use restricted maximum likelihood.
119	6858x	Type XtWX_log_det = XtWX_D.log().sum();
120	8634x	neg_log_lik = (x_matrix.rows() - x_matrix.cols()) / 2.0 * log(2.0 * M_PI) +
121	6266x	sum_log_det +
122	12532x	XtWX_log_det / 2.0 +
123	8634x	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	478x	neg_log_lik = x_matrix.rows() / 2.0 * log(2.0 * M_PI) +
127		sum_log_det +
128	418x	0.5 * (epsilonTilde * epsilonTilde).sum();
129		}
130
131		// Report quantities to R.
132	6564x	REPORT(beta);
133
134		// We already compute the inverse of XtWX here because we already did the
135		// matrix decomposition above.
136	7216x	matrix<Type> Identity(XtWX.rows(), XtWX.cols());
137	7216x	Identity.setIdentity();
138	7216x	matrix<Type> beta_vcov = XtWX_decomposition.solve(Identity);
139	6564x	REPORT(beta_vcov);
140
141		// Also return the decomposition components L and D.
142	7216x	matrix<Type> XtWX_L(XtWX.rows(), XtWX.cols());
143	7216x	XtWX_L = XtWX_decomposition.matrixL();
144	6564x	REPORT(XtWX_L);
145	6564x	REPORT(XtWX_D);
146
147		// normalized residual
148	6564x	REPORT(epsilonTilde);
149		// inverse square root of diagonal of covariance
150	6564x	REPORT(diag_cov_inv_sqrt);
151	7216x	matrix<Type> covariance_lower_chol = chols_group.get_default_chol();
152	6564x	REPORT(covariance_lower_chol);
153
154	7868x	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	65x	void expect_equal_matrix(const T& target, const T& current)
32		{
33	65x	int nrow = target.rows();
34	65x	int ncol = target.cols();
35
36	!	expect_true(nrow == current.rows());
37	!	expect_true(ncol == current.cols());
38
39	410x	for (int i = 0; i < nrow; i++) {
40	1206x	for (int j = 0; j < ncol; j++) {
41	!	expect_equal(target(i, j), current(i, j));
42		}
43		}
44		}
45
46		template <class T>
47	9x	void expect_equal_vector(const T& target, const T& current)
48		{
49	9x	int n = target.size();
50	!	expect_true(n == current.size());
51
52	54x	for (int i = 0; i < n; i++) {
53	!	expect_equal(target(i), current(i));
54		}
55		}
56
57		#endif

1		#include "testthat-helpers.h"
2		#include "kr_comp.h"
3
4	2x	context("KR covariance directions") {
5	6x	test_that("directions reconstruct positive, semidefinite and indefinite W without truncation") {
6	1x	matrix<double> rotation(3, 3);
7	1x	rotation << 1, 2, -1, 0, 1, 2, 2, -1, 1;
8	5x	for (double last : {0.5, 0.0, -0.5, 1e-12}) {
9	4x	vector<double> values(3);
10	4x	values << 2, 1, last;
11	4x	matrix<double> w = rotation * values.matrix().asDiagonal() * rotation.transpose();
12	4x	kr_directions directions(w);
13	4x	matrix<double> reconstructed = directions.B * directions.signs.matrix().asDiagonal() * directions.B.transpose();
14		expect_true(directions.B.cols() == 3);
15		expect_true((reconstructed - w).norm() / w.norm() < 1e-12);
16		if (last < 0) expect_true(directions.signs.minCoeff() == -1);
17		}
18		}
19	6x	test_that("pattern directions equal first-derivative combinations with no second derivatives") {
20	1x	vector<double> theta(3);
21	1x	theta << 0.1, 0.2, 0.4;
22	1x	derivatives_nonspatial<double> cache(theta, 2, "us", false);
23	1x	matrix<double> w = matrix<double>::Identity(3, 3);
24	1x	w(0, 1) = w(1, 0) = 0.3;
25	1x	kr_directions directions(w);
26	1x	matrix<double> dist(0, 0);
27	6x	for (std::vector<int> visits : {std::vector<int>{0, 1}, std::vector<int>{1}}) {
28	2x	kr_pattern pattern(&cache, visits, dist, directions.B);
29	2x	int m = visits.size();
30	8x	for (int ell = 0; ell < 3; ++ell) {
31	6x	matrix<double> expected = matrix<double>::Zero(m, m);
32	24x	for (int h = 0; h < 3; ++h) {
33	18x	expected += directions.B(h, ell) * pattern.inverse_d1.block(h * m, 0, m, m);
34		}
35		expect_true((expected - pattern.inverse_directions.block(ell * m, 0, m, m)).norm() < 1e-12);
36		}
37		}
38		expect_true(cache.sigmad2_cache.empty());
39		}
40		}