R/draws.R

In insightsengineering/rbmi: Reference Based Multiple Imputation

#' @title Fit the base imputation model and get parameter estimates
#'
#' @description `draws` fits the base imputation model to the observed outcome data
#' according to the given multiple imputation methodology.
#' According to the user's method specification, it returns either draws from the posterior distribution of the
#' model parameters as required for Bayesian multiple imputation or frequentist parameter estimates from the
#' original data and bootstrapped or leave-one-out datasets as required for conditional mean imputation.
#' The purpose of the imputation model is to estimate model parameters
#' in the absence of intercurrent events (ICEs) handled using reference-based imputation methods.
#' For this reason, any observed outcome data after ICEs, for which reference-based imputation methods are
#' specified, are removed and considered as missing for the purpose of estimating the imputation model, and for
#' this purpose only. The imputation model is a mixed model for repeated measures (MMRM) that is valid
#' under a missing-at-random (MAR) assumption.
#' It can be fit using maximum likelihood (ML) or restricted ML (REML) estimation,
#' a Bayesian approach, or an approximate Bayesian approach according to the user's method specification.
#' The ML/REML approaches and the approximate Bayesian approach support several possible covariance structures,
#' while the Bayesian approach based on MCMC sampling supports only an unstructured covariance structure.
#' In any case the covariance matrix can be assumed to be the same or different across each group.
#'
#' @name draws
#' @param data A `data.frame` containing the data to be used in the model. See details.
#' @param data_ice A `data.frame` that specifies the information related
#' to the ICEs and the imputation strategies. See details.
#' @param vars A `vars` object as generated by [set_vars()]. See details.
#' @param method A `method` object as generated by either [method_bayes()],
#' [method_approxbayes()], [method_condmean()] or [method_bmlmi()].
#' It specifies the multiple imputation methodology to be used. See details.
#' @param ncores A single numeric specifying the number of cores to use in creating the draws object.
#' Note that this parameter is ignored for [method_bayes()] (Default = 1). Can also be a cluster object
#' generated by [`make_rbmi_cluster()`]
#' @param quiet Logical, if `TRUE` will suppress printing of progress information that is printed to
#' the console.
#'
#' @details
#'
#' `draws` performs the first step of the multiple imputation (MI) procedure: fitting the
#' base imputation model. The goal is to estimate the parameters of interest needed
#' for the imputation phase (i.e. the regression coefficients and the covariance matrices
#' from a MMRM model).
#'
#' The function distinguishes between the following methods:
#' - Bayesian MI based on MCMC sampling: `draws` returns the draws
#' from the posterior distribution of the parameters using a Bayesian approach based on
#' MCMC sampling. This method can be specified by using `method = method_bayes()`.
#' - Approximate Bayesian MI based on bootstrapping: `draws` returns
#' the draws from the posterior distribution of the parameters using an approximate Bayesian approach,
#' where the sampling from the posterior distribution is simulated by fitting the MMRM model
#' on bootstrap samples of the original dataset. This method can be specified by using
#' `method = method_approxbayes()]`.
#' - Conditional mean imputation with bootstrap re-sampling: `draws` returns the
#' MMRM parameter estimates from the original dataset and from `n_samples` bootstrap samples.
#' This method can be specified by using `method = method_condmean()` with
#' argument `type = "bootstrap"`.
#' - Conditional mean imputation with jackknife re-sampling: `draws` returns the
#' MMRM parameter estimates from the original dataset and from each leave-one-subject-out sample.
#' This method can be specified by using `method = method_condmean()` with
#' argument `type = "jackknife"`.
#' - Bootstrapped Maximum Likelihood MI: `draws` returns the MMRM parameter estimates from
#' a given number of bootstrap samples needed to perform random imputations of the bootstrapped samples.
#' This method can be specified by using `method = method_bmlmi()`.
#'
#' Bayesian MI based on MCMC sampling has been proposed in Carpenter, Roger, and Kenward (2013) who first introduced
#' reference-based imputation methods. Approximate Bayesian MI is discussed in Little and Rubin (2002).
#' Conditional mean imputation methods are discussed in Wolbers et al (2022).
#' Bootstrapped Maximum Likelihood MI is described in Von Hippel & Bartlett (2021).
#'
#' The argument `data` contains the longitudinal data. It must have at least the following variables:
#' - `subjid`: a factor vector containing the subject ids.
#' - `visit`: a factor vector containing the visit the outcome was observed on.
#' - `group`: a factor vector containing the group that the subject belongs to.
#' - `outcome`: a numeric vector containing the outcome variable. It might contain missing values.
#' Additional baseline or time-varying covariates must be included in `data`.
#'
#' `data` must have one row per visit per subject. This means that incomplete
#' outcome data must be set as `NA` instead of having the related row missing. Missing values
#' in the covariates are not allowed.
#' If `data` is incomplete
#' then the [expand_locf()] helper function can be used to insert any missing rows using
#' Last Observation Carried Forward (LOCF) imputation to impute the covariates values.
#' Note that LOCF is generally not a principled imputation method and should only be used when appropriate
#' for the specific covariate.
#'
#' Please note that there is no special provisioning for the baseline outcome values. If you do not want baseline
#' observations to be included in the model as part of the response variable then these should be removed in advance
#' from the outcome variable in `data`. At the same time if you want to include the baseline outcome as covariate in
#' the model, then this should be included as a separate column of `data` (as any other covariate).
#'
#' Character covariates will be explicitly
#' cast to factors. If you use a custom analysis function that requires specific reference
#' levels for the character covariates (for example in the computation of the least square means
#' computation) then you are advised
#' to manually cast your character covariates to factor in advance of running [draws()].
#'
#' The argument `data_ice` contains information about the occurrence of ICEs. It is a
#' `data.frame` with 3 columns:
#' - **Subject ID**: a character vector containing the ids of the subjects that experienced
#'   the ICE. This column must be named as specified in `vars$subjid`.
#' - **Visit**: a character vector containing the first visit after the occurrence of the ICE
#'   (i.e. the first visit affected by the ICE).
#'   The visits must be equal to one of the levels of `data[[vars$visit]]`.
#'   If multiple ICEs happen for the same subject, then only the first non-MAR visit should be used.
#'   This column must be named as specified in `vars$visit`.
#' - **Strategy**: a character vector specifying the imputation strategy to address the ICE for this subject.
#'   This column must be named as specified in `vars$strategy`.
#'   Possible imputation strategies are:
#'   - `"MAR"`: Missing At Random.
#'   - `"CIR"`: Copy Increments in Reference.
#'   - `"CR"`: Copy Reference.
#'   - `"JR"`: Jump to Reference.
#'   - `"LMCF"`: Last Mean Carried Forward.
#' For explanations of these imputation strategies, see Carpenter, Roger, and Kenward (2013), Cro et al (2021),
#' and Wolbers et al (2022).
#' Please note that user-defined imputation strategies can also be set.
#'
#' The `data_ice` argument is necessary at this stage since (as explained in Wolbers et al (2022)), the model is fitted
#' after removing the observations which are incompatible with the imputation model, i.e.
#' any observed data on or after `data_ice[[vars$visit]]` that are addressed with an imputation
#' strategy different from MAR are excluded for the model fit. However such observations
#' will not be discarded from the data in the imputation phase
#' (performed with the function ([impute()]). To summarize, **at this stage only pre-ICE data
#' and post-ICE data that is after ICEs for which MAR imputation is specified are used**.
#'
#' If the `data_ice` argument is omitted, or if a subject doesn't have a record within `data_ice`, then it is
#' assumed that all of the relevant subject's data is pre-ICE and as such all missing
#' visits will be imputed under the MAR assumption and all observed data will be used to fit the base imputation model.
#' Please note that the ICE visit cannot be updated via the `update_strategy` argument
#' in [impute()]; this means that subjects who didn't have a record in `data_ice` will always have their
#' missing data imputed under the MAR assumption even if their strategy is updated.
#'
#' The `vars` argument is a named list that specifies the names of key variables within
#' `data` and `data_ice`. This list is created by [set_vars()] and contains the following named elements:
#' - `subjid`: name of the column in `data` and `data_ice` which contains the subject ids variable.
#' - `visit`: name of the column in `data` and `data_ice` which contains the visit variable.
#' - `group`: name of the column in `data` which contains the group variable.
#' - `outcome`: name of the column in `data` which contains the outcome variable.
#' - `covariates`: vector of characters which contains the covariates to be included
#'   in the model (including interactions which are specified as `"covariateName1*covariateName2"`).
#'   If no covariates are provided the default model specification of `outcome ~ 1 + visit + group` will be used.
#'   Please note that the `group*visit` interaction
#'   is **not** included in the model by default.
#' - `strata`: covariates used as stratification variables in the bootstrap sampling.
#'   By default only the `vars$group` is set as stratification variable.
#'   Needed only for `method_condmean(type = "bootstrap")` and `method_approxbayes()`.
#' - `strategy`: name of the column in `data_ice` which contains the subject-specific imputation strategy.
#'
#' In our experience, Bayesian MI (`method = method_bayes()`) with a relatively low number of
#' samples (e.g. `n_samples` below 100) frequently triggers STAN warnings about R-hat such as
#' "The largest R-hat is X.XX, indicating chains have not mixed".  In many instances, this warning
#' might be spurious, i.e. standard diagnostics analysis of the MCMC samples do not indicate any
#' issues and results look reasonable. Increasing the number of samples to e.g. above 150 usually
#' gets rid of the warning.
#'
#' @inherit as_draws return
#'
#'@seealso [method_bayes()], [method_approxbayes()], [method_condmean()], [method_bmlmi()] for setting `method`.
#'@seealso [set_vars()] for setting `vars`.
#'@seealso [expand_locf()] for expanding `data` in case of missing rows.
#'
#' For more details see the quickstart vignette:
#' \code{vignette("quickstart", package = "rbmi")}.
#'
#' @references
#'
#' James R Carpenter, James H Roger, and Michael G Kenward. Analysis of longitudinal trials with protocol deviation: a
#' framework for relevant, accessible assumptions, and inference via multiple imputation. Journal of Biopharmaceutical
#' Statistics, 23(6):1352–1371, 2013.
#'
#' Suzie Cro, Tim P Morris, Michael G Kenward, and James R Carpenter. Sensitivity analysis for clinical trials with
#' missing continuous outcome data using controlled multiple imputation: a practical guide. Statistics in
#' Medicine, 39(21):2815–2842, 2020.
#'
#' Roderick J. A. Little and Donald B. Rubin. Statistical Analysis with Missing Data, Second Edition. John Wiley & Sons,
#' Hoboken, New Jersey, 2002. \[Section 10.2.3\]
#'
#' Marcel Wolbers, Alessandro Noci, Paul Delmar, Craig Gower-Page, Sean Yiu, Jonathan W. Bartlett.
#' Standard and reference-based
#' conditional mean imputation. \url{https://arxiv.org/abs/2109.11162}, 2022.
#'
#' Von Hippel, Paul T and Bartlett, Jonathan W.
#' Maximum likelihood multiple imputation: Faster imputations and consistent standard errors without
#' posterior draws. 2021.
#'
#' @export
draws <- function(data, data_ice = NULL, vars, method, ncores = 1, quiet = FALSE) {
    UseMethod("draws", method)
}


#' @rdname draws
#' @export
draws.approxbayes <- function(data, data_ice = NULL, vars, method, ncores = 1, quiet = FALSE) {
    longdata <- longDataConstructor$new(data, vars)
    longdata$set_strategies(data_ice)
    x <- get_draws_mle(
        longdata = longdata,
        method = method,
        sample_stack = get_bootstrap_stack(longdata, method),
        use_samp_ids = FALSE,
        first_sample_orig = FALSE,
        ncores = ncores,
        n_target_samples = method$n_samples,
        failure_limit = ceiling(method$threshold * method$n_samples),
        quiet = quiet
    )
    return(x)
}



#' @rdname draws
#' @export
draws.condmean <- function(data, data_ice = NULL, vars, method, ncores = 1, quiet = FALSE) {
    longdata <- longDataConstructor$new(data, vars)
    longdata$set_strategies(data_ice)

    sample_opts <- list(
        longdata = longdata,
        method = method,
        ncores = ncores,
        first_sample_orig = TRUE,
        use_samp_ids = TRUE,
        quiet = quiet
    )

    if (method$type == "bootstrap") {
        extra_opts <- list(
            n_target_samples = method$n_samples,
            failure_limit = ceiling(method$threshold * method$n_samples),
            sample_stack = get_bootstrap_stack(longdata, method)
        )
    } else if (method$type == "jackknife") {
        extra_opts <- list(
            n_target_samples = length(longdata$ids),
            failure_limit = 0,
            sample_stack = get_jackknife_stack(longdata)
        )
    } else {
        stop("Invalid Method type")
    }
    sample_opts <- append(sample_opts, extra_opts)
    x <- do.call(get_draws_mle, sample_opts)
    return(x)
}


#' @rdname draws
#' @export
draws.bmlmi <- function(data, data_ice = NULL, vars, method, ncores = 1, quiet = FALSE) {
    longdata <- longDataConstructor$new(data, vars)
    longdata$set_strategies(data_ice)
    method$n_samples <- method$B
    x <- get_draws_mle(
        longdata = longdata,
        method = method,
        sample_stack = get_bootstrap_stack(longdata, method),
        n_target_samples = method$n_samples,
        failure_limit = ceiling(method$threshold * method$n_samples),
        use_samp_ids = TRUE,
        first_sample_orig = FALSE,
        ncores = ncores,
        quiet = quiet
    )
    x$method$n_samples <- NULL
    return(x)
}


#' Fit the base imputation model on bootstrap samples
#'
#' @description
#' Fit the base imputation model using a ML/REML approach on a given number of bootstrap samples as
#' specified by `method$n_samples`. Returns the parameter estimates from the model fit.
#'
#' @param longdata R6 `longdata` object containing all relevant input data information.
#' @param method A `method` object as generated by either
#' [method_approxbayes()] or [method_condmean()] with argument `type = "bootstrap"`.
#' @param sample_stack A stack object containing the subject ids to be used on each mmrm iteration.
#' @param n_target_samples Number of samples needed to be created
#' @param first_sample_orig Logical. If `TRUE` the function returns `method$n_samples + 1` samples where
#' the first sample contains the parameter estimates from the original dataset and `method$n_samples`
#' samples contain the parameter estimates from bootstrap samples.
#' If `FALSE` the function returns `method$n_samples` samples containing the parameter estimates from
#' bootstrap samples.
#' @param use_samp_ids Logical. If `TRUE`, the sampled subject ids are returned. Otherwise
#' the subject ids from the original dataset are returned. These values are used to tell [impute()]
#' what subjects should be used to derive the imputed dataset.
#' @param failure_limit Number of failed samples that are allowed before throwing an error
#' @param ncores Number of processes to parallelise the job over
#' @param quiet Logical, If `TRUE` will suppress printing of progress information that is printed to
#' the console.
#'
#' @details
#'
#' This function takes a `Stack` object which contains multiple lists of patient ids. The function
#' takes this Stack and pulls a set ids and then constructs a dataset just consisting of these
#' patients (i.e. potentially a bootstrap or a jackknife sample).
#'
#' The function then fits a MMRM model to this dataset to create a sample object. The function
#' repeats this process until `n_target_samples` have been reached. If more than `failure_limit`
#' samples fail to converge then the function throws an error.
#'
#' After reaching the desired number of samples the function generates and returns a draws object.
#'
#' @inherit as_draws return
get_draws_mle <- function(
    longdata,
    method,
    sample_stack,
    n_target_samples,
    first_sample_orig,
    use_samp_ids,
    failure_limit = 0,
    ncores = 1,
    quiet = FALSE
) {

    max_sample_attempts <- n_target_samples + failure_limit
    assert_that(
        length(sample_stack$stack) >= max_sample_attempts,
        msg = sprintf(
            "Only %s samples in the sample_stack, %s are required",
            length(sample_stack$stack),
            max_sample_attempts
        )
    )

    time_taken <- system.time({
        initial_sample <- get_mmrm_sample(
            ids = longdata$ids,
            longdata = longdata,
            method = method
        )
    })

    if (initial_sample$failed) {
        stop("Fitting MMRM to original dataset failed")
    }

    if (!quiet) {
        cat(
            sprintf(
                "\nEstimated running time (assuming single core) is %s seconds\n\n",
                round(time_taken[[3]] * n_target_samples, 2)
            )
        )
    }


    cl <- make_rbmi_cluster(ncores, objects = list("longdata" = longdata, "method" = method))

    # If the user provided the clusters object directly then do not close it on completion
    if (!is(ncores, "cluster")){
        on.exit(
            { if (!is.null(cl)) parallel::stopCluster(cl) },
            add = TRUE,
            after = FALSE
        )
    }

    # Encapsulate arguments into a single function on `ids` and handle parallelisation
    par_get_mmrm_sample <- function(ids) {
        par_lapply(
            cl,
            get_mmrm_sample,
            ids,
            longdata = longdata,
            method = method
        )
    }

    samples <- list()
    n_failed_samples <- 0
    logger <- progressLogger$new(n_target_samples, quiet = quiet)

    while (length(samples) < n_target_samples) {
        ids <- sample_stack$pop(min(ncores, n_target_samples - length(samples)))
        new_samples <- par_get_mmrm_sample(ids)
        isfailure <- vapply(new_samples, function(x) x$failed, logical(1))
        new_samples_keep <- new_samples[!isfailure]
        n_failed_samples <- n_failed_samples + sum(isfailure)
        if (n_failed_samples > failure_limit) {
            if (!is.null(method$type)) {
                if (method$type == "jackknife") {
                    ids_fail <- ids[isfailure][[1]]
                    ids_jack <- longdata$ids[!longdata$ids %in% ids_fail][[1]]
                    msg <- "MMRM failed to fit to the data after removing subject '%s'"
                    stop(sprintf(msg, ids_jack))
                }
            }
            msg <- "More than %s failed fits. Try using a simpler covariance structure"
            stop(sprintf(msg, failure_limit))
        }
        logger$add(length(new_samples_keep))
        samples <- append(samples, new_samples_keep)
    }

    assert_that(
        length(samples) == n_target_samples,
        msg = "Incorrect number of samples were produced"
    )

    if (first_sample_orig) {
        samples <- append(list(initial_sample), samples)
    }

    if (!use_samp_ids) {
        for (i in seq_along(samples)) {
            samples[[i]]$ids <- longdata$ids
        }
    }

    ret <- as_draws(
        method = method,
        samples = sample_list(samples),
        data = longdata,
        formula = longdata$formula,
        n_failures = n_failed_samples
    )
    return(ret)
}



#' Fit MMRM and returns parameter estimates
#'
#' @description
#' `get_mmrm_sample` fits the base imputation model using a ML/REML approach.
#' Returns the parameter estimates from the fit.
#'
#' @param ids vector of characters containing the ids of the subjects.
#' @param longdata R6 `longdata` object containing all relevant input data information.
#' @param method A `method` object as generated by either
#' [method_approxbayes()] or [method_condmean()].
#'
#' @inherit sample_single return
get_mmrm_sample <- function(ids, longdata, method) {

    vars <- longdata$vars
    dat <- longdata$get_data(ids, nmar.rm = TRUE, na.rm = TRUE)
    model_df <- as_model_df(
        dat = dat,
        frm = longdata$formula
    )

    sample <- fit_mmrm(
        designmat = model_df[, -1, drop = FALSE],
        outcome = as.data.frame(model_df)[, 1],
        subjid = dat[[vars$subjid]],
        visit = dat[[vars$visit]],
        group = dat[[vars$group]],
        cov_struct = method$covariance,
        REML = method$REML,
        same_cov = method$same_cov
    )

    if (sample$failed) {
        ret <- sample_single(
            ids = ids,
            failed = TRUE
        )
    } else {
        ret <- sample_single(
            ids = ids,
            failed = FALSE,
            beta = sample$beta,
            sigma = sample$sigma,
            theta = sample$theta
        )
    }
    return(ret)
}


#' Set to NA outcome values that would be MNAR if they were missing
#' (i.e. which occur after an ICE handled using a reference-based imputation strategy)
#'
#' @param longdata R6 `longdata` object containing all relevant input data information.
#'
#' @return
#' A `data.frame` containing `longdata$get_data(longdata$ids)`, but MNAR outcome
#' values are set to `NA`.
extract_data_nmar_as_na <- function(longdata) {
    # remove non-MAR data
    data <- longdata$get_data(longdata$ids, nmar.rm = FALSE, na.rm = FALSE)
    is_mar <- unlist(longdata$is_mar)
    data[!is_mar, longdata$vars$outcome] <- NA
    return(data)
}


#' @rdname draws
#' @export
draws.bayes <- function(data, data_ice = NULL, vars, method, ncores = 1, quiet = FALSE) {

    longdata <- longDataConstructor$new(data, vars)
    longdata$set_strategies(data_ice)

    data2 <- extract_data_nmar_as_na(longdata)

    # compute design matrix
    model_df <- as_model_df(
        dat = data2,
        frm = longdata$formula
    )

    # scale input data
    scaler <- scalerConstructor$new(model_df)
    model_df_scaled <- scaler$scale(model_df)

    fit <- fit_mcmc(
        designmat = model_df_scaled[, -1, drop = FALSE],
        outcome = model_df_scaled[, 1, drop = TRUE],
        group = data2[[vars$group]],
        visit = data2[[vars$visit]],
        subjid = data2[[vars$subjid]],
        method = method,
        quiet = quiet
    )

    # set names of covariance matrices
    fit$samples$sigma <- lapply(
        fit$samples$sigma,
        function(sample_cov) {
            lvls <- levels(data2[[vars$group]])
            sample_cov <- ife(
                method$same_cov == TRUE,
                rep(sample_cov, length(lvls)),
                sample_cov
            )
            setNames(sample_cov, lvls)
        }
    )

    # unscale samples
    samples <- mapply(
        function(x, y) list("beta" = x, "sigma" = y),
        lapply(fit$samples$beta, scaler$unscale_beta),
        lapply(fit$samples$sigma, function(covs) lapply(covs, scaler$unscale_sigma)),
        SIMPLIFY = FALSE
    )

    # set ids associated to each sample
    samples <- lapply(
        samples,
        function(x) {
            sample_single(
                ids = longdata$ids,
                beta = x$beta,
                sigma = x$sigma,
                failed = FALSE
            )
        }
    )

    result <- as_draws(
        method = method,
        samples = sample_list(samples),
        data = longdata,
        fit = fit$fit,
        formula = longdata$formula,
        n_failures = 0
    )

    return(result)
}




#' Print `draws` object
#'
#' @param x A `draws` object generated by [draws()].
#' @param ... not used.
#' @export
print.draws <- function(x, ...) {

    frm <- as.character(x$formula)
    frm_str <- sprintf("%s ~ %s", frm[[2]], frm[[3]])

    meth <- switch(
        class(x$method)[[2]],
        "approxbayes" = "Approximate Bayes",
        "condmean" = "Conditional Mean",
        "bayes" = "Bayes"
    )

    method <- x$method

    control_args <- if ("control" %in% names(method)) {
        control <- method$control
        method <- method[!(names(method) == "control")]
        format_method_descriptions(control)
    } else {
        character()
    }
    meth_args <- format_method_descriptions(method)

    n_samp <- length(x$samples)
    n_samp_string <- ife(
        has_class(x$method, "condmean"),
        sprintf("1 + %s", n_samp - 1),
        as.character(n_samp)
    )

    string <- c(
        "",
        "Draws Object",
        "------------",
        sprintf("Number of Samples: %s", n_samp_string),
        sprintf("Number of Failed Samples: %s", x$n_failures),
        sprintf("Model Formula: %s", frm_str),
        sprintf("Imputation Type: %s", class(x)[[2]]),
        "Method:",
        sprintf("    name: %s", meth),
        meth_args,
        ife(
            length(control_args),
            c(
                "Controls:",
                control_args
            ),
            NULL
        ),
        ""
    )

    cat(string, sep = "\n")
    return(invisible(x))
}



#' Creates a `draws` object
#'
#' @description
#' Creates a `draws` object which is the final output of a call to [draws()].
#'
#' @param method A `method` object as generated by either [method_bayes()],
#' [method_approxbayes()], [method_condmean()] or [method_bmlmi()].
#' @param samples A list of `sample_single` objects. See [sample_single()].
#' @param data R6 `longdata` object containing all relevant input data information.
#' @param formula Fixed effects formula object used for the model specification.
#' @param n_failures Absolute number of failures of the model fit.
#' @param fit If `method_bayes()` is chosen, returns the MCMC Stan fit object. Otherwise `NULL`.
#'
#' @return
#' A `draws` object which is a named list containing the following:
#' - `data`: R6 `longdata` object containing all relevant input data information.
#' - `method`: A `method` object as generated by either [method_bayes()],
#' [method_approxbayes()] or [method_condmean()].
#' - `samples`: list containing the estimated parameters of interest.
#'   Each element of `samples` is a named list containing the following:
#'   - `ids`: vector of characters containing the ids of the subjects included in the original dataset.
#'   - `beta`: numeric vector of estimated regression coefficients.
#'   - `sigma`: list of estimated covariance matrices (one for each level of `vars$group`).
#'   - `theta`: numeric vector of transformed covariances.
#'   - `failed`: Logical. `TRUE` if the model fit failed.
#'   - `ids_samp`: vector of characters containing the ids of the subjects included in the given sample.
#' - `fit`: if `method_bayes()` is chosen, returns the MCMC Stan fit object. Otherwise `NULL`.
#' - `n_failures`: absolute number of failures of the model fit.
#' Relevant only for `method_condmean(type = "bootstrap")`, `method_approxbayes()` and `method_bmlmi()`.
#' - `formula`: fixed effects formula object used for the model specification.
#'
as_draws <- function(
    method,
    samples,
    data,
    formula,
    n_failures = NULL,
    fit = NULL
) {
    x <- list(
        data = data,
        method = method,
        samples = samples,
        fit = fit,
        n_failures = n_failures,
        formula = formula
    )

    next_class <- switch(class(x$method)[[2]],
        "approxbayes" = "random",
        "condmean" = "condmean",
        "bayes" = "random",
        "bmlmi" = "random"
    )

    class(x) <- c("draws", next_class, "list")
    return(x)
}


#' Validate `draws` object
#'
#' @param x A `draws` object generated by [as_draws()].
#' @param ... Not used.
#' @export
validate.draws <- function(x, ...) {
    assert_that(
        has_class(x$data, "longdata"),
        has_class(x$method, "method"),
        has_class(x$samples, "sample_list"),
        validate(x$samples),
        is.null(x$n_failures) | is.numeric(x$n_failures),
        is.null(x$fit) | has_class(x$fit, "stanfit"),
        has_class(x$formula, "formula")
    )
}






#' R6 Class for printing current sampling progress
#'
#' @description
#'
#' Object is initalised with total number of iterations that are expected to occur.
#' User can then update the object with the `add` method to indicate how many more iterations
#' have just occurred.
#' Every time `step` * 100 % of iterations have occurred a message is printed to the console.
#' Use the `quiet` argument to prevent the object from printing anything at all
#'
#' @import R6
progressLogger <- R6::R6Class(
    classname = "progressLogger",
    public = list(
        #' @field step real, percentage of iterations to allow before printing the
        #' progress to the console
        step = NULL,

        #' @field step_current integer, the total number of iterations completed since
        #' progress was last printed to the console
        step_current = 0,

        #' @field n integer, the current number of completed iterations
        n = 0,

        #' @field n_max integer, total number of expected iterations to be completed
        #' acts as the denominator for calculating progress percentages
        n_max = NULL,

        #' @field quiet logical holds whether or not to print anything
        quiet = FALSE,

        #' @description
        #' Create progressLogger object
        #' @param n_max integer, sets field `n_max`
        #' @param quiet logical, sets field `quiet`
        #' @param step real, sets field `step`
        initialize = function(n_max, quiet = FALSE, step = 0.1) {
            self$step <- step
            self$n_max <- n_max
            self$quiet <- quiet
        },

        #' @description
        #' Records that `n` more iterations have been completed
        #' this will add that number to the current step count (`step_current`) and will
        #' print a progress message to the log if the step limit (`step`) has been reached.
        #' This function will do nothing if `quiet` has been set to `TRUE`
        #' @param n the number of successfully complete iterations since `add()` was last called
        add = function(n) {
            if (self$quiet) {
                return(invisible())
            }
            self$n <- self$n + n
            self$step_current <- self$step_current + n / self$n_max
            if (self$step_current >= self$step) {
                self$print_progress()
                self$step_current <- 0
            }
        },

        #' @description
        #' method to print the current state of progress
        print_progress = function() {
            cat(
                sprintf("Progress: %3.0f%%\n", self$n * 100 / self$n_max)
            )
        }
    )
)