jmpost coverage - 96.01%

Files
Source

#' Row Numbers of Data with Missing Variables
#'
#' @param df (`data.frame`)\cr input data.
#' @param formula (`formula` or `NULL`)\cr which variables to inspect for missingness, if `NULL`
#'   all variables are considered.
#'
#' @returns Numeric vector specifying which rows contain at least 1 missing observation
#'   in any of the inspected variables.
#'
#' @keywords internal
get_missing_rownumbers <- function(df, formula = NULL) {
    if (is.null(formula)) {
        formula <- ~ .
    }
    mdf <- stats::model.frame(formula, data = df, na.action = stats::na.pass)
    which(!stats::complete.cases(mdf))
}

#' Remove Rows with Missing Variables
#'
#' Removes any rows from a data set that contain missing values in the inspected
#' variables. Allows users to specify which variables to inspect for missing values
#' based on either a formula or a character vector of variable names.
#'
#' @param data (`data.frame`)\cr input data.
#' @param formula (`formula` or `NULL`)\cr which variables to inspect for missingness.
#' @param extra_vars (`character`)\cr additional variables to inspect for missingness.
#'
#' @returns The `data` after removing observations that contain missing values in the required variables.
#'   Note that additional variables not listed in `formula` or `extra_vars` are not dropped and may
#'   still contain missing values.
#'
#' @keywords internal
remove_missing_rows <- function(data, formula, extra_vars = NULL) {
    if (!is.null(extra_vars)) {
        extra_vars <- paste(extra_vars, collapse = " + ")
        formula_update_string <- paste0(". ~ . + ", extra_vars)
        formula <- stats::update(formula, formula_update_string)
    }

    missing_rows <- get_missing_rownumbers(data, formula)

    if (length(missing_rows) == 0) {
        return(data)
    }
    message(sprintf(
        "Note that %d observations were removed as one of more required variables contained missing values",
        length(missing_rows)
    ))
    data_reduced <- data[-missing_rows, ]
    rownames(data_reduced) <- NULL
    data_reduced
}

#' Replicate Single Values in a List
#'
#' @param initial_values (`list`)\cr initial values with names.
#' @param sizes (`list`)\cr each size corresponds to an element in `initial_values`,
#'   matched by the names. An attribute `array` must be attached to each element,
#'   see [replace_with_lookup()].
#'
#' @returns A named list of values, with any single values in the `initial_values` list
#'   replicated according to the corresponding values in the `sizes` list.
#'   Even when the size is 1, the value is passed as an `array` if the corresponding
#'   attribute is `TRUE` in `sizes`.
#'
#' @note The resulting list has the same names as the original lists.
#'
#' @keywords internal
expand_initial_values <- function(initial_values, sizes) {
    assert_that(
        is.list(initial_values),
        msg = "`initial_values` must be a list"
    )
    assert_that(
        is.list(sizes),
        msg = "`sizes` must be a list"
    )
    assert_that(
        all(names(sizes) %in% names(initial_values)),
        all(names(initial_values) %in% names(sizes)),
        msg = "`initial_values` and `sizes` must have identical names"
    )

    for (name in names(initial_values)) {
        # Check for single values and replicate them according to sizes.
        if (length(initial_values[[name]]) == 1) {
            initial_values[[name]] <- rep(initial_values[[name]], sizes[[name]])
        }
        # Check for array handling.
        needs_array <- attr(sizes[[name]], "array")
        assert_that(
            is.flag(needs_array),
            msg = "each sizes element must have array flag attribute"
        )
        if (needs_array) {
            initial_values[[name]] <- as.array(initial_values[[name]])
        }
    }

    # Check that each element of initial_values has the same length as specified in sizes.
    for (name in names(initial_values)) {
        assert_that(
            length(initial_values[[name]]) == sizes[[name]],
            msg = "length of element in `initial_values` does not match specified size"
        )
    }

    initial_values
}

#' Replace Character Size by Looked Up Numbers
#'
#' @param sizes (`list`)\cr may include character elements that correspond to
#'   names in the data list.
#' @param data (`list`)\cr data containing numeric values.
#'
#' @returns A list of sizes with character elements in `sizes`
#'   replaced by their corresponding numeric values in `data`.
#'
#' @details An attribute `array` for each returned list element indicates
#'   whether the parameter needs to be handled
#'   as an array. This is the case when the size is larger than 1, or when
#'   the size was looked up in the `data`, because in that case it is flexible
#'   and hence is handled as an array in the Stan code.
#'
#' @note Each element in the final list of sizes must be a single number.
#'
#' @keywords internal
replace_with_lookup <- function(sizes, data) {

    assert_that(is.list(sizes), msg = "`sizes` must be a list")
    assert_that(is.list(data), msg = "`data` must be a list")

    for (idx in seq_along(sizes)) {
        val <- sizes[[idx]]
        if (is.character(val)) {
            assert_that(
                length(val) == 1,
                msg = "character elements of `sizes` must be strings"
            )
            assert_that(
                val %in% names(data),
                msg = sprintf("`%s` is not available in `data`", val)
            )
            new_val <- data[[val]]
            assert_that(
                is.number(new_val),
                msg = "Selected values from data must be single numbers"
            )
            sizes[[idx]] <- structure(new_val, array = TRUE)
        } else {
            assert_that(
                is.number(val),
                msg = "Existing values in sizes must be single numbers"
            )
            sizes[[idx]] <- structure(val, array = val > 1)
        }
    }
    sizes
}

#' Obtain Median and Credible Intervals from MCMC samples
#'
#' @param samples (`matrix`)\cr with samples in rows and parameters in columns.
#' @param level (`number`)\cr credibility level to use for the credible intervals.
#'
#' @returns A `data.frame` with columns `median`, `lower` and `upper`.
#' @keywords internal
#' samples_median_ci(samples)
samples_median_ci <- function(samples, level = 0.95) {
    assert_that(is.matrix(samples))
    assert_that(is.number(level), level < 1, level > 0)

    samples_median <- apply(samples, MARGIN = 2L, FUN = stats::median)
    probs <- c((1 - level) / 2, (1 + level) / 2)
    samples_ci <- t(apply(samples, MARGIN = 2L, FUN = stats::quantile, probs = probs))
    colnames(samples_ci) <- c("lower", "upper")
    as.data.frame(cbind(
        median = samples_median,
        samples_ci
    ))
}



#' `decorated_render`
#'
#' Simple wrapper around [jinjar::render()] that provides some additional default
#' variables about the system (avoids each call to jinjar having to specify them)
#' @param ... Arguments passed onto [jinjar::render()]
#' @returns See [jinjar::render()]
#' @keywords internal
decorated_render <- function(...) {
    jinjar::render(
        ...,
        machine_double_eps = sqrt(.Machine$double.eps),
        machine_double_neg_eps = sqrt(.Machine$double.neg.eps)
    )
}


is_windows <- function() {
    sysname <- Sys.info()["sysname"]
    return(sysname == "Windows")
}

#' `validate_time_grid`
#'
#' Validate that the provided time grid is:
#' - finite
#' - numeric
#' - non-missing
#' - sorted
#' - unique
#'
#' @param time_grid (`numeric`)\cr A vector of times which quantities will be
#' evaluated at.
#'
#' @keywords internal
validate_time_grid <- function(time_grid) {
    assert_that(
        !any(is.na(time_grid)),
        is.numeric(time_grid),
        !is.null(time_grid),
        !is.unsorted(time_grid),
        !any(duplicated(time_grid)),
        all(is.finite(time_grid)),
        msg = "`time_grid` needs to be finite, sorted, unique valued numeric vector"
    )
    invisible(time_grid)
}



#' `expand_subjects`
#'
#' This function checks and expands a given subjects vector.
#' The input vector must be unique and contain only values
#' as specified by `all_subjects`
#'
#' @param subjects (`character` or `NULL`)\cr Character vector representing the subjects.
#' If NULL, it will be set to the value of `all_subjects`.
#' @param all_subjects (`character`)\cr Character vector representing all possible subjects.
#' @return Returns the expanded `subjects` vector.
#' @keywords internal
expand_subjects <- function(subjects, all_subjects) {
    assert_that(
        is.character(all_subjects),
        msg = "`all_subjects` must be a character vector"
    )
    if (is.null(subjects)) {
        subjects <- unique(all_subjects)
    }
    assert_that(
        is.character(subjects),
        all(subjects %in% all_subjects),
        !any(duplicated(subjects)),
        msg = "`subjects` should be a unique character vector containing only values from the original df"
    )
    return(subjects)
}



#' Decompose subjects into Relevant Components
#'
#' This function takes in a character vector or list of subjects and decomposes it into a
#' structured format.
#'
#' The primary use of this function is to correctly setup indexing variables for
#' predicting survival quantities (see [SurvivalQuantities()])
#'
#' @param subjects (`character` or `list`)\cr subject identifiers. If `NULL` will be set to `all_subjects`.
#'
#' @param all_subjects (`character`)\cr the set of allowable subject identifiers.
#' Will cause an error if any value of `subjects` is not in this vector.
#'
#' @return A list containing three components:
#' - `groups`: (`list`)\cr each element of the list is a character vector
#' specifying which subjects belong to a given "group" where the "group" is the element name
#' - `unique_values`: (`character`)\cr vector of the unique subjects within `subjects`
#' - `indexes`: (`list`)\cr each element is a named and is a numeric index vector
#' that maps the values of `grouped` to `unique_values`
#' @examples
#' \dontrun{
#' result <- decompose_subjects(c("A", "B"), c("A", "B", "C", "D"))
#' result <- decompose_subjects(
#'     list("g1" = c("A", "B"), "g2" = c("B", "C")),
#'     c("A", "B", "C", "D")
#' )
#' }
#' @seealso [expand_subjects()], [SurvivalQuantities()]
#' @keywords internal
decompose_subjects <- function(subjects, all_subjects) {
    if (is.character(subjects) || is.null(subjects)) {
        subjects <- expand_subjects(subjects, all_subjects)
        names(subjects) <- subjects
        subjects <- as.list(subjects)
    }
    subjects <- lapply(
        subjects,
        expand_subjects,
        all_subjects = all_subjects
    )
    assert_that(
        is.list(subjects),
        length(unique(names(subjects))) == length(subjects),
        all(vapply(subjects, is.character, logical(1)))
    )
    subjects_vec_unordered <- unique(unlist(subjects))
    subjects_vec <- subjects_vec_unordered[order(subjects_vec_unordered)]
    subjects_lookup <- stats::setNames(seq_along(subjects_vec), subjects_vec)
    subjects_index <- lapply(
        subjects,
        \(x) {
            z <- subjects_lookup[x]
            names(z) <- NULL
            z
        }
    )
    list(
        groups = subjects,
        unique_values = subjects_vec,
        indexes = subjects_index
    )
}

is_cmdstanr_available <- function() {
    requireNamespace("cmdstanr", quietly = TRUE)
}


is_connection <- function(obj) {
    inherits(obj, "connection")
}

#' @include generics.R
NULL

# STAN_BLOCKS ----

#' List of Stan Blocks
#'
#' @description
#' A list with 1 element per standard Stan program blocks.
#' This object is  mostly used internally as a reference for
#' what blocks are expected to exist within a given Stan program.
#'
#' @export
STAN_BLOCKS <- list(
    functions = "functions",
    data = "data",
    transformed_data = "transformed data",
    parameters = "parameters",
    transformed_parameters = "transformed parameters",
    model = "model",
    generated_quantities = "generated quantities"
)

# add_missing_stan_blocks ----

#' Add Missing Stan Blocks
#'
#' @param x (`list`)\cr list of Stan code blocks
#' @param stan_blocks (`list`)\cr reference list of stan blocks.
#'
#' @return Amended list `x` such that all blocks in the global variable
#' `STAN_BLOCKS` are contained.
#'
#' @keywords internal
add_missing_stan_blocks <- function(x, stan_blocks = STAN_BLOCKS) {
    for (block in names(stan_blocks)) {
        if (is.null(x[[block]])) {
            x[[block]] <- ""
        }
    }
    return(x)
}

# StanModule-class ----

#' `StanModule` Object and Constructor Function
#'
#' @param x (`string`)\cr file path to a Stan program or a character vector
#' of Stan code to be parsed.
#' @param ... additional arguments passed to the constructor.
#'
#' @slot functions (`character`)\cr the `functions` block.
#' @slot data (`character`)\cr the `data` block.
#' @slot transformed_data (`character`)\cr the `transformed_data` block.
#' @slot parameters (`character`)\cr the `parameters` block.
#' @slot transformed_parameters (`character`)\cr the `transformed_parameters` block.
#' @slot model (`character`)\cr the `model` block.
#' @slot generated_quantities (`character`)\cr the `generated_quantities` block.
#'
#' @exportClass StanModule
#' @export StanModule
#' @family StanModule
.StanModule <- setClass(
    Class = "StanModule",
    slots = list(
        functions = "character",
        data = "character",
        transformed_data = "character",
        parameters = "character",
        transformed_parameters = "character",
        model = "character",
        generated_quantities = "character"
    )
)

# StanModule-constructors ----

#' @rdname StanModule-class
StanModule <- function(
    x = "",
    ...
) {
    assert_that(
        is.character(x),
        length(x) == 1,
        msg = "`x` must be a length 1 character vector"
    )
    code <- read_stan(x)
    code_fragments <- as_stan_fragments(code)

    if (paste0(unlist(code_fragments), collapse = "") == "" && paste0(x, collaspe = "") != "") {
        warning("Non-empty input resulted in an empty StanModule object. Is the input correct?")
    }

    .StanModule(
        functions = code_fragments$functions,
        data = code_fragments$data,
        transformed_data = code_fragments$transformed_data,
        parameters = code_fragments$parameters,
        transformed_parameters = code_fragments$transformed_parameters,
        model = code_fragments$model,
        generated_quantities = code_fragments$generated_quantities,
        ...
    )
}

# as.character-StanModule ----

#' `StanModule` -> `character`
#' @param x ([`StanModule`])\cr A stan program
#' @param ... Not Used.
#' @description
#' Converts a [`StanModule`] object into a valid Stan program file where each
#' line of the returned `character` vector represents a line of the program
#' @family StanModule
#' @export
as.character.StanModule <- function(x, ...) {
    as_stan_file(
        functions = x@functions,
        data = x@data,
        transformed_data = x@transformed_data,
        parameters = x@parameters,
        transformed_parameters = x@transformed_parameters,
        model = x@model,
        generated_quantities = x@generated_quantities
    )
}



# merge-StanModule,StanModule ----

#' @param stan_blocks (`list`)\cr reference list of stan blocks.
#' @rdname merge
setMethod(
    f = "merge",
    signature = c("StanModule", "StanModule"),
    definition = function(x, y, stan_blocks = STAN_BLOCKS, ...) {
        stan_fragments <- lapply(
            names(stan_blocks),
            \(par) {
                if (all(slot(y, par) == "")) {
                    return(slot(x, par))
                }
                if (all(slot(x, par) == "")) {
                    return(slot(y, par))
                }
                return(c(slot(x, par), slot(y, par)))
            }
        )
        names(stan_fragments) <- names(stan_blocks)
        stan_code <- do.call(as_stan_file, stan_fragments)
        StanModule(
            x = stan_code
        )
    }
)

# compileStanModel-StanModule,character ----

#' @rdname compileStanModel
compileStanModel.StanModule <- function(object) {
    exe_dir <- getOption("jmpost.cache_dir")
    if (!dir.exists(exe_dir)) {
        dir.create(exe_dir, recursive = TRUE)
    }
    stan_code <- as.character(object)
    hash_name <- digest::digest(stan_code, "md5")
    exe_name <- paste0(
        "model_",
        hash_name,
        if (is_windows()) ".exe" else ""
    )
    exe_file <- file.path(exe_dir, exe_name)
    source_file <- cmdstanr::write_stan_file(
        code = stan_code,
        dir = exe_dir,
        basename = sprintf("model_%s.stan", hash_name)
    )

    # Suppress "model executable is up to date" message as
    # users are not in control of the cache so this message is meaningless
    withCallingHandlers(
        {
            x <- cmdstanr::cmdstan_model(
                stan_file = source_file,
                exe_file = exe_file,
                quiet = TRUE
            )
        },
        message = function(m) {
            if (m$message == "Model executable is up to date!\n") {
                invokeRestart("muffleMessage")
            }
        }
    )
    invisible(x)
}



# as.list-StanModule ----

#' `StanModule` -> `list`
#' @description
#' Returns a named list where each element of the list corresponds
#' to a Stan modelling block e.g. `data`, `model`, etc.
#' @param x ([`StanModule`])\cr A Stan Module
#' @param stan_blocks (`list`)\cr reference list of stan blocks.
#' @param ... Not Used.
#' @family StanModule
#' @export
as.list.StanModule <- function(x, stan_blocks = STAN_BLOCKS, ...) {
    string <- as.character(x)
    li <- as_stan_fragments(string)
    for (block in names(stan_blocks)) {
        li[[block]] <- paste(li[[block]], collapse = "\n")
    }
    return(li)
}

# is_file ----

#' Is String a Valid File?
#'
#' A utility function to check if a string is a valid file or not.
#' Used to help address short comings of [file.exists()] that will return `TRUE`
#' for a directory as well as a file.
#'
#' @param filename (`string`)\cr file name.
#'
#' @keywords internal
is_file <- function(filename = NULL) {
    if (is.null(filename)) {
        return(FALSE)
    }
    assert_that(
        is.character(filename),
        length(filename) == 1,
        msg = "`filename` must be a length 1 character"
    )
    if (nchar(filename) > 1000) {
        return(FALSE)
    }
    if (is.na(filename)) {
        return(FALSE)
    }
    return(file.exists(filename) & !dir.exists(filename))
}

# read_stan ----

#' Stan Code as Character
#'
#' @param string Character, either the absolute path of a stan file, or the name of the stan
#' file in the package directory or the stan code as a string.
read_stan <- function(string) {
    local_inst_file <- file.path("inst", "stan", string)
    system_file <- system.file(file.path("stan", string), package = "jmpost")
    local_file <- string
    files <- c(local_file, local_inst_file, system_file)
    for (fi in files) {
        if (is_file(fi)) {
            string <- readLines(fi)
            break
        }
    }
    string <- paste0(string, collapse = "\n")
    return(string)
}

# as_stan_file ----

#' Merging Code Blocks into Stan Code Character Vector
#'
#' @param functions (`character`)\cr code block.
#' @param data (`character`)\cr code block.
#' @param transformed_data (`character`)\cr code block.
#' @param parameters (`character`)\cr code block.
#' @param transformed_parameters (`character`)\cr code block.
#' @param model (`character`)\cr code block.
#' @param generated_quantities (`character`)\cr code block.
#' @param stan_blocks (`list`)\cr reference list of stan blocks.
#'
#' @return Character vector of the complete Stan code.
#'
#' @keywords internal
as_stan_file <- function(
    functions = "",
    data = "",
    transformed_data = "",
    parameters = "",
    transformed_parameters = "",
    model = "",
    generated_quantities = "",
    stan_blocks = STAN_BLOCKS
) {
    block_strings <- lapply(
        names(stan_blocks),
        function(id) {
            char <- get(id)
            if (any(nchar(char) >= 1)) {
                return(sprintf("%s {\n%s\n}\n\n", stan_blocks[[id]], paste(char, collapse = "\n")))
            } else {
                return("")
            }
        }
    )
    return(paste0(block_strings, collapse = ""))
}

# as_stan_fragments ----

#' Conversion of Character Vector into Stan Code Block List
#'
#' @param x (`character`)\cr the single Stan code vector.
#' @param stan_blocks (`list`)\cr reference list of stan blocks.
#'
#' @return A list with the Stan code blocks.
#'
#' @details
#' Function only works if code is in format
#' ```
#' data {
#'     <code>
#' }
#' model {
#'     <code>
#' }
#' ```
#' That is to say we do not support code in inline format i.e.
#' ```
#' data { <code> }
#' model { <code> }
#' ```
#'
#' @keywords internal
as_stan_fragments <- function(x, stan_blocks = STAN_BLOCKS) {
    code <- unlist(stringr::str_split(x, "\n"))

    errmsg <- paste(
        "There were problems parsing the `%s` block.",
        "Please consult the `Formatting Stan Files` section of the",
        "`Extending jmpost` vignette"
    )

    # Check to see if any block openings exist that have code on the same line
    # e.g.  `data { int i;}`. This is unsupported so we throw an error
    for (block in stan_blocks) {
        regex <- sprintf("^\\s*%s\\s*\\{\\s*[^\\s-]+", block)
        if (any(grepl(regex, code, perl = TRUE))) {
            stop(sprintf(errmsg, block))
        }
    }

    # We first look to identify the opening of a block e.g.  `data {`
    # We then regard all lines that follow as belonging to that block
    # until we see another block being opened e.g. `model{`
    results <- list()
    target <- NULL
    for (line in code) {
        for (block in names(stan_blocks)) {
            regex <- sprintf("^\\s*%s\\s*\\{\\s*$", stan_blocks[[block]])
            if (stringr::str_detect(line, regex)) {
                target <- block
                line <- NULL
                break
            }
        }
        if (!is.null(target)) {
            # This is memory inefficient but given the relatively small size of
            # stan files its regarded as a acceptable simplification to ease the
            # code burden
            results[[target]] <- c(results[[target]], line)
        }
    }

    # Loop over each block to remove trailing "}".
    for (block in names(results)) {
        block_length <- length(results[[block]])
        # The following processing is only required if the block actually has content
        if (block_length == 1 && results[[block]] == "") {
            next
        }
        has_removed_char <- FALSE
        # Walk backwards to find the closing `}` that corresponds to the `<block> {`
        for (index in rev(seq_len(block_length))) {
            line <- results[[block]][[index]]
            # This code will exit the for loop as soon as it hits the closing `}`
            # thus if we ever see a line that ends in text/numbers it means
            # somethings gone wrong
            if (stringr::str_detect(line, "[\\w\\d]+\\s*$")) {
                stop(sprintf(errmsg, block))
            }
            if (stringr::str_detect(line, "\\}\\s*$")) {
                new_line <- stringr::str_replace(line, "\\s*\\}\\s*$", "")
                # If the line is now blank after removing the closing `}` then drop the line
                keep_offset <- if (nchar(new_line) == 0) -1 else 0
                # Only keep lines from the start of the block to the closing `}`
                # this is to ensure we drop blank lines that were between the end
                # of the block and the start of the next
                keep_range <- seq_len(index + keep_offset)
                results[[block]][[index]] <- new_line
                results[[block]] <- results[[block]][keep_range]
                has_removed_char <- TRUE
                break
            }
        }
        # If we haven't actually removed a closing `}` then something has gone wrong...
        if (!has_removed_char) {
            stop(sprintf(errmsg, block))
        }
    }

    # Add any missing blocks back in
    for (block in names(stan_blocks)) {
        if (is.null(results[[block]])) {
            results[[block]] <- ""
        }
    }
    results
}

#' `StanModule` -> Printable `Character`
#'
#' Converts [`StanModule`] object into a printable string.
#' @param object ([`StanModule`])\cr A stan program
#' @family StanModule
#' @param indent (`numeric`)\cr how much white space to prefix the print string with.
#' @keywords internal
#' @export
as_print_string.StanModule <- function(object, indent = 1, ...) {
    slots <- names(getSlots("StanModule"))
    components <- Filter(\(block) paste(slot(object, block), collapse = "") != "", slots)
    template <- c(
        "StanModule Object with components:",
        paste("    ", components)
    )
    pad <- rep(" ", indent) |> paste(collapse = "")
    template_padded <- paste(pad, template)
    sprintf(
        paste(template_padded, collapse = "\n")
    )
}

#' @rdname show-object
#' @export
setMethod(
    f = "show",
    signature = "StanModule",
    definition = function(object) {
        string <- as_print_string(object)
        cat("\n", string, "\n\n")
    }
)

#' @include LongitudinalModel.R
#' @include StanModule.R
#' @include generics.R
#' @include ParameterList.R
#' @include Parameter.R
#' @include Link.R
NULL

# LongitudinalGSF-class ----

#' `LongitudinalGSF`
#'
#' This class extends the general [`LongitudinalModel`] class for using the
#' Generalized Stein-Fojo (GSF) model for the longitudinal outcome.
#'
#' @section Available Links:
#' - [`linkDSLD()`]
#' - [`linkTTG()`]
#' - [`linkIdentity()`]
#' - [`linkGrowth()`]
#' @exportClass LongitudinalGSF
.LongitudinalGSF <- setClass(
    Class = "LongitudinalGSF",
    contains = "LongitudinalModel"
)

# LongitudinalGSF-constructors ----

#' @rdname LongitudinalGSF-class
#'
#' @param mu_bsld (`Prior`)\cr for the mean baseline value `mu_bsld`.
#' @param mu_ks (`Prior`)\cr for the mean shrinkage rate `mu_ks`.
#' @param mu_kg (`Prior`)\cr for the mean growth rate `mu_kg`.
#' @param mu_phi (`Prior`)\cr for the mean proportion of cells affected by the treatment `mu_phi`.
#'
#' @param omega_bsld (`Prior`)\cr for the baseline value standard deviation `omega_bsld`.
#' @param omega_ks (`Prior`)\cr for the shrinkage rate standard deviation `omega_ks`.
#' @param omega_kg (`Prior`)\cr for the growth rate standard deviation `omega_kg`.
#' @param omega_phi (`Prior`)\cr for the standard deviation of the proportion of cells
#' affected by the treatment `omega_phi`.
#'
#' @param sigma (`Prior`)\cr for the variance of the longitudinal values `sigma`.
#'
#' @param centred (`logical`)\cr whether to use the centred parameterization.
#'
#' @importFrom stats qlogis
#' @export
LongitudinalGSF <- function(

    mu_bsld = prior_normal(log(60), 1),
    mu_ks = prior_normal(log(0.5), 1),
    mu_kg = prior_normal(log(0.3), 1),
    mu_phi = prior_normal(qlogis(0.5), 1),

    omega_bsld = prior_lognormal(log(0.2), 1),
    omega_ks = prior_lognormal(log(0.2), 1),
    omega_kg = prior_lognormal(log(0.2), 1),
    omega_phi = prior_lognormal(log(0.2), 1),

    sigma = prior_lognormal(log(0.1), 1),

    centred = FALSE
) {

    gsf_model <- StanModule(decorated_render(
        .x = read_stan("lm-gsf/model.stan"),
        centred = centred
    ))

    # Apply constraints
    omega_bsld <- set_limits(omega_bsld, lower = 0)
    omega_ks <- set_limits(omega_ks, lower = 0)
    omega_kg <- set_limits(omega_kg, lower = 0)
    omega_phi <- set_limits(omega_phi, lower = 0)
    sigma <- set_limits(sigma, lower = 0)


    parameters <- list(
        Parameter(name = "lm_gsf_mu_bsld", prior = mu_bsld, size = "n_studies"),
        Parameter(name = "lm_gsf_mu_ks", prior = mu_ks, size = "n_arms"),
        Parameter(name = "lm_gsf_mu_kg", prior = mu_kg, size = "n_arms"),
        Parameter(name = "lm_gsf_mu_phi", prior = mu_phi, size = "n_arms"),

        Parameter(name = "lm_gsf_omega_bsld", prior = omega_bsld, size = 1),
        Parameter(name = "lm_gsf_omega_ks", prior = omega_ks, size = 1),
        Parameter(name = "lm_gsf_omega_kg", prior = omega_kg, size = 1),
        Parameter(name = "lm_gsf_omega_phi", prior = omega_phi, size = 1),

        Parameter(name = "lm_gsf_sigma", prior = sigma, size = 1)
    )

    assert_flag(centred)
    parameters_extra <- if (centred) {
        list(
            Parameter(
                name = "lm_gsf_psi_bsld",
                prior = prior_init_only(prior_lognormal(median(mu_bsld), median(omega_bsld))),
                size = "n_subjects"
            ),
            Parameter(
                name = "lm_gsf_psi_ks",
                prior = prior_init_only(prior_lognormal(median(mu_ks), median(omega_ks))),
                size = "n_subjects"
            ),
            Parameter(
                name = "lm_gsf_psi_kg",
                prior = prior_init_only(prior_lognormal(median(mu_kg), median(omega_kg))),
                size = "n_subjects"
            ),
            Parameter(
                name = "lm_gsf_psi_phi_logit",
                prior = prior_init_only(prior_normal(median(mu_phi), median(omega_phi))),
                size = "n_subjects"
            )
        )
    } else {
        list(
            Parameter(name = "lm_gsf_eta_tilde_bsld", prior = prior_std_normal(), size = "n_subjects"),
            Parameter(name = "lm_gsf_eta_tilde_ks", prior = prior_std_normal(), size = "n_subjects"),
            Parameter(name = "lm_gsf_eta_tilde_kg", prior = prior_std_normal(), size = "n_subjects"),
            Parameter(name = "lm_gsf_eta_tilde_phi", prior = prior_std_normal(), size = "n_subjects")
        )
    }
    parameters <- append(parameters, parameters_extra)

    x <- LongitudinalModel(
        name = "Generalized Stein-Fojo",
        stan = merge(
            gsf_model,
            StanModule("lm-gsf/functions.stan")
        ),
        parameters = do.call(ParameterList, parameters)
    )
    .LongitudinalGSF(x)
}



#' @export
enableGQ.LongitudinalGSF <- function(object, ...) {
    StanModule("lm-gsf/quantities.stan")
}


#' @export
enableLink.LongitudinalGSF <- function(object, ...) {
    object@stan <- merge(
        object@stan,
        StanModule("lm-gsf/link.stan")
    )
    object
}

#' @export
linkDSLD.LongitudinalGSF <- function(prior = prior_normal(0, 2), model, ...) {
    LinkComponent(
        key = "link_dsld",
        stan = StanModule("lm-gsf/link_dsld.stan"),
        prior = prior
    )
}

#' @export
linkTTG.LongitudinalGSF <- function(prior = prior_normal(0, 2), model, ...) {
    LinkComponent(
        key = "link_ttg",
        stan = StanModule("lm-gsf/link_ttg.stan"),
        prior = prior
    )
}

#' @export
linkIdentity.LongitudinalGSF <- function(prior = prior_normal(0, 2), model, ...) {
    LinkComponent(
        key = "link_identity",
        stan = StanModule("lm-gsf/link_identity.stan"),
        prior = prior
    )
}

#' @export
linkGrowth.LongitudinalGSF <- function(prior = prior_normal(0, 2), model, ...) {
    LinkComponent(
        key = "link_growth",
        stan = StanModule("lm-gsf/link_growth.stan"),
        prior = prior
    )
}


#' @rdname getPredictionNames
#' @export
getPredictionNames.LongitudinalGSF <- function(object, ...) {
    c("b", "s", "g", "phi")
}

#' @include generics.R
#' @include utilities.R
NULL

#' Re-used documentation for `DataLongitudinal`
#'
#' @param object ([`DataLongitudinal`]) \cr Longitudinal Data.
#' @param x ([`DataLongitudinal`]) \cr Longitudinal Data.
#' @param ... Not Used.
#'
#' @name DataLongitudinal-Shared
#' @keywords internal
NULL


# DataLongitudinal-class ----

#' Longitudinal Data Object and Constructor Function
#'
#' The [`DataLongitudinal`] class handles the processing of the longitudinal data for fitting a Joint Model.
#'
#'
#' @slot data (`data.frame`)\cr See Arguments for details; Note that
#'   observations that contain missing values in the required variables are removed.
#' @slot formula (`formula`)\cr See Arguments for details
#' @slot threshold (`numeric`)\cr See Arguments for details
#'
#' @family DataObjects
#' @family DataLongitudinal
#' @export DataLongitudinal
#' @exportClass DataLongitudinal
.DataLongitudinal <- setClass(
    Class = "DataLongitudinal",
    representation = list(
        data = "data.frame",
        formula = "formula",
        threshold = "numeric_or_NULL"
    )
)


#' @param data (`data.frame`)\cr containing the observed longitudinal data.
#' @param formula (`formula`)\cr of the form `outcome ~ time`, and cannot contain any additional covariates.
#' @param threshold (`numeric`)\cr cut-off value to be used to declare an observation as censored
#'   (below detection limit).
#' @rdname DataLongitudinal-class
DataLongitudinal <- function(data, formula, threshold = NULL) {
    .DataLongitudinal(
        data = remove_missing_rows(data, formula),
        formula = formula,
        threshold = threshold
    )
}

setValidity(
    "DataLongitudinal",
    function(object) {
        if (!length(object@formula) == 3) {
            return("`formula` should be a 2 sided formula")
        }
        if (!length(object@formula[[3]]) == 1) {
            return("the RHS of `formula` should only have 1 value")
        }
        if (!length(object@threshold) <= 1) {
            return("`threshold` must be of length 1 or `NULL`")
        }
        vars <- extractVariableNames(object)
        vars$threshold <- NULL
        vars$frm <- NULL
        for (i in unlist(vars)) {
            if (! i %in% names(object@data)) {
                return(sprintf("Variable `%s` is not in data", i))
            }
        }
        return(TRUE)
    }
)




#' @rdname harmonise
harmonise.DataLongitudinal <- function(object, subject_var, subject_ord, ...) {
    data <- as.data.frame(object)
    vars <- extractVariableNames(object)
    assert_string(subject_var, na.ok = FALSE)
    assert_character(subject_ord, any.missing = FALSE)
    assert_that(
        subject_var %in% names(data),
        msg = sprintf("Subject variable `%s` not found in `longitudinal`", subject_var)
    )
    assert_that(
        all(data[[subject_var]] %in% subject_ord),
        msg = "There are subjects in `longitudinal` that are not present in `subjects`"
    )
    assert_that(
        all(subject_ord %in% data[[subject_var]]),
        msg = "There are subjects in `subjects` that are not present in `longitudinal`"
    )
    data[[subject_var]] <- factor(
        as.character(data[[subject_var]]),
        levels = subject_ord
    )
    data_re_ord <- order(
        data[[subject_var]],
        data[[vars$time]],
        data[[vars$outcome]]
    )
    data_ord <- data[data_re_ord, ]
    DataLongitudinal(
        data = data_ord,
        formula = object@formula,
        threshold = object@threshold
    )
}



#' `DataLongitudinal` -> `data.frame`
#'
#' @inheritParams DataLongitudinal-Shared
#'
#' @description
#' Converts a [`DataLongitudinal`] object into a `data.frame`.
#' The subject variable is cast to factor.
#' @family DataLongitudinal
#' @export
as.data.frame.DataLongitudinal <- function(x, ...) {
    x <- x@data
    rownames(x) <- NULL
    x
}




#' @inheritParams DataLongitudinal-Shared
#' @inherit extractVariableNames description title
#'
#' @returns
#' A list with the following named elements:
#' - `subject` (`character`)\cr The name of the variable containing the subject identifier
#' - `frm` (`formula`)\cr of the form `outcome ~ time`
#' - `time` (`character`)\cr The name of the variable containing the outcome time
#' - `outcome` (`character`)\cr The name of the variable containing the outcome values
#' - `threshold` (`numeric`)\cr cut-off value to be used to declare an observation as censored
#'   (below detection limit).
#' @family DataLongitudinal
#' @family extractVariableNames
extractVariableNames.DataLongitudinal <- function(object) {
    list(
        frm = object@formula,
        time = as.character(object@formula[[3]]),
        outcome = as.character(object@formula[[2]]),
        threshold = object@threshold
    )
}


#' @rdname as_stan_list.DataObject
#' @family DataLongitudinal
#' @export
as_stan_list.DataLongitudinal <- function(object, subject_var, ...) {

    df <- as.data.frame(object)
    vars <- extractVariableNames(object)

    assert_factor(df[[subject_var]])

    mat_sld_index <- stats::model.matrix(
        stats::as.formula(paste("~ -1 + ", subject_var)),
        data = df
    ) |>
        t()

    adj_threshold <- if (is.null(vars$threshold)) {
        -999999
    } else {
        vars$threshold
    }

    index_obs <- which(df[[vars$outcome]] >= adj_threshold)
    index_cen <- which(df[[vars$outcome]] < adj_threshold)

    sparse_mat_inds_all_y <- rstan::extract_sparse_parts(mat_sld_index)
    sparse_mat_inds_obs_y <- rstan::extract_sparse_parts(mat_sld_index[, index_obs])
    sparse_mat_inds_cens_y <- rstan::extract_sparse_parts(mat_sld_index[, index_cen])

    model_data <- list(
        n_tumour_all = nrow(df),

        # Number of individuals and tumour assessments.
        n_tumour_obs = length(index_obs),
        n_tumour_cens = length(index_cen),

        # Index vectors
        subject_tumour_index = as.numeric(df[[subject_var]]),
        subject_tumour_index_obs = index_obs,
        subject_tumour_index_cens = index_cen,

        tumour_value = df[[vars$outcome]],
        tumour_time = df[[vars$time]],
        tumour_value_lloq = adj_threshold,

        # Sparse matrix parameters
        # Matrix of individuals x observed tumour assessments.
        n_mat_inds_obs_y = c(
            length(sparse_mat_inds_obs_y$w),
            length(sparse_mat_inds_obs_y$v),
            length(sparse_mat_inds_obs_y$u)
        ),
        w_mat_inds_obs_y = sparse_mat_inds_obs_y$w,
        v_mat_inds_obs_y = sparse_mat_inds_obs_y$v,
        u_mat_inds_obs_y = sparse_mat_inds_obs_y$u,

        # Matrix of individuals x censored tumour assessments.
        n_mat_inds_cens_y = c(
            length(sparse_mat_inds_cens_y$w),
            length(sparse_mat_inds_cens_y$v),
            length(sparse_mat_inds_cens_y$u)
        ),
        w_mat_inds_cens_y = sparse_mat_inds_cens_y$w,
        v_mat_inds_cens_y = sparse_mat_inds_cens_y$v,
        u_mat_inds_cens_y = sparse_mat_inds_cens_y$u,

        # Matrix of all individuals tumour assessments
        n_mat_inds_all_y = c(
            length(sparse_mat_inds_all_y$w),
            length(sparse_mat_inds_all_y$v),
            length(sparse_mat_inds_all_y$u)
        ),
        w_mat_inds_all_y = sparse_mat_inds_all_y$w,
        v_mat_inds_all_y = sparse_mat_inds_all_y$v,
        u_mat_inds_all_y = sparse_mat_inds_all_y$u

    )

    return(model_data)
}

#' @rdname as_stan_list.DataObject
#' @export
as.list.DataLongitudinal <- function(x, ...) {
    as_stan_list(x, ...)
}


#' `DataLongitudinal` -> Printable `Character`
#'
#' Converts [`DataLongitudinal`] object into a printable string.
#' @inheritParams DataLongitudinal-Shared
#' @family DataLongitudinal
#' @param indent (`numeric`)\cr how much white space to prefix the print string with.
#' @keywords internal
#' @export
as_print_string.DataLongitudinal <- function(object, indent = 1, ...) {
    template <- c(
        "Longitudinal-Data Object:",
        "    # of Rows     = %d",
        "    # of Columns  = %d",
        "    # of Cen-Obvs = %d",
        "    Formula       = %s"
    )
    pad <- rep(" ", indent) |> paste(collapse = "")
    template_padded <- paste(pad, template)
    vars <- extractVariableNames(object)
    sprintf(
        paste(template_padded, collapse = "\n"),
        nrow(object@data),
        ncol(object@data),
        sum(object@data[[vars$outcome]] < vars$threshold),
        Reduce(paste, deparse(vars$frm))
    )
}

#' @rdname show-object
#' @export
setMethod(
    f = "show",
    signature = "DataLongitudinal",
    definition = function(object) {
        string <- as_print_string(object)
        cat("\n", string, "\n\n")
    }
)

#' @include StanModule.R
#' @include LongitudinalModel.R
#' @include ParameterList.R
#' @include LinkComponent.R
#' @include Prior.R
NULL


#' `Link` Function Arguments
#'
#' This exists just to contain all the common arguments for [`Link`] methods.
#'
#' @param x ([`Link`])\cr a link object.
#' @param object ([`Link`])\cr a link object.
#' @param ... Not Used.
#'
#' @name Link-Shared
#' @keywords internal
NULL



#' `Link`
#'
#' @slot components (`list`)\cr a list of [`LinkComponent`] or [`PromiseLinkComponent`] objects.
#' @slot resolved (`logical`)\cr indicates if all the `components` have been resolved.
#'
#' @param ... ([`LinkComponent`] or [`PromiseLinkComponent`])\cr
#' an arbitrary number of link components.
#'
#' @description
#' Simple container class to enable the use of multiple link components in a joint model.
#' Note that the constructor of this object is idempotent e.g. `Link(Link(x)) == Link(x)`
#'
#' @examples
#' Link(
#'     linkDSLD(),
#'     linkTTG()
#' )
#'
#' @family Link
#' @name Link-class
#' @exportClass Link
.Link <- setClass(
    Class = "Link",
    slots = list(
        components = "list",
        resolved = "logical"
    )
)

#' @rdname Link-class
#' @export
Link <- function(...) {
    components <- list(...)

    # If the input is already a Link object, return it (e.g. implement
    # a constructor that is idempotent)
    if (length(components) == 1 && is(components[[1]], "Link")) {
        return(components[[1]])
    }

    .Link(
        components = components,
        resolved = !any(vapply(components, \(x) is(x, "PromiseLinkComponent"), logical(1)))
    )
}


#' Resolve any promises
#'
#' Loops over all components and ensures that any [`PromiseLinkComponent`] objects
#' are resolved to [`LinkComponent`] objects.
#'
#' @param object ([`Link`])\cr a link object.
#' @param model ([`LongitudinalModel`])\cr the model object.
#' @param ... Not Used.
#'
#' @export
resolvePromise.Link <- function(object, model, ...) {
    if (length(object) == 0) {
        return(object)
    }
    assert_that(
        is(model, "LongitudinalModel"),
        msg = "model must be of class `LongitudinalModel`"
    )
    do.call(Link, lapply(object@components, resolvePromise, model = model))
}


setValidity(
    Class = "Link",
    method = function(object) {

        for (i in object@components) {
            if (!(is(i, "LinkComponent") || is(i, "PromiseLinkComponent"))) {
                return("All components must be of class `LinkComponent` or `PromiseLinkComponent`")
            }
        }

        contains_promise <- any(
            vapply(
                object@components,
                \(x) is(x, "PromiseLinkComponent"),
                logical(1)
            )
        )
        if (contains_promise & object@resolved) {
            return("Object cannot be resolved if it contains promises")
        }

        if (length(object@resolved) > 1) {
            return("The `resolved` slot must be a logical scalar")
        }
        return(TRUE)
    }
)




#' `Link` -> `StanModule`
#'
#' Converts a [`Link`] object to a [`StanModule`] object
#'
#' @inheritParams Link-Shared
#'
#' @family Link
#' @family as.StanModule
#' @export
as.StanModule.Link <- function(object, ...) {

    if (length(object@components) == 0) {
        return(StanModule("base/link_none.stan"))
    }

    keys <- vapply(
        object@components,
        function(x) x@key,
        character(1)
    )

    base_stan <- StanModule(
        decorated_render(
            .x = read_stan("base/link.stan"),
            items = as.list(keys)
        )
    )

    stan_list <- lapply(
        object@components,
        as.StanModule
    )

    stan <- Reduce(
        merge,
        append(base_stan, stan_list)
    )
    return(stan)
}




#' `Link` -> `list`
#'
#' @inheritParams Link-Shared
#'
#' @description
#' Returns a named list where each element of the list corresponds
#' to a Stan modelling block e.g. `data`, `model`, etc.
#'
#' @family Link
#' @export
as.list.Link <- function(x, ...) {
    as.list(as.StanModule(x, ...))
}



#' @export
#' @rdname getParameters
getParameters.Link <- function(object, ...) {
    parameters_list <- lapply(
        object@components,
        getParameters,
        ...
    )
    Reduce(
        merge,
        parameters_list
    )
}


#' @rdname initialValues
#' @export
initialValues.Link <- function(object, ...) {
    unlist(
        lapply(object@components, initialValues),
        recursive = FALSE
    )
}


#' `Link` -> `list`
#'
#' @inheritParams Link-Shared
#'
#' @description
#' Returns the number of link components within the [`Link`] object
#'
#' @family Link
#' @export
length.Link <- function(x) {
    length(x@components)
}


#' @export
as_print_string.Link <- function(object, ...) {
    if (length(object) == 0) {
        return("\nNo Link")
    }

    strings <- vapply(object@components, as_print_string, character(1))

    paste(
        c(
            "\nLink with the following components/parameters:",
            paste0("    ", strings)
        ),
        collapse = "\n"
    )
}


#' @rdname show-object
#' @export
setMethod(
    f = "show",
    signature = "Link",
    definition = function(object) {
        cat(paste0(as_print_string(object), "\n"))
    }
)

#' @include Grid.R
#' @include GridFixed.R
#' @include generics.R
NULL

#' @rdname Grid-Dev
.GridGrouped <- setClass(
    "GridGrouped",
    contains = "Grid",
    slots = c(
        "groups" = "list",
        "times" = "numeric_or_NULL"
    )
)

#' @rdname Grid-Functions
#' @export
GridGrouped <- function(groups, times = NULL) {
    .GridGrouped(
        groups = groups,
        times = times
    )
}


setValidity(
    "GridGrouped",
    function(object) {
        if (!all(vapply(object@groups, is.character, logical(1)))) {
            return("Each element of `groups` must be a character vector")
        }
        gnames <- names(object@groups)
        gnames <- gnames[!is.na(gnames) & gnames != ""]
        if (length(gnames) != length(object@groups)) {
            return("Each element of `groups` must be named")
        }
        return(TRUE)
    }
)


#' @rdname Quant-Dev
#' @export
as.QuantityGenerator.GridGrouped <- function(object, data, ...) {
    assert_class(data, "DataJoint")
    data_list <- as.list(data)
    subjects_unique <- unique(unlist(object@groups))
    assert_that(
        all(subjects_unique %in% names(data_list$subject_to_index))
    )
    as.QuantityGenerator(
        GridFixed(
            times = object@times,
            subjects = subjects_unique
        ),
        data = data
    )
}


#' @rdname Quant-Dev
#' @export
as.QuantityCollapser.GridGrouped <- function(object, data, ...) {
    assert_class(data, "DataJoint")
    data_list <- as.list(data)
    assert_that(
        all(unique(unlist(object@groups)) %in% names(data_list$subject_to_index))
    )

    validate_time_grid(object@times)

    group_grid <- expand.grid(
        group = names(object@groups),
        time = object@times,
        stringsAsFactors = FALSE
    )

    generator <- as.QuantityGenerator(object, data)

    select_indexes <- mapply(
        function(group, time) {
            correct_subject <- generator@subjects %in% object@groups[[group]]
            correct_time <- generator@times == time
            seq_along(correct_time)[correct_subject & correct_time]
        },
        group_grid$group,
        group_grid$time,
        SIMPLIFY = FALSE
    )
    names(select_indexes) <- NULL

    QuantityCollapser(
        times = group_grid$time,
        groups = group_grid$group,
        indexes = select_indexes
    )
}

#' @export
as.list.GridGrouped <- function(x, ...) {
    x@groups
}

#' @rdname coalesceGridTime
#' @export
coalesceGridTime.GridGrouped <- function(object, times, ...) {
    if (is.null(object@times)) {
        object <- GridGrouped(
            groups = object@groups,
            times = times
        )
    }
    object
}

#' @include Grid.R
#' @include generics.R
NULL


#' @rdname Grid-Dev
.GridPrediction <- setClass(
    "GridPrediction",
    contains = "Grid",
    slots = c(
        "times" = "numeric_or_NULL",
        "newdata" = "data.frame",
        "params" = "list"
    )
)
#' @rdname Grid-Functions
#' @export
GridPrediction <- function(times = NULL, newdata, params = list()) {
    .GridPrediction(
        times = times,
        params = params,
        newdata = newdata
    )
}
setValidity(
    "GridPrediction",
    function(object) {
        for (param in names(object@params)) {
            if (length(object@params[[param]]) != 1) {
                return(sprintf("Parameter '%s' must be length 1", param))
            }
        }
        if (length(object@params) != length(names(object@params))) {
            return("All elements of `params` must be named")
        }
        if (!is.null(object@newdata[["..new_subject.."]])) {
            return("`newdata` must not contain a column named '..new_subject..'")
        }
        return(TRUE)
    }
)


#' @rdname Quant-Dev
#' @export
as.QuantityGenerator.GridPrediction <- function(object, data, ...) {

    assert_class(data, "DataJoint")
    data_list <- as.list(data)
    validate_time_grid(object@times)

    n_times <- length(object@times)
    n_obs <- nrow(object@newdata)
    newdata <- object@newdata
    newdata[["..new_subject.."]] <- sprintf(
        "new_subject_%i",
        seq_len(nrow(newdata))
    )

    QuantityGeneratorPrediction(
        times = rep(object@times, each = n_obs),
        newdata = replicate(newdata, n = n_times, simplify = FALSE) |> dplyr::bind_rows(),
        params = object@params
    )
}

#' @rdname Quant-Dev
#' @export
as.QuantityCollapser.GridPrediction <- function(object, data, ...) {
    assert_class(data, "DataJoint")
    data_list <- as.list(data)
    generator <- as.QuantityGenerator(object, data)
    QuantityCollapser(
        times = generator@times,
        groups = generator@newdata[["..new_subject.."]],
        indexes = as.list(seq_along(generator@times))
    )
}


#' @export
as.list.GridPrediction <- function(x, data, ...) {
    stop("`as.list()` is not implemented for `GridPrediction` objects")
}


#' @rdname coalesceGridTime
#' @export
coalesceGridTime.GridPrediction <- function(object, times, ...) {
    if (is.null(object@times)) {
        object <- GridPrediction(
            times = times,
            newdata = object@newdata,
            params = object@params
        )
    }
    object
}

#' @include DataSurvival.R
#' @include DataLongitudinal.R
#' @include DataSubject.R
NULL


#' Re-used documentation for `DataJoint`
#'
#' @param object ([`DataJoint`]) \cr Survival and Longitudinal Data.
#' @param x ([`DataJoint`]) \cr Survival and Longitudinal Data.
#' @param ... Not Used.
#'
#' @name DataJoint-Shared
#' @keywords internal
NULL

setClassUnion("DataLongitudinal_or_NULL", c("DataLongitudinal", "NULL"))
setClassUnion("DataSurvival_or_NULL", c("DataSurvival", "NULL"))


# DataJoint-class ----

#' @title
#' Joint Data Object and Constructor Function
#'
#' @description
#' The `DataJoint` class handles combining data from a [`DataSurvival`] object and a
#' [`DataLongitudinal`] object.
#'
#' @slot subject (`DataSubject`)\cr See Argument for details.
#' @slot survival (`DataSurvival`)\cr See Argument for details.
#' @slot longitudinal (`DataLongitudinal`)\cr See Argument for details.
#'
#' @family DataObjects
#' @family DataJoint
#' @export DataJoint
#' @exportClass DataJoint
.DataJoint <- setClass(
    Class = "DataJoint",
    representation = list(
        subject = "DataSubject",
        survival = "DataSurvival_or_NULL",
        longitudinal = "DataLongitudinal_or_NULL"
    )
)

#' @param subject (`DataSubject`)\cr object created by [DataSubject()].
#' @param survival (`DataSurvival`)\cr object created by [DataSurvival()].
#' @param longitudinal (`DataLongitudinal`)\cr object created by [DataLongitudinal()].
#' @rdname DataJoint-class
DataJoint <- function(subject, survival = NULL, longitudinal = NULL) {

    subject_suited <- harmonise(subject)
    vars <- extractVariableNames(subject)
    subject_var <- vars$subject
    subject_ord <- levels(as.data.frame(subject_suited)[[vars$subject]])

    survival_suited <- harmonise(
        survival,
        subject_var = subject_var,
        subject_ord = subject_ord
    )

    longitudinal_suited <- harmonise(
        longitudinal,
        subject_var = subject_var,
        subject_ord = subject_ord
    )

    .DataJoint(
        subject = subject_suited,
        survival = survival_suited,
        longitudinal = longitudinal_suited
    )
}



setValidity(
    Class = "DataJoint",
    method = function(object) {
        vars <- extractVariableNames(object@subject)
        subject_var <- vars$subject
        subject_ord <- as.character(as.data.frame(object@subject)[[vars$subject]])
        if (!is.null(object@survival)) {
            survival_df <- as.data.frame(object@survival)
            if (!subject_var %in% names(survival_df)) {
                return(sprintf("Unable to find `%s` in `survival`", sujbect_var))
            }
            if (!all(survival_df[[subject_var]] %in% subject_ord)) {
                return("There are subjects in `survival` that are not in `subject`")
            }
            if (!nrow(survival_df) == length(unique(survival_df[[subject_var]]))) {
                return("There are duplicate subjects in `survival`")
            }
        }
        if (!is.null(object@longitudinal)) {
            long_df <- as.data.frame(object@longitudinal)
            if (!subject_var %in% names(long_df)) {
                return(sprintf("Unable to find `%s` in `longitudinal`", sujbect_var))
            }
            if (!all(long_df[[subject_var]] %in% subject_ord)) {
                return("There are subjects in `longitudinal` that are not in `subject`")
            }
        }
        subject_df <- as.data.frame(object@subject)
        if (!subject_var %in% names(subject_df)) {
            return(sprintf("Unable to find `%s` in `subject`", sujbect_var))
        }
        if (!nrow(subject_df) == length(unique(subject_df[[subject_var]]))) {
            return("There are duplicate subjects in `subject`")
        }
        return(TRUE)
    }
)


# DataJoint-as.list ----




#' Data Object -> `list`
#'
#' @param object (`DataSubject` or `DataLongitudinal` or `DataSurvival`) \cr
#' data object to convert to a `list`.
#' @param x (`DataSubject` or `DataLongitudinal` or `DataSurvival`) \cr
#' data object to convert to a `list`.
#' @param subject_var (`character`) \cr the name of the variable
#' containing the subject identifier.
#' @param ... not used.
#'
#' @description
#' Coerces a data object into a `list` of data components required
#' for fitting a [`JointModel`]. See the "Extending jmpost" vignette for more details.
#'
#' @name as_stan_list.DataObject
#' @family as_stan_list
#' @family DataJoint
#' @export
as_stan_list.DataJoint <- function(object, ...) {
    vars <- extractVariableNames(object@subject)
    subject_var <- vars$subject
    as_stan_list(object@subject) |>
        append(as_stan_list(object@survival)) |>
        append(as_stan_list(
            object@longitudinal,
            subject_var = subject_var
        ))
}

#' @rdname as_stan_list.DataObject
#' @export
as.list.DataJoint <- function(x, ...) {
    as_stan_list(x, ...)
}



#' Subsetting `DataJoint` as a `data.frame`
#'
#' @param x (`DataJoint`) \cr object created by [DataJoint()].
#' @param subjects (`character` or `list`)\cr subjects that you wish to subset the `data.frame`
#' to contain. See details.
#' @param ... Not used.
#'
#' @description
#'
#' Coerces the object into a `data.frame` containing just event times and status
#' filtering for specific subjects If `subjects` is a list then an additional variable `group` will be added
#' onto the dataset specifying which group the row belongs to.
#'
#' @examples
#' \dontrun{
#' subjects <- c("SUB1", "SUB2", "SUB3", "SUB4")
#' subset(x, subjects)
#'
#' groups <- list(
#'     "g1" = c("SUB1", "SUB3", "SUB4"),
#'     "g2" = c("SUB2", "SUB3")
#' )
#' subset(x, groups)
#' }
#' @family DataJoint
#' @export
subset.DataJoint <- function(x, subjects, ...) {
    data <- as.list(x)
    dat <- data.frame(
        time = data[["event_times"]],
        event = as.numeric(seq_along(data[["event_times"]]) %in% data[["subject_event_index"]]),
        subject = names(data[["subject_to_index"]])
    )
    subset_and_add_grouping(dat, subjects)
}



#' `subset_and_add_grouping`
#'
#' @param dat (`data.frame`) \cr must have a column called `subject` which corresponds to the
#' values passed to `groupings`.
#' @param groupings (`character` or `list`)\cr subjects that you wish to subset the dataset
#' to contain. If `groupings` is a list then an additional variable `group` will be added
#' onto the dataset specifying which group the row belongs to.
#'
#' @details
#' Example of usage
#' ```
#' subjects <- c("SUB1", "SUB2", "SUB3", "SUB4")
#' subset_and_add_grouping(dat, subjects)
#'
#' groups <- list(
#'     "g1" = c("SUB1", "SUB3", "SUB4"),
#'     "g2" = c("SUB2", "SUB3")
#' )
#' subset_and_add_grouping(dat, groups)
#' ```
#'
#' @keywords internal
subset_and_add_grouping <- function(dat, groupings) {
    groupings <- decompose_subjects(groupings, dat$subject)$groups
    dat_subset_list <- lapply(
        seq_along(groupings),
        \(i) {
            dat_reduced <- dat[dat$subject %in% groupings[[i]], , drop = FALSE]
            dat_reduced[["group"]] <- names(groupings)[[i]]
            dat_reduced
        }
    )
    x <- Reduce(rbind, dat_subset_list)
    row.names(x) <- NULL
    x
}


#' Extract Observed Longitudinal Values
#'
#' Utility function to extract the observed longitudinal values from a [`DataJoint`] object
#' @param object ([`DataJoint`])\cr data used to fit a [`JointModel`].
#' @return A data.frame with the following columns
#' - `subject` (`character`)\cr The subject identifier
#' - `time` (`numeric`)\cr The time at which the observation occurred
#' - `Yob` (`numeric`)\cr The observed value
#' @keywords internal
extract_observed_values <- function(object) {
    assert_class(object, "DataJoint")
    data <- as.list(object)
    x <- data.frame(
        subject = names(data$subject_to_index)[data$subject_tumour_index],
        time = data$tumour_time,
        Yob = data$tumour_value
    )
    row.names(x) <- NULL
    x
}

#' @rdname show-object
#' @export
setMethod(
    f = "show",
    signature = "DataJoint",
    definition = function(object) {
        string_survival <- if (is.null(object@survival)) {
            "      Survival-Data Object:\n          Not Specified"
        } else {
            as_print_string(object@survival, indent = 5)
        }

        string_longitudinal <- if (is.null(object@longitudinal)) {
            "      Longitudinal-Data Object:\n          Not Specified"
        } else {
            as_print_string(object@longitudinal, indent = 5)
        }

        template <- c(
            "Joint-Data Object Containing:",
            "",
            as_print_string(object@subject, indent = 5),
            "",
            string_survival,
            "",
            string_longitudinal
        )
        cat("\n", paste(template, collapse = "\n"), "\n\n")
    }
)


#' Re-used documentation for `Quantities`
#'
#' @param x ([`Quantities`]) \cr generated quantities.
#' @param object ([`Quantities`]) \cr generated quantities.
#' @param type (`character`)\cr sets the `type` variable.
#' @param conf.level (`numeric`) \cr confidence level of the interval.
#' @param ... not used.
#'
#' @keywords internal
#' @name Quantities-Shared
NULL


#' Generated Quantities Container
#'
#' A simple wrapper around a `matrix` to store required metadata
#'
#' @param quantities (`matrix`)\cr of generated quantities.
#' @param times (`numeric`)\cr labels specifying which time point the quantity was generated at.
#' @param groups (`character`)\cr labels for which group the quantity belongs to.
#'
#' @slot quantities (`matrix`)\cr See Arguments for details.
#' @slot times (`numeric`)\cr See Arguments for details.
#' @slot groups (`character`)\cr See Arguments for details.
#'
#' @details
#' Each row of the matrix represents a sample and each column represents a distinct quantity.
#' As such the number of columns in the matrix should equal the length of `times` and `groups`
#' which provide metadata for who the quantity belongs to and at what time point it was generated at.
#'
#' @keywords internal
#' @name Quantities-class
#' @family Quantities
.Quantities <- setClass(
    "Quantities",
    slots = list(
        "quantities" = "matrix",
        "times" = "numeric",
        "groups" = "character"
    )
)
#' @rdname Quantities-class
Quantities <- function(quantities, times, groups) {
    .Quantities(
        quantities = quantities,
        times = times,
        groups = groups
    )
}

setValidity(
    Class = "Quantities",
    method = function(object) {
        if (length(object@times) != length(object@groups)) {
            return("Length of `times` must be equal to the length of `groups`")
        }
        if (length(object@times) != ncol(object@quantities)) {
            return("Length of `times` must be equal to the number of columns in `quantities`")
        }
        return(TRUE)
    }
)



#' @export
dim.Quantities <- function(x) {
    dim(x@quantities)
}


#' `Quantities` -> `data.frame`
#'
#' @inheritParams Quantities-Shared
#'
#' @keywords internal
#' @family Quantities
#' @export
as.data.frame.Quantities <- function(x, ...) {
    data.frame(
        group = rep(x@groups, each = nrow(x@quantities)),
        time = rep(x@times, each = nrow(x@quantities)),
        values = as.vector(x@quantities)
    )
}


#' summary
#'
#' @description
#' This method returns a summary statistic `data.frame` of the quantities. Note that this
#' is just an internal utility method in order to share common code between
#' [LongitudinalQuantities] and [SurvivalQuantities]
#'
#' @inheritParams Quantities-Shared
#'
#' @returns
#' A `data.frame` with the following variables:
#' - `median` (`numeric`) \cr the median value of the quantity.
#' - `lower` (`numeric`) \cr the lower CI value of the quantity.
#' - `upper` (`numeric`) \cr the upper CI value of the quantity.
#' - `time` (`numeric`) \cr the time point which the quantity is for.
#' - `group` (`character`) \cr which group the quantity belongs to.
#' - `type` (`character`) \cr what type of quantity is it.
#'
#' @keywords internal
#' @family Quantities
#' @export
summary.Quantities <- function(object, conf.level = 0.95, ...) {
    quantities_summarised <- samples_median_ci(
        object@quantities,
        level = conf.level
    )

    quantities_summarised$group <- object@groups
    quantities_summarised$time <- object@times
    quantities_summarised[, c("group", "time", "median", "lower", "upper")]
}



#' Create Grouped Quantities
#'
#' This function takes a matrix of quantity samples and aggregates them by calculating the
#' pointwise average.
#'
#' @param quantities_raw (`matrix`)\cr of samples with 1 row per sample and 1 column per
#' distinct quantity.
#'
#' @param collapser ([`QuantityCollapser`])\cr specifies which columns to combine together.
#'
#' @details
#' This function essentially implements the group wise average by collapsing multiple columns
#' together based on the specification provided by the `QuantityCollapser` object.
#' The [Grid-Dev] page provides an example of what this function implements
#'
#' @keywords internal
collapse_quantities <- function(quantities_raw, collapser) {
    assert_class(quantities_raw, "matrix")
    assert_class(collapser, "QuantityCollapser")

    quantities <- matrix(
        NA,
        nrow = nrow(quantities_raw),
        ncol = length(collapser)
    )

    for (idx in seq_len(length(collapser))) {
        quantities[, idx] <- quantities_raw[
            ,
            collapser@indexes[[idx]],
            drop = FALSE
        ] |> rowMeans()
    }

    return(quantities)
}

#' Extract Survival Quantities
#'
#' Utility function to extract generated quantities from a [cmdstanr::CmdStanGQ] object.
#' Multiple quantities are generated by default so this is a convenience function to extract
#' the desired ones and return them them as a user friendly [posterior::draws_matrix] object
#'
#' @param gq (`CmdStanGQ`) \cr a [cmdstanr::CmdStanGQ] object created by [generateQuantities()].
#' @param type (`character`)\cr quantity to be generated.
#' Must be one of `surv`, `haz`, `loghaz`, `cumhaz`, `lm_identity`.
#' @keywords internal
extract_quantities <- function(gq, type = c("surv", "haz", "loghaz", "cumhaz", "lm_identity")) {
    type <- match.arg(type)
    assert_class(gq, "CmdStanGQ")
    meta <- switch(type,
        surv = list("log_surv_fit_at_time_grid", exp),
        cumhaz = list("log_surv_fit_at_time_grid", \(x) -x),
        haz = list("log_haz_fit_at_time_grid", exp),
        loghaz = list("log_haz_fit_at_time_grid", identity),
        lm_identity = list("y_fit_at_time_grid", identity)
    )
    result <- gq$draws(meta[[1]], format = "draws_matrix")
    result_transformed <- meta[[2]](result)
    cnames <- colnames(result_transformed)
    colnames(result_transformed) <- gsub(meta[[1]], "quantity", cnames)
    result_transformed
}


#' `Quantities` -> Printable `Character`
#'
#' Converts [`Quantities`] object into a printable string.
#' @inheritParams Quantities-Shared
#' @family Quantities
#' @keywords internal
#' @export
as_print_string.Quantities <- function(object, indent = 1, ...) {
    template <- c(
        "Quantities Object:",
        "    # of samples    = %d",
        "    # of quantities = %d"
    )
    pad <- rep(" ", indent) |> paste(collapse = "")
    template_padded <- paste(pad, template)
    sprintf(
        paste(template_padded, collapse = "\n"),
        nrow(object),
        ncol(object)
    )
}


#' @rdname show-object
#' @export
setMethod(
    f = "show",
    signature = "Quantities",
    definition = function(object) {
        string <- as_print_string(object)
        cat("\n", string, "\n\n")
    }
)


#' `SimSurvival` Function Arguments
#'
#' The documentation lists all the conventional arguments for [`SimSurvival`]
#' constructors.
#'
#' @param time_max (`number`)\cr the maximum time to simulate to.
#' @param time_step (`number`)\cr the time interval between evaluating the log-hazard function.
#' @param lambda_censor (`number`)\cr the censoring rate.
#' @param beta_cont (`number`)\cr the continuous covariate coefficient.
#' @param beta_cat (`numeric`)\cr the categorical covariate coefficients.
#' @param loghazard (`function`)\cr the log hazard function.
#' @param name (`character`)\cr the name of the object.
#' @param ... Not Used.
#'
#' @section Hazard Evaluation:
#'
#' Event times are simulated by sampling a cumulative hazard limit from a \eqn{U(0, 1)} distribution
#' for
#' each subject and then counting how much hazard they've been exposed to by evaluating the
#' log-hazard function at a set interval. The `time_max` argument sets the upper bound for the
#' number of time points to evaluate the log-hazard function at with subjects who have not had an
#' event being censored at `time_max`. The `time_step` argument sets the interval at which to
#' evaluate the log-hazard function. Setting smaller values for `time_step` will increase the
#' precision of the simulation at the cost of increased computation time. Likewise, setting large
#' values for `time_max` will minimize the number of censored subjects at the cost of
#' increased computation time.
#'
#' @name SimSurvival-Shared
#' @keywords internal
NULL


#' Abstract Simulation Class for Survival Data
#'
#' @inheritParams SimSurvival-Shared
#' @inheritSection SimSurvival-Shared Hazard Evaluation
#'
#' @slot time_max (`numeric`)\cr See arguments.
#' @slot time_step (`numeric`)\cr See arguments.
#' @slot lambda_censor (`numeric`)\cr See arguments.
#' @slot beta_cont (`numeric`)\cr See arguments.
#' @slot beta_cat (`numeric`)\cr See arguments.
#' @slot loghazard (`function`)\cr See arguments.
#' @slot name (`character`)\cr See arguments.
#'
#' @family SimSurvival
#' @exportClass SimSurvival
#' @name SimSurvival-class
.SimSurvival <- setClass(
    "SimSurvival",
    slots = c(
        time_max = "numeric",
        time_step = "numeric",
        lambda_censor = "numeric",
        beta_cont = "numeric",
        beta_cat = "numeric",
        loghazard = "function",
        name = "character"
    )
)

#' @rdname SimSurvival-class
#' @export
SimSurvival <- function(
    time_max = 2000,
    time_step = 1,
    lambda_censor = 1 / 3000,
    beta_cont = 0.2,
    beta_cat = c("A" = 0, "B" = -0.4, "C" = 0.2),
    loghazard,
    name = "SimSurvival"
) {
    .SimSurvival(
        time_max = time_max,
        time_step = time_step,
        lambda_censor = lambda_censor,
        beta_cont = beta_cont,
        beta_cat = beta_cat,
        loghazard = loghazard,
        name = name
    )
}


#' @rdname show-object
#' @export
setMethod(
    f = "show",
    signature = "SimSurvival",
    definition = function(object) {
        x <- sprintf("\nA %s Object\n\n", as_print_string(object))
        cat(x)
        return(object)
    }
)

#' @rdname as_print_string
as_print_string.SimSurvival <- function(object) {
    return(object@name)
}


#' Construct Time Intervals
#'
#' @param object (`SimSurvival`)\cr the survival simulation object to create evaluation points for.
#'
#' @return A `tibble` with `lower`, `upper`, `time`, `eval` and `width`.
#' @keywords internal
hazardWindows.SimSurvival <- function(object) {
    times <- seq(0, object@time_max, object@time_step)
    bound_lower <- times[-length(times)]
    bound_upper <- times[-1]
    bound_width <- bound_upper - bound_lower
    mid_point <- bound_upper - (bound_width / 2)
    tibble::tibble(
        lower = bound_lower,
        upper = bound_upper,
        midpoint = mid_point,
        width = bound_width
    )
}

#' @rdname sampleSubjects
#' @export
sampleSubjects.SimSurvival <- function(object, subjects_df) {
    subjects_df |>
        dplyr::mutate(cov_cont = stats::rnorm(dplyr::n())) |>
        dplyr::mutate(cov_cat = factor(
            sample(names(object@beta_cat), replace = TRUE, size = dplyr::n()),
            levels = names(object@beta_cat)
        )) |>
        dplyr::mutate(log_haz_cov = .data$cov_cont * object@beta_cont + object@beta_cat[.data$cov_cat]) |>
        dplyr::mutate(survival = stats::runif(dplyr::n())) |>
        dplyr::mutate(chazard_limit = -log(.data$survival)) |>
        dplyr::mutate(time_cen = stats::rexp(dplyr::n(), object@lambda_censor))
}


#' @rdname sampleObservations
#' @export
sampleObservations.SimSurvival <- function(object, times_df) {

    assert_that(
        all(times_df$time >= 0),
        msg = "All time points must be greater than or equal to 0"
    )

    os_dat_chaz <- times_df |>
        dplyr::mutate(log_bl_haz = object@loghazard(.data$midpoint)) |>
        # Fix to avoid issue with log(0) = NaN values
        dplyr::mutate(log_bl_haz = dplyr::if_else(.data$midpoint == 0, -999, .data$log_bl_haz)) |>
        dplyr::mutate(hazard_instant = exp(.data$log_bl_haz + .data$log_haz_cov + .data$log_haz_link)) |>
        # Reset Inf values to large number to avoid NaN issues downstream
        # This is suitable as Hazard limits tend to be in the range of -10 to 10 so large numbers
        # are essentially equivalent to infinity for simulation purposes
        dplyr::mutate(hazard_instant = dplyr::if_else(.data$hazard_instant == Inf, 999, .data$hazard_instant)) |>
        dplyr::mutate(hazard_instant = dplyr::if_else(.data$hazard_instant == -Inf, -999, .data$hazard_instant)) |>
        dplyr::mutate(hazard_interval = .data$hazard_instant * .data$width) |>
        dplyr::group_by(.data$subject) |>
        dplyr::mutate(chazard = cumsum(.data$hazard_interval)) |>
        dplyr::ungroup()

    os_had_event <- os_dat_chaz |>
        dplyr::filter(.data$chazard >= .data$chazard_limit) |>
        dplyr::group_by(.data$subject) |>
        dplyr::slice(1) |>
        dplyr::ungroup() |>
        dplyr::mutate(event = 1)

    os_had_censor <- os_dat_chaz |>
        dplyr::filter(!.data$subject %in% os_had_event$subject) |>
        dplyr::group_by(.data$subject) |>
        dplyr::slice(dplyr::n()) |>
        dplyr::ungroup() |>
        dplyr::mutate(event = 0)

    if (!(nrow(os_had_censor) == 0)) {
        message(sprintf("INFO: %i subjects did not die before max(times)", nrow(os_had_censor)))
    }

    os_dat_complete <- os_had_event |>
        dplyr::bind_rows(os_had_censor) |>
        dplyr::mutate(real_time = .data$time) |>
        dplyr::mutate(event = dplyr::if_else(.data$real_time <= .data$time_cen, .data$event, 0)) |>
        dplyr::mutate(time = dplyr::if_else(.data$real_time <= .data$time_cen, .data$real_time, .data$time_cen)) |>
        dplyr::arrange(.data$subject)

    os_dat_complete[, c("subject", "study", "arm", "time", "event", "cov_cont", "cov_cat")]
}


#' Simulate Survival Data from a Weibull Proportional Hazard Model
#'
#' @param lambda (`number`)\cr the scale parameter.
#' @param gamma (`number`)\cr the shape parameter.
#'
#' @inheritParams SimSurvival-Shared
#' @inheritSection SimSurvival-Shared Hazard Evaluation
#'
#' @family SimSurvival
#'
#' @export
SimSurvivalWeibullPH <- function(
    lambda,
    gamma,
    time_max = 2000,
    time_step = 1,
    lambda_censor = 1 / 3000,
    beta_cont = 0.2,
    beta_cat = c("A" = 0, "B" = -0.4, "C" = 0.2)
) {
    SimSurvival(
        time_max = time_max,
        time_step = time_step,
        lambda_censor = lambda_censor,
        beta_cont = beta_cont,
        beta_cat = beta_cat,
        loghazard = function(time) {
            log(lambda) + log(gamma) + (gamma - 1) * log(time)
        },
        name = "SimSurvivalWeibullPH"
    )
}

#' Simulate Survival Data from a Log-Logistic Proportional Hazard Model
#'
#' @param a (`number`)\cr the scale parameter.
#' @param b (`number`)\cr the shape parameter.
#'
#' @inheritParams SimSurvival-Shared
#' @inheritSection SimSurvival-Shared Hazard Evaluation
#'
#' @family SimSurvival
#' @export
SimSurvivalLogLogistic <- function(
    a,
    b,
    time_max = 2000,
    time_step = 1,
    lambda_censor = 1 / 3000,
    beta_cont = 0.2,
    beta_cat = c("A" = 0, "B" = -0.4, "C" = 0.2)
) {
    SimSurvival(
        time_max = time_max,
        time_step = time_step,
        lambda_censor = lambda_censor,
        beta_cont = beta_cont,
        beta_cat = beta_cat,
        loghazard = function(time) {
            c1 <- - log(a) + log(b) + (b - 1) * (- log(a) + log(time))
            c2 <- log(1 + (time / a)^b)
            return(c1 - c2)
        },
        name = "SimSurvivalLogLogistic"
    )
}



#' Simulate Survival Data from a Exponential Proportional Hazard Model
#'
#' @param lambda (`number`)\cr the rate parameter.
#'
#' @inheritParams SimSurvival-Shared
#' @inheritSection SimSurvival-Shared Hazard Evaluation
#'
#' @family SimSurvival
#'
#' @export
SimSurvivalExponential <- function(
    lambda,
    time_max = 2000,
    time_step = 1,
    lambda_censor = 1 / 3000,
    beta_cont = 0.2,
    beta_cat = c("A" = 0, "B" = -0.4, "C" = 0.2)
) {
    SimSurvival(
        time_max = time_max,
        time_step = time_step,
        lambda_censor = lambda_censor,
        beta_cont = beta_cont,
        beta_cat = beta_cat,
        loghazard = function(time) {
            log(lambda)
        },
        name = "SimSurvivalExponential"
    )
}


#' @include DataJoint.R
#' @include Quantities.R
NULL


#' Re-used documentation for `LongitudinalQuantities`
#'
#' @param x ([`LongitudinalQuantities`]) \cr longitudinal quantities.
#' @param object ([`LongitudinalQuantities`]) \cr longitudinal quantities.
#' @param ... not used.
#'
#' @keywords internal
#' @name LongitudinalQuantities-Shared
NULL


#' `LongitudinalQuantities` Object & Constructor Function
#'
#' Constructor function to generate a `LongitudinalQuantities` object.
#'
#' @details
#' Note that unlike [`SurvivalQuantities`], [`LongitudinalQuantities`] does not support
#' group aggregation.
#'
#' @slot quantities (`Quantities`)\cr The sampled quantities. Should contain 1 element per
#' element of `group`
#'
#' @slot data (`DataJoint`)\cr Survival and Longitudinal Data.
#'
#' @family LongitudinalQuantities
#' @name LongitudinalQuantities-class
#' @export LongitudinalQuantities
.LongitudinalQuantities <- setClass(
    "LongitudinalQuantities",
    slots = c(
        "quantities" = "Quantities",
        "data" = "DataJoint"
    )
)

#' @param object ([`JointModelSamples`]) \cr samples as drawn from a Joint Model.
#'
#' @param grid (`Grid`) \cr object that specifies which subjects and time points to calculate the
#' quantities for. See [Grid-Functions].
#' @rdname LongitudinalQuantities-class
LongitudinalQuantities <- function(
    object,
    grid
) {
    assert_class(object, "JointModelSamples")
    assert_class(grid, "Grid")
    assert_that(
        !is(grid, "GridPrediction"),
        msg = "`GridPrediction` objects are not supported for `LongitudinalQuantities`"
    )

    time_grid <- seq(
        from = 0,
        to = max(as.list(object@data)[["tumour_time"]]),
        length = 201
    )

    grid <- coalesceGridTime(grid, time_grid)

    gq <- generateQuantities(
        object,
        generator = as.QuantityGenerator(grid, object@data),
        type = "longitudinal"
    )

    quantities_raw <- extract_quantities(gq, type = "lm_identity")

    collapser <- as.QuantityCollapser(grid, object@data)
    quantities <- collapse_quantities(quantities_raw, collapser)

    .LongitudinalQuantities(
        quantities = Quantities(
            quantities,
            groups = collapser@groups,
            times = collapser@times
        ),
        data = object@data
    )
}



#' `as.data.frame`
#'
#' @param x ([`LongitudinalQuantities`]) \cr longitudinal quantities.
#' @param ... not used.
#' @family LongitudinalQuantities
#' @export
as.data.frame.LongitudinalQuantities <- function(x, ...) {
    as.data.frame(x@quantities)
}



#' summary
#'
#' @description
#' This method returns a `data.frame` of the longitudinal quantities.
#'
#' @param conf.level (`numeric`) \cr confidence level of the interval.
#' @inheritParams LongitudinalQuantities-Shared
#'
#' @family LongitudinalQuantities
#' @family summary
#' @export
summary.LongitudinalQuantities <- function(
    object,
    conf.level = 0.95,
    ...
) {
    summary(object@quantities, conf.level = conf.level)
}


#' Longitudinal Plot
#'
#' Internal plotting function to create longitudinal plots
#' This function predominately exists to extract core logic into its own function
#' to enable easier unit testing.
#'
#' @param data (`data.frame`)\cr summary statistics for longitudinal
#' value estimates. See details.
#' @param data_obs (`data.frame`)\cr real observed values to be
#' overlaid for reference.  See details.
#' @param add_ci (`logical`)\cr Should confidence intervals be added? Default = `TRUE`.
#' @details
#'
#' ## `data`
#' Should contain the following columns:
#' - `time` (`numeric`) \cr time point for the summary statistic.
#' - `group` (`character`) \cr the group in which the observation belongs to.
#' - `median` (`numeric`) \cr the median value for the summary statistic.
#' - `upper` (`numeric`) \cr the upper 95% CI for the summary statistic.
#' - `lower` (`numeric`) \cr the lower 95% CI for the summary statistic.
#'
#' ## `data_obs`
#' Should contain the following columns:
#' - `time` (`numeric`) \cr the time at which the observed value occurred.
#' - `Yob` (`numeric`) \cr the real observed value.
#' - `group` (`character`) \cr which group the event belongs to, should correspond to
#' values in `data$group`.
#' @keywords internal
longitudinal_plot <- function(
    data,
    data_obs = NULL,
    add_ci = FALSE
) {
    p <- ggplot() +
        geom_line(aes(x = .data$time, y = .data$median), data = data) +
        xlab(expression(t)) +
        ylab(expression(y)) +
        facet_wrap(~group) +
        theme_bw()

    if (add_ci) {
        p <- p + geom_ribbon(
            aes(x = .data$time, ymin = .data$lower, ymax = .data$upper),
            data = data,
            alpha = 0.3
        )
    }

    if (!is.null(data_obs)) {
        p <- p + geom_point(aes(x = .data$time, y = .data$Yob), data = data_obs)
    }
    return(p)
}


#' Automatic Plotting for `LongitudinalQuantities`
#'
#' @param conf.level (`numeric`) \cr confidence level of the interval. If values of `FALSE`,
#' `NULL` or `0` are provided then confidence regions will not be added to the plot.
#' @inheritParams LongitudinalQuantities-Shared
#'
#' @family LongitudinalQuantities
#' @family autoplot
#' @export
autoplot.LongitudinalQuantities <- function(object, conf.level = 0.95, ...) {
    include_ci <- !is.null(conf.level) && is.numeric(conf.level) && conf.level > 0
    # If CI aren't needed supply a default 0.95 to summary function as it needs
    # a value to be specified to work.
    conf.level <- if (include_ci) conf.level else 0.95
    data_sum <- summary(object, conf.level = conf.level)
    data_obs <- extract_observed_values(object@data)
    assert_that(
        "group" %in% names(data_sum),
        "subject" %in% names(data_obs)
    )
    data_obs$group <- data_obs$subject
    data_obs <- data_obs[data_obs$group %in% data_sum$group, ]
    longitudinal_plot(
        data = data_sum,
        data_obs = data_obs,
        add_ci = include_ci
    )
}



#' @rdname show-object
#' @export
setMethod(
    f = "show",
    signature = "LongitudinalQuantities",
    definition = function(object) {
        template <- c(
            "LongitudinalQuantities Object:",
            "    # of samples    = %d",
            "    # of quantities = %d"
        )
        string <- sprintf(
            paste(template, collapse = "\n"),
            nrow(object@quantities),
            ncol(object@quantities)
        )
        cat("\n", string, "\n\n")
    }
)

#' @include Grid.R
#' @include generics.R
NULL

#' @rdname Grid-Dev
.GridObserved <- setClass(
    "GridObserved",
    contains = "Grid",
    slots = c(
        "subjects" = "character_or_NULL"
    )
)


#' @rdname Grid-Functions
#' @export
GridObserved <- function(subjects = NULL) {
    .GridObserved(
        subjects = subjects
    )
}

#' @rdname Quant-Dev
#' @export
as.QuantityGenerator.GridObserved <- function(object, data, ...) {
    assert_class(data, "DataJoint")
    data_list <- as.list(data)
    subjects <- unlist(as.list(object, data = data), use.names = FALSE)
    unique_visits <- tapply(data_list$tumour_time, data_list$subject_tumour_index, unique)
    subject_visits <- unique_visits[data_list$subject_to_index[subjects]]
    visit_lengths <- vapply(subject_visits, length, numeric(1))
    QuantityGeneratorSubject(
        times = unlist(subject_visits, use.names = FALSE),
        subjects = rep(subjects, visit_lengths)
    )
}

#' @rdname Quant-Dev
#' @export
as.QuantityCollapser.GridObserved <- function(object, data, ...) {
    generator <- as.QuantityGenerator(object, data)

    QuantityCollapser(
        times = generator@times,
        groups = generator@subjects,
        indexes = as.list(seq_along(generator@times))
    )
}

#' @export
as.list.GridObserved <- function(x, data, ...) {
    subjects_to_list(x@subjects, data)
}

#' @include JointModel.R
#' @include SurvivalQuantities.R
NULL

setOldClass("CmdStanMCMC")

# JointModelSamples-class ----

#' `JointModelSamples`
#'
#' Contains samples from a [`JointModel`].
#'
#' @slot model ([`JointModel`])\cr the model that the samples were drawn from.
#' @slot data ([`DataJoint`])\cr the data that the model was fitted on.
#' @slot results ([`cmdstanr::CmdStanMCMC`])\cr the STAN samples.
#'
#' @aliases JointModelSamples
#' @export
.JointModelSamples <- setClass(
    "JointModelSamples",
    slots = c(
        model = "JointModel",
        data = "DataJoint",
        results = "CmdStanMCMC"
    )
)


#' @rdname generateQuantities
#' @param generator (`QuantityGenerator`)\cr object that specifies which subjects and time points
#' to calculate the quantities at
#' @param type (`character`)\cr type of quantities to be generated, must be either "survival" or
#' "longitudinal".
#' @export
generateQuantities.JointModelSamples <- function(object, generator, type, ...) {

    data <- as_stan_list(object@data) |>
        append(as_stan_list(object@model@parameters)) |>
        append(as_stan_list(generator, data = object@data, model = object@model))

    stanobj <- as.StanModule(object, generator = generator, type = type)
    model <- compileStanModel(stanobj)

    devnull <- utils::capture.output(
        results <- model$generate_quantities(
            data = data,
            fitted_params = object@results
        )
    )
    return(results)
}


#' `JointModelSamples` -> `StanModule`
#'
#' Converts a `JointModelSamples` object into a `StanModule` object ensuring
#' that the resulting `StanModule` object is able to generate post sampling
#' quantities.
#'
#' @inheritParams generateQuantities
#' @export
as.StanModule.JointModelSamples <- function(object, generator, type, ...) {
    assert_that(
        is(generator, "QuantityGenerator"),
        length(type) == 1,
        type %in% c("survival", "longitudinal")
    )

    quant_stanobj <- read_stan("base/quantities.stan") |>
        decorated_render(
            include_gq_longitudinal_idv = (type == "longitudinal") & is(generator, "QuantityGeneratorSubject"),
            include_gq_longitudinal_pop = (type == "longitudinal") & is(generator, "QuantityGeneratorPopulation"),
            include_gq_survival_idv = (type == "survival") & is(generator, "QuantityGeneratorSubject"),
            include_gq_survival_pred = (type == "survival") & is(generator, "QuantityGeneratorPrediction")
        ) |>
        StanModule()

    stanobj <- Reduce(
        merge,
        list(
            as.StanModule(object@model),
            enableGQ(object@model),
            quant_stanobj
        )
    )
    stanobj
}



#' `JointModelSamples` -> Printable `Character`
#'
#' Converts [`JointModelSamples`] object into a printable string.
#' @param object ([`JointModelSamples`])\cr samples as drawn from a [`JointModel`].
#' @family JointModelSamples
#' @param indent (`numeric`)\cr how much white space to prefix the print string with.
#' @keywords internal
#' @export
as_print_string.JointModelSamples <- function(object, indent = 1, ...) {
    sizes <- vapply(
        as.CmdStanMCMC(object)$metadata()[["stan_variable_sizes"]],
        \(x) {
            if (length(x) == 1 && x == 1) return("")
            paste0("[", paste(x, collapse = ", "), "]")
        },
        character(1)
    )
    variable_string <- paste0(
        "        ",
        as.CmdStanMCMC(object)$metadata()[["stan_variables"]],
        sizes
    )
    template <- c(
        "JointModelSamples Object with:",
        "",
        "    # of samples per chain = %d",
        "    # of chains            = %d",
        "",
        "    Variables:",
        variable_string[order(variable_string)]
    )
    pad <- rep(" ", indent) |> paste(collapse = "")
    template_padded <- paste(pad, template)
    sprintf(
        paste(template_padded, collapse = "\n"),
        as.CmdStanMCMC(object)$metadata()$iter_sampling,
        as.CmdStanMCMC(object)$num_chains()
    )
}

#' @rdname show-object
#' @export
setMethod(
    f = "show",
    signature = "JointModelSamples",
    definition = function(object) {
        string <- as_print_string(object)
        cat("\n", string, "\n\n")
    }
)


#' @rdname as.CmdStanMCMC
#' @export
as.CmdStanMCMC.JointModelSamples <- function(object, ...) {
    return(object@results)
}

#' @include generics.R
#' @include Link.R
#' @include LongitudinalModel.R
#' @include SurvivalModel.R
#' @include StanModule.R
#' @include StanModel.R
#' @include Parameter.R
NULL



#' @rdname getParameters
#' @export
getParameters.default <- function(object, ...) {
    if (missing(object) || is.null(object)) {
        return(NULL)
    }
    stop(sprintf("No default method implemented for getParameters(<%s>)", typeof(object)))
}


# merge ----

## merge-StanModel,NULL ----

#' @rdname merge
setMethod(
    "merge",
    signature = c("StanModel", "NULL"),
    definition = function(x, y, ...) x@stan
)

## merge-NULL,StanModel ----

#' @rdname merge
setMethod(
    "merge",
    signature = c("NULL", "StanModel"),
    definition = function(x, y, ...) y@stan
)

## merge-StanModule,NULL ----

#' @rdname merge
setMethod(
    "merge",
    signature = c("StanModule", "NULL"),
    definition = function(x, y, ...) x
)

## merge-NULL,StanModule ----

#' @rdname merge
setMethod(
    "merge",
    signature = c("NULL", "StanModule"),
    definition = function(x, y, ...) y
)

## merge-ParameterList,NULL ----

#' @rdname merge
setMethod(
    "merge",
    signature = c("ParameterList", "NULL"),
    definition = function(x, y, ...) x
)

## merge-NULL,ParameterList ----

#' @rdname merge
setMethod(
    "merge",
    signature = c("NULL", "ParameterList"),
    definition = function(x, y, ...) y
)

## merge-NULL,NULL ----

#' @rdname merge
setMethod(
    "merge",
    signature = c("NULL", "NULL"),
    definition = function(x, y, ...) NULL
)

#' @include generics.R
#' @include Grid.R
NULL

#' @rdname Quant-Dev
.QuantityGeneratorPrediction <- setClass(
    "QuantityGeneratorPrediction",
    contains = "QuantityGenerator",
    slots = c(
        "times" = "numeric",
        "newdata" = "data.frame",
        "params" = "list"
    )
)


#' @rdname Quant-Dev
QuantityGeneratorPrediction <- function(times, newdata = NULL, params = NULL) {
    .QuantityGeneratorPrediction(
        times = times,
        newdata = newdata,
        params = params
    )
}
setValidity(
    "QuantityGeneratorPrediction",
    function(object) {
        if (length(object@times) != nrow(object@newdata)) {
            return("Length of `times` and `newdata` must be equal")
        }
        return(TRUE)
    }
)




#' @rdname as_stan_list.QuantityGenerator
#' @export
as_stan_list.QuantityGeneratorPrediction <- function(object, data, model, ...) {
    assert_that(
        is(data, "DataJoint")
    )
    ret <- list()
    data_list <- as_stan_list(data)

    ret[["gq_times"]] <- object@times
    ret[["gq_n_quant"]] <- length(object@times)

    # Get a list of which longutidunal parameters need to be defined based
    # on the selected longitudinal model
    par_names <- getPredictionNames(model@longitudinal)
    for (nam in par_names) {
        assert_that(
            nam %in% names(object@params),
            msg = sprintf("Parameter '%s' not found in `params`", nam)
        )
        assert_that(
            is.numeric(object@params[[nam]]),
            msg = sprintf("Parameter '%s' must be numeric", nam)
        )
        assert_that(
            length(object@params[[nam]]) == 1,
            msg = sprintf("Parameter '%s' must be length 1", nam)
        )
    }

    par_vals <- object@params[par_names]
    if (length(par_vals) == 0) {
        par_vals <- 0
        par_names <- "null_model"
    }

    ret[["gq_n_par"]] <- length(par_names)

    # Replicate the longitudinal parameters so the same parameter values are used
    # for all observations that are being predicted
    ret[["gq_link_function_inputs"]] <- matrix(
        rep(unlist(par_vals), each = ret[["gq_n_quant"]]),
        ncol = ret[["gq_n_par"]],
        nrow = ret[["gq_n_quant"]]
    )

    # Create design matrix from new data ensuring that it has the same
    # structure as the original design matrix
    ret[["gq_os_cov_design"]] <- mirror_design_matrix(
        data@survival,
        object@newdata
    )

    # dummy pop indexes in order for stan code to actualy compile. In this setting
    # this matrix isn't actually used so doesn't matter what these values are
    # but don't want to have to burden individual longitudinal models with the
    # conditional logic to check if they are generating population quantities or not
    ret[["gq_long_pop_arm_index"]] <- rep(1, ret[["gq_n_quant"]])
    ret[["gq_long_pop_study_index"]] <- rep(1, ret[["gq_n_quant"]])


    # Sanity checks
    assert_that(
        nrow(ret[["gq_os_cov_design"]]) == ret[["gq_n_quant"]],
        ncol(ret[["gq_os_cov_design"]]) == data_list[["p_os_cov_design"]],
        all(!is.na(ret[["gq_link_function_inputs"]]))
    )

    return(ret)
}



#' Re-used documentation for `DataSubject`
#'
#' @param object ([`DataSubject`]) \cr subject-level data.
#' @param x ([`DataSubject`]) \cr subject-level data.
#' @param ... Not Used.
#'
#' @name DataSubject-Shared
#' @keywords internal
NULL



#' Subject Data Object and Constructor Function
#'
#' The [`DataSubject`] class handles the processing of the subject data for
#' fitting a [`JointModel`].
#'
#' @slot data (`data.frame`)\cr the subject-level data.
#' @slot subject (`character`)\cr the name of the variable containing the subject identifier.
#' @slot arm (`character`)\cr the name of the variable containing the arm identifier.
#' @slot study (`character`)\cr the name of the variable containing the study identifier.
#'
#' @family DataObjects
#' @family DataSubject
#' @exportClass DataSubject
#' @export DataSubject
.DataSubject <- setClass(
    Class = "DataSubject",
    representation = list(
        data = "data.frame",
        subject = "character",
        arm = "character",
        study = "character"
    )
)


#' @param data (`data.frame`)\cr the subject-level data.
#' @param subject (`character`)\cr the name of the variable containing the subject identifier.
#' @param arm (`character`)\cr the name of the variable containing the arm identifier.
#' @param study (`character`)\cr the name of the variable containing the study identifier.
#' @rdname DataSubject-class
#' @export
DataSubject <- function(data, subject, arm, study) {
    vars <- c(subject, arm, study)
    vars_frm_chr <- paste0("~ ", paste(vars, collapse = " + "))
    .DataSubject(
        data = remove_missing_rows(data, stats::as.formula(vars_frm_chr)),
        subject = subject,
        arm = arm,
        study = study
    )
}

setValidity(
    "DataSubject",
    method = function(object) {
        if (length(object@subject) != 1) {
            return("`subject` must be a length 1 character")
        }
        if (length(object@arm) != 1) {
            return("`arm` must be a length 1 character")
        }
        if (length(object@study) != 1) {
            return("`study` must be a length 1 character")
        }
        if (nrow(object@data) == 0) {
            return("`data` should not have 0 rows")
        }
        sbj <- object@data[[object@subject]]
        if (!(is(sbj, "character") | is(sbj, "factor"))) {
            return("`data[[subject]]` should be of type character or factor")
        }
    }
)



#' @inheritParams DataSubject-Shared
#' @inherit extractVariableNames description title
#'
#' @returns
#' A list with the following named elements:
#' - `subject` (`character`)\cr the name of the variable containing the subject identifier.
#' - `arm` (`character`)\cr the name of the variable containing the arm identifier.
#' - `study` (`character`) \cr the name of the variable containing the study identifier.
#' @family DataSubject
#' @family extractVariableNames
#' @export
#' @keywords internal
extractVariableNames.DataSubject <- function(object) {
    list(
        subject = object@subject,
        arm = object@arm,
        study = object@study
    )
}


#' @rdname as_stan_list.DataObject
#' @family DataSubject
#' @export
as_stan_list.DataSubject <- function(object, ...) {
    df <- as.data.frame(harmonise(object))
    vars <- extractVariableNames(object)

    unique_arm_study_combos <- unique(
        data.frame(
            arm = as.numeric(df[[vars$arm]]),
            study = as.numeric(df[[vars$study]])
        )
    )

    list(
        n_subjects = nrow(df),
        n_studies = length(unique(df[[vars$study]])),
        n_arms = length(unique(df[[vars$arm]])),
        subject_study_index = as.numeric(df[[vars$study]]),
        subject_arm_index = as.numeric(df[[vars$arm]]),
        subject_to_index = stats::setNames(
            seq_len(nlevels(df[[vars$subject]])),
            levels(df[[vars$subject]])
        ),
        arm_to_index = stats::setNames(
            seq_len(nlevels(df[[vars$arm]])),
            levels(df[[vars$arm]])
        ),
        study_to_index = stats::setNames(
            seq_len(nlevels(df[[vars$study]])),
            levels(df[[vars$study]])
        ),
        pop_arm_index = unique_arm_study_combos$arm,
        pop_study_index = unique_arm_study_combos$study
    )
}

#' @rdname as_stan_list.DataObject
#' @export
as.list.DataSubject <- function(x, ...) {
    as_stan_list(x, ...)
}


#' `DataSubject` -> `data.frame`
#'
#' @inheritParams DataSubject-Shared
#'
#' @description
#' Converts a [`DataSubject`] object into a `data.frame`.
#' The subject variable is cast to factor.
#' @family DataSubject
#' @export
as.data.frame.DataSubject <- function(x, ...) {
    x <- x@data
    rownames(x) <- NULL
    x
}



#' @rdname harmonise
harmonise.DataSubject <- function(object, ...) {
    data <- as.data.frame(object)
    vars <- extractVariableNames(object)
    assert_that(
        vars$subject %in% names(data),
        vars$arm %in% names(data),
        vars$study %in% names(data)
    )
    assert_character(
        as.character(data[[vars$subject]]),
        any.missing = FALSE,
        unique = TRUE
    )
    data[[vars$subject]] <- factor(data[[vars$subject]])
    data[[vars$arm]] <- factor(data[[vars$arm]])
    data[[vars$study]] <- factor(data[[vars$study]])
    data <- data[order(data[[vars$subject]]), ]
    DataSubject(
        data = data,
        subject = object@subject,
        arm = object@arm,
        study = object@study
    )
}

#' `DataSubject` -> Printable `Character`
#'
#' Converts [`DataSubject`] object into a printable string.
#' @inheritParams DataSubject-Shared
#' @family DataSubject
#' @keywords internal
#' @param indent (`numeric`)\cr how much white space to prefix the print string with.
#' @export
as_print_string.DataSubject <- function(object, indent = 1, ...) {
    template <- c(
        "Subject-Data Object:",
        "    # of Subjects = %d",
        "    # of Studies  = %d",
        "    # of Arms     = %d"
    )
    pad <- rep(" ", indent) |> paste(collapse = "")
    template_padded <- paste(pad, template)
    vars <- extractVariableNames(object)
    sprintf(
        paste(template_padded, collapse = "\n"),
        nrow(object@data),
        length(unique(object@data[[vars$study]])),
        length(unique(object@data[[vars$arm]]))
    )
}

#' @rdname show-object
#' @export
setMethod(
    f = "show",
    signature = "DataSubject",
    definition = function(object) {
        string <- as_print_string(object)
        cat("\n", string, "\n\n")
    }
)



#' Re-used documentation for Brier Score components
#'
#' @param t (`numeric`)\cr timepoints to calculate the desired quantity at.
#' @param times (`numeric`)\cr observed times.
#' @param events (`numeric`)\cr event indicator for `times`. Either 1 for an event or 0 for censor.
#' @param event_offset (`logical`)\cr If `TRUE` then \eqn{G(T_i)} is evaluated at \eqn{G(T_i-)}.
#' Setting this as `TRUE` mirrors the implementation of the `{pec}` package.
#' @param maintain_cen_order (`logical`)\cr If `TRUE` then, in the case of ties,
#' censor times are always considered
#' to have occurred after the event times when calculating the "reverse Kaplan-Meier" for the
#' IPCW estimates. Setting this to `TRUE` mirrors the implementation of the `{prodlim}`
#' package.
#' @param ... not used.
#'
#' @name Brier-Score-Shared
#' @keywords internal
NULL




#' Reverse Kaplan-Meier
#'
#' @inheritParams Brier-Score-Shared
#' @description
#' Calculates the survival estimates of the censoring distribution
#' using the Kaplan-Meier estimate.
#' This is primarily used in the calculation of the IPCW estimates.
#'
#' @details
#' With regards to ties between censor and event times; the standard
#' approach is to regard events as occurring before censors. However,
#' when modelling the censoring distribution we are regarding the
#' censors as "events" so which should come first in the case of ties?
#'
#' The `reverse_km_event_first()` function maintains the rule
#' that events always come first even if we are regarding the censors
#' as "events". This matches the implementation of
#' `prodlim::prodlim(..., reverse = TRUE)`.
#'
#' The `reverse_km_cen_first()` function provides the alternative
#' implementation assuming that in the case of ties the censor "events"
#' come before the event "censors". This is essentially a thin wrapper
#' around `survival::survfit(Surv(time, 1 - event), ...)`
#'
#' @name reverse_km
#' @keywords internal
reverse_km_event_first <- function(t, times, events) {
    assert_numeric(t, any.missing = FALSE, finite = TRUE)
    assert_numeric(times, any.missing = FALSE, finite = TRUE)
    assert_numeric(events, any.missing = FALSE, finite = TRUE)
    assert_that(
        length(times) == length(events),
        all(events == 1 | events == 0)
    )
    events_cen <- 1 - events
    ord <- order(times, events_cen)
    times <- times[ord]
    events <- events[ord]
    events_cen <- events_cen[ord]
    cs_events <- cumsum(events)
    cs_censor <- cumsum(events_cen)

    g_times <- unique(times)
    g_n_events_cen <- tapply(events_cen, times, sum)
    g_cs_events_cen <- tapply(cs_censor, times, max)
    g_cs_events <- tapply(cs_events, times, max)
    g_is_cen <- tapply(events_cen, times, max)

    nrisk <- length(times) - g_cs_events - g_cs_events_cen + g_n_events_cen
    surv_interval <- 1 - g_n_events_cen / nrisk
    surv_interval <- ifelse(nrisk == 0, 1, surv_interval)
    surv <- cumprod(surv_interval)

    ct <- g_times[which(g_is_cen == 1)]
    sv <- surv[which(g_is_cen == 1)]

    names(surv) <- NULL
    names(ct) <- NULL
    names(sv) <- NULL
    list(
        t = t,
        surv = c(1, sv)[findInterval(t, ct) + 1]
    )
}


#' @rdname reverse_km
reverse_km_cen_first <- function(t, times, events) {
    assert_numeric(t, any.missing = FALSE, finite = TRUE)
    assert_numeric(times, any.missing = FALSE, finite = TRUE)
    assert_numeric(events, any.missing = FALSE, finite = TRUE)
    assert_that(
        length(times) == length(events),
        all(events == 1 | events == 0)
    )
    dat <- data.frame(
        times = times,
        events = events,
        cen_events = 1 - events
    )
    mod <- survival::survfit(
        survival::Surv(times, cen_events) ~ 1,
        data = dat
    )
    preds <- summary(mod, times = t[order(t)], extend = TRUE)$surv

    assert_that(
        length(preds) == length(t)
    )
    list(
        t = t,
        surv = preds[match_order(t)]
    )
}

#' Match Order
#'
#' @param x (`numeric`)\cr a vector for which we want to generate an index so other
#' vectors can be put into the same sort order
#'
#' Assuming we have a vector that is sorted then this function
#' will return the index vector to convert that sorted vector
#' into the same sort order as the input vector `x`.
#' For example let `x = 8, 7, 9 , 7`. If sorted we would get
#' `x_sort = 7, 7, 8, 9`. So in order to convert `x_sort`
#' back into `x` we'd need an index vector of `3, 1, 4, 2`.
#' This function is used to determine that corresponding
#' index vector for an arbitrarily sorted vector `x`.
#'
#' There is no specific handling of ties. It is assuming that in the case
#' of ties for `x` that the vector you are re-indexing also has tied values
#' thus the specific tied element selection does not matter.
#' @keywords internal
match_order <- function(x) {
    order(order(x))
}


#' Brier Score
#'
#' @inheritParams Brier-Score-Shared
#'
#' @description
#' Implements the Brier Score as detailed in \insertCite{blanche2015}{jmpost}
#'
#' @details
#' - `bs_get_squared_dist()` - implements the squared distance part
#' of the formula.
#' - `bs_get_weights()` - implements the IPCW weighting
#'
#' @references
#' \insertAllCited{}
#'
#' @keywords internal
brier_score <- function(
    t,
    times,
    events,
    pred_mat,
    maintain_cen_order = TRUE,
    event_offset = TRUE
) {

    square_diff_mat <- bs_get_squared_dist(
        t = t,
        times = times,
        events = events,
        pred_mat = pred_mat
    )

    weight_mat <- bs_get_weights(
        t = t,
        times = times,
        events = events,
        event_offset = event_offset,
        maintain_cen_order = maintain_cen_order
    )

    # the following is a computational shortcut for diag(A %*% B)
    # as we don't want to compute the off-diagonal entries of the
    # matrix multiplication
    x <- colSums(weight_mat * square_diff_mat)
    names(x) <- t
    x / length(times)
}


#' @rdname brier_score
bs_get_squared_dist <- function(t, times, events, pred_mat) {

    assert_numeric(
        times,
        finite = TRUE,
        any.missing = FALSE
    )
    assert_numeric(
        t,
        finite = TRUE,
        any.missing = FALSE
    )
    assert_numeric(
        events,
        finite = TRUE,
        any.missing = FALSE,
        lower = 0,
        upper = 1
    )
    assert_matrix(
        pred_mat,
        any.missing = FALSE,
        nrows = length(times),
        ncols = length(t)
    )
    assert_that(
        length(events) == length(times)
    )


    expected_mat <- mapply(
        \(ti, event) (ti <= t) * event * 1,
        ti = times,
        event = events,
        SIMPLIFY = FALSE
    ) |>
        unlist() |>
        matrix(ncol = length(t), byrow = TRUE)

    assert_that(
        nrow(pred_mat) == nrow(expected_mat),
        ncol(pred_mat) == ncol(expected_mat)
    )

    (expected_mat - pred_mat)^2
}


#' @rdname brier_score
bs_get_weights <- function(
    t,
    times,
    events,
    event_offset = TRUE,
    maintain_cen_order = TRUE
) {
    assert_numeric(
        times,
        finite = TRUE,
        any.missing = FALSE
    )
    assert_numeric(
        t,
        finite = TRUE,
        any.missing = FALSE
    )
    assert_numeric(
        events,
        finite = TRUE,
        any.missing = FALSE,
        lower = 0,
        upper = 1
    )
    assert_flag(event_offset, na.ok = FALSE, null.ok = FALSE)
    assert_flag(maintain_cen_order, na.ok = FALSE, null.ok = FALSE)
    n_col <- length(t)
    n_row <- length(times)

    reverse_km <- if (maintain_cen_order) reverse_km_event_first else reverse_km_cen_first
    offset <- if (event_offset) -.Machine$double.eps^(1 / 2) else 0

    censor_dist_t <- reverse_km(t, times, events)
    weight_mat_t <- 1 / matrix(
        rep(censor_dist_t$surv, n_row),
        nrow = n_row,
        byrow = TRUE
    )

    censor_dist_ti <- reverse_km(times + offset, times, events)
    weight_mat_ti <- 1 / matrix(
        rep(censor_dist_ti$surv, n_col),
        nrow = n_row
    )

    indicator_mat_t <- mapply(
        \(ti) (ti > t) * 1,
        ti = times,
        SIMPLIFY = FALSE
    ) |>
        unlist() |>
        matrix(ncol = n_col, byrow = TRUE)

    indicator_mat_ti <- mapply(
        \(ti, event) (ti <= t) * event * 1,
        ti = times,
        event = events,
        SIMPLIFY = FALSE
    ) |>
        unlist() |>
        matrix(ncol = n_col, byrow = TRUE)

    assert_that(
        all(indicator_mat_t + indicator_mat_ti <= 1)
    )

    weight_mat_t[indicator_mat_t == 0] <- 0
    weight_mat_ti[indicator_mat_ti == 0] <- 0

    (weight_mat_t + weight_mat_ti)
}


#' @include StanModule.R
#' @include generics.R
#' @include ParameterList.R
#' @include DataJoint.R
#' @include DataSurvival.R
#' @include DataLongitudinal.R
#' @include LongitudinalModel.R
#' @include SurvivalModel.R
#' @include Link.R
#' @include constants.R
NULL


#' Re-used documentation for `JointModel-Shared`
#'
#' @param object ([`JointModel`]) \cr Joint model specification.
#' @param x ([`JointModel`]) \cr Joint model specification.
#' @param ... Not Used.
#'
#' @name JointModel-Shared
#' @keywords internal
NULL

setClassUnion("LongitudinalModel_OR_NULL", c("LongitudinalModel", "NULL"))
setClassUnion("SurvivalModel_OR_NULL", c("SurvivalModel", "NULL"))

# JointModel-class ----

#' Joint Model Object and Constructor Function
#'
#' @slot longitudinal ([`LongitudinalModel`] or `NULL`)\cr the longitudinal model.
#' @slot survival ([`SurvivalModel`] or `NULL`)\cr the survival model.
#' @slot link (`Link`)\cr the link.
#' @slot parameters (`ParameterList`)\cr the parameter specification.
#'
#' @family JointModel
#' @export JointModel
#' @exportClass JointModel
.JointModel <- setClass(
    Class = "JointModel",
    slots = list(
        longitudinal = "LongitudinalModel_OR_NULL",
        survival = "SurvivalModel_OR_NULL",
        link = "Link",
        parameters = "ParameterList"
    )
)

#' @param longitudinal ([`LongitudinalModel`] or `NULL`)\cr the longitudinal model.
#' @param survival ([`SurvivalModel`] or `NULL`)\cr the survival model.
#' @param link (`Link`)\cr the link.
#' @rdname JointModel-class
JointModel <- function(
    longitudinal = NULL,
    survival = NULL,
    link = Link()
) {
    link <- resolvePromise(Link(link), longitudinal)

    if (length(link) > 0) {
        longitudinal <- enableLink(longitudinal)
    }

    parameters <- Reduce(
        merge,
        list(
            getParameters(longitudinal),
            getParameters(survival),
            getParameters(link)
        )
    )

    .JointModel(
        longitudinal = longitudinal,
        survival = survival,
        link = link,
        parameters = parameters
    )
}


#' @export
enableGQ.JointModel <- function(object, ...) {
    merge(
        enableGQ(object@survival),
        enableGQ(object@longitudinal)
    )
}


#' `JointModel` -> `StanModule`
#'
#' Converts a [`JointModel`] object to a [`StanModule`] object
#'
#' @inheritParams JointModel-Shared
#' @family JointModel
#' @family as.StanModule
#' @export
as.StanModule.JointModel <- function(object, ...) {

    base_model <- read_stan("base/base.stan")

    stan_full <- decorated_render(
        .x = base_model,
        longitudinal = add_missing_stan_blocks(as.list(object@longitudinal)),
        survival = add_missing_stan_blocks(as.list(object@survival)),
        link = add_missing_stan_blocks(as.list(object@link)),
        priors = add_missing_stan_blocks(as.list(object@parameters))
    )
    # Unresolved Jinja code within the longitudinal / Survival / Link
    # models won't be resolved by the above call to `decorated_render`.
    # Instead they it will just be inserted into the template asis. Thus
    # we run `decorated_render` again to resolve any lingering Jinja code
    # Main example being models that don't have any Jinja code but still
    # use the `decorated_render` constants `machine_double_eps`.
    stan_full <- decorated_render(.x = stan_full)

    x <- merge(
        StanModule("base/functions.stan"),
        StanModule(stan_full)
    )

    return(x)
}



#' `JointModel` -> `character`
#'
#' Renders a [`JointModel`] object to a stan program
#'
#' @inheritParams JointModel-Shared
#' @family JointModel
#' @export
as.character.JointModel <- function(x, ...) {
    as.character(as.StanModule(x))
}


# write_stan-JointModel ----

#' @rdname write_stan
#' @export
write_stan.JointModel <- function(object, destination, ...) {
    if (is_connection(destination)) {
        return(writeLines(as.character(object), con = destination))
    }
    fi <- file(destination, open = "w")
    writeLines(as.character(object), con = fi)
    close(fi)
}


# compileStanModel-JointModel ----

#' @rdname compileStanModel
#' @export
compileStanModel.JointModel <- function(object) {
    object |>
        as.StanModule() |>
        compileStanModel() |>
        invisible()
}


# sampleStanModel-JointModel ----

#' @rdname sampleStanModel
#'
#' @param data (`DataJoint` or `list`)\cr input data.
#' @export
sampleStanModel.JointModel <- function(object, data, ...) {

    assert_class(data, "DataJoint")

    if (!is.null(object@survival)) {
        assert_that(
            !is.null(data@survival),
            msg = "`DataSurvival` can't be missing if a `SurvivalModel` has been specified"
        )
    }
    if (!is.null(object@longitudinal)) {
        assert_that(
            !is.null(data@longitudinal),
            msg = "`DataLongitudinal` can't be missing if a `LongitudinalModel` has been specified"
        )
    }

    args <- list(...)

    args[["data"]] <- append(
        as_stan_list(data),
        as_stan_list(object@parameters)
    )

    args[["chains"]] <- if ("chains" %in% names(args)) {
        args[["chains"]]
    } else {
        # Magic constant from R/constants.R
        CMDSTAN_DEFAULT_CHAINS
    }

    initial_values <- if ("init" %in% names(args)) {
        args[["init"]]
    } else {
        initialValues(object, n_chains = args[["chains"]])
    }

    args[["init"]] <- ensure_initial_values(
        initial_values,
        args[["data"]],
        object@parameters
    )

    model <- compileStanModel(object)

    results <- do.call(
        model$sample,
        args
    )

    .JointModelSamples(
        model = object,
        data = data,
        results = results
    )
}


#' Ensure that initial values are correctly specified
#'
#' @param initial_values (`list`)\cr A list of lists containing the initial values
#' must be 1 list per desired chain. All elements should have identical names
#' @param data (`list`)\cr specifies the size to expand each of our initial values to be.
#' That is elements of size 1 in `initial_values` will be expanded to be the same
#' size as the corresponding element in `data` by broadcasting the value.
#' @param parameters ([`ParameterList`])\cr the parameters object
#'
#' @details
#' This function is mostly a thin wrapper around `expand_initial_values` to
#' enable easier unit testing.
#'
#' @keywords internal
ensure_initial_values <- function(initial_values, data, parameters) {
    if (is.function(initial_values)) {
        return(initial_values)
    }

    assert_class(data, "list")
    assert_class(parameters, "ParameterList")
    assert_class(initial_values, "list")

    values_sizes <- size(parameters)
    values_sizes_complete <- replace_with_lookup(
        values_sizes,
        data
    )
    lapply(
        initial_values,
        expand_initial_values,
        sizes = values_sizes_complete
    )
}



#' @rdname initialValues
#' @export
initialValues.JointModel <- function(object, n_chains, ...) {
    initialValues(object@parameters, n_chains)
}


pad_with_white_space <- function(x, pad = 4) {
    padding <- paste0(rep(" ", each = pad), collapse = "")
    x_sep <- x |>
        strsplit(split = "\n") |>
        unlist()
    x_padded <- paste(padding, x_sep) |>
        paste(collapse = "\n")
    return(x_padded)
}


#' @rdname show-object
#' @export
setMethod(
    f = "show",
    signature = "JointModel",
    definition = function(object) {
        survival_string <- if (is.null(object@survival)) {
            "\n     Not Specified\n"
        } else {
            as_print_string(object@survival) |> pad_with_white_space()
        }

        longitudinal_string <- if (is.null(object@longitudinal)) {
            "\n     Not Specified\n"
        } else {
            as_print_string(object@longitudinal) |> pad_with_white_space()
        }

        link_string <- as_print_string(object@link) |> pad_with_white_space()

        string <- "\nA Joint Model with:\n\n  Survival:%s\n  Longitudinal:%s\n  Link:%s\n"
        cat(sprintf(
            string,
            survival_string,
            longitudinal_string,
            link_string
        ))
    }
)

#' Simulating Joint Longitudinal and Time-to-Event Data
#'
#' @param design (`list`)\cr a list of [`SimGroup`] objects. See details.
#' @param longitudinal ([`SimLongitudinal`])\cr object specifying how to simulate the longitudinal data
#' @param survival ([`SimSurvival`])\cr object specifying how to simulate the survival data
#' @param .silent (`flag`)\cr whether to suppress info messages
#'
#' @slot longitudinal (`data.frame`)\cr the simulated longitudinal data.
#' @slot survival (`data.frame`)\cr the simulated survival data.
#'
#' @details
#'
#' The `design` argument is used to specify how many distinct groups should be simulated
#' including key information such as the number of subjects within the group as well as
#' which treatment arm and study the group belongs to. The `design` argument should be a
#' list of [`SimGroup`] objects e.g.
#' ```
#' design = list(
#'     SimGroup(n = 50, study = "Study-1", arm = "Arm-A"),
#'     SimGroup(n = 50, study = "Study-1", arm = "Arm-B")
#' )
#' ```
#'
#' @name SimJointData-class
#' @exportClass SimJointData
.SimJointData <- setClass(
    "SimJointData",
    slots = list(
        longitudinal = "data.frame",
        survival = "data.frame"
    )
)



#' @rdname SimJointData-class
#' @export
SimJointData <- function(
    design = list(
        SimGroup(n = 50, study = "Study-1", arm = "Arm-A"),
        SimGroup(n = 50, study = "Study-1", arm = "Arm-B")
    ),
    longitudinal,
    survival,
    .silent = FALSE
) {

    assert(
        all(vapply(design, \(x) is(x, "SimGroup"), logical(1))),
        msg = "All elements of `design` must be of class `SimGroup`"
    )

    hazard_evaluation_info <- hazardWindows(survival)

    n_group <- vapply(design, function(x) x@n, numeric(1))
    arms <- vapply(design, function(x) x@arm, character(1))
    studies <- vapply(design, function(x) x@study, character(1))
    n_subjects <- sum(n_group)
    n_times <- length(hazard_evaluation_info$midpoint)

    sprintf_string <- paste0("subject_%0", ceiling(log(n_subjects, 10)) + 1, "i")

    baseline <- dplyr::tibble(subject = sprintf(sprintf_string, seq_len(n_subjects))) |>
        dplyr::mutate(arm = factor(rep(arms, times = n_group), levels = unique(arms))) |>
        dplyr::mutate(study = factor(rep(studies, times = n_group), levels = unique(studies)))

    os_baseline <- sampleSubjects(survival, subjects_df = baseline)
    lm_baseline <- sampleSubjects(longitudinal, subjects_df = baseline)

    lm_dat_no_obvs <- lapply(
        longitudinal@times,
        \(time) {
            baseline[["time"]] <- time
            baseline
        }
    ) |>
        dplyr::bind_rows() |>
        dplyr::left_join(lm_baseline, by = c("subject", "study", "arm"))

    lm_dat <- sampleObservations(longitudinal, lm_dat_no_obvs)


    hazard_eval_df <- dplyr::tibble(
        subject = rep(lm_baseline$subject, each = n_times),
        arm = rep(lm_baseline$arm, each = n_times),
        study = rep(lm_baseline$study, each = n_times),
        midpoint = rep(as.double(hazard_evaluation_info$midpoint), times = n_subjects),
        time = rep(as.double(hazard_evaluation_info$upper), times = n_subjects),
        width = rep(as.double(hazard_evaluation_info$width), times = n_subjects)
    )

    lm_link_dat <- sampleObservations(
        longitudinal,
        dplyr::left_join(hazard_eval_df, lm_baseline, by = c("subject", "study", "arm"))
    )[, c("subject", "study", "arm", "log_haz_link", "time", "width", "midpoint")]

    os_eval_df <- lm_link_dat |>
        dplyr::left_join(os_baseline, by = c("subject", "study", "arm"))

    withCallingHandlers(
        os_dat <- sampleObservations(survival, os_eval_df),
        message = function(e) {
            if (!.silent) message(e)
            invokeRestart("muffleMessage")
        }
    )

    lm_dat2 <- lm_dat |>
        dplyr::left_join(dplyr::select(os_dat, "subject", os_time = "time"), by = "subject") |>
        dplyr::mutate(observed = (.data$time <= .data$os_time)) |>
        dplyr::arrange(dplyr::pick(c("subject", "time")))

    assert_that(
        length(unique(os_dat$subject)) == length(os_dat$subject),
        length(os_dat$subject) == n_subjects,
        all(os_dat$time >= 0),
        all(os_dat$event %in% c(0, 1)),
        msg = "Assumptions for the Survival data are not met (please report this issue)"
    )

    assert_that(
        nrow(lm_dat2) == n_subjects * length(longitudinal@times),
        length(unique(lm_dat2$subject)) == n_subjects,
        msg = "Assumptions for the Longitudinal data are not met (please report this issue)"
    )

    return(
        .SimJointData(
            survival = os_dat,
            longitudinal = lm_dat2[, c("subject", "arm", "study", "time", "sld", "observed")]
        )
    )
}

#' @rdname show-object
#' @export
setMethod(
    f = "show",
    signature = "SimJointData",
    definition = function(object) {
        x <- sprintf("\nA SimJointData Object\n\n")
        cat(x)
        return(object)
    }
)


#' Quantity Grid Specification
#'
#' @param subjects (`character` or `NULL`)\cr vector of subjects to extract quantities for.
#' If `NULL` will default to all subjects within the dataset.
#'
#' @param times (`numeric` or `NULL`)\cr vector of time points to extract quantities at.
#' If `NULL` will default to 201 evenly spaced timepoints between 0 and either the max
#' observation time (for [`LongitudinalQuantities`]) or max event time (for [`SurvivalQuantities`]).
#'
#' @param groups (`list`)\cr named list of subjects to extract quantities for. See Group Specification.
#'
#' @param spec (`list`)\cr named list of subjects to extract quantities for. The names of each
#' element should be the required subjects with the element itself being a numeric vector of timepoints
#' to generate the quantity at.
#'
#' @param length.out (`numeric`)\cr number of evenly spaced timepoints to generate quantities at.
#'
#' @param newdata (`data.frame`) \cr new data to generate quantities for. Must contain the same columns
#' and factor levels of the original data used in the [`DataSurvival`] object.
#'
#' @param params (`list`)\cr named list of parameters to fix the longitudinal model parameters at when
#' predicting survival quantities. See [`getPredictionNames()`] for the required parameters.
#'
#' @description
#' These functions are used to specify which subjects and timepoints should be generated
#' when calculating quantities via [`SurvivalQuantities`] and [`LongitudinalQuantities`].
#'
#' @details
#'
#' - `GridFixed()` is used to specify a fixed set of timepoints to generate quantities at for
#' all the specified subjects.
#'
#' - `GridGrouped()` is similar to `GridFixed()` but allows for groupwise averaging
#' (see Group Specification).
#'
#' - `GridObserved()` generates quantities at the observed longitudinal timepoints for each
#' subject.
#'
#' - `GridManual()` allows for individual timepoint specification for each subject.
#'
#' - `GridEven()` generates quantities for each subject at N evenly spaced timepoints
#' between each subjects first and last longitudinal observations.
#'
#' - `GridEvent()` generates one quantity for each subject at their event/censor time
#' as indicated by the `time` variable in the survival dataset.
#'
#' - `GridPopulation()` generates longitudinal model quantities based on the population parameters at the
#' specified time points. Generates 1 set of quantities for each distinct combination of `arm`
#' and `study` within the [`DataSubject`] object provided to the [`JointModel`].
#'
#' - `GridPrediction()` generates survival quantities based on any user-defined values at the
#' specified time points. This is useful for generating quantities for a new dataset
#' on specific longitudinal model parameters. See [`getPredictionNames()`] to determine which
#' longitudinal model parameters need to be defined for a given longitudinal model.
#'
#' @section Group Specification:
#' For `GridGrouped()`, `groups` must be a named list of character vectors. Each element of the list
#' must be a character vector of the subjects that will form the group where the element name
#' is the corresponding name of the group. For example if the goal was to create two groups
#' named `Group-1` and `Group-2` which are composed of the subjects `sub-1`, `sub-2` and
#' `sub-3`, `sub-4` respectively then this would be specified as:
#' ```
#' GridGrouped(
#'     groups = list(
#'         "Group-1" = c("sub-1", "sub-2"),
#'         "Group-2" = c("sub-3", "sub-4")
#'     )
#' )
#' ```
#' @seealso [`SurvivalQuantities`], [`LongitudinalQuantities`]
#' @name Grid-Functions
NULL






#' Grid Developer Notes
#'
#' @description
#' Developer details for implementing / extending `Grid` objects for defining
#' generated quantities for [`SurvivalQuantities`] and [`LongitudinalQuantities`].
#'
#' @slot subjects (`character` or `NULL`)\cr vector of subjects to extract quantities for.
#' If `NULL` will default to all subjects within the dataset.
#' @slot times (`numeric` or `NULL`)\cr vector of time points to extract quantities at.
#' If `NULL` will default to 201 evenly spaced timepoints between 0 and either the max
#' @slot groups (`list`)\cr named list of subjects to extract quantities for. See details.
#'
#' @details
#' All grid classes must inherit from the abstract `Grid` class.
#' All grid classes must provide `as.QuantityGenerator(object, data)` and
#' `as.QuantityCollapser(object, data)` methods where `data` is a [`DataJoint`] object.
#' These methods must return a `QuantityGenerator` and `QuantityCollapser` object respectively.
#' The `QuantityGenerator` object specifies unique subject/timepoint combinations that samples
#' should be generated at.
#' The `QuantityCollapser` object specifies how to combine these generated samples
#' to form the desired quantities.
#' As an example say we want to generate grouped samples for the groups `Group-1` and `Group-2`
#' which consist of the subjects `sub-1`, `sub-2` and `sub-3`, `sub-4` respectively at two time points
#' `10` and `20`. We can achieve this as follows:
#' ```
#' QuantityGenerator(
#'     times = c(10, 10, 10, 10, 20, 20, 20, 20),
#'     subjects = c("sub-1" "sub-2", "sub-3", "sub-4", "sub-1" "sub-2", "sub-3", "sub-4")
#' )
#' QuantityCollapser(
#'     times = c(10, 20, 10 , 20),
#'     groups = c("Group-1", "Group-1", "Group-2", "Group-2"),
#'     indexes = list(c(1, 2), c(5, 6), c(3, 4), c(7, 8))
#' )
#' ```
#' For population based quantities use the `arms` and `studies` arguments of `QuantityGenerator`
#' instead of `subjects`.
#'
#' @inheritSection Grid-Functions Group Specification
#'
#' @keywords internal
#' @seealso `Quant-Dev`
#' @name Grid-Dev
NULL



#' Quantity Developer Notes
#'
#' @description
#' Developer details for `QuantityX` objects/methods. This page just outlines the arguments
#' and slots of these objects/methods. For the full implementation details please see [Grid-Dev]
#'
#' @slot times (`numeric`)\cr See Arguments for details.
#' @slot subjects (`character`)\cr See Arguments for details.
#' @slot groups (`character`)\cr See Arguments for details.
#' @slot indexes (`list`)\cr See Arguments for details.
#'
#' @param times (`numeric`)\cr vector of time points to extract quantities at.
#' @param subjects (`character`)\cr vector of subjects to extract quantities for.
#' @param groups (`character`)\cr vector of labels to apply to the generated quantities.
#' @param indexes (`list`)\cr list of indexes that specify which observations from a
#' `QuantityGenerator` should be combined to form the desired quantities.
#'
#' @param object (`Grid`)\cr object to convert to a `QuantityGenerator` or `QuantityCollapser`.
#' @param data (`DataJoint`)\cr Survival and Longitudinal Data.
#' @param ... Not currently used.
#'
#' @details
#' The `as.QuantityGenerator` must return a `QuantityGenerator` object.
#' The `as.QuantityCollapser` must return a `QuantityCollapser` object.
#' @keywords internal
#' @name Quant-Dev
NULL


#' @rdname Grid-Dev
.Grid <- setClass("Grid")


#' @rdname Quant-Dev
.QuantityGenerator <- setClass(
    "QuantityGenerator"
)
#' `QuantityGenerator` -> `list`
#' @description
#' Converts a `QuantityGenerator` object to a list containing the required input data for a stan
#' model.
#' @param object (`QuantityGenerator`)\cr object to convert to a list.
#' @param data (`DataJoint`)\cr Survival and Longitudinal Data.
#' @param ... Not currently used.
#' @keywords internal
as_stan_list.QuantityGenerator <- function(object, data, ...) {
    stop("as_stan_list.QuantityGenerator not implemented")
}



#' @rdname Quant-Dev
.QuantityCollapser <- setClass(
    "QuantityCollapser",
    slots = c(
        "times" = "numeric",
        "groups" = "character",
        "indexes" = "list"
    )
)
#' @rdname Quant-Dev
QuantityCollapser <- function(times, groups, indexes) {
    .QuantityCollapser(
        times = times,
        groups = groups,
        indexes = indexes
    )
}

setValidity(
    "QuantityCollapser",
    function(object) {
        if (
            length(object@times) != length(object@groups) ||
                length(object@times) != length(object@indexes)
        ) {
            return("Length of `times`, `groups`, and `indexes` must be equal")
        }
    }
)

#' @export
length.QuantityCollapser <- function(x) {
    length(x@indexes)
}


#' Expand and Validate Subjects
#'
#' @param subjects (`character`)\cr vector of subjects that should exist in `data`
#' @param data (`DataJoint`)\cr Survival and Longitudinal Data.
#'
#' @description
#' If `subjects` is `NULL` this will return a named list of all subjects in `data`.
#' Else it will return `subjects` as a named list ensuring that all subjects exist in `data`.
#'
#' @keywords internal
subjects_to_list <- function(subjects = NULL, data) {
    data_list <- as.list(data)
    subjects_exp <- if (is.null(subjects)) {
        subs <- as.list(names(data_list$subject_to_index))
        names(subs) <- names(data_list$subject_to_index)
        subs
    } else {
        subs <- as.list(subjects)
        names(subs) <- subjects
        subs
    }
    subjects_exp_vec <- unlist(subjects_exp, use.names = FALSE)
    assert_that(
        identical(subjects_exp_vec, unique(subjects_exp_vec)),
        msg = "All subject names must be unique"
    )
    assert_that(
        all(subjects_exp_vec %in% names(data_list$subject_to_index)),
        msg = "Not all subjects exist within the data object"
    )
    subjects_exp
}


#' @include Grid.R
#' @include generics.R
NULL


#' @rdname Grid-Dev
.GridManual <- setClass(
    "GridManual",
    contains = "Grid",
    slots = c(
        "spec" = "list"
    )
)


#' @rdname Grid-Functions
#' @export
GridManual <- function(spec) {
    .GridManual(
        spec = spec
    )
}


setValidity(
    "GridManual",
    function(object) {
        subject_names <- names(object@spec)
        subject_names_valid <- subject_names[!is.na(subject_names) & subject_names != ""]
        if (length(subject_names_valid) != length(object@spec)) {
            return("Each element of `subjects` must be named")
        }
        for (times in object@spec) {
            if (!is.numeric(times)) {
                return("Each element of `spec` must be a numeric vector")
            }
            if (length(times) != length(unique(times))) {
                return("Each time vector per subject must be unique")
            }
        }
        return(TRUE)
    }
)


#' @rdname Quant-Dev
#' @export
as.QuantityGenerator.GridManual <- function(object, data, ...) {
    assert_class(data, "DataJoint")
    data_list <- as.list(data)
    assert_that(
        all(names(object@spec) %in% names(data_list$subject_to_index)),
        msg = "All subject names must be in the `DataSubject` object"
    )
    lens <- vapply(object@spec, length, numeric(1))
    QuantityGeneratorSubject(
        times = unlist(object@spec, use.names = FALSE),
        subjects = rep(names(object@spec), lens)
    )
}


#' @rdname Quant-Dev
#' @export
as.QuantityCollapser.GridManual <- function(object, data, ...) {
    generator <- as.QuantityGenerator(object, data)
    QuantityCollapser(
        times = generator@times,
        groups = generator@subjects,
        indexes = as.list(seq_along(generator@times))
    )
}


#' @export
as.list.GridManual <- function(x, data, ...) {
    subs <- as.list(names(x@spec))
    names(subs) <- names(x@spec)
    subs
}

#' @include LongitudinalModel.R
#' @include StanModule.R
#' @include generics.R
#' @include ParameterList.R
#' @include Parameter.R
#' @include Link.R
NULL


#' `LongitudinalSteinFojo`
#'
#' This class extends the general [`LongitudinalModel`] class for using the
#' Stein-Fojo model for the longitudinal outcome.
#'
#' @section Available Links:
#' - [`linkDSLD()`]
#' - [`linkTTG()`]
#' - [`linkIdentity()`]
#' - [`linkGrowth()`]
#' @exportClass LongitudinalSteinFojo
.LongitudinalSteinFojo <- setClass(
    Class = "LongitudinalSteinFojo",
    contains = "LongitudinalModel"
)


#' @rdname LongitudinalSteinFojo-class
#'
#' @param mu_bsld (`Prior`)\cr for the mean baseline value `mu_bsld`.
#' @param mu_ks (`Prior`)\cr for the mean shrinkage rate `mu_ks`.
#' @param mu_kg (`Prior`)\cr for the mean growth rate `mu_kg`.
#'
#' @param omega_bsld (`Prior`)\cr for the baseline value standard deviation `omega_bsld`.
#' @param omega_ks (`Prior`)\cr for the shrinkage rate standard deviation `omega_ks`.
#' @param omega_kg (`Prior`)\cr for the growth rate standard deviation `omega_kg`.
#'
#' @param sigma (`Prior`)\cr for the variance of the longitudinal values `sigma`.
#'
#' @param centred (`logical`)\cr whether to use the centred parameterization.
#'
#' @export
LongitudinalSteinFojo <- function(

    mu_bsld = prior_normal(log(60), 1),
    mu_ks = prior_normal(log(0.5), 1),
    mu_kg = prior_normal(log(0.3), 1),

    omega_bsld = prior_lognormal(log(0.2), 1),
    omega_ks = prior_lognormal(log(0.2), 1),
    omega_kg = prior_lognormal(log(0.2), 1),

    sigma = prior_lognormal(log(0.1), 1),

    centred = FALSE
) {

    sf_model <- StanModule(decorated_render(
        .x = read_stan("lm-stein-fojo/model.stan"),
        centred = centred
    ))

    # Apply constriants
    omega_bsld <- set_limits(omega_bsld, lower = 0)
    omega_ks <- set_limits(omega_ks, lower = 0)
    omega_kg <- set_limits(omega_kg, lower = 0)
    sigma <- set_limits(sigma, lower = 0)

    parameters <- list(
        Parameter(name = "lm_sf_mu_bsld", prior = mu_bsld, size = "n_studies"),
        Parameter(name = "lm_sf_mu_ks", prior = mu_ks, size = "n_arms"),
        Parameter(name = "lm_sf_mu_kg", prior = mu_kg, size = "n_arms"),

        Parameter(name = "lm_sf_omega_bsld", prior = omega_bsld, size = 1),
        Parameter(name = "lm_sf_omega_ks", prior = omega_ks, size = 1),
        Parameter(name = "lm_sf_omega_kg", prior = omega_kg, size = 1),

        Parameter(name = "lm_sf_sigma", prior = sigma, size = 1)
    )

    assert_flag(centred)
    parameters_extra <- if (centred) {
        list(
            Parameter(
                name = "lm_sf_psi_bsld",
                prior = prior_init_only(prior_lognormal(median(mu_bsld), median(omega_bsld))),
                size = "n_subjects"
            ),
            Parameter(
                name = "lm_sf_psi_ks",
                prior = prior_init_only(prior_lognormal(median(mu_ks), median(omega_ks))),
                size = "n_subjects"
            ),
            Parameter(
                name = "lm_sf_psi_kg",
                prior = prior_init_only(prior_lognormal(median(mu_kg), median(omega_kg))),
                size = "n_subjects"
            )
        )
    } else {
        list(
            Parameter(name = "lm_sf_eta_tilde_bsld", prior = prior_std_normal(), size = "n_subjects"),
            Parameter(name = "lm_sf_eta_tilde_ks", prior = prior_std_normal(), size = "n_subjects"),
            Parameter(name = "lm_sf_eta_tilde_kg", prior = prior_std_normal(), size = "n_subjects")
        )
    }
    parameters <- append(parameters, parameters_extra)

    x <- LongitudinalModel(
        name = "Stein-Fojo",
        stan = merge(
            sf_model,
            StanModule("lm-stein-fojo/functions.stan")
        ),
        parameters = do.call(ParameterList, parameters)
    )
    .LongitudinalSteinFojo(x)
}



#' @export
enableGQ.LongitudinalSteinFojo <- function(object, ...) {
    StanModule("lm-stein-fojo/quantities.stan")
}

#' @export
enableLink.LongitudinalSteinFojo <- function(object, ...) {
    object@stan <- merge(
        object@stan,
        StanModule("lm-stein-fojo/link.stan")
    )
    object
}

#' @export
linkDSLD.LongitudinalSteinFojo <- function(prior = prior_normal(0, 2), model, ...) {
    LinkComponent(
        key = "link_dsld",
        stan = StanModule("lm-stein-fojo/link_dsld.stan"),
        prior = prior
    )
}

#' @export
linkTTG.LongitudinalSteinFojo <- function(prior = prior_normal(0, 2), model, ...) {
    LinkComponent(
        key = "link_ttg",
        stan = StanModule("lm-stein-fojo/link_ttg.stan"),
        prior = prior
    )
}

#' @export
linkIdentity.LongitudinalSteinFojo <- function(prior = prior_normal(0, 2), model, ...) {
    LinkComponent(
        key = "link_identity",
        stan = StanModule("lm-stein-fojo/link_identity.stan"),
        prior = prior
    )
}

#' @export
linkGrowth.LongitudinalSteinFojo <- function(prior = prior_normal(0, 2), model, ...) {
    LinkComponent(
        key = "link_growth",
        stan = StanModule("lm-stein-fojo/link_growth.stan"),
        prior = prior
    )
}

#' @rdname getPredictionNames
#' @export
getPredictionNames.LongitudinalSteinFojo <- function(object, ...) {
    c("b", "s", "g")
}

#' @include generics.R
#' @include utilities.R
NULL

#' Re-used documentation for `DataSurvival`
#'
#' @param object ([`DataSurvival`]) \cr Survival Data.
#' @param x ([`DataSurvival`]) \cr Survival Data.
#' @param ... Not Used.
#'
#' @name DataSurvival-Shared
#' @keywords internal
NULL



# DataSurvival-class ----

#' Survival Data Object and Constructor Function
#'
#' The [`DataSurvival`] class handles the processing of the survival data
#' for fitting a [`JointModel`].
#'
#' @slot data (`data.frame`)\cr See Arguments for details.
#' @slot formula (`formula`)\cr See Arguments for details.
#'
#' @family DataObjects
#' @family DataSurvival
#' @exportClass DataSurvival
#' @export DataSurvival
.DataSurvival <- setClass(
    Class = "DataSurvival",
    representation = list(
        data = "data.frame",
        formula = "formula"
    )
)

#' @param data (`data.frame`)\cr the observed time-to-event data.
#' @param formula (`formula`)\cr of the form `Surv(time, event) ~ cov1 + cov2 + ...`.
#'   See [survival::Surv()] for more details, though note that this package only supports right censoring.
#' @rdname DataSurvival-class
DataSurvival <- function(data, formula) {
    .DataSurvival(
        data = remove_missing_rows(data, formula),
        formula = formula
    )
}

setValidity(
    "DataSurvival",
    method = function(object) {
        dat <- object@data
        x <- extractVariableNames(object)
        dnames <- names(dat)
        if (nrow(object@data) == 0) {
            return("`data` should not have 0 rows")
        }
        if (length(x$frm) != 3) {
            return("`formula` should be a 2 sided formula")
        }
        LHS <- as.character(x$frm[[2]][[1]])
        if (!(identical(LHS, "Surv") || identical(LHS, c("::", "survival", "Surv")))) {
            return("The LHS of `formula` should be a call to survival::Surv()")
        }
        for (v in c(x$time, x$event)) {
            if (! v %in% dnames) {
                return(sprintf("Variable %s is not in `data`", x$subject))
            }
        }
        return(TRUE)
    }
)






#' @inheritParams DataSurvival-Shared
#' @inherit extractVariableNames description title
#'
#' @returns
#' A list with the following named elements:
#' - `frm` (`formula`)\cr a symbolic description of the survival model to be fitted
#' - `time` (`character`)\cr  The name of the variable containing the event time
#' - `event` (`character`) \cr  The name of the variable containing the event status
#' @export
#' @family DataSurvival
#' @family extractVariableNames
#' @keywords internal
extractVariableNames.DataSurvival <- function(object) {
    list(
        frm = object@formula,
        time = as.character(object@formula[[2]][[2]]),
        event = as.character(object@formula[[2]][[3]])
    )
}


#' `DataSurvival` -> `data.frame`
#'
#' @inheritParams DataSurvival-Shared
#'
#' @description
#' Converts a [`DataSurvival`] object into a `data.frame`.
#' The subject variable is cast to factor.
#' @family DataSurvival
#' @export
as.data.frame.DataSurvival <- function(x, ...) {
    x <- x@data
    rownames(x) <- NULL
    x
}



#' @rdname as_stan_list.DataObject
#' @family DataSurvival
#' @export
as_stan_list.DataSurvival <- function(object, ...) {
    df <- as.data.frame(object)
    vars <- extractVariableNames(object)

    design_mat <- stats::model.matrix(vars$frm, data = df)
    remove_index <- grep("(Intercept)", colnames(design_mat), fixed = TRUE)
    design_mat <- design_mat[, -remove_index, drop = FALSE]
    rownames(design_mat) <- NULL

    # Parameters for efficient integration of hazard function -> survival function
    gh_parameters <- statmod::gauss.quad(
        n = getOption("jmpost.gauss_quad_n"),
        kind = "legendre"
    )

    model_data <- list(
        n_subject_event = sum(df[[vars$event]]),
        subject_event_index = which(df[[vars$event]] == 1),
        event_times = df[[vars$time]],
        p_os_cov_design = ncol(design_mat),
        os_cov_design = design_mat,
        n_nodes = length(gh_parameters$nodes),
        nodes = gh_parameters$nodes,
        weights = gh_parameters$weights
    )
    return(model_data)
}

#' @rdname as_stan_list.DataObject
#' @export
as.list.DataSurvival <- function(x, ...) {
    as_stan_list(x, ...)
}

#' @rdname harmonise
harmonise.DataSurvival <- function(object, subject_var, subject_ord, ...) {

    data <- as.data.frame(object)

    assert_string(subject_var, na.ok = FALSE)
    assert_character(subject_ord, any.missing = FALSE)
    assert_that(
        subject_var %in% names(data),
        msg = sprintf("Subject variable `%s` not found in `survival`", subject_var)
    )
    assert_that(
        all(data[[subject_var]] %in% subject_ord),
        msg = "There are subjects in `survival` that are not present in `subjects`"
    )
    assert_that(
        all(subject_ord %in% data[[subject_var]]),
        msg = "There are subjects in `subjects` that are not present in `survival`"
    )

    data[[subject_var]] <- factor(
        as.character(data[[subject_var]]),
        levels = subject_ord
    )

    data_ord <- data[order(data[[subject_var]]), ]

    DataSurvival(
        data = data_ord,
        formula = object@formula
    )
}



#' `DataSurvival` -> Printable `Character`
#'
#' Converts [`DataSurvival`] object into a printable string.
#' @inheritParams DataSurvival-Shared
#' @family DataSurvival
#' @param indent (`numeric`)\cr how much white space to prefix the print string with.
#' @keywords internal
#' @export
as_print_string.DataSurvival <- function(object, indent = 1, ...) {
    template <- c(
        "Survival-Data Object:",
        "    # of Rows     = %d",
        "    # of Columns  = %d",
        "    # of Events   = %d",
        "    Formula       = %s"
    )
    pad <- rep(" ", indent) |> paste(collapse = "")
    template_padded <- paste(pad, template)
    vars <- extractVariableNames(object)
    sprintf(
        paste(template_padded, collapse = "\n"),
        nrow(object@data),
        ncol(object@data),
        sum(object@data[[vars$event]]),
        Reduce(paste, deparse(vars$frm))
    )
}

#' @rdname show-object
#' @export
setMethod(
    f = "show",
    signature = "DataSurvival",
    definition = function(object) {
        string <- as_print_string(object)
        cat("\n", string, "\n\n")
    }
)

#' Build design matrix for prediction data
#'
#' @description
#' This function takes a `DataSurvival` object and a `data.frame` object and generates
#' a design matrix for the `data.frame` that has the identical structure to the
#' design matrix of the `DataSurvival` object.
#'
#' This is used for predicting new data using a model that was trained on a different
#' original data source
#'
#' @param olddata ([`DataSurvival`]) \cr The original data to be used as a template for the new data
#' @param newdata ([`data.frame`]) \cr The new data to be used to generate the design matrix
#' @importFrom stats .checkMFClasses terms delete.response model.frame model.matrix
#' @importFrom survival coxph
#' @keywords internal
mirror_design_matrix <- function(olddata, newdata) {
    frm <- as_formula(olddata)
    # Dummy model to generate a bunch of meta information that we can use to
    # re-construct the design matrix
    model <- coxph(data = as.data.frame(olddata), formula = frm)
    model_terms <- delete.response(terms(model))
    model_frame <- model.frame(
        model_terms,
        newdata,
        xlev = model$xlevels
    )
    if (
        !is.null(data_classes <- attr(model_terms, "dataClasses"))) {
        .checkMFClasses(data_classes, model_frame)
    }
    design_mat <- model.matrix(
        model_terms,
        model_frame,
        contrasts.arg = model$contrasts
    )
    remove_index <- grep("(Intercept)", colnames(design_mat), fixed = TRUE)
    design_mat <- design_mat[, -remove_index, drop = FALSE]
    rownames(design_mat) <- NULL
    design_mat
}


#' @export
as_formula.DataSurvival <- function(x, ...) {
    vars <- extractVariableNames(x)
    vars$frm
}


#' @include DataJoint.R
#' @include Quantities.R
NULL



#' Re-used documentation for `SurvivalQuantities`
#'
#' @param object ([`SurvivalQuantities`]) \cr survival quantities.
#' @param x ([`SurvivalQuantities`]) \cr survival quantities.
#' @param ... not used.
#'
#' @keywords internal
#' @name SurvivalQuantities-Shared
NULL



#' `SurvivalQuantities` Object & Constructor Function
#'
#' Constructor function to generate a `SurvivalQuantities` object.
#'
#' @slot quantities (`Quantities`)\cr The sampled quantities. Should contain 1 element per
#' element of `group`
#' @slot groups (`list`)\cr See argument section for details
#' @slot type (`character`)\cr See See argument section for details
#' @slot time_grid (`numeric`)\cr See argument section for details
#' @slot data ([`DataJoint`])\cr The data that the Joint Model was fitted to to produce
#' the samples/quantities
#'
#' @section Group Specification:
#' If `groups` is a character vector of subject IDs then the survival quantities will
#' only be calculated for those specific subjects.
#'
#' If `groups` is a list then any elements with more than 1 subject ID will be grouped together
#' and their quantities will be calculated by taking a point-wise average.
#' For example: `groups = list("g1" = c("sub1", "sub2"), "g2" = c("sub3", "sub4"))` would result
#' in 2 groups being created whose values are the pointwise average
#' of `c("sub1", "sub2")` and `c("sub3", "sub4")` respectively.
#'
#' If `groups=NULL` then all subjects from original dataset will be selected
#'
#' @family SurvivalQuantities
#' @name SurvivalQuantities-class
#' @export SurvivalQuantities
.SurvivalQuantities <- setClass(
    Class = "SurvivalQuantities",
    slots = c(
        "quantities" = "Quantities",
        "grid" = "Grid",
        "type" = "character",
        "data" = "DataJoint"
    )
)

#' @param object ([`JointModelSamples`]) \cr samples as drawn from a Joint Model.
#'
#' @param grid (`Grid`) \cr object that specifies which subjects and time points to calculate the
#' quantities for. See [Grid-Functions].
#'
#' @param type (`character`)\cr quantity to be generated.
#' Must be one of `surv`, `haz`, `loghaz`, `cumhaz`.
#'
#' @rdname SurvivalQuantities-class
SurvivalQuantities <- function(
    object,
    grid,
    type = c("surv", "haz", "loghaz", "cumhaz")
) {
    type <- match.arg(type)
    assert_class(object, "JointModelSamples")
    assert_class(grid, "Grid")
    assert_that(
        !is(grid, "GridPopulation"),
        msg = "GridPopulation objects are not supported for `SurvivalQuantities`"
    )

    time_grid <- seq(
        from = 0,
        to = max(as.list(object@data)[["event_times"]]),
        length = 201
    )

    grid <- coalesceGridTime(grid, time_grid)

    generator <- as.QuantityGenerator(grid, data = object@data)

    assert_that(
        all(generator@times >= 0),
        msg = "Time points must be greater than or equal to 0"
    )

    gq <- generateQuantities(
        object,
        generator = generator,
        type = "survival"
    )

    quantities_raw <- extract_quantities(gq, type)
    collapser <- as.QuantityCollapser(grid, object@data)
    quantities <- collapse_quantities(quantities_raw, collapser)

    .SurvivalQuantities(
        quantities = Quantities(
            quantities,
            groups = collapser@groups,
            times = collapser@times
        ),
        grid = grid,
        data = object@data,
        type = type
    )
}



#' `as.data.frame`
#'
#' @param x ([`SurvivalQuantities`]) \cr longitudinal quantities.
#' @param ... not used.
#' @family SurvivalQuantities
#' @export
as.data.frame.SurvivalQuantities <- function(x, ...) {
    as.data.frame(x@quantities)
}



#' summary
#'
#' @description
#' This method returns a `data.frame` of the longitudinal quantities.
#'
#' @param conf.level (`numeric`) \cr confidence level of the interval.
#' @inheritParams SurvivalQuantities-Shared
#'
#' @family SurvivalQuantities
#' @family summary
#' @export
summary.SurvivalQuantities <- function(
    object,
    conf.level = 0.95,
    ...
) {
    summary(object@quantities, conf.level = conf.level)
}


#' Automatic Plotting for `SurvivalQuantities``
#'
#' @inheritParams SurvivalQuantities-Shared
#' @param add_km (`logical`) \cr if `TRUE` Kaplan-Meier curves will be added to the plot for
#' each group/subject.
#' @param add_wrap (`logical`) \cr if `TRUE` will apply a [ggplot2::facet_wrap()] to the plot
#' by each group/subject.
#' @param conf.level (`numeric`) \cr confidence level of the interval. If values of `FALSE`,
#' `NULL` or `0` are provided then confidence regions will not be added to the plot
#' @param ... not used.
#'
#' @family SurvivalQuantities
#' @family autoplot
#' @export
autoplot.SurvivalQuantities <- function(
    object,
    conf.level = 0.95,
    add_km = FALSE,
    add_wrap = TRUE,
    ...
) {
    include_ci <- !is.null(conf.level) && is.numeric(conf.level) && conf.level > 0
    # If CI aren't needed supply a default 0.95 to summary function as it needs
    # a value to be specified to work
    conf.level <- if (include_ci) conf.level else 0.95
    assert_that(
        is.flag(add_km),
        length(conf.level) == 1,
        conf.level < 1,
        is.flag(add_wrap)
    )
    kmdf <- if (add_km) {
        subset(
            object@data,
            as.list(object@grid, data = object@data)
        )
    } else {
        NULL
    }
    all_fit_df <- summary(object, conf.level = conf.level)
    label <- switch(
        object@type,
        "surv" = expression(S(t)),
        "cumhaz" = expression(H(t)),
        "haz" = expression(h(t)),
        "loghaz" = expression(log(h(t)))
    )
    survival_plot(
        data = all_fit_df,
        add_ci = include_ci,
        add_wrap = add_wrap,
        kmdf = kmdf,
        y_label = label
    )
}




#' Survival Plot
#'
#' Internal plotting function to create survival plots with KM curve overlays
#' This function predominately exists to extract core logic into its own function
#' to enable easier unit testing.
#'
#' @param data (`data.frame`)\cr summary statistics for a survival
#' curve to be plotted. See details.
#' @param add_ci (`logical`)\cr should confidence intervals be added? Default = `TRUE`.
#' @param add_wrap (`logical`)\cr should the plots be wrapped by `data$group`? Default = `TRUE`.
#' @param kmdf (`data.frame` or `NULL`)\cr event times and status used to plot
#' overlaying KM curves. If `NULL` no KM curve will be plotted. See details.
#' @param y_label (`character` or `expression`) \cr label to display on the y-axis.
#' Default = `expression(S(t))`.
#' @param x_label (`character` or `expression`) \cr label to display on the x-axis.
#'
#' @details
#'
#' ## `data`
#' Should contain the following columns:
#' - `time` (`numeric`) \cr time point for the summary statistic.
#' - `group` (`character`) \cr the group in which the observation belongs to.
#' - `median` (`numeric`) \cr the median value for the summary statistic.
#' - `upper` (`numeric`) \cr the upper 95% CI for the summary statistic.
#' - `lower` (`numeric`) \cr the lower 95% CI for the summary statistic.
#'
#' ## `kmdf`
#' Should contain the following columns:
#' - `time` (`numeric`) \cr the time at which an event occurred.
#' - `event` (`numeric`) \cr 1/0 status indicator for the event.
#' - `group` (`character`) \cr which group the event belongs to, should correspond to values in `data$group`.
#' @keywords internal
survival_plot <- function(
    data,
    add_ci = TRUE,
    add_wrap = TRUE,
    kmdf = NULL,
    y_label = expression(S(t)),
    x_label = expression(t)
) {
    assert_that(
        is.flag(add_ci),
        is.flag(add_wrap),
        is.expression(y_label) || is.character(y_label),
        is.expression(x_label) || is.character(x_label),
        is.null(kmdf) | is.data.frame(kmdf)
    )

    p <- ggplot() +
        xlab(x_label) +
        ylab(y_label) +
        theme_bw()

    if (add_wrap) {
        p <- p + facet_wrap(~group)
        aes_ci <- aes(x = .data$time, ymin = .data$lower, ymax = .data$upper)
        aes_line <- aes(x = .data$time, y = .data$median)
        aes_km <- aes(time = .data$time, status = .data$event)
    } else {
        aes_ci <- aes(
            x = .data$time,
            ymin = .data$lower,
            ymax = .data$upper,
            fill = .data$group,
            group = .data$group
        )
        aes_line <- aes(
            x = .data$time,
            y = .data$median,
            colour = .data$group,
            group = .data$group
        )
        aes_km <- aes(
            time = .data$time,
            status = .data$event,
            group = .data$group,
            colour = .data$group
        )
    }
    p <- p + geom_line(aes_line, data = data)
    if (add_ci) {
        p <- p + geom_ribbon(aes_ci, data = data, alpha = 0.3)
    }
    if (!is.null(kmdf)) {
        p <- p +
            ggplot2.utils::geom_km(aes_km, data = kmdf) +
            ggplot2.utils::geom_km_ticks(aes_km, data = kmdf)
    }
    p
}


#' @rdname show-object
#' @export
setMethod(
    f = "show",
    signature = "SurvivalQuantities",
    definition = function(object) {
        template <- c(
            "SurvivalQuantities Object:",
            "    # of samples    = %d",
            "    # of quantities = %d",
            "    Type            = %s"
        )
        string <- sprintf(
            paste(template, collapse = "\n"),
            nrow(object@quantities),
            ncol(object@quantities),
            object@type
        )
        cat("\n", string, "\n\n")
    }
)



#' `brierScore`
#'
#' @description
#' Derives the Brier Scores (using Inverse Probability of Censoring Weighting)
#' for the Survival estimates as detailed in \insertCite{blanche2015}{jmpost}.
#'
#' @inheritParams SurvivalQuantities-Shared
#' @inheritParams Brier-Score-Shared
#'
#' @family brierScore
#' @family SurvivalQuantities
#' @references
#' \insertAllCited{}
#' @export
brierScore.SurvivalQuantities <- function(
    object,
    maintain_cen_order = TRUE,
    event_offset = TRUE,
    ...
) {
    assert_that(
        object@type == "surv",
        msg = paste(
            "Brier Score can only be calculated when the survival quantities were",
            "generated with `type = 'surv'`",
            collapse = " "
        )
    )
    assert_that(
        is(object@grid, "GridFixed"),
        msg = paste(
            "Brier Score can only be calculated when the survival quantities were",
            "generated with `grid = GridFixed()`",
            collapse = " "
        )
    )

    sdat <- summary(object)
    times <- unique(as.QuantityGenerator(object@grid, object@data)@times)
    times <- times[order(times)]
    assert_that(
        nrow(sdat) == length(times) * length(unique(sdat$group))
    )

    subject_col <- extractVariableNames(object@data@subject)$subject
    time_col <- extractVariableNames(object@data@survival)$time
    event_col <- extractVariableNames(object@data@survival)$event
    groups <- as.character(object@data@survival@data[[subject_col]])
    orig_times <- object@data@survival@data[[time_col]]
    events <- object@data@survival@data[[event_col]]

    pred_mat <- matrix(
        ncol = length(times),
        nrow = length(unique(sdat$group))
    )
    for (i in seq_along(times)) {
        pred_mat[, i] <- sdat[sdat["time"] == times[i], "median"]
        assert_that(
            all(groups == sdat[sdat["time"] == times[i], "group"])
        )
    }
    brier_score(
        t = times,
        times = orig_times,
        events = events,
        pred_mat = 1 - pred_mat,
        maintain_cen_order = maintain_cen_order,
        event_offset = event_offset
    )
}

#' @include LongitudinalModel.R
#' @include Link.R
NULL

# LongitudinalRandomSlope-class ----

#' `LongitudinalRandomSlope`
#'
#' This class extends the general [`LongitudinalModel`] class for using the
#' random slope linear model for the longitudinal outcome.
#'
#' @section Available Links:
#' - [`linkDSLD()`]
#' - [`linkIdentity()`]
#' @exportClass LongitudinalRandomSlope
.LongitudinalRandomSlope <- setClass(
    Class = "LongitudinalRandomSlope",
    contains = "LongitudinalModel"
)


# LongitudinalRandomSlope-constructors ----

#' @rdname LongitudinalRandomSlope-class
#'
#' @param intercept (`Prior`)\cr for the `intercept`.
#' @param slope_mu (`Prior`)\cr for the population slope `slope_mu`.
#' @param slope_sigma (`Prior`)\cr for the random slope standard deviation `slope_sigma`.
#' @param sigma (`Prior`)\cr for the variance of the longitudinal values `sigma`.
#'
#' @export
LongitudinalRandomSlope <- function(
    intercept = prior_normal(30, 10),
    slope_mu = prior_normal(1, 3),
    slope_sigma = prior_lognormal(0, 1.5),
    sigma = prior_lognormal(0, 1.5)
) {

    stan <- StanModule(
        x = "lm-random-slope/model.stan"
    )

    # Apply constriants
    sigma <- set_limits(sigma, lower = 0)
    slope_sigma <- set_limits(slope_sigma, lower = 0)

    .LongitudinalRandomSlope(
        LongitudinalModel(
            name = "Random Slope",
            stan = stan,
            parameters = ParameterList(
                Parameter(name = "lm_rs_intercept", prior = intercept, size = "n_studies"),
                Parameter(name = "lm_rs_slope_mu", prior = slope_mu, size = "n_arms"),
                Parameter(name = "lm_rs_slope_sigma", prior = slope_sigma, size = 1),
                Parameter(name = "lm_rs_sigma", prior = sigma, size = 1),
                Parameter(
                    name = "lm_rs_ind_rnd_slope",
                    prior = prior_init_only(prior_normal(median(slope_mu), median(slope_sigma))),
                    size = "n_subjects"
                )
            )
        )
    )
}


#' @export
enableGQ.LongitudinalRandomSlope <- function(object, ...) {
    StanModule("lm-random-slope/quantities.stan")
}

#' @export
enableLink.LongitudinalRandomSlope <- function(object, ...) {
    object@stan <- merge(
        object@stan,
        StanModule("lm-random-slope/link.stan")
    )
    object
}


#' @export
linkDSLD.LongitudinalRandomSlope <- function(prior = prior_normal(0, 2), model, ...) {
    LinkComponent(
        key = "link_dsld",
        stan = StanModule("lm-random-slope/link_dsld.stan"),
        prior = prior
    )
}

#' @export
linkIdentity.LongitudinalRandomSlope <- function(prior = prior_normal(0, 2), model, ...) {
    LinkComponent(
        key = "link_identity",
        stan = StanModule("lm-random-slope/link_identity.stan"),
        prior = prior
    )
}

#' @export
linkGrowth.LongitudinalRandomSlope <- function(prior = prior_normal(0, 2), model, ...) {
    LinkComponent(
        key = "link_growth",
        stan = StanModule("lm-random-slope/link_growth.stan"),
        prior = prior
    )
}

#' @rdname getPredictionNames
#' @export
getPredictionNames.LongitudinalRandomSlope <- function(object, ...) {
    c("intercept", "slope")
}

#' @include Grid.R
#' @include generics.R
NULL

#' @rdname Grid-Dev
.GridEvent <- setClass(
    "GridEvent",
    contains = "Grid",
    slots = c(
        "subjects" = "character_or_NULL"
    )
)


#' @rdname Grid-Functions
#' @export
GridEvent <- function(subjects = NULL) {
    .GridEvent(
        subjects = subjects
    )
}

#' @rdname Quant-Dev
#' @export
as.QuantityGenerator.GridEvent <- function(object, data, ...) {
    assert_class(data, "DataJoint")
    assert_that(
        !is.null(data@survival),
        msg = "Survival data must have been provided to `DataJoint()` in order to use `GridEvent()`"
    )
    data_list <- as.list(data)
    subjects <- unlist(as.list(object, data = data), use.names = FALSE)
    event_times <- data_list$event_times[data_list$subject_to_index[subjects]]
    QuantityGeneratorSubject(
        times = event_times,
        subjects = subjects
    )
}

#' @rdname Quant-Dev
#' @export
as.QuantityCollapser.GridEvent <- function(object, data, ...) {
    generator <- as.QuantityGenerator(object, data)
    QuantityCollapser(
        times = generator@times,
        groups = generator@subjects,
        indexes = as.list(seq_along(generator@times))
    )
}

#' @export
as.list.GridEvent <- function(x, data, ...) {
    subjects_to_list(x@subjects, data)
}

# "missing" = no argument provided
# "NULL" = explicit NULL
setClassUnion("empty", c("missing", "NULL"))
setClassUnion("numeric_or_NULL", c("numeric", "NULL"))
setClassUnion("character_or_NULL", c("character", "NULL"))

# merge ----

#' `merge`
#'
#' Merge two `StanModule` or `ParameterList` objects.
#'
#' @param x first module.
#' @param y second module.
#' @param ... additional arguments.
#'
#' @export
# Needs to be S4 for multiple dispatch !
setGeneric(
    name = "merge",
    def = function(x, y, ...) standardGeneric("merge")
)


# show ----

#' Printing of Different Classes
#'
#' These methods print objects of different classes.
#'
#' @name show
#' @aliases show
#'
#' @param object what to print.
#'
#' @export
NULL


# write_stan ----

#' `write_stan`
#'
#' Write the Stan code for a Stan module.
#'
#' @param object the module.
#' @param destination (`character` or  `connection`)\cr Where to write stan code to.
#' @param ... Additional arguments
#'
#' @export
write_stan <- function(object, destination, ...) {
    UseMethod("write_stan")
}

# compileStanModel ----

#' `compileStanModel`
#'
#' Compile the Stan module.
#'
#' @param object the module.
#'
#' @export
compileStanModel <- function(object) {
    UseMethod("compileStanModel")
}


# sampleStanModel ----

#' `sampleStanModel`
#'
#' Sample from a Stan Module.
#'
#' @param object the module.
#' @param ... additional arguments.
#'
#' @export
sampleStanModel <- function(object, ...) {
    UseMethod("sampleStanModel")
}


# as.StanModule ----

#' `as.StanModule`
#'
#' Converts an object into a [`StanModule`].
#'
#' @param object what to convert.
#' @param ... additional options.
#' @family as.StanModule
#' @keywords internal
as.StanModule <- function(object, ...) {
    UseMethod("as.StanModule")
}



#' `getParameters`
#'
#' Extract any modelling parameters as a [`ParameterList`] object
#' from a model.
#'
#' @param object where to obtain the parameters from.
#' @param ... additional options.
#'
#' @export
getParameters <- function(object, ...) {
    UseMethod("getParameters")
}


# extractVariableNames ----

#' Extract Mapping to Standardised Variable Names
#'
#' @description
#' Extract a `list` that maps the variable names in a user-defined
#' `data.frame` to standardised values.
#'
#' @param object the data object.
#' @family extractVariableNames
#' @keywords internal
extractVariableNames <- function(object) {
    UseMethod("extractVariableNames")
}



# initialValues ----

#' `initialValues`
#'
#' Obtain the `list` of initial values to be passed to the Stan sampler.
#'
#' @param object where to get the initial values from.
#' @param n_chains the number of initial values to generate. See details.
#' @param ... Not currently used.
#'
#' @details
#' There are multiple ways of specifying initial values to Stan, see the `init` argument
#' in [cmdstanr::model-method-sample] for full details. Within this package we supply
#' initial values via a list of lists where each inner list contains the initial values
#' for a single chain. As such the `n_chains` argument specifies the number of inner lists
#' to generate.
#'
#' See the Vignette for further details of how to specify initial values.
#'
#' @export
initialValues <- function(object, ...) {
    UseMethod("initialValues")
}


# size ----

#' `size`
#'
#' Obtain the `list` of parameter sizes.
#'
#' @param object where to get the parameter sizes from.
#'
#' @keywords internal
size <- function(object) {
    UseMethod("size")
}


# generateQuantities ----

#' `generateQuantities`
#'
#' Obtain the generated quantities from a Stan Model.
#'
#' @param object object to obtain generated quantities from
#' @param ... additional options.
#'
#' @export
generateQuantities <- function(object, ...) {
    UseMethod("generateQuantities")
}


#' Prepare Data Object
#'
#' @param object (`DataSubject` or `DataLongitudinal` or `DataSurvival`) \cr data object to "harmonise"
#' @param subject_var (`character`) \cr the name of the variable containing the subject identifier.
#' @param subject_ord (`character`) \cr the expected levels (in order) of the subject identifier.
#' @param ... not used.
#'
#' @details
#' This utility function prepares the datasets in the data objects in order to ensure they
#' are consistent and compatible with each other.
#'
#' In particular it ensures that the `subject` variable, as specified by `DataSubject`,
#' is available in `DataLongitudinal` and `DataSurvival` and that all levels are present
#' in all 3 data objects.
#'
#' It also sorts the datasets to ensure that indexes are consistent e.g. index 1 for
#' `DataSubject@data` corresponds to the same subject as index 1 for `DataSurvival@data`.
#' For `DataLongitudinal` the data is additionally sorted by time and outcome value.
#'
#' @seealso [`DataJoint`], [`DataSurvival`], [`DataSubject`], [`DataLongitudinal`]
#'
#' @keywords internal
#' @return Returns the original object but with the data standardised (see details)
harmonise <- function(object, ...) {
    UseMethod("harmonise")
}


#' @rdname harmonise
harmonise.default <- function(object, ...) {
    NULL
}


#' `as_stan_list`
#'
#' @description
#' Extracts a list of data elements from an object to be used as input
#' to a Stan Model
#'
#' @param object to be converted.
#' @param ... additional options.
#'
#' @family as_stan_list
#' @export
as_stan_list <- function(object, ...) {
    UseMethod("as_stan_list")
}

#' @rdname as_stan_list
#' @export
as_stan_list.default <- function(object, ...) {
    NULL
}


#' `as_print_string`
#'
#' @description
#' Returns the character representation of an object which is suitable
#' for printing to the console
#'
#' @param object to be converted to string.
#' @param ... additional options.
#'
#' @family as_print_string
#' @keywords internal
as_print_string <- function(object, ...) {
    UseMethod("as_print_string")
}

#' Show an Object
#'
#' Prints an object to the console.
#'
#' @param object Object to be printed
#'
#' @name show-object
NULL



#' `brierScore`
#'
#' @description
#' Returns the Brier Score for a given model
#'
#' @param object to calculate Brier Score for.
#' @param ... additional options.
#'
#' @family brierScore
#' @export
brierScore <- function(object, ...) {
    UseMethod("brierScore")
}



#' Generate Simulated Observations
#'
#' @param object (`SimLongitudinal` or `SimSurvival`) \cr object to generate observations from.
#' @param times_df (`data.frame`) \cr the times at which to generate observations. See details.
#'
#' @details
#' The `times_df` argument should be a `data.frame` as created by `sampleSubjects` but
#' replicated for each time point at which observations are to be generated. That is if you want
#' to generate observations for times `c(0, 1, 2, 3)` then `times_df` should be created as:
#' ```
#' subject_dat <- sampleSubjects(object, ...)
#' times_df <- tidyr::expand_grid(
#'     subject_dat,
#'     time = c(0, 1, 2, 3)
#' )
#' ```
#'
#' @export
sampleObservations <- function(object, times_df) {
    UseMethod("sampleObservations")
}


#' Generate Simulated Subjects
#'
#' @param object (`SimLongitudinal` or `SimSurvival`) \cr object to generate subjects from.
#' @param subjects_df (`data.frame`) \cr the subjects to generate observations for. See details.
#'
#' @details
#' The `subjects_df` argument should be a `data.frame` with 1 row per desired subject to create
#' with the following columns:
#' - `study` (`factor`) the study identifier.
#' - `arm` (`factor`) the treatment arm identifier.
#' - `subject` (`character`) the subject identifier.
#'
#' This method takes care of generating all the individual subject data required for the
#' [`sampleObservations`] method to generate the observations.
#' @export
sampleSubjects <- function(object, subjects_df) {
    UseMethod("sampleSubjects")
}


#' Generate time windows for evaluating a hazard function
#'
#' @param object (`SurvivalModel`) \cr object to generate time windows for.
#' @param ... Not used.
#'
hazardWindows <- function(object, ...) {
    UseMethod("hazardWindows")
}


#' Coerce to `CmdStanMCMC`
#'
#' @param object to be converted
#' @param ... additional options
#'
#' @description
#' Coerces an object to a [`cmdstanr::CmdStanMCMC`] object
#'
#' @export
as.CmdStanMCMC <- function(object, ...) {
    UseMethod("as.CmdStanMCMC")
}

#' @rdname Quant-Dev
#' @export
as.QuantityGenerator <- function(object, ...) {
    UseMethod("as.QuantityGenerator")
}

#' @rdname Quant-Dev
#' @export
as.QuantityCollapser <- function(object, ...) {
    UseMethod("as.QuantityCollapser")
}


#' Coalesce Time
#'
#' @param object ([`Grid`]) \cr object to coalesce time for.
#' @param times (`numeric`) \cr the times to coalesce to.
#' @param ... Not used
#'
#' Method used to replace NULL times on grid objects (if appropriate)
#'
#' @keywords internal
coalesceGridTime <- function(object, times, ...) {
    UseMethod("coalesceGridTime")
}
#' @export
coalesceGridTime.default <- function(object, times, ...) {
    object
}


#' Resolve a Promise
#'
#' @param object (`ANY`)\cr an object to resolve.
#' @param ... (`ANY`)\cr additional arguments.
#'
#' If `object` is not a promise will just return itself else will resolve the promise
#' and return the promised object.
#'
#' @export
resolvePromise <- function(object, ...) {
    UseMethod("resolvePromise")
}

#' @rdname resolvePromise
#' @export
resolvePromise.default <- function(object, ...) {
    object
}

#' Enable Link Generic
#'
#' @param object ([`LongitudinalModel`])\cr to enable link for.
#' @param ... Not used.
#'
#' Optional hook method that is called on a [`LongitudinalModel`] only if a link method
#' is provided to [`JointModel`]. This can be used to allow the model to include any
#' optional stan code that is only required if there are links present.
#'
#' @return [`LongitudinalModel`] object
#'
#' @export
enableLink <- function(object, ...) {
    UseMethod("enableLink")
}
#' @export
enableLink.default <- function(object, ...) {
    object
}


#' Enable Generated Quantities Generic
#'
#' @param object ([`StanModel`])\cr to enable generated quantities for.
#' @param ... Not used.
#'
#' Optional hook method that is called on a [`StanModel`] if attempting to use
#' either [`LongitudinalQuantities`] or [`SurvivalQuantities`]
#'
#' @return [`StanModule`] object
#'
#' @export
enableGQ <- function(object, ...) {
    UseMethod("enableGQ")
}
#' @export
enableGQ.default <- function(object, ...) {
    StanModule()
}



#' Get Prediction Names
#'
#' Utility function that returns the names of the required parameters for predicting
#' survival quantities with [`GridPrediction`].
#'
#' @param object (`LongitudinalModel`) \cr A longitudinal model object
#' @param ... Not used.
#' @export
getPredictionNames <- function(object, ...) {
    UseMethod("getPredictionNames")
}

#' @rdname getPredictionNames
getPredictionNames.default <- function(object, ...) {
    NULL
}

#' As Formula
#'
#' Utility wrapper function to convert an object to a formula.
#' @param x (`ANY`) \cr object to convert to a formula.
#' @param ... Not used.
#' @export
as_formula <- function(x, ...) {
    UseMethod("as_formula")
}

#' @importFrom stats as.formula
#' @export
as_formula.default <- function(x, ...) {
    as.formula(x, ...)
}


#' Set Constraints
#'
#' Applies constraints to a prior distribution to ensure any sampled numbers
#' from the distribution fall within the constraints
#'
#' @param object (`Prior`)\cr a prior distribution to apply constraints to
#' @param lower (`numeric`)\cr lower constraint boundary
#' @param upper (`numeric`)\cr upper constraint boundary
#'
#' @export
set_limits <- function(object, lower = -Inf, upper = Inf) {
    UseMethod("set_limits")
}



#' Promise
#'
#' Abstract class for promise objects to inherit off of
#' @aliases Promise
#' @exportClass Promise
.Promise <- setClass("Promise")


#' Promise of a `LongitudinalModel`
#'
#' An object that promises to resolve to a [`LongitudinalModel`] object.
#'
#' @exportClass PromiseLongitudinalModel
.PromiseLongitudinalModel <- setClass(
    "PromiseLongitudinalModel",
    contains = "Promise"
)

#' @rdname PromiseLongitudinalModel-class
#' @export
PromiseLongitudinalModel <- function() {
    .PromiseLongitudinalModel()
}




#' Promise of a `LinkComponent`
#'
#' An object that promises to resolve to a [`LinkComponent`] object.
#' Inheriting from [`Promise`] and [`LinkComponent`].
#'
#' @slot fun (`function`) \cr a function that returns a `LinkComponent`. See details.
#'
#' @param fun (`function`) \cr a function that returns a `LinkComponent`. See details.
#' @inheritParams LinkComponent
#'
#' @details
#'
#' The `fun` slot should be a function of signature `function(prior, model)` and should return
#' a [`LinkComponent`] object. An error will be thrown if the returned [`LinkComponent`] object
#' does not have the same `key` slot value as the original `PromiseLinkComponent`.
#'
#' @exportClass PromiseLinkComponent
.PromiseLinkComponent <- setClass(
    "PromiseLinkComponent",
    contains = c("Promise", "LinkComponent"),
    slots = list(
        fun = "function"
    )
)

#' @rdname PromiseLinkComponent-class
#' @export
PromiseLinkComponent <- function(fun, prior, key) {
    .PromiseLinkComponent(
        fun = fun,
        stan = StanModule(),
        parameters = ParameterList(Parameter(name = key, prior = prior, size = 1)),
        key = key
    )
}


#' @export
as.StanModule.PromiseLinkComponent <- function(object, model, ...) {
    resolved_object <- resolvePromise(object, model = model)
    as.StanModule(resolved_object, ...)
}

#' Resolve a `PromiseLinkComponent`
#'
#' Resolves a [`PromiseLinkComponent`] object to a [`LinkComponent`] object.
#' An error will be thrown if the returned [`LinkComponent`] object
#' does not have the same `key` slot value as the original [`PromiseLinkComponent`].
#'
#' @param object ([`PromiseLinkComponent`]) \cr the promise to resolve
#' @param model ([`LongitudinalModel`]) \cr the model to resolve the promise with
#' @param ... Not used.
#' @return ([`LinkComponent`]) \cr the resolved `LinkComponent` object
#'
#' @export
resolvePromise.PromiseLinkComponent <- function(object, model, ...) {
    x <- object@fun(
        prior = object@parameters@parameters[[1]]@prior,
        model = model
    )
    assert_that(
        is(x, "LinkComponent"),
        msg = "Resolved `PromiseLinkComponent` did not produce a `LinkComponent` object"
    )
    assert_that(
        names(object) == names(x),
        msg = paste(
            "Resolved `PromiseLinkComponent` did not produce a `LinkComponent` object",
            "with the same key as the promise"
        )
    )
    x
}

#' @include StanModel.R
NULL

# LongitudinalModel-class ----

#' `LongitudinalModel`
#'
#' This class extends the general [`StanModel`] class to comprise the longitudinal
#' model specification.
#'
#' @exportClass LongitudinalModel
.LongitudinalModel <- setClass(
    Class = "LongitudinalModel",
    contains = "StanModel"
)

# LongitudinalModel-constructors ----

#' @rdname LongitudinalModel-class
#'
#' @inheritParams stanmodel_arguments
#'
#' @export
LongitudinalModel <- function(
    stan = StanModule(),
    parameters = ParameterList(),
    name = "<Unnamed>",
    ...
) {
    base_stan <- read_stan("base/longitudinal.stan")

    stan_full <- decorated_render(
        .x = base_stan,
        stan = add_missing_stan_blocks(as.list(stan))
    )

    .LongitudinalModel(
        StanModel(
            stan = StanModule(stan_full),
            parameters = parameters,
            name = name,
            ...
        )
    )
}

#' @export
as_print_string.LongitudinalModel <- function(object, ...) {
    string <- sprintf(
        "\n%s Longitudinal Model with parameters:\n%s\n\n",
        object@name,
        paste("   ", as_print_string(object@parameters)) |> paste(collapse = "\n")
    )
    return(string)
}




#' jmpost settings
#'
#' @description
#' Define settings that modify the behaviour of the `jmpost` package
#'
#' Each of the following are the name of options that can be set via:
#' ```
#' options(<option_name> = <value>)
#' ```
#'
#' ## `jmpost.prior_shrinkage`
#'
#' Default = `0.5`
#'
#' By default all initial values are drawn as random sample from the respective prior
#' distribution with a shrinkage factor towards the mean. That is:
#' ```
#' initial_value = prior_mean * prior_shrinkage + (1 - prior_shrinkage) * prior_sample
#' ```
#' This setting controls the shrinkage factor. A value of 0 means no shrinkage (i.e.
#' pure random draw) whilst a value of 1 means the initial value is just the mean.
#'
#' ## `jmpost.cache_dir`
#'
#' Default = `tempfile()`
#'
#' Directory to store compiled stan models in. If not set, a temporary directory is used for
#' the given R session. Can also be set via the environment variable `JMPOST_CACHE_DIR`.
#'
#'
#'
#' ## `jmpost.gauss_quad_n`
#'
#' Default = 15
#'
#' In most cases the survival function of the joint model does not have a closed form
#' and as such it is calculated by integrating the hazard function. `jmpost` estimates this
#' via Gaussian Quadrature, in particular it uses [`statmod::gauss.quad`] with
#' `kind = "legendre"` to create the nodes and weights.
#'
#' This option specifies the `n` argument in the call to [`statmod::gauss.quad`]. In general
#' higher values of `n` lead to better accuracy of the approximation but at the cost of
#' increased computational time.
#'
#' @examples
#' \dontrun{
#' options(jmpost.prior_shrinkage = 0.5)
#' }
#' @name jmpost-settings
set_options <- function() {

    cache_dir <- Sys.getenv("JMPOST_CACHE_DIR")

    if (cache_dir == "" || is.null(cache_dir)) {
        cache_dir <- tempfile()
    }

    current_opts <- names(options())
    jmpost_opts <- list(
        jmpost.cache_dir = cache_dir,
        jmpost.prior_shrinkage = 0.5,
        jmpost.gauss_quad_n = 15
    )
    for (opt in names(jmpost_opts)) {
        if (!opt %in% current_opts) {
            options(jmpost_opts[opt])
        }
    }
}

#' @include Grid.R
#' @include generics.R
NULL


#' @rdname Grid-Dev
.GridPopulation <- setClass(
    "GridPopulation",
    contains = "Grid",
    slots = c(
        "times" = "numeric_or_NULL"
    )
)

#' @rdname Grid-Functions
#' @export
GridPopulation <- function(times = NULL) {
    .GridPopulation(
        times = times
    )
}

#' @rdname Quant-Dev
#' @export
as.QuantityGenerator.GridPopulation <- function(object, data, ...) {

    assert_class(data, "DataJoint")
    data_list <- as.list(data)
    validate_time_grid(object@times)

    n_times <- length(object@times)
    n_quant <- length(data_list$pop_study_index)

    QuantityGeneratorPopulation(
        times = rep(object@times, each = n_quant),
        arms = rep(names(data_list$arm_to_index)[data_list$pop_arm_index], n_times),
        studies = rep(names(data_list$study_to_index)[data_list$pop_study_index], n_times)
    )
}

#' @rdname Quant-Dev
#' @export
as.QuantityCollapser.GridPopulation <- function(object, data, ...) {
    assert_class(data, "DataJoint")
    data_list <- as.list(data)
    generator <- as.QuantityGenerator(object, data)
    QuantityCollapser(
        times = generator@times,
        groups = sprintf(
            "arm=%s; study=%s",
            generator@arms,
            generator@studies
        ),
        indexes = as.list(seq_along(generator@times))
    )
}


#' @export
as.list.GridPopulation <- function(x, data, ...) {
    stop("`as.list()` is not implemented for `GridPopulation` objects")
}


#' @rdname coalesceGridTime
#' @export
coalesceGridTime.GridPopulation <- function(object, times, ...) {
    if (is.null(object@times)) {
        object <- GridPopulation(
            times = times
        )
    }
    object
}

#' @include SurvivalModel.R
NULL

# SurvivalExponential-class ----

#' `SurvivalExponential`
#'
#' This class extends the general [`SurvivalModel`] class for using the
#' exponential survival model.
#'
#' @exportClass SurvivalExponential
.SurvivalExponential <- setClass(
    Class = "SurvivalExponential",
    contains = "SurvivalModel"
)

# SurvivalExponential-constructors ----

#' @rdname SurvivalExponential-class
#'
#' @param lambda (`Prior`)\cr for the exponential rate `lambda`.
#' @param beta (`Prior`)\cr for covariates coefficients `beta`.
#'
#' @export
SurvivalExponential <- function(
    lambda = prior_gamma(2, 5),
    beta = prior_normal(0, 2)
) {
    .SurvivalExponential(
        SurvivalModel(
            name = "Exponential",
            stan = StanModule("sm-exponential/model.stan"),
            parameters = ParameterList(
                Parameter(name = "sm_exp_lambda", prior = lambda, size = 1),
                Parameter(name = "beta_os_cov", prior = beta, size = "p_os_cov_design")
            )
        )
    )
}

#' @include generics.R
#' @include Parameter.R
#' @include Prior.R
NULL


#' ParameterList-Shared
#' @param object (`ParameterList`) \cr A List of [`Parameter`] Objects.
#' @param x (`ParameterList`) \cr A List of [`Parameter`] Objects.
#' @param ... Not Used.
#' @keywords internal
#' @name ParameterList-Shared
NULL


# ParameterList-class ----

#' `ParameterList`
#'
#' This class extends the general [`list`] type for containing [`Parameter`]
#' specifications.
#'
#'

#' @slot parameters (`list`) \cr a list of [`Parameter`] objects
#' @family ParameterList
#' @export ParameterList
#' @exportClass ParameterList
.ParameterList <- setClass(
    Class = "ParameterList",
    slots = c(
        parameters = "list"
    )
)
#' @param ... (`Parameter`)\cr which parameter specifications to include.
#' @rdname ParameterList-class
ParameterList <- function(...) {
    .ParameterList(parameters = list(...))
}

# ParameterList-validity ----

setValidity(
    Class = "ParameterList",
    method = function(object) {
        is_parameters <- vapply(object@parameters, function(x) is(x, "Parameter"), logical(1))
        if (!all(is_parameters)) {
            return("all elements must be of class 'Parameter'")
        }
        return(TRUE)
    }
)



# as.StanModule-ParameterList ----

#' `ParameterList` -> `StanModule`
#'
#' Converts a [`ParameterList`] object to a [`StanModule`] object
#'
#' @inheritParams ParameterList-Shared
#'
#' @family ParameterList
#' @family as.StanModule
#' @export
as.StanModule.ParameterList <- function(object, ...) {
    stan_modules <- lapply(
        object@parameters,
        as.StanModule
    )
    assert_that(
        all(vapply(stan_modules, inherits, logical(1), "StanModule"))
    )
    Reduce(merge, stan_modules)
}



#' `ParameterList` -> `list`
#'
#' Converts a ParameterList object to a list of parameter data values
#' for a Stan model.
#'
#' @inheritParams ParameterList-Shared
#'
#' @family as_stan_list
#' @family ParameterList
#' @export
as_stan_list.ParameterList <- function(object, ...) {
    stan_lists <- lapply(
        object@parameters,
        as_stan_list
    )
    assert_that(
        all(vapply(stan_lists, is.list, logical(1)))
    )
    Reduce(append, stan_lists)
}


# merge-ParameterList,ParameterList ----

#' @rdname merge
setMethod(
    f = "merge",
    signature = c(x = "ParameterList", y = "ParameterList"),
    definition = function(x, y) {
        parameters <- append(x@parameters, y@parameters)
        do.call(ParameterList, parameters)
    }
)

# as.list-ParameterList ----

#' `ParameterList` -> `list`
#' @description
#' Returns a named list where each element of the list corresponds
#' to a Stan modelling block e.g. `data`, `model`, etc.
#' @inheritParams ParameterList-Shared
#' @family ParameterList
#' @export
as.list.ParameterList <- function(x, ...) {
    as.list(as.StanModule(x))
}




#' Parameter-List Getter Functions
#' @description
#' Getter functions for the slots of a [`ParameterList`] object
#' @inheritParams ParameterList-Shared
#' @family ParameterList
#' @param n_chains (`integer`) \cr the number of chains.
#' @name ParameterList-Getter-Methods
NULL


#' @describeIn ParameterList-Getter-Methods The parameter-list's parameter names
#' @export
names.ParameterList <- function(x) {
    vapply(x@parameters, names, character(1))
}


#' @describeIn ParameterList-Getter-Methods The parameter-list's parameter initial values
#' @export
initialValues.ParameterList <- function(object, n_chains, ...) {
    # Generate initial values as a list of lists. This is to ensure it is in the required
    # format as specified by cmdstanr see the `init` argument of
    # `help("model-method-sample", "cmdstanr")` for more details
    lapply(
        seq_len(n_chains),
        \(i) {
            vals <- lapply(object@parameters, initialValues)
            name <- vapply(object@parameters, names, character(1))
            names(vals) <- name
            vals
        }
    )
}


#' @describeIn ParameterList-Getter-Methods The parameter-list's parameter dimensionality
#' @export
size.ParameterList <- function(object) {
    x <- lapply(object@parameters, size)
    names(x) <- names(object)
    return(x)
}


#' `ParameterList` -> Printable `Character`
#'
#' Converts [`ParameterList`] object into a printable string.
#' @inheritParams ParameterList-Shared
#' @family ParameterList
#' @keywords internal
#' @export
as_print_string.ParameterList <- function(object, ...) {
    x <- vapply(object@parameters, as.character, character(1))
    if (length(x) == 0) {
        x <- "<No Parameters>"
    }
    return(x)
}



#' @rdname show-object
#' @export
setMethod(
    f = "show",
    signature = "ParameterList",
    definition = function(object) {
        chrs <- as_print_string(object)
        string <- paste("   ", chrs) |> paste(collapse = "\n")
        x <- sprintf("\nParameterList Object:\n%s\n\n", string)
        cat(x)
        return(object)
    }
)



#' @include SimLongitudinal.R
#' @include generics.R
NULL

#' Simulate Longitudinal Data from a GSF Model
#'
#' @param times (`numeric`)\cr the times to generate observations at.
#' @param sigma (`number`)\cr the variance of the longitudinal values.
#' @param mu_s (`numeric`)\cr the mean shrinkage rates.
#' @param mu_g (`numeric`)\cr the mean growth rates.
#' @param mu_b (`numeric`)\cr the mean baseline values.
#' @param mu_phi (`numeric`)\cr the mean proportion of cells affected by the treatment
#' @param omega_b (`number`)\cr the baseline value standard deviation.
#' @param omega_s (`number`)\cr the shrinkage rate standard deviation.
#' @param omega_g (`number`)\cr the growth rate standard deviation.
#' @param omega_phi (`number`)\cr for the standard deviation of the proportion of cells
#' affected by the treatment `omega_phi`.
#' @param link_dsld (`number`)\cr the link coefficient for the derivative contribution.
#' @param link_ttg (`number`)\cr the link coefficient for the time-to-growth contribution.
#' @param link_identity (`number`)\cr the link coefficient for the SLD Identity contribution.
#' @param link_growth (`number`)\cr the link coefficient for the growth parameter contribution.
#'
#' @slot sigma (`numeric`)\cr See arguments.
#' @slot mu_s (`numeric`)\cr See arguments.
#' @slot mu_g (`numeric`)\cr See arguments.
#' @slot mu_b (`numeric`)\cr See arguments.
#' @slot mu_phi (`numeric`)\cr See arguments.
#' @slot omega_b (`numeric`)\cr See arguments.
#' @slot omega_s (`numeric`)\cr See arguments.
#' @slot omega_g (`numeric`)\cr See arguments.
#' @slot omega_phi (`numeric`)\cr See arguments.
#' @slot link_dsld (`numeric`)\cr See arguments.
#' @slot link_ttg (`numeric`)\cr See arguments.
#' @slot link_identity (`numeric`)\cr See arguments.
#' @slot link_growth (`numeric`)\cr See arguments.
#' @family SimLongitudinal
#' @name SimLongitudinalGSF-class
#' @exportClass SimLongitudinalGSF
.SimLongitudinalGSF <- setClass(
    "SimLongitudinalGSF",
    contains = "SimLongitudinal",
    slots = c(
        sigma = "numeric",
        mu_s = "numeric",
        mu_g = "numeric",
        mu_b = "numeric",
        mu_phi = "numeric",
        omega_b = "numeric",
        omega_s = "numeric",
        omega_g = "numeric",
        omega_phi = "numeric",
        link_dsld = "numeric",
        link_ttg = "numeric",
        link_identity = "numeric",
        link_growth = "numeric"
    )
)

#' @rdname SimLongitudinalGSF-class
#' @export
SimLongitudinalGSF <- function(
    times = c(-100, -50, 0, 50, 100, 150, 250, 350, 450, 550) / 365,
    sigma = 0.01,
    mu_s = log(c(0.6, 0.4)),
    mu_g = log(c(0.25, 0.35)),
    mu_b = log(60),
    mu_phi = qlogis(c(0.4, 0.6)),
    omega_b = 0.2,
    omega_s = 0.2,
    omega_g = 0.2,
    omega_phi = 0.2,
    link_dsld = 0,
    link_ttg = 0,
    link_identity = 0,
    link_growth = 0
) {
    .SimLongitudinalGSF(
        times = times,
        sigma = sigma,
        mu_s = mu_s,
        mu_g = mu_g,
        mu_b = mu_b,
        mu_phi = mu_phi,
        omega_b = omega_b,
        omega_s = omega_s,
        omega_g = omega_g,
        omega_phi = omega_phi,
        link_dsld = link_dsld,
        link_ttg = link_ttg,
        link_identity = link_identity,
        link_growth = link_growth
    )
}


setValidity(
    "SimLongitudinalGSF",
    function(object) {
        par_lengths <- c(
            length(object@mu_s),
            length(object@mu_g),
            length(object@mu_phi)
        )
        if (length(unique(par_lengths)) != 1) {
            return("The parameters `mu_s`, `mu_g` and `mu_phi` must have the same length.")
        }

        len_1_pars <- c(
            "sigma", "omega_b", "omega_s", "omega_g", "omega_phi",
            "link_dsld", "link_ttg", "link_identity", "link_growth"
        )
        for (par in len_1_pars) {
            if (length(slot(object, par)) != 1) {
                return(sprintf("The `%s` parameter must be a length 1 numeric.", par))
            }
        }

        return(TRUE)
    }
)

#' @rdname as_print_string
as_print_string.SimLongitudinalGSF <- function(object) {
    return("SimLongitudinalGSF")
}

#' @rdname sampleObservations
#' @export
sampleObservations.SimLongitudinalGSF <- function(object, times_df) {
    times_df |>
        dplyr::mutate(mu_sld = gsf_sld(.data$time, .data$psi_b, .data$psi_s, .data$psi_g, .data$psi_phi)) |>
        dplyr::mutate(dsld = gsf_dsld(.data$time, .data$psi_b, .data$psi_s, .data$psi_g, .data$psi_phi)) |>
        dplyr::mutate(ttg = gsf_ttg(.data$time, .data$psi_b, .data$psi_s, .data$psi_g, .data$psi_phi)) |>
        dplyr::mutate(sld = stats::rnorm(dplyr::n(), .data$mu_sld, .data$mu_sld * object@sigma)) |>
        dplyr::mutate(
            log_haz_link =
                (object@link_dsld * .data$dsld) +
                (object@link_ttg * .data$ttg) +
                (object@link_identity * .data$mu_sld) +
                (object@link_growth * log(.data$psi_g))
        )
}


#' @rdname sampleSubjects
#' @export
sampleSubjects.SimLongitudinalGSF <- function(object, subjects_df) {
    assert_that(
        is.factor(subjects_df$study),
        is.factor(subjects_df$arm),
        length(levels(subjects_df$study)) == length(object@mu_b),
        length(levels(subjects_df$arm)) == length(object@mu_s),
        length(levels(subjects_df$arm)) == length(object@mu_g),
        length(levels(subjects_df$arm)) == length(object@mu_phi)
    )

    res <- subjects_df |>
        dplyr::distinct(.data$subject, .data$arm, .data$study) |>
        dplyr::mutate(study_idx = as.numeric(.data$study)) |>
        dplyr::mutate(arm_idx = as.numeric(.data$arm)) |>
        dplyr::mutate(psi_b = stats::rlnorm(dplyr::n(), object@mu_b[.data$study_idx], object@omega_b)) |>
        dplyr::mutate(psi_s = stats::rlnorm(dplyr::n(), object@mu_s[.data$arm_idx], object@omega_s)) |>
        dplyr::mutate(psi_g = stats::rlnorm(dplyr::n(), object@mu_g[.data$arm_idx], object@omega_g)) |>
        dplyr::mutate(psi_phi_logit = stats::rnorm(
            dplyr::n(),
            object@mu_phi[.data$arm_idx],
            object@omega_phi
        )) |>
        dplyr::mutate(psi_phi = stats::plogis(.data$psi_phi_logit))

    res[, c("subject", "arm", "study", "psi_b", "psi_s", "psi_g", "psi_phi")]
}




## sim_lm_gsf ----

#' Generalized Stein-Fojo Functionals
#'
#' @param time (`numeric`)\cr time grid.
#' @param b (`number`)\cr baseline.
#' @param s (`number`)\cr shrinkage.
#' @param g (`number`)\cr growth.
#' @param phi (`number`)\cr shrinkage proportion.
#'
#' @returns The function results.
#'
#' @keywords internal
gsf_sld <- function(time, b, s, g, phi) {
    phi <- dplyr::if_else(time >= 0, phi, 0)
    b * (phi * exp(-s * time) + (1 - phi) * exp(g * time))
}


#' @rdname gsf_sld
gsf_ttg <- function(time, b, s, g, phi) {
    t1 <- (log(s * phi / (g * (1 - phi))) / (g + s))
    t1[t1 <= 0] <- 0
    return(t1)
}


#' @rdname gsf_sld
gsf_dsld <- function(time, b, s, g, phi) {
    phi <- dplyr::if_else(time >= 0, phi, 0)
    t1 <- (1 - phi) * g * exp(g * time)
    t2 <- phi * s * exp(-s * time)
    return(b * (t1 - t2))
}


#' @include SimLongitudinal.R
#' @include generics.R
NULL

#' Simulate Longitudinal Data from a Stein-Fojo Model
#'
#' @param times (`numeric`)\cr the times to generate observations at.
#' @param sigma (`number`)\cr the variance of the longitudinal values.
#' @param mu_s (`numeric`)\cr the mean shrinkage rates for the two treatment arms.
#' @param mu_g (`numeric`)\cr the mean growth rates for the two treatment arms.
#' @param mu_b (`numeric`)\cr the mean baseline values for the two treatment arms.
#' @param omega_b (`number`)\cr the baseline value standard deviation.
#' @param omega_s (`number`)\cr the shrinkage rate standard deviation.
#' @param omega_g (`number`)\cr the growth rate standard deviation.
#' @param link_dsld (`number`)\cr the link coefficient for the derivative contribution.
#' @param link_ttg (`number`)\cr the link coefficient for the time-to-growth contribution.
#' @param link_identity (`number`)\cr the link coefficient for the SLD Identity contribution.
#' @param link_growth (`number`)\cr the link coefficient for the growth parameter contribution.
#'
#' @slot sigma (`numeric`)\cr See arguments.
#' @slot mu_s (`numeric`)\cr See arguments.
#' @slot mu_g (`numeric`)\cr See arguments.
#' @slot mu_b (`numeric`)\cr See arguments.
#' @slot omega_b (`numeric`)\cr See arguments.
#' @slot omega_s (`numeric`)\cr See arguments.
#' @slot omega_g (`numeric`)\cr See arguments.
#' @slot link_dsld (`numeric`)\cr See arguments.
#' @slot link_ttg (`numeric`)\cr See arguments.
#' @slot link_identity (`numeric`)\cr See arguments.
#' @slot link_growth (`numeric`)\cr See arguments.
#'
#' @family SimLongitudinal
#' @name SimLongitudinalSteinFojo-class
#' @exportClass SimLongitudinalSteinFojo
.SimLongitudinalSteinFojo <- setClass(
    "SimLongitudinalSteinFojo",
    contains = "SimLongitudinal",
    slots = c(
        sigma = "numeric",
        mu_s = "numeric",
        mu_g = "numeric",
        mu_b = "numeric",
        omega_b = "numeric",
        omega_s = "numeric",
        omega_g = "numeric",
        link_dsld = "numeric",
        link_ttg = "numeric",
        link_identity = "numeric",
        link_growth = "numeric"
    )
)

#' @rdname SimLongitudinalSteinFojo-class
#' @export
SimLongitudinalSteinFojo <- function(
    times = c(-100, -50, 0, 50, 100, 150, 250, 350, 450, 550) / 365,
    sigma = 0.01,
    mu_s = log(c(0.6, 0.4)),
    mu_g = log(c(0.25, 0.35)),
    mu_b = log(60),
    omega_b = 0.2,
    omega_s = 0.2,
    omega_g = 0.2,
    link_dsld = 0,
    link_ttg = 0,
    link_identity = 0,
    link_growth = 0
) {
    .SimLongitudinalSteinFojo(
        times = times,
        sigma = sigma,
        mu_s = mu_s,
        mu_g = mu_g,
        mu_b = mu_b,
        omega_b = omega_b,
        omega_s = omega_s,
        omega_g = omega_g,
        link_dsld = link_dsld,
        link_ttg = link_ttg,
        link_identity = link_identity,
        link_growth = link_growth
    )
}


setValidity(
    "SimLongitudinalSteinFojo",
    function(object) {
        par_lengths <- c(
            length(object@mu_s),
            length(object@mu_g)
        )
        if (length(unique(par_lengths)) != 1) {
            return("The parameters `mu_s` and `mu_g` must have the same length.")
        }
        len_1_pars <- c(
            "sigma", "omega_b", "omega_s", "omega_g",
            "link_dsld", "link_ttg", "link_identity",
            "link_growth"
        )
        for (par in len_1_pars) {
            if (length(slot(object, par)) != 1) {
                return(sprintf("The `%s` parameter must be a length 1 numeric.", par))
            }
        }
        return(TRUE)
    }
)

#' @rdname as_print_string
as_print_string.SimLongitudinalSteinFojo <- function(object) {
    return("SimLongitudinalSteinFojo")
}

#' @rdname sampleObservations
#' @export
sampleObservations.SimLongitudinalSteinFojo <- function(object, times_df) {
    times_df |>
        dplyr::mutate(mu_sld = sf_sld(.data$time, .data$psi_b, .data$psi_s, .data$psi_g)) |>
        dplyr::mutate(dsld = sf_dsld(.data$time, .data$psi_b, .data$psi_s, .data$psi_g)) |>
        dplyr::mutate(ttg = sf_ttg(.data$time, .data$psi_b, .data$psi_s, .data$psi_g)) |>
        dplyr::mutate(sld = stats::rnorm(dplyr::n(), .data$mu_sld, .data$mu_sld * object@sigma)) |>
        dplyr::mutate(
            log_haz_link =
                (object@link_dsld * .data$dsld) +
                (object@link_ttg * .data$ttg) +
                (object@link_identity * .data$mu_sld) +
                (object@link_growth * log(.data$psi_g))
        )
}


#' @rdname sampleSubjects
#' @export
sampleSubjects.SimLongitudinalSteinFojo <- function(object, subjects_df) {
    assert_that(
        is.factor(subjects_df$study),
        is.factor(subjects_df$arm),
        length(levels(subjects_df$study)) == length(object@mu_b),
        length(levels(subjects_df$arm)) == length(object@mu_s)
    )

    res <- subjects_df |>
        dplyr::distinct(.data$subject, .data$arm, .data$study) |>
        dplyr::mutate(study_idx = as.numeric(.data$study)) |>
        dplyr::mutate(arm_idx = as.numeric(.data$arm)) |>
        dplyr::mutate(psi_b = stats::rlnorm(dplyr::n(), object@mu_b[.data$study_idx], object@omega_b)) |>
        dplyr::mutate(psi_s = stats::rlnorm(dplyr::n(), object@mu_s[.data$arm_idx], object@omega_s)) |>
        dplyr::mutate(psi_g = stats::rlnorm(dplyr::n(), object@mu_g[.data$arm_idx], object@omega_g))

    res[, c("subject", "arm", "study", "psi_b", "psi_s", "psi_g")]
}


#' Stein-Fojo Functionals
#'
#' @param time (`numeric`)\cr time grid.
#' @param b (`number`)\cr baseline.
#' @param s (`number`)\cr shrinkage.
#' @param g (`number`)\cr growth.
#'
#' @returns The function results.
#' @keywords internal
sf_sld <- function(time, b, s, g) {
    s <- dplyr::if_else(time >= 0, s, 0)
    b * (exp(-s * time) + exp(g * time) - 1)
}


#' @rdname sf_sld
sf_ttg <- function(time, b, s, g) {
    t1 <- (log(s) - log(g)) / (g + s)
    t1[t1 <= 0] <- 0
    return(t1)
}


#' @rdname sf_sld
sf_dsld <- function(time, b, s, g) {
    s <- dplyr::if_else(time >= 0, s, 0)
    t1 <- g * exp(g * time)
    t2 <- s * exp(-s * time)
    return(b * (t1 - t2))
}


#' @include SimLongitudinal.R
#' @include generics.R
NULL

#' Simulate Longitudinal Data from a Claret-Bruno Model
#'
#' @param times (`numeric`)\cr the times to generate observations at.
#' @param sigma (`number`)\cr the variance of the longitudinal values.
#'
#' @param mu_b (`numeric`)\cr the mean population baseline sld value.
#' @param mu_g (`numeric`)\cr the mean population growth rate.
#' @param mu_c (`numeric`)\cr the mean population resistance rate.
#' @param mu_p (`numeric`)\cr the mean population growth inhibition.
#'
#' @param omega_b (`number`)\cr the population standard deviation for the baseline sld value.
#' @param omega_g (`number`)\cr the population standard deviation for the growth rate.
#' @param omega_c (`number`)\cr the population standard deviation for the resistance rate.
#' @param omega_p (`number`)\cr the population standard deviation for the growth inhibition.
#'
#' @param link_dsld (`number`)\cr the link coefficient for the derivative contribution.
#' @param link_ttg (`number`)\cr the link coefficient for the time-to-growth contribution.
#' @param link_identity (`number`)\cr the link coefficient for the SLD Identity contribution.
#' @param link_growth (`number`)\cr the link coefficient for the growth parameter contribution.
#'
#' @slot sigma (`numeric`)\cr See arguments.
#'
#' @slot mu_b (`numeric`)\cr See arguments.
#' @slot mu_g (`numeric`)\cr See arguments.
#' @slot mu_c (`numeric`)\cr See arguments.
#' @slot mu_p (`numeric`)\cr See arguments.
#'
#' @slot omega_b (`numeric`)\cr See arguments.
#' @slot omega_g (`numeric`)\cr See arguments.
#' @slot omega_c (`numeric`)\cr See arguments.
#' @slot omega_p (`numeric`)\cr See arguments.
#'
#' @slot link_dsld (`numeric`)\cr See arguments.
#' @slot link_ttg (`numeric`)\cr See arguments.
#' @slot link_identity (`numeric`)\cr See arguments.
#' @slot link_growth (`numeric`)\cr See arguments.
#'
#' @family SimLongitudinal
#' @name SimLongitudinalClaretBruno-class
#' @exportClass SimLongitudinalClaretBruno
.SimLongitudinalClaretBruno <- setClass(
    "SimLongitudinalClaretBruno",
    contains = "SimLongitudinal",
    slots = c(
        sigma = "numeric",
        mu_b = "numeric",
        mu_g = "numeric",
        mu_c = "numeric",
        mu_p = "numeric",
        omega_b = "numeric",
        omega_g = "numeric",
        omega_c = "numeric",
        omega_p = "numeric",
        link_dsld = "numeric",
        link_ttg = "numeric",
        link_identity = "numeric",
        link_growth = "numeric"
    )
)

#' @rdname SimLongitudinalClaretBruno-class
#' @export
SimLongitudinalClaretBruno <- function(
    times = c(-100, -50, 0, 50, 100, 150, 250, 350, 450, 550) / 365,
    sigma = 0.01,
    mu_b = log(60),
    mu_g = log(c(0.9, 1.1)),
    mu_c = log(c(0.25, 0.35)),
    mu_p = log(c(1.5, 2)),
    omega_b = 0.2,
    omega_g = 0.2,
    omega_c = 0.2,
    omega_p = 0.2,
    link_dsld = 0,
    link_ttg = 0,
    link_identity = 0,
    link_growth = 0
) {
    .SimLongitudinalClaretBruno(
        times = times,
        sigma = sigma,
        mu_b = mu_b,
        mu_g = mu_g,
        mu_c = mu_c,
        mu_p = mu_p,
        omega_b = omega_b,
        omega_g = omega_g,
        omega_c = omega_c,
        omega_p = omega_p,
        link_dsld = link_dsld,
        link_ttg = link_ttg,
        link_identity = link_identity,
        link_growth = link_growth
    )
}


setValidity(
    "SimLongitudinalClaretBruno",
    function(object) {
        par_lengths <- c(
            length(object@mu_g),
            length(object@mu_c),
            length(object@mu_p)
        )
        if (length(unique(par_lengths)) != 1) {
            return("The parameters `mu_g`, `mu_c` & `mu_p` must have the same length.")
        }
        len_1_pars <- c(
            "sigma", "omega_b", "omega_g", "omega_c", "omega_p",
            "link_dsld", "link_ttg", "link_identity",
            "link_growth"
        )
        for (par in len_1_pars) {
            if (length(slot(object, par)) != 1) {
                return(sprintf("The `%s` parameter must be a length 1 numeric.", par))
            }
        }
        return(TRUE)
    }
)

#' @rdname as_print_string
as_print_string.SimLongitudinalClaretBruno <- function(object) {
    return("SimLongitudinalClaretBruno")
}

#' @rdname sampleObservations
#' @export
sampleObservations.SimLongitudinalClaretBruno <- function(object, times_df) {
    times_df |>
        dplyr::mutate(mu_sld = clbr_sld(.data$time, .data$ind_b, .data$ind_g, .data$ind_c, .data$ind_p)) |>
        dplyr::mutate(dsld = clbr_dsld(.data$time, .data$ind_b, .data$ind_g, .data$ind_c, .data$ind_p)) |>
        dplyr::mutate(ttg = clbr_ttg(.data$time, .data$ind_b, .data$ind_g, .data$ind_c, .data$ind_p)) |>
        dplyr::mutate(sld = stats::rnorm(dplyr::n(), .data$mu_sld, .data$mu_sld * object@sigma)) |>
        dplyr::mutate(
            log_haz_link =
                (object@link_dsld * .data$dsld) +
                (object@link_ttg * .data$ttg) +
                (object@link_identity * .data$mu_sld) +
                (object@link_growth * log(.data$ind_g))
        )
}


#' @rdname sampleSubjects
#' @export
sampleSubjects.SimLongitudinalClaretBruno <- function(object, subjects_df) {
    assert_that(
        is.factor(subjects_df$study),
        is.factor(subjects_df$arm),
        length(levels(subjects_df$study)) == length(object@mu_b),
        length(levels(subjects_df$arm)) == length(object@mu_g),
        length(levels(subjects_df$arm)) == length(object@mu_c),
        length(levels(subjects_df$arm)) == length(object@mu_p)
    )

    res <- subjects_df |>
        dplyr::distinct(.data$subject, .data$arm, .data$study) |>
        dplyr::mutate(study_idx = as.numeric(.data$study)) |>
        dplyr::mutate(arm_idx = as.numeric(.data$arm)) |>
        dplyr::mutate(ind_b = stats::rlnorm(dplyr::n(), object@mu_b[.data$study_idx], object@omega_b)) |>
        dplyr::mutate(ind_g = stats::rlnorm(dplyr::n(), object@mu_g[.data$arm_idx], object@omega_g)) |>
        dplyr::mutate(ind_c = stats::rlnorm(dplyr::n(), object@mu_c[.data$arm_idx], object@omega_c)) |>
        dplyr::mutate(ind_p = stats::rlnorm(dplyr::n(), object@mu_p[.data$arm_idx], object@omega_p))

    res[, c("subject", "arm", "study", "ind_b", "ind_g", "ind_c", "ind_p")]
}


#' Claret-Bruno Functionals
#'
#' @param t (`numeric`)\cr time grid.
#' @param b (`number`)\cr baseline sld.
#' @param g (`number`)\cr growth rate.
#' @param c (`number`)\cr resistance rate.
#' @param p (`number`)\cr growth inhibition.
#'
#' @returns The function results.
#' @keywords internal
clbr_sld <- function(t, b, g, c, p) {
    p <- ifelse(t >= 0, p, 0)
    b * exp((g * t) - (p / c) * (1 - exp(-c * t)))
}

#' @rdname clbr_sld
clbr_ttg <- function(t, b, g, c, p) {
    log(p / g) / c
}

#' @rdname clbr_sld
clbr_dsld <- function(t, b, g, c, p) {
    lt0 <- b * g * exp(g * t)
    gt0 <- (g - p * exp(-c * t)) * clbr_sld(t, b, g, c, p)
    ifelse(t >= 0, gt0, lt0)
}

#' @include LongitudinalModel.R
#' @include StanModule.R
#' @include generics.R
#' @include ParameterList.R
#' @include Parameter.R
#' @include Link.R
NULL


#' `LongitudinalClaretBruno`
#'
#' This class extends the general [`LongitudinalModel`] class for using the
#' Claret-Bruno model for the longitudinal outcome.
#'
#' @section Available Links:
#' - [`linkDSLD()`]
#' - [`linkTTG()`]
#' - [`linkIdentity()`]
#' - [`linkGrowth()`]
#' @exportClass LongitudinalClaretBruno
.LongitudinalClaretBruno <- setClass(
    Class = "LongitudinalClaretBruno",
    contains = "LongitudinalModel"
)


#' @rdname LongitudinalClaretBruno-class
#'
#' @param mu_b (`Prior`)\cr for the mean population baseline sld value.
#' @param mu_g (`Prior`)\cr for the mean population growth rate.
#' @param mu_c (`Prior`)\cr for the mean population resistance rate.
#' @param mu_p (`Prior`)\cr for the mean population growth inhibition
#'
#' @param omega_b (`Prior`)\cr for the population standard deviation for the baseline sld value.
#' @param omega_g (`Prior`)\cr for the population standard deviation for the growth rate.
#' @param omega_c (`Prior`)\cr for the population standard deviation for the resistance rate.
#' @param omega_p (`Prior`)\cr for the population standard deviation for the growth inhibition.
#'
#' @param sigma (`Prior`)\cr for the variance of the longitudinal values.
#'
#' @param centred (`logical`)\cr whether to use the centred parameterization.
#'
#' @export
LongitudinalClaretBruno <- function(

    mu_b = prior_normal(log(60), 0.5),
    mu_g = prior_normal(log(1), 0.5),
    mu_c = prior_normal(log(0.4), 0.5),
    mu_p = prior_normal(log(2), 0.5),

    omega_b = prior_lognormal(log(0.2), 0.5),
    omega_g = prior_lognormal(log(0.2), 0.5),
    omega_c = prior_lognormal(log(0.2), 0.5),
    omega_p = prior_lognormal(log(0.2), 0.5),

    sigma = prior_lognormal(log(0.1), 0.5),

    centred = FALSE
) {

    sf_model <- StanModule(decorated_render(
        .x = read_stan("lm-claret-bruno/model.stan"),
        centred = centred
    ))

    # Apply constraints
    omega_b <- set_limits(omega_b, lower = 0)
    omega_g <- set_limits(omega_g, lower = 0)
    omega_c <- set_limits(omega_c, lower = 0)
    omega_p <- set_limits(omega_p, lower = 0)
    sigma <- set_limits(sigma, lower = 0)


    parameters <- list(
        Parameter(name = "lm_clbr_mu_b", prior = mu_b, size = "n_studies"),
        Parameter(name = "lm_clbr_mu_g", prior = mu_g, size = "n_arms"),
        Parameter(name = "lm_clbr_mu_c", prior = mu_c, size = "n_arms"),
        Parameter(name = "lm_clbr_mu_p", prior = mu_p, size = "n_arms"),

        Parameter(name = "lm_clbr_omega_b", prior = omega_b, size = 1),
        Parameter(name = "lm_clbr_omega_g", prior = omega_g, size = 1),
        Parameter(name = "lm_clbr_omega_c", prior = omega_c, size = 1),
        Parameter(name = "lm_clbr_omega_p", prior = omega_p, size = 1),

        Parameter(name = "lm_clbr_sigma", prior = sigma, size = 1)
    )

    assert_flag(centred)
    parameters_extra <- if (centred) {
        list(
            Parameter(
                name = "lm_clbr_ind_b",
                prior = prior_init_only(prior_lognormal(median(mu_b), median(omega_b))),
                size = "n_subjects"
            ),
            Parameter(
                name = "lm_clbr_ind_g",
                prior = prior_init_only(prior_lognormal(median(mu_g), median(omega_g))),
                size = "n_subjects"
            ),
            Parameter(
                name = "lm_clbr_ind_c",
                prior = prior_init_only(prior_lognormal(median(mu_c), median(omega_c))),
                size = "n_subjects"
            ),
            Parameter(
                name = "lm_clbr_ind_p",
                prior = prior_init_only(prior_lognormal(median(mu_p), median(omega_p))),
                size = "n_subjects"
            )
        )
    } else {
        list(
            Parameter(name = "lm_clbr_eta_b", prior = prior_std_normal(), size = "n_subjects"),
            Parameter(name = "lm_clbr_eta_g", prior = prior_std_normal(), size = "n_subjects"),
            Parameter(name = "lm_clbr_eta_c", prior = prior_std_normal(), size = "n_subjects"),
            Parameter(name = "lm_clbr_eta_p", prior = prior_std_normal(), size = "n_subjects")
        )
    }
    parameters <- append(parameters, parameters_extra)

    x <- LongitudinalModel(
        name = "Claret-Bruno",
        stan = merge(
            sf_model,
            StanModule("lm-claret-bruno/functions.stan")
        ),
        parameters = do.call(ParameterList, parameters)
    )
    .LongitudinalClaretBruno(x)
}



#' @export
enableLink.LongitudinalClaretBruno <- function(object, ...) {
    object@stan <- merge(
        object@stan,
        StanModule("lm-claret-bruno/link.stan")
    )
    object
}

#' @export
linkDSLD.LongitudinalClaretBruno <- function(prior = prior_normal(0, 2), model, ...) {
    LinkComponent(
        key = "link_dsld",
        stan = StanModule("lm-claret-bruno/link_dsld.stan"),
        prior = prior
    )
}

#' @export
linkTTG.LongitudinalClaretBruno <- function(prior = prior_normal(0, 2), model, ...) {
    LinkComponent(
        key = "link_ttg",
        stan = StanModule("lm-claret-bruno/link_ttg.stan"),
        prior = prior
    )
}

#' @export
linkIdentity.LongitudinalClaretBruno <- function(prior = prior_normal(0, 2), model, ...) {
    LinkComponent(
        key = "link_identity",
        stan = StanModule("lm-claret-bruno/link_identity.stan"),
        prior = prior
    )
}

#' @export
linkGrowth.LongitudinalClaretBruno <- function(prior = prior_normal(0, 2), model, ...) {
    LinkComponent(
        key = "link_growth",
        stan = StanModule("lm-claret-bruno/link_growth.stan"),
        prior = prior
    )
}

#' @rdname getPredictionNames
#' @export
getPredictionNames.LongitudinalClaretBruno <- function(object, ...) {
    c("b", "g", "c", "p")
}

#' @export
enableGQ.LongitudinalClaretBruno <- function(object, ...) {
    StanModule("lm-claret-bruno/quantities.stan")
}


#' @include SimLongitudinal.R
#' @include generics.R
NULL

#' Simulate Longitudinal Data from a Random Slope Model
#'
#' @param times (`numeric`)\cr the times to generate observations at.
#' @param intercept (`number`)\cr the mean baseline value for each study.
#' @param slope_mu (`numeric`)\cr the population slope for each treatment arm.
#' @param slope_sigma (`number`)\cr the random slope standard deviation.
#' @param sigma (`number`)\cr the variance of the longitudinal values.
#' @param link_dsld (`number`)\cr the link coefficient for the DSLD contribution.
#' @param link_identity (`number`)\cr the link coefficient for the identity contribution.
#'
#' @slot intercept (`numeric`)\cr See arguments.
#' @slot slope_mu (`numeric`)\cr See arguments.
#' @slot slope_sigma (`numeric`)\cr See arguments.
#' @slot sigma (`numeric`)\cr See arguments.
#' @slot link_dsld (`numeric`)\cr See arguments.
#' @slot link_identity (`numeric`)\cr See arguments.
#'
#' @family SimLongitudinal
#' @name SimLongitudinalRandomSlope-class
#' @exportClass SimLongitudinalRandomSlope
.SimLongitudinalRandomSlope <- setClass(
    "SimLongitudinalRandomSlope",
    contains = "SimLongitudinal",
    slots = c(
        intercept = "numeric",
        slope_mu = "numeric",
        slope_sigma = "numeric",
        sigma = "numeric",
        link_dsld = "numeric",
        link_identity = "numeric"
    )
)

#' @rdname SimLongitudinalRandomSlope-class
#' @export
SimLongitudinalRandomSlope <- function(
    times = c(-100, -50, 0, 50, 100, 150, 250, 350, 450, 550),
    intercept = 50,
    slope_mu = c(0.01, 0.03),
    slope_sigma = 0.5,
    sigma = 2,
    link_dsld = 0,
    link_identity = 0
) {
    .SimLongitudinalRandomSlope(
        times = times,
        intercept = intercept,
        slope_mu = slope_mu,
        slope_sigma = slope_sigma,
        sigma = sigma,
        link_dsld = link_dsld,
        link_identity = link_identity
    )
}

#' @rdname as_print_string
as_print_string.SimLongitudinalRandomSlope <- function(object) {
    return("SimLongitudinalRandomSlope")
}

#' @rdname sampleObservations
#' @export
sampleObservations.SimLongitudinalRandomSlope <- function(object, times_df) {
    times_df |>
        dplyr::mutate(err = stats::rnorm(dplyr::n(), 0, object@sigma)) |>
        dplyr::mutate(sld_mu = .data$intercept + .data$slope_ind * .data$time) |>
        dplyr::mutate(sld = .data$sld_mu + .data$err) |>
        dplyr::mutate(
            log_haz_link =
                object@link_dsld * .data$slope_ind +
                object@link_identity * .data$sld_mu
        )
}

#' @rdname sampleSubjects
#' @export
sampleSubjects.SimLongitudinalRandomSlope <- function(object, subjects_df) {
    assert_that(
        is.factor(subjects_df[["study"]]),
        is.factor(subjects_df[["arm"]])
    )

    assert_that(
        length(object@slope_mu) == length(unique(subjects_df[["arm"]])),
        msg = "`length(slope_mu)` should be equal to the number of unique arms"
    )

    assert_that(
        length(object@intercept) == length(unique(subjects_df[["study"]])),
        msg = "`length(intercept)` should be equal to the number of unique studies"
    )

    assert_that(
        nrow(subjects_df) == length(unique(subjects_df[["subject"]])),
        msg = "The number of rows in `subjects_df` should be equal to the number of unique subjects"
    )

    subjects_df |>
        dplyr::mutate(intercept = object@intercept[as.numeric(.data$study)]) |>
        dplyr::mutate(slope_ind = stats::rnorm(
            n = dplyr::n(),
            mean = object@slope_mu[as.numeric(.data$arm)],
            sd = object@slope_sigma
        ))
}

#' @include generics.R
#' @include Prior.R
NULL


#' `Parameter` Function Arguments
#'
#' The documentation lists all the conventional arguments for [`Parameter`]
#' constructors.
#'
#' @param x ([`Parameter`])\cr a prior Distribution
#' @param object ([`Parameter`])\cr a prior Distribution
#' @param ... Not Used.
#'
#' @name Parameter-Shared
#' @keywords internal
NULL


setClassUnion(name = "numeric_OR_character", c("numeric", "character"))

# Parameter-class ----

#' `Parameter`
#'
#' Stores the name, the prior distribution and the size of a parameter.
#' If `size` is a string then this indicates the name of the variable
#' within the stan data object that specifies the size of this parameter.
#'
#' @slot name (`string`)\cr of the parameter.
#' @slot prior (`Prior`)\cr for the parameter.
#' @slot size (`numeric` or `string`)\cr dimension of the parameter.
#'
#' @family Parameter
#' @exportClass Parameter
#' @export Parameter
.Parameter <- setClass(
    Class = "Parameter",
    slots = list(
        "name" = "character",
        "prior" = "Prior",
        "size" = "numeric_OR_character"
    )
)
#' @param prior (`Prior`)\cr for the parameter.
#' @param name (`string`)\cr of the parameter.
#' @param size (`numeric` or `string`)\cr dimension of the parameter.
#' @rdname Parameter-class
Parameter <- function(prior, name, size = 1) {
    .Parameter(
        prior = prior,
        name = name,
        size = size
    )
}
setValidity(
    Class = "Parameter",
    method = function(object) {
        if (!length(object@name) == 1) {
            return("Name must be a length 1 character vector")
        }
        if (is.character(object@size)) {
            if (!length(object@size) == 1) {
                return("Size must be a numeric vector or length 1 character vector")
            }
        }
        return(TRUE)
    }
)


#' `Parameter` -> `StanModule`
#'
#' Converts a [`Parameter`] object to a [`StanModule`] object
#'
#' @inheritParams Parameter-Shared
#'
#' @family Parameter
#' @family as.StanModule
#' @export
as.StanModule.Parameter <- function(object, ...) {
    as.StanModule(object@prior, name = object@name)
}


#' `Parameter` -> `list`
#'
#' Converts a Parameter object to a list of parameter data values
#' for a Stan model.
#'
#' @inheritParams Parameter-Shared
#'
#' @family as_stan_list
#' @family Parameter
#' @export
as_stan_list.Parameter <- function(object, ...) {
    as_stan_list(object@prior, name = object@name)
}


#' Parameter Getter Functions
#'
#' @param x (`Paramater`) \cr A model parameter
#' @param object (`Paramater`) \cr A model parameter
#' @param ... Not used.
#'
#' @description
#' Getter functions for the slots of a [`Parameter`] object
#' @family Parameter
#' @name Parameter-Getter-Methods
NULL

#' @describeIn Parameter-Getter-Methods The parameter's name
#' @export
names.Parameter <- function(x) x@name

#' @describeIn Parameter-Getter-Methods The parameter's initial values
#' @export
initialValues.Parameter <- function(object, ...) initialValues(object@prior)

#' @describeIn Parameter-Getter-Methods The parameter's dimensionality
#' @export
size.Parameter <- function(object) object@size


#' `Parameter` -> `Character`
#'
#' Converts a [`Parameter`] object to a character vector
#' @inheritParams Parameter-Shared
#' @family Parameter
#' @export
as.character.Parameter <- function(x, ...) {
    paste0(x@name, " ~ ", as.character(x@prior))
}


#' @rdname show-object
#' @export
setMethod(
    f = "show",
    signature = "Parameter",
    definition = function(object) {
        x <- sprintf("\nParameter Object:\n   %s\n\n", as.character(object))
        cat(x)
        return(object)
    }
)


#' @include Grid.R
#' @include generics.R
NULL


#' @rdname Grid-Dev
.GridEven <- setClass(
    "GridEven",
    contains = "Grid",
    slots = c(
        "subjects" = "character",
        "length.out" = "numeric"
    )
)


#' @rdname Grid-Functions
#' @export
GridEven <- function(subjects = NULL, length.out = 30) {
    .GridEven(
        subjects = subjects,
        length.out = length.out
    )
}


setValidity(
    "GridEven",
    function(object) {
        if (length(object@length.out) != 1 || all(object@length.out <= 0)) {
            return("The `length.out` argument must be a positive integer")
        }
    }
)


#' @rdname Quant-Dev
#' @export
as.QuantityGenerator.GridEven <- function(object, data, ...) {
    assert_class(data, "DataJoint")
    data_list <- as.list(data)
    subjects <- unlist(as.list(object, data = data), use.names = FALSE)
    assert_that(
        all(subjects %in% names(data_list$subject_to_index)),
        msg = "All subject names must be in the `DataSubject` object"
    )

    spec <- lapply(
        subjects,
        \(sub) {
            subject_index <- data_list$subject_to_index[[sub]]
            time_index <- data_list$subject_tumour_index == subject_index
            subject_times <- data_list$tumour_time[time_index]
            seq(min(subject_times), max(subject_times), length.out = object@length.out)
        }
    )
    names(spec) <- subjects

    as.QuantityGenerator(
        GridManual(spec),
        data = data
    )
}


#' @rdname Quant-Dev
#' @export
as.QuantityCollapser.GridEven <- function(object, data, ...) {
    generator <- as.QuantityGenerator(object, data)
    QuantityCollapser(
        times = generator@times,
        groups = generator@subjects,
        indexes = as.list(seq_along(generator@times))
    )
}


#' @export
as.list.GridEven <- function(x, data, ...) {
    subjects_to_list(x@subjects, data)
}

#' @include generics.R
#' @include Grid.R
NULL


#' @rdname Quant-Dev
.QuantityGeneratorSubject <- setClass(
    "QuantityGeneratorSubject",
    contains = "QuantityGenerator",
    slots = c(
        "times" = "numeric",
        "subjects" = "character_or_NULL"
    )
)


#' @rdname Quant-Dev
QuantityGeneratorSubject <- function(times, subjects = NULL) {
    .QuantityGeneratorSubject(
        times = times,
        subjects = subjects
    )
}


setValidity(
    "QuantityGeneratorSubject",
    function(object) {
        if (length(object@times) != length(object@subjects)) {
            return("Length of `times` and `subjects` must be equal")
        }
        return(TRUE)
    }
)


#' @rdname as_stan_list.QuantityGenerator
#' @export
as_stan_list.QuantityGeneratorSubject <- function(object, data, ...) {
    assert_that(
        is(data, "DataJoint")
    )
    ret <- list()
    data_list <- as.list(data)
    ret[["gq_subject_index"]] <- data_list$subject_to_index[as.character(object@subjects)]
    ret[["gq_n_quant"]] <- length(object@subjects)
    ret[["gq_times"]] <- object@times

    # dummy pop indexes in order for stan code to actualy compile. In this setting
    # this matrix isn't actually used so doesn't matter what these values are
    # but don't want to have to burden individual longitudinal models with the
    # conditional logic to check if they are generating population quantities or not
    ret[["gq_long_pop_arm_index"]] <- rep(1, ret[["gq_n_quant"]])
    ret[["gq_long_pop_study_index"]] <- rep(1, ret[["gq_n_quant"]])

    # Sanity checks
    assert_that(
        length(ret[["gq_times"]]) == ret[["gq_n_quant"]],
        all(object@subjects %in% names(data_list$subject_to_index))
    )
    return(ret)
}

#' @include Grid.R
#' @include generics.R
NULL


#' @rdname Grid-Dev
.GridFixed <- setClass(
    "GridFixed",
    contains = "Grid",
    slots = c(
        "subjects" = "character_or_NULL",
        "times" = "numeric_or_NULL"
    )
)

#' @rdname Grid-Functions
#' @export
GridFixed <- function(subjects = NULL, times = NULL) {
    .GridFixed(
        subjects = subjects,
        times = times
    )
}

#' @rdname Quant-Dev
#' @export
as.QuantityGenerator.GridFixed <- function(object, data, ...) {

    assert_class(data, "DataJoint")
    data_list <- as.list(data)
    subjects <- unlist(as.list(object, data = data), use.names = FALSE)

    validate_time_grid(object@times)
    subject_times <- expand.grid(
        subject = subjects,
        time = object@times,
        stringsAsFactors = FALSE
    )

    QuantityGeneratorSubject(
        times = subject_times$time,
        subjects = subject_times$subject
    )
}

#' @rdname Quant-Dev
#' @export
as.QuantityCollapser.GridFixed <- function(object, data, ...) {
    generator <- as.QuantityGenerator(object, data)
    QuantityCollapser(
        times = generator@times,
        groups = generator@subjects,
        indexes = as.list(seq_along(generator@times))
    )
}


#' @export
as.list.GridFixed <- function(x, data, ...) {
    subjects_to_list(x@subjects, data)
}

#' @rdname coalesceGridTime
#' @export
coalesceGridTime.GridFixed <- function(object, times, ...) {
    if (is.null(object@times)) {
        object <- GridFixed(
            subjects = object@subjects,
            times = times
        )
    }
    object
}

#' @include generics.R
NULL



#' Standard Links
#'
#' @param prior ([`Prior`]) \cr A [`Prior`] object.
#' @param model ([`LongitudinalModel`]) \cr A [`LongitudinalModel`] object.
#' @param ... Not used.
#'
#' @description
#' These functions are used to enable the use of the corresponding link function between
#' the survival and longitudinal models in a joint model. Note that the exact implementation
#' of the link function is model specific, see
#' \code{vignette("Statistical Specifications", package = "jmpost")} for more details.
#'
#' @name standard-link-user
NULL




#' @describeIn standard-link-user No link (fit the survival and longitudinal models independently)
#' @export
linkNone <- function() {
    Link()
}


#' @describeIn standard-link-user Time to growth link
#' @export
linkTTG <- function(prior, model = PromiseLongitudinalModel(), ...) {
    UseMethod("linkTTG", model)
}
#' @export
linkTTG.PromiseLongitudinalModel <- function(prior = prior_normal(0, 2), model, ...) {
    PromiseLinkComponent(fun = linkTTG, prior = prior, key = "link_ttg")
}
#' @export
linkTTG.default <- function(prior, model, ...) {
    stop(sprintf("Method `linkTTG` is not available for `%s`", class(model)[[1]]))
}




#' @describeIn standard-link-user Derivative of the SLD over time link
#' @export
linkDSLD <- function(prior, model = PromiseLongitudinalModel(), ...) {
    UseMethod("linkDSLD", model)
}
#' @export
linkDSLD.PromiseLongitudinalModel <- function(prior = prior_normal(0, 2), model, ...) {
    PromiseLinkComponent(fun = linkDSLD, prior = prior, key = "link_dsld")
}
#' @export
linkDSLD.default <- function(prior, model, ...) {
    stop(sprintf("Method `linkDSLD` is not available for `%s`", class(model)[[1]]))
}




#' @describeIn standard-link-user Current SLD value link
#' @export
linkIdentity <- function(prior, model = PromiseLongitudinalModel(), ...) {
    UseMethod("linkIdentity", model)
}
#' @export
linkIdentity.PromiseLongitudinalModel <- function(prior = prior_normal(0, 2), model, ...) {
    PromiseLinkComponent(fun = linkIdentity, prior = prior, key = "link_identity")
}
#' @export
linkIdentity.default <- function(prior, model, ...) {
    stop(sprintf("Method `linkIdentity` is not available for `%s`", class(model)[[1]]))
}


#' @describeIn standard-link-user Growth Parameter link
#' @export
linkGrowth <- function(prior, model = PromiseLongitudinalModel(), ...) {
    UseMethod("linkGrowth", model)
}
#' @export
linkGrowth.PromiseLongitudinalModel <- function(prior = prior_normal(0, 2), model, ...) {
    PromiseLinkComponent(fun = linkGrowth, prior = prior, key = "link_growth")
}
#' @export
linkGrowth.default <- function(prior, model, ...) {
    stop(sprintf("Method `linkGrowth` is not available for `%s`", class(model)[[1]]))
}


.onLoad <- function(libname, pkgname) {
    set_options()
}

.onAttach <- function(libname, pkgname) {
    if (!is_cmdstanr_available()) {
        packageStartupMessage(
            "jmpost uses cmdstanr for compiling and sampling from models, but it does not seem to be installed.\n",
            "To install:\n",
            "install.packages(\"cmdstanr\", repos = c(\"https://stan-dev.r-universe.dev/\", getOption(\"repos\")))"
        )
    } else if (is.null(cmdstanr::cmdstan_version(error_on_NA = FALSE))) {
        possible_paths <- unique(c(
            cmdstanr::cmdstan_default_install_path(),
            Sys.getenv("CMDSTAN"),
            Sys.getenv("CMDSTAN_PATH"),
            "/root/.cmdstan",
            "~/.cmdstan"
        ))
        possible_paths <- possible_paths[dir.exists(possible_paths)]

        if (length(possible_paths)) {
            for (try_path in possible_paths) {
                new_path <- tryCatch(
                    suppressMessages(cmdstanr::set_cmdstan_path(try_path)),
                    warning = function(w) NULL,
                    error = function(e) NULL
                )
            }
            if (!is.null(new_path)) {
                packageStartupMessage("CmdStan path set to: ", new_path)
            }
        } else {
            packageStartupMessage("jmpost could not identify CmdStan path. Please use cmdstanr::set_cmdstan_path()")
        }
    }
    return(invisible(NULL))
}


# This only exists to silence the false positive R CMD CHECK warning about
# importing but not using the posterior package. posterior is a dependency
# of rcmdstan that we use a lot implicitly. Also we link to their documentation
# pages in ours
.never_run <- function() {
    posterior::as_draws()
}

#' @include StanModel.R
NULL

# SurvivalModel-class ----

#' `SurvivalModel`
#'
#' This class extends the general [`StanModel`] class to comprise the survival
#' model specification.
#'
#' @exportClass SurvivalModel
.SurvivalModel <- setClass(
    Class = "SurvivalModel",
    contains = "StanModel"
)

# SurvivalModel-constructors ----

#' @rdname SurvivalModel-class
#'
#' @inheritParams stanmodel_arguments
#'
#' @export
SurvivalModel <- function(
    stan = StanModule(),
    parameters = ParameterList(),
    name = "<Unnamed>",
    ...
) {
    base_stan <- read_stan("base/survival.stan")
    stan_full <- decorated_render(
        .x = base_stan,
        stan = add_missing_stan_blocks(as.list(stan))
    )
    .SurvivalModel(
        StanModel(
            name = name,
            stan = StanModule(stan_full),
            parameters = parameters,
            ...
        )
    )
}

#' @export
as_print_string.SurvivalModel <- function(object, ...) {
    string <- sprintf(
        "\n%s Survival Model with parameters:\n%s\n\n",
        object@name,
        paste("   ", as_print_string(object@parameters)) |> paste(collapse = "\n")
    )
    return(string)
}

#' @include StanModule.R
#' @include LongitudinalModel.R
#' @include ParameterList.R
#' @include generics.R
#' @include Prior.R
NULL




#' `LinkComponent` Function Arguments
#'
#' This exists to contain all the common arguments for [`LinkComponent`] methods.
#'
#' @param stan (`StanModule`)\cr Stan code.
#' @param x ([`LinkComponent`])\cr a link component.
#' @param object ([`LinkComponent`])\cr a link component.
#' @param ... Not Used.
#'
#' @name LinkComponent-Shared
#' @keywords internal
NULL




#' `LinkComponent`
#'
#' @slot stan (`StanModule`)\cr See Arguments.
#' @slot name (`character`)\cr See Arguments.
#' @slot parameters (`ParameterList`)\cr The parameter specification.
#'
#' @param stan (`StanModule`)\cr Stan code. See Details.
#' @param prior (`Prior`)\cr The prior for the scaling coeficient.
#' @param key (`character`)\cr Link identifier. See Details.
#'
#' @details
#'
#' This object provides key information needed to construct a link contribution in the
#' survival model based on the parameters of the longitudinal model.
#'
#' Each link component defines a stan function of the longitudinal model parameters which is
#' multiplied by a model coefficient and added to the survival models hazard function.
#'
#' For full details about the specification of a `LinkComponent` please see
#' \code{vignette("extending-jmpost", package = "jmpost")}.
#'
#' @inheritParams stanmodel_arguments
#' @family LinkComponent
#' @name LinkComponent-class
#' @exportClass Link
.LinkComponent <- setClass(
    Class = "LinkComponent",
    slots = list(
        "stan" = "StanModule",
        "parameters" = "ParameterList",
        "key" = "character"
    )
)


#' @rdname LinkComponent-class
#' @export
LinkComponent <- function(stan, prior, key, ...) {
    .LinkComponent(
        stan = stan,
        key = key,
        parameters = ParameterList(Parameter(name = key, prior = prior, size = 1)),
        ...
    )
}




#' @family LinkComponent
#' @rdname getParameters
#' @export
getParameters.LinkComponent <- function(object, ...) {
    object@parameters
}


#' @family LinkComponent
#' @rdname initialValues
#' @export
initialValues.LinkComponent <- function(object, n_chains, ...) {
    initialValues(object@parameters, n_chains)
}




#' `LinkComponent` -> `StanModule`
#'
#' Converts a [`LinkComponent`] object to a [`StanModule`] object
#'
#' @inheritParams LinkComponent-Shared
#'
#' @family LinkComponent
#' @family as.StanModule
#' @export
as.StanModule.LinkComponent <- function(object, ...) {
    object@stan
}



#' `LinkComponent` -> `list`
#'
#' @inheritParams LinkComponent-Shared
#'
#' @description
#' Returns a named list where each element of the list corresponds
#' to a Stan modelling block e.g. `data`, `model`, etc.
#'
#' @family LinkComponent
#' @export
as.list.LinkComponent <- function(x, ...) {
    stan <- as.StanModule(x, ...)
    as.list(stan)
}

#' @family LinkComponent
#' @export
as_print_string.LinkComponent <- function(object, ...) {
    as_print_string(object@parameters)
}


#' @rdname show-object
#' @export
setMethod(
    f = "show",
    signature = "LinkComponent",
    definition = function(object) {
        cat(
            paste0(
                "\nLinkComponent with parameter:\n    ",
                as_print_string(object), "\n\n"
            )
        )
    }
)

#' @family LinkComponent
#' @export
names.LinkComponent <- function(x, ...) {
    names(x@parameters)
}

#' @include generics.R
#' @include Grid.R
NULL

#' @rdname Quant-Dev
.QuantityGeneratorPopulation <- setClass(
    "QuantityGeneratorPopulation",
    contains = "QuantityGenerator",
    slots = c(
        "times" = "numeric",
        "studies" = "character_or_NULL",
        "arms" = "character_or_NULL"
    )
)


#' @rdname Quant-Dev
QuantityGeneratorPopulation <- function(times, studies = NULL, arms = NULL) {
    .QuantityGeneratorPopulation(
        times = times,
        studies = studies,
        arms = arms
    )
}


setValidity(
    "QuantityGeneratorPopulation",
    function(object) {
        if (length(object@times) != length(object@arms)) {
            return("Length of `times` and `arms` must be equal")
        }
        if (length(object@times) != length(object@studies)) {
            return("Length of `times` and `studies` must be equal")
        }
        return(TRUE)
    }
)


#' @rdname as_stan_list.QuantityGenerator
#' @export
as_stan_list.QuantityGeneratorPopulation <- function(object, data, ...) {
    assert_that(
        is(data, "DataJoint")
    )
    ret <- list()
    data_list <- as.list(data)
    ret[["gq_times"]] <- object@times
    ret[["gq_n_quant"]] <- length(object@arms)
    ret[["gq_long_pop_arm_index"]] <- data_list$arm_to_index[object@arms]
    ret[["gq_long_pop_study_index"]] <- data_list$study_to_index[object@studies]

    # Sanity checks
    assert_that(
        length(ret[["gq_long_pop_arm_index"]]) == length(ret[["gq_long_pop_study_index"]]),
        length(ret[["gq_long_pop_study_index"]]) == length(ret[["gq_times"]]),
        length(ret[["gq_long_pop_study_index"]]) == ret[["gq_n_quant"]],
        all(!is.na(ret[["gq_long_pop_arm_index"]])),
        all(!is.na(ret[["gq_long_pop_study_index"]]))
    )
    return(ret)
}

#' @include SurvivalModel.R
NULL

# SurvivalLogLogistic-class ----

#' `SurvivalLogLogistic`
#'
#' This class extends the general [`SurvivalModel`] class for using the
#' log-logistic survival model.
#'
#' @exportClass SurvivalLogLogistic
.SurvivalLogLogistic <- setClass(
    Class = "SurvivalLogLogistic",
    contains = "SurvivalModel"
)

# SurvivalLogLogistic-constructors ----

#' @rdname SurvivalLogLogistic-class
#'
#' @param a (`Prior`)\cr Prior distribution for the scale parameter `a`.
#' @param b (`Prior`)\cr Prior distribution for the shape parameter `b`.
#' @param beta (`Prior`)\cr Prior distribution for covariates coefficients `beta`.
#'
#' @export
SurvivalLogLogistic <- function(
    a = prior_lognormal(log(0.1), 5),
    b = prior_gamma(2, 5),
    beta = prior_normal(0, 2)
) {
    .SurvivalLogLogistic(
        SurvivalModel(
            name = "Log-Logistic",
            stan = StanModule("sm-loglogistic/model.stan"),
            parameters = ParameterList(
                Parameter(name = "sm_loglogis_a", prior = a, size = 1),
                Parameter(name = "sm_loglogis_b", prior = b, size = 1),
                Parameter(name = "beta_os_cov", prior = beta, size = "p_os_cov_design")
            )
        )
    )
}


#' Define Simulation Group
#'
#' Specifies a simulation group to be used by [`SimJointData()`].
#'
#' @param n (`numeric`)\cr number of subjects in the group.
#' @param arm (`character`)\cr treatment arm.
#' @param study (`character`)\cr study name.
#'
#' @slot n (`numeric`)\cr See arguments.
#' @slot arm (`character`)\cr See arguments.
#' @slot study (`character`)\cr See arguments.
#'
#' @examples
#' SimGroup(n = 50, arm = "Arm-A", study = "Study-1")
#' @name SimGroup-class
#' @exportClass SimGroup
.SimGroup <- setClass(
    "SimGroup",
    slots = c(
        n = "numeric",
        arm = "character",
        study = "character"
    )
)

#' @export
#' @rdname SimGroup-class
SimGroup <- function(n, arm, study) {
    .SimGroup(
        n = n,
        arm = arm,
        study = study
    )
}


setValidity(
    "SimGroup",
    function(object) {
        if (length(object@n) != 1) {
            return("`n` must be a length 1 integer")
        }
        if (length(object@arm) != 1) {
            return("`arm` must be a length 1 string")
        }
        if (length(object@study) != 1) {
            return("`study` must be a length 1 string")
        }
        if (any(object@n < 1) | any(object@n %% 1 != 0)) {
            return("`n` must be positive integer")
        }
        return(TRUE)
    }
)


#' Abstract Simulation Class for Longitudinal Data
#'
#' @param times (`numeric`) the times to generate observations at.
#'
#' @description
#' This class exists to be extended by other classes that simulate longitudinal data.
#' It is not intended to be used directly.
#' @name SimLongitudinal-class
#' @family SimLongitudinal
#' @exportClass SimLongitudinal
.SimLongitudinal <- setClass(
    "SimLongitudinal",
    slots = list(
        times = "numeric"
    )
)

#' @rdname SimLongitudinal-class
#' @export
SimLongitudinal <- function(times = seq(0, 100, 50)) {
    .SimLongitudinal(times = times)
}


#' @rdname show-object
#' @export
setMethod(
    f = "show",
    signature = "SimLongitudinal",
    definition = function(object) {
        x <- sprintf("\nA %s Object\n\n", as_print_string(object))
        cat(x)
        return(object)
    }
)

#' @rdname as_print_string
as_print_string.SimLongitudinal <- function(object) {
    return("SimLongitudinal")
}

#' @include StanModule.R
#' @include ParameterList.R
#' @include generics.R
NULL

#' `StanModel` Function Arguments
#'
#' The documentation lists all the conventional arguments for wrappers around
#' [StanModel()].
#'
#' @param stan (`StanModule`)\cr code containing the Stan code specification.
#' @param parameters (`ParameterList`)\cr the parameter specification.
#' @param parameter (`ParameterList`)\cr the (single) parameter specification.
#' @param name (`character`)\cr display name for the model object.
#' @param ... additional arguments for [StanModel()].
#'
#' @name stanmodel_arguments
#' @keywords internal
NULL

# StanModel-class ----


#' Stan Model Object and Constructor Function
#'
#' @slot stan (`StanModule`)\cr See Arguments.
#' @slot parameters (`ParameterList`)\cr See Arguments.
#' @slot name (`character`)\cr display name for the model object.
#'
#' @export StanModel
#' @exportClass StanModel
#' @family StanModel
.StanModel <- setClass(
    Class = "StanModel",
    slots = list(
        "stan" = "StanModule",
        "parameters" = "ParameterList",
        "name" = "character"
    )
)

# StanModel-constructor ----

#' @inheritParams stanmodel_arguments
#' @rdname StanModel-class
StanModel <- function(stan, parameters, name = "<Unnamed>") {
    .StanModel(
        stan = stan,
        parameters = parameters,
        name = name
    )
}

# as.list-StanModel ----

#' `StanModel` -> `list`
#' @description
#' Returns a named list where each element of the list corresponds
#' to a Stan modelling block e.g. `data`, `model`, etc.
#' @param x ([`StanModel`])\cr A Stan Model
#' @param ... Not Used.
#' @family StanModel
#' @export
as.list.StanModel <- function(x, ...) {
    as.list(x@stan)
}

# getParameters-StanModel ----

#' @rdname getParameters
#' @export
getParameters.StanModel <- function(object, ...) object@parameters


#' @export
as_print_string.StanModel <- function(object, ...) {
    string <- sprintf(
        "\n%s Model Object with parameters:\n%s\n\n",
        object@name,
        paste("   ", as_print_string(object@parameters)) |> paste(collapse = "\n")
    )
    return(string)
}

#' @rdname show-object
#' @export
setMethod(
    f = "show",
    signature = "StanModel",
    definition = function(object) {
        cat(as_print_string(object))
    }
)

#' @rdname initialValues
#' @export
initialValues.StanModel <- function(object, n_chains, ...) {
    initialValues(object@parameters, n_chains)
}

#' @include SurvivalModel.R
NULL

# SurvivalWeibullPH-class ----

#' `SurvivalWeibullPH`
#'
#' This class extends the general [`SurvivalModel`] class for using the
#' Weibull proportional hazards survival model.
#'
#' @exportClass SurvivalWeibullPH
.SurvivalWeibullPH <- setClass(
    Class = "SurvivalWeibullPH",
    contains = "SurvivalModel"
)

# SurvivalWeibullPH-constructors ----

#' @rdname SurvivalWeibullPH-class
#'
#' @param lambda (`Prior`)\cr for the scale `lambda`.
#' @param gamma (`Prior`)\cr for the shape `gamma`.
#' @param beta (`Prior`)\cr for covariates coefficients `beta`.
#'
#' @export
SurvivalWeibullPH <- function(
    lambda = prior_gamma(2, 0.5),
    gamma = prior_gamma(2, 0.5),
    beta = prior_normal(0, 2)
) {
    .SurvivalWeibullPH(
        SurvivalModel(
            name = "Weibull-PH",
            stan = StanModule(x = "sm-weibull-ph/model.stan"),
            parameters = ParameterList(
                Parameter(name = "sm_weibull_ph_lambda", prior = lambda, size = 1),
                Parameter(name = "sm_weibull_ph_gamma", prior = gamma, size = 1),
                Parameter(name = "beta_os_cov", prior = beta, size = "p_os_cov_design")
            )
        )
    )
}

1		#' @include LongitudinalModel.R
2		#' @include StanModule.R
3		#' @include generics.R
4		#' @include ParameterList.R
5		#' @include Parameter.R
6		#' @include Link.R
7		NULL
8
9		# LongitudinalGSF-class ----
10
11		#' `LongitudinalGSF`
12		#'
13		#' This class extends the general [`LongitudinalModel`] class for using the
14		#' Generalized Stein-Fojo (GSF) model for the longitudinal outcome.
15		#'
16		#' @section Available Links:
17		#' - [`linkDSLD()`]
18		#' - [`linkTTG()`]
19		#' - [`linkIdentity()`]
20		#' - [`linkGrowth()`]
21		#' @exportClass LongitudinalGSF
22		.LongitudinalGSF <- setClass(
23		Class = "LongitudinalGSF",
24		contains = "LongitudinalModel"
25		)
26
27		# LongitudinalGSF-constructors ----
28
29		#' @rdname LongitudinalGSF-class
30		#'
31		#' @param mu_bsld (`Prior`)\cr for the mean baseline value `mu_bsld`.
32		#' @param mu_ks (`Prior`)\cr for the mean shrinkage rate `mu_ks`.
33		#' @param mu_kg (`Prior`)\cr for the mean growth rate `mu_kg`.
34		#' @param mu_phi (`Prior`)\cr for the mean proportion of cells affected by the treatment `mu_phi`.
35		#'
36		#' @param omega_bsld (`Prior`)\cr for the baseline value standard deviation `omega_bsld`.
37		#' @param omega_ks (`Prior`)\cr for the shrinkage rate standard deviation `omega_ks`.
38		#' @param omega_kg (`Prior`)\cr for the growth rate standard deviation `omega_kg`.
39		#' @param omega_phi (`Prior`)\cr for the standard deviation of the proportion of cells
40		#' affected by the treatment `omega_phi`.
41		#'
42		#' @param sigma (`Prior`)\cr for the variance of the longitudinal values `sigma`.
43		#'
44		#' @param centred (`logical`)\cr whether to use the centred parameterization.
45		#'
46		#' @importFrom stats qlogis
47		#' @export
48		LongitudinalGSF <- function(
49
50		mu_bsld = prior_normal(log(60), 1),
51		mu_ks = prior_normal(log(0.5), 1),
52		mu_kg = prior_normal(log(0.3), 1),
53		mu_phi = prior_normal(qlogis(0.5), 1),
54
55		omega_bsld = prior_lognormal(log(0.2), 1),
56		omega_ks = prior_lognormal(log(0.2), 1),
57		omega_kg = prior_lognormal(log(0.2), 1),
58		omega_phi = prior_lognormal(log(0.2), 1),
59
60		sigma = prior_lognormal(log(0.1), 1),
61
62		centred = FALSE
63		) {
64
65	21x	gsf_model <- StanModule(decorated_render(
66	21x	.x = read_stan("lm-gsf/model.stan"),
67	21x	centred = centred
68		))
69
70		# Apply constraints
71	21x	omega_bsld <- set_limits(omega_bsld, lower = 0)
72	21x	omega_ks <- set_limits(omega_ks, lower = 0)
73	21x	omega_kg <- set_limits(omega_kg, lower = 0)
74	21x	omega_phi <- set_limits(omega_phi, lower = 0)
75	21x	sigma <- set_limits(sigma, lower = 0)
76
77
78	21x	parameters <- list(
79	21x	Parameter(name = "lm_gsf_mu_bsld", prior = mu_bsld, size = "n_studies"),
80	21x	Parameter(name = "lm_gsf_mu_ks", prior = mu_ks, size = "n_arms"),
81	21x	Parameter(name = "lm_gsf_mu_kg", prior = mu_kg, size = "n_arms"),
82	21x	Parameter(name = "lm_gsf_mu_phi", prior = mu_phi, size = "n_arms"),
83
84	21x	Parameter(name = "lm_gsf_omega_bsld", prior = omega_bsld, size = 1),
85	21x	Parameter(name = "lm_gsf_omega_ks", prior = omega_ks, size = 1),
86	21x	Parameter(name = "lm_gsf_omega_kg", prior = omega_kg, size = 1),
87	21x	Parameter(name = "lm_gsf_omega_phi", prior = omega_phi, size = 1),
88
89	21x	Parameter(name = "lm_gsf_sigma", prior = sigma, size = 1)
90		)
91
92	21x	assert_flag(centred)
93	21x	parameters_extra <- if (centred) {
94	4x	list(
95	4x	Parameter(
96	4x	name = "lm_gsf_psi_bsld",
97	4x	prior = prior_init_only(prior_lognormal(median(mu_bsld), median(omega_bsld))),
98	4x	size = "n_subjects"
99		),
100	4x	Parameter(
101	4x	name = "lm_gsf_psi_ks",
102	4x	prior = prior_init_only(prior_lognormal(median(mu_ks), median(omega_ks))),
103	4x	size = "n_subjects"
104		),
105	4x	Parameter(
106	4x	name = "lm_gsf_psi_kg",
107	4x	prior = prior_init_only(prior_lognormal(median(mu_kg), median(omega_kg))),
108	4x	size = "n_subjects"
109		),
110	4x	Parameter(
111	4x	name = "lm_gsf_psi_phi_logit",
112	4x	prior = prior_init_only(prior_normal(median(mu_phi), median(omega_phi))),
113	4x	size = "n_subjects"
114		)
115		)
116		} else {
117	17x	list(
118	17x	Parameter(name = "lm_gsf_eta_tilde_bsld", prior = prior_std_normal(), size = "n_subjects"),
119	17x	Parameter(name = "lm_gsf_eta_tilde_ks", prior = prior_std_normal(), size = "n_subjects"),
120	17x	Parameter(name = "lm_gsf_eta_tilde_kg", prior = prior_std_normal(), size = "n_subjects"),
121	17x	Parameter(name = "lm_gsf_eta_tilde_phi", prior = prior_std_normal(), size = "n_subjects")
122		)
123		}
124	21x	parameters <- append(parameters, parameters_extra)
125
126	21x	x <- LongitudinalModel(
127	21x	name = "Generalized Stein-Fojo",
128	21x	stan = merge(
129	21x	gsf_model,
130	21x	StanModule("lm-gsf/functions.stan")
131		),
132	21x	parameters = do.call(ParameterList, parameters)
133		)
134	21x	.LongitudinalGSF(x)
135		}
136
137
138
139		#' @export
140		enableGQ.LongitudinalGSF <- function(object, ...) {
141	7x	StanModule("lm-gsf/quantities.stan")
142		}
143
144
145		#' @export
146		enableLink.LongitudinalGSF <- function(object, ...) {
147	4x	object@stan <- merge(
148	4x	object@stan,
149	4x	StanModule("lm-gsf/link.stan")
150		)
151	4x	object
152		}
153
154		#' @export
155		linkDSLD.LongitudinalGSF <- function(prior = prior_normal(0, 2), model, ...) {
156	6x	LinkComponent(
157	6x	key = "link_dsld",
158	6x	stan = StanModule("lm-gsf/link_dsld.stan"),
159	6x	prior = prior
160		)
161		}
162
163		#' @export
164		linkTTG.LongitudinalGSF <- function(prior = prior_normal(0, 2), model, ...) {
165	4x	LinkComponent(
166	4x	key = "link_ttg",
167	4x	stan = StanModule("lm-gsf/link_ttg.stan"),
168	4x	prior = prior
169		)
170		}
171
172		#' @export
173		linkIdentity.LongitudinalGSF <- function(prior = prior_normal(0, 2), model, ...) {
174	3x	LinkComponent(
175	3x	key = "link_identity",
176	3x	stan = StanModule("lm-gsf/link_identity.stan"),
177	3x	prior = prior
178		)
179		}
180
181		#' @export
182		linkGrowth.LongitudinalGSF <- function(prior = prior_normal(0, 2), model, ...) {
183	1x	LinkComponent(
184	1x	key = "link_growth",
185	1x	stan = StanModule("lm-gsf/link_growth.stan"),
186	1x	prior = prior
187		)
188		}
189
190
191		#' @rdname getPredictionNames
192		#' @export
193		getPredictionNames.LongitudinalGSF <- function(object, ...) {
194	3x	c("b", "s", "g", "phi")
195		}

1		#' @include Grid.R
2		#' @include GridFixed.R
3		#' @include generics.R
4		NULL
5
6		#' @rdname Grid-Dev
7		.GridGrouped <- setClass(
8		"GridGrouped",
9		contains = "Grid",
10		slots = c(
11		"groups" = "list",
12		"times" = "numeric_or_NULL"
13		)
14		)
15
16		#' @rdname Grid-Functions
17		#' @export
18		GridGrouped <- function(groups, times = NULL) {
19	13x	.GridGrouped(
20	13x	groups = groups,
21	13x	times = times
22		)
23		}
24
25
26		setValidity(
27		"GridGrouped",
28		function(object) {
29		if (!all(vapply(object@groups, is.character, logical(1)))) {
30		return("Each element of `groups` must be a character vector")
31		}
32		gnames <- names(object@groups)
33		gnames <- gnames[!is.na(gnames) & gnames != ""]
34		if (length(gnames) != length(object@groups)) {
35		return("Each element of `groups` must be named")
36		}
37		return(TRUE)
38		}
39		)
40
41
42		#' @rdname Quant-Dev
43		#' @export
44		as.QuantityGenerator.GridGrouped <- function(object, data, ...) {
45	17x	assert_class(data, "DataJoint")
46	17x	data_list <- as.list(data)
47	17x	subjects_unique <- unique(unlist(object@groups))
48	17x	assert_that(
49	17x	all(subjects_unique %in% names(data_list$subject_to_index))
50		)
51	17x	as.QuantityGenerator(
52	17x	GridFixed(
53	17x	times = object@times,
54	17x	subjects = subjects_unique
55		),
56	17x	data = data
57		)
58		}
59
60
61		#' @rdname Quant-Dev
62		#' @export
63		as.QuantityCollapser.GridGrouped <- function(object, data, ...) {
64	8x	assert_class(data, "DataJoint")
65	8x	data_list <- as.list(data)
66	8x	assert_that(
67	8x	all(unique(unlist(object@groups)) %in% names(data_list$subject_to_index))
68		)
69
70	8x	validate_time_grid(object@times)
71
72	8x	group_grid <- expand.grid(
73	8x	group = names(object@groups),
74	8x	time = object@times,
75	8x	stringsAsFactors = FALSE
76		)
77
78	8x	generator <- as.QuantityGenerator(object, data)
79
80	8x	select_indexes <- mapply(
81	8x	function(group, time) {
82	490x	correct_subject <- generator@subjects %in% object@groups[[group]]
83	490x	correct_time <- generator@times == time
84	490x	seq_along(correct_time)[correct_subject & correct_time]
85		},
86	8x	group_grid$group,
87	8x	group_grid$time,
88	8x	SIMPLIFY = FALSE
89		)
90	8x	names(select_indexes) <- NULL
91
92	8x	QuantityCollapser(
93	8x	times = group_grid$time,
94	8x	groups = group_grid$group,
95	8x	indexes = select_indexes
96		)
97		}
98
99		#' @export
100		as.list.GridGrouped <- function(x, ...) {
101	1x	x@groups
102		}
103
104		#' @rdname coalesceGridTime
105		#' @export
106		coalesceGridTime.GridGrouped <- function(object, times, ...) {
107	10x	if (is.null(object@times)) {
108	2x	object <- GridGrouped(
109	2x	groups = object@groups,
110	2x	times = times
111		)
112		}
113	10x	object
114		}

1		#' @include Grid.R
2		#' @include generics.R
3		NULL
4
5
6		#' @rdname Grid-Dev
7		.GridPrediction <- setClass(
8		"GridPrediction",
9		contains = "Grid",
10		slots = c(
11		"times" = "numeric_or_NULL",
12		"newdata" = "data.frame",
13		"params" = "list"
14		)
15		)
16		#' @rdname Grid-Functions
17		#' @export
18		GridPrediction <- function(times = NULL, newdata, params = list()) {
19	7x	.GridPrediction(
20	7x	times = times,
21	7x	params = params,
22	7x	newdata = newdata
23		)
24		}
25		setValidity(
26		"GridPrediction",
27		function(object) {
28		for (param in names(object@params)) {
29		if (length(object@params[[param]]) != 1) {
30		return(sprintf("Parameter '%s' must be length 1", param))
31		}
32		}
33		if (length(object@params) != length(names(object@params))) {
34		return("All elements of `params` must be named")
35		}
36		if (!is.null(object@newdata[["..new_subject.."]])) {
37		return("`newdata` must not contain a column named '..new_subject..'")
38		}
39		return(TRUE)
40		}
41		)
42
43
44		#' @rdname Quant-Dev
45		#' @export
46		as.QuantityGenerator.GridPrediction <- function(object, data, ...) {
47
48	6x	assert_class(data, "DataJoint")
49	6x	data_list <- as.list(data)
50	6x	validate_time_grid(object@times)
51
52	6x	n_times <- length(object@times)
53	6x	n_obs <- nrow(object@newdata)
54	6x	newdata <- object@newdata
55	6x	newdata[["..new_subject.."]] <- sprintf(
56	6x	"new_subject_%i",
57	6x	seq_len(nrow(newdata))
58		)
59
60	6x	QuantityGeneratorPrediction(
61	6x	times = rep(object@times, each = n_obs),
62	6x	newdata = replicate(newdata, n = n_times, simplify = FALSE) \|> dplyr::bind_rows(),
63	6x	params = object@params
64		)
65		}
66
67		#' @rdname Quant-Dev
68		#' @export
69		as.QuantityCollapser.GridPrediction <- function(object, data, ...) {
70	2x	assert_class(data, "DataJoint")
71	2x	data_list <- as.list(data)
72	2x	generator <- as.QuantityGenerator(object, data)
73	2x	QuantityCollapser(
74	2x	times = generator@times,
75	2x	groups = generator@newdata[["..new_subject.."]],
76	2x	indexes = as.list(seq_along(generator@times))
77		)
78		}
79
80
81		#' @export
82		as.list.GridPrediction <- function(x, data, ...) {
83	!	stop("`as.list()` is not implemented for `GridPrediction` objects")
84		}
85
86
87		#' @rdname coalesceGridTime
88		#' @export
89		coalesceGridTime.GridPrediction <- function(object, times, ...) {
90	4x	if (is.null(object@times)) {
91	!	object <- GridPrediction(
92	!	times = times,
93	!	newdata = object@newdata,
94	!	params = object@params
95		)
96		}
97	4x	object
98		}

1		#' @include Grid.R
2		#' @include generics.R
3		NULL
4
5		#' @rdname Grid-Dev
6		.GridObserved <- setClass(
7		"GridObserved",
8		contains = "Grid",
9		slots = c(
10		"subjects" = "character_or_NULL"
11		)
12		)
13
14
15		#' @rdname Grid-Functions
16		#' @export
17		GridObserved <- function(subjects = NULL) {
18	2x	.GridObserved(
19	2x	subjects = subjects
20		)
21		}
22
23		#' @rdname Quant-Dev
24		#' @export
25		as.QuantityGenerator.GridObserved <- function(object, data, ...) {
26	4x	assert_class(data, "DataJoint")
27	4x	data_list <- as.list(data)
28	4x	subjects <- unlist(as.list(object, data = data), use.names = FALSE)
29	4x	unique_visits <- tapply(data_list$tumour_time, data_list$subject_tumour_index, unique)
30	4x	subject_visits <- unique_visits[data_list$subject_to_index[subjects]]
31	4x	visit_lengths <- vapply(subject_visits, length, numeric(1))
32	4x	QuantityGeneratorSubject(
33	4x	times = unlist(subject_visits, use.names = FALSE),
34	4x	subjects = rep(subjects, visit_lengths)
35		)
36		}
37
38		#' @rdname Quant-Dev
39		#' @export
40		as.QuantityCollapser.GridObserved <- function(object, data, ...) {
41	2x	generator <- as.QuantityGenerator(object, data)
42
43	2x	QuantityCollapser(
44	2x	times = generator@times,
45	2x	groups = generator@subjects,
46	2x	indexes = as.list(seq_along(generator@times))
47		)
48		}
49
50		#' @export
51		as.list.GridObserved <- function(x, data, ...) {
52	4x	subjects_to_list(x@subjects, data)
53		}

1		#' @include JointModel.R
2		#' @include SurvivalQuantities.R
3		NULL
4
5		setOldClass("CmdStanMCMC")
6
7		# JointModelSamples-class ----
8
9		#' `JointModelSamples`
10		#'
11		#' Contains samples from a [`JointModel`].
12		#'
13		#' @slot model ([`JointModel`])\cr the model that the samples were drawn from.
14		#' @slot data ([`DataJoint`])\cr the data that the model was fitted on.
15		#' @slot results ([`cmdstanr::CmdStanMCMC`])\cr the STAN samples.
16		#'
17		#' @aliases JointModelSamples
18		#' @export
19		.JointModelSamples <- setClass(
20		"JointModelSamples",
21		slots = c(
22		model = "JointModel",
23		data = "DataJoint",
24		results = "CmdStanMCMC"
25		)
26		)
27
28
29		#' @rdname generateQuantities
30		#' @param generator (`QuantityGenerator`)\cr object that specifies which subjects and time points
31		#' to calculate the quantities at
32		#' @param type (`character`)\cr type of quantities to be generated, must be either "survival" or
33		#' "longitudinal".
34		#' @export
35		generateQuantities.JointModelSamples <- function(object, generator, type, ...) {
36
37	31x	data <- as_stan_list(object@data) \|>
38	31x	append(as_stan_list(object@model@parameters)) \|>
39	31x	append(as_stan_list(generator, data = object@data, model = object@model))
40
41	30x	stanobj <- as.StanModule(object, generator = generator, type = type)
42	30x	model <- compileStanModel(stanobj)
43
44	30x	devnull <- utils::capture.output(
45	30x	results <- model$generate_quantities(
46	30x	data = data,
47	30x	fitted_params = object@results
48		)
49		)
50	30x	return(results)
51		}
52
53
54		#' `JointModelSamples` -> `StanModule`
55		#'
56		#' Converts a `JointModelSamples` object into a `StanModule` object ensuring
57		#' that the resulting `StanModule` object is able to generate post sampling
58		#' quantities.
59		#'
60		#' @inheritParams generateQuantities
61		#' @export
62		as.StanModule.JointModelSamples <- function(object, generator, type, ...) {
63	38x	assert_that(
64	38x	is(generator, "QuantityGenerator"),
65	38x	length(type) == 1,
66	38x	type %in% c("survival", "longitudinal")
67		)
68
69	38x	quant_stanobj <- read_stan("base/quantities.stan") \|>
70	38x	decorated_render(
71	38x	include_gq_longitudinal_idv = (type == "longitudinal") & is(generator, "QuantityGeneratorSubject"),
72	38x	include_gq_longitudinal_pop = (type == "longitudinal") & is(generator, "QuantityGeneratorPopulation"),
73	38x	include_gq_survival_idv = (type == "survival") & is(generator, "QuantityGeneratorSubject"),
74	38x	include_gq_survival_pred = (type == "survival") & is(generator, "QuantityGeneratorPrediction")
75		) \|>
76	38x	StanModule()
77
78	38x	stanobj <- Reduce(
79	38x	merge,
80	38x	list(
81	38x	as.StanModule(object@model),
82	38x	enableGQ(object@model),
83	38x	quant_stanobj
84		)
85		)
86	38x	stanobj
87		}
88
89
90
91		#' `JointModelSamples` -> Printable `Character`
92		#'
93		#' Converts [`JointModelSamples`] object into a printable string.
94		#' @param object ([`JointModelSamples`])\cr samples as drawn from a [`JointModel`].
95		#' @family JointModelSamples
96		#' @param indent (`numeric`)\cr how much white space to prefix the print string with.
97		#' @keywords internal
98		#' @export
99		as_print_string.JointModelSamples <- function(object, indent = 1, ...) {
100	17x	sizes <- vapply(
101	17x	as.CmdStanMCMC(object)$metadata()[["stan_variable_sizes"]],
102	17x	\(x) {
103	91x	if (length(x) == 1 && x == 1) return("")
104	154x	paste0("[", paste(x, collapse = ", "), "]")
105		},
106	17x	character(1)
107		)
108	17x	variable_string <- paste0(
109		" ",
110	17x	as.CmdStanMCMC(object)$metadata()[["stan_variables"]],
111	17x	sizes
112		)
113	17x	template <- c(
114	17x	"JointModelSamples Object with:",
115		"",
116	17x	" # of samples per chain = %d",
117	17x	" # of chains = %d",
118		"",
119	17x	" Variables:",
120	17x	variable_string[order(variable_string)]
121		)
122	17x	pad <- rep(" ", indent) \|> paste(collapse = "")
123	17x	template_padded <- paste(pad, template)
124	17x	sprintf(
125	17x	paste(template_padded, collapse = "\n"),
126	17x	as.CmdStanMCMC(object)$metadata()$iter_sampling,
127	17x	as.CmdStanMCMC(object)$num_chains()
128		)
129		}
130
131		#' @rdname show-object
132		#' @export
133		setMethod(
134		f = "show",
135		signature = "JointModelSamples",
136		definition = function(object) {
137	17x	string <- as_print_string(object)
138	17x	cat("\n", string, "\n\n")
139		}
140		)
141
142
143		#' @rdname as.CmdStanMCMC
144		#' @export
145		as.CmdStanMCMC.JointModelSamples <- function(object, ...) {
146	81x	return(object@results)
147		}

1		#' Row Numbers of Data with Missing Variables
2		#'
3		#' @param df (`data.frame`)\cr input data.
4		#' @param formula (`formula` or `NULL`)\cr which variables to inspect for missingness, if `NULL`
5		#' all variables are considered.
6		#'
7		#' @returns Numeric vector specifying which rows contain at least 1 missing observation
8		#' in any of the inspected variables.
9		#'
10		#' @keywords internal
11		get_missing_rownumbers <- function(df, formula = NULL) {
12	455x	if (is.null(formula)) {
13	!	formula <- ~ .
14		}
15	455x	mdf <- stats::model.frame(formula, data = df, na.action = stats::na.pass)
16	455x	which(!stats::complete.cases(mdf))
17		}
18
19		#' Remove Rows with Missing Variables
20		#'
21		#' Removes any rows from a data set that contain missing values in the inspected
22		#' variables. Allows users to specify which variables to inspect for missing values
23		#' based on either a formula or a character vector of variable names.
24		#'
25		#' @param data (`data.frame`)\cr input data.
26		#' @param formula (`formula` or `NULL`)\cr which variables to inspect for missingness.
27		#' @param extra_vars (`character`)\cr additional variables to inspect for missingness.
28		#'
29		#' @returns The `data` after removing observations that contain missing values in the required variables.
30		#' Note that additional variables not listed in `formula` or `extra_vars` are not dropped and may
31		#' still contain missing values.
32		#'
33		#' @keywords internal
34		remove_missing_rows <- function(data, formula, extra_vars = NULL) {
35	454x	if (!is.null(extra_vars)) {
36	1x	extra_vars <- paste(extra_vars, collapse = " + ")
37	1x	formula_update_string <- paste0(". ~ . + ", extra_vars)
38	1x	formula <- stats::update(formula, formula_update_string)
39		}
40
41	454x	missing_rows <- get_missing_rownumbers(data, formula)
42
43	454x	if (length(missing_rows) == 0) {
44	452x	return(data)
45		}
46	2x	message(sprintf(
47	2x	"Note that %d observations were removed as one of more required variables contained missing values",
48	2x	length(missing_rows)
49		))
50	2x	data_reduced <- data[-missing_rows, ]
51	2x	rownames(data_reduced) <- NULL
52	2x	data_reduced
53		}
54
55		#' Replicate Single Values in a List
56		#'
57		#' @param initial_values (`list`)\cr initial values with names.
58		#' @param sizes (`list`)\cr each size corresponds to an element in `initial_values`,
59		#' matched by the names. An attribute `array` must be attached to each element,
60		#' see [replace_with_lookup()].
61		#'
62		#' @returns A named list of values, with any single values in the `initial_values` list
63		#' replicated according to the corresponding values in the `sizes` list.
64		#' Even when the size is 1, the value is passed as an `array` if the corresponding
65		#' attribute is `TRUE` in `sizes`.
66		#'
67		#' @note The resulting list has the same names as the original lists.
68		#'
69		#' @keywords internal
70		expand_initial_values <- function(initial_values, sizes) {
71	28x	assert_that(
72	28x	is.list(initial_values),
73	28x	msg = "`initial_values` must be a list"
74		)
75	28x	assert_that(
76	28x	is.list(sizes),
77	28x	msg = "`sizes` must be a list"
78		)
79	28x	assert_that(
80	28x	all(names(sizes) %in% names(initial_values)),
81	28x	all(names(initial_values) %in% names(sizes)),
82	28x	msg = "`initial_values` and `sizes` must have identical names"
83		)
84
85	28x	for (name in names(initial_values)) {
86		# Check for single values and replicate them according to sizes.
87	181x	if (length(initial_values[[name]]) == 1) {
88	178x	initial_values[[name]] <- rep(initial_values[[name]], sizes[[name]])
89		}
90		# Check for array handling.
91	181x	needs_array <- attr(sizes[[name]], "array")
92	181x	assert_that(
93	181x	is.flag(needs_array),
94	181x	msg = "each sizes element must have array flag attribute"
95		)
96	181x	if (needs_array) {
97	101x	initial_values[[name]] <- as.array(initial_values[[name]])
98		}
99		}
100
101		# Check that each element of initial_values has the same length as specified in sizes.
102	28x	for (name in names(initial_values)) {
103	181x	assert_that(
104	181x	length(initial_values[[name]]) == sizes[[name]],
105	181x	msg = "length of element in `initial_values` does not match specified size"
106		)
107		}
108
109	28x	initial_values
110		}
111
112		#' Replace Character Size by Looked Up Numbers
113		#'
114		#' @param sizes (`list`)\cr may include character elements that correspond to
115		#' names in the data list.
116		#' @param data (`list`)\cr data containing numeric values.
117		#'
118		#' @returns A list of sizes with character elements in `sizes`
119		#' replaced by their corresponding numeric values in `data`.
120		#'
121		#' @details An attribute `array` for each returned list element indicates
122		#' whether the parameter needs to be handled
123		#' as an array. This is the case when the size is larger than 1, or when
124		#' the size was looked up in the `data`, because in that case it is flexible
125		#' and hence is handled as an array in the Stan code.
126		#'
127		#' @note Each element in the final list of sizes must be a single number.
128		#'
129		#' @keywords internal
130		replace_with_lookup <- function(sizes, data) {
131
132	21x	assert_that(is.list(sizes), msg = "`sizes` must be a list")
133	21x	assert_that(is.list(data), msg = "`data` must be a list")
134
135	21x	for (idx in seq_along(sizes)) {
136	130x	val <- sizes[[idx]]
137	130x	if (is.character(val)) {
138	70x	assert_that(
139	70x	length(val) == 1,
140	70x	msg = "character elements of `sizes` must be strings"
141		)
142	70x	assert_that(
143	70x	val %in% names(data),
144	70x	msg = sprintf("`%s` is not available in `data`", val)
145		)
146	70x	new_val <- data[[val]]
147	70x	assert_that(
148	70x	is.number(new_val),
149	70x	msg = "Selected values from data must be single numbers"
150		)
151	69x	sizes[[idx]] <- structure(new_val, array = TRUE)
152		} else {
153	60x	assert_that(
154	60x	is.number(val),
155	60x	msg = "Existing values in sizes must be single numbers"
156		)
157	59x	sizes[[idx]] <- structure(val, array = val > 1)
158		}
159		}
160	19x	sizes
161		}
162
163		#' Obtain Median and Credible Intervals from MCMC samples
164		#'
165		#' @param samples (`matrix`)\cr with samples in rows and parameters in columns.
166		#' @param level (`number`)\cr credibility level to use for the credible intervals.
167		#'
168		#' @returns A `data.frame` with columns `median`, `lower` and `upper`.
169		#' @keywords internal
170		#' samples_median_ci(samples)
171		samples_median_ci <- function(samples, level = 0.95) {
172	38x	assert_that(is.matrix(samples))
173	38x	assert_that(is.number(level), level < 1, level > 0)
174
175	38x	samples_median <- apply(samples, MARGIN = 2L, FUN = stats::median)
176	38x	probs <- c((1 - level) / 2, (1 + level) / 2)
177	38x	samples_ci <- t(apply(samples, MARGIN = 2L, FUN = stats::quantile, probs = probs))
178	38x	colnames(samples_ci) <- c("lower", "upper")
179	38x	as.data.frame(cbind(
180	38x	median = samples_median,
181	38x	samples_ci
182		))
183		}
184
185
186
187		#' `decorated_render`
188		#'
189		#' Simple wrapper around [jinjar::render()] that provides some additional default
190		#' variables about the system (avoids each call to jinjar having to specify them)
191		#' @param ... Arguments passed onto [jinjar::render()]
192		#' @returns See [jinjar::render()]
193		#' @keywords internal
194		decorated_render <- function(...) {
195	343x	jinjar::render(
196		...,
197	343x	machine_double_eps = sqrt(.Machine$double.eps),
198	343x	machine_double_neg_eps = sqrt(.Machine$double.neg.eps)
199		)
200		}
201
202
203		is_windows <- function() {
204	52x	sysname <- Sys.info()["sysname"]
205	52x	return(sysname == "Windows")
206		}
207
208		#' `validate_time_grid`
209		#'
210		#' Validate that the provided time grid is:
211		#' - finite
212		#' - numeric
213		#' - non-missing
214		#' - sorted
215		#' - unique
216		#'
217		#' @param time_grid (`numeric`)\cr A vector of times which quantities will be
218		#' evaluated at.
219		#'
220		#' @keywords internal
221		validate_time_grid <- function(time_grid) {
222	80x	assert_that(
223	80x	!any(is.na(time_grid)),
224	80x	is.numeric(time_grid),
225	80x	!is.null(time_grid),
226	80x	!is.unsorted(time_grid),
227	80x	!any(duplicated(time_grid)),
228	80x	all(is.finite(time_grid)),
229	80x	msg = "`time_grid` needs to be finite, sorted, unique valued numeric vector"
230		)
231	76x	invisible(time_grid)
232		}
233
234
235
236		#' `expand_subjects`
237		#'
238		#' This function checks and expands a given subjects vector.
239		#' The input vector must be unique and contain only values
240		#' as specified by `all_subjects`
241		#'
242		#' @param subjects (`character` or `NULL`)\cr Character vector representing the subjects.
243		#' If NULL, it will be set to the value of `all_subjects`.
244		#' @param all_subjects (`character`)\cr Character vector representing all possible subjects.
245		#' @return Returns the expanded `subjects` vector.
246		#' @keywords internal
247		expand_subjects <- function(subjects, all_subjects) {
248	38x	assert_that(
249	38x	is.character(all_subjects),
250	38x	msg = "`all_subjects` must be a character vector"
251		)
252	38x	if (is.null(subjects)) {
253	3x	subjects <- unique(all_subjects)
254		}
255	38x	assert_that(
256	38x	is.character(subjects),
257	38x	all(subjects %in% all_subjects),
258	38x	!any(duplicated(subjects)),
259	38x	msg = "`subjects` should be a unique character vector containing only values from the original df"
260		)
261	30x	return(subjects)
262		}
263
264
265
266		#' Decompose subjects into Relevant Components
267		#'
268		#' This function takes in a character vector or list of subjects and decomposes it into a
269		#' structured format.
270		#'
271		#' The primary use of this function is to correctly setup indexing variables for
272		#' predicting survival quantities (see [SurvivalQuantities()])
273		#'
274		#' @param subjects (`character` or `list`)\cr subject identifiers. If `NULL` will be set to `all_subjects`.
275		#'
276		#' @param all_subjects (`character`)\cr the set of allowable subject identifiers.
277		#' Will cause an error if any value of `subjects` is not in this vector.
278		#'
279		#' @return A list containing three components:
280		#' - `groups`: (`list`)\cr each element of the list is a character vector
281		#' specifying which subjects belong to a given "group" where the "group" is the element name
282		#' - `unique_values`: (`character`)\cr vector of the unique subjects within `subjects`
283		#' - `indexes`: (`list`)\cr each element is a named and is a numeric index vector
284		#' that maps the values of `grouped` to `unique_values`
285		#' @examples
286		#' \dontrun{
287		#' result <- decompose_subjects(c("A", "B"), c("A", "B", "C", "D"))
288		#' result <- decompose_subjects(
289		#' list("g1" = c("A", "B"), "g2" = c("B", "C")),
290		#' c("A", "B", "C", "D")
291		#' )
292		#' }
293		#' @seealso [expand_subjects()], [SurvivalQuantities()]
294		#' @keywords internal
295		decompose_subjects <- function(subjects, all_subjects) {
296	13x	if (is.character(subjects) \|\| is.null(subjects)) {
297	7x	subjects <- expand_subjects(subjects, all_subjects)
298	4x	names(subjects) <- subjects
299	4x	subjects <- as.list(subjects)
300		}
301	10x	subjects <- lapply(
302	10x	subjects,
303	10x	expand_subjects,
304	10x	all_subjects = all_subjects
305		)
306	8x	assert_that(
307	8x	is.list(subjects),
308	8x	length(unique(names(subjects))) == length(subjects),
309	8x	all(vapply(subjects, is.character, logical(1)))
310		)
311	8x	subjects_vec_unordered <- unique(unlist(subjects))
312	8x	subjects_vec <- subjects_vec_unordered[order(subjects_vec_unordered)]
313	8x	subjects_lookup <- stats::setNames(seq_along(subjects_vec), subjects_vec)
314	8x	subjects_index <- lapply(
315	8x	subjects,
316	8x	\(x) {
317	22x	z <- subjects_lookup[x]
318	22x	names(z) <- NULL
319	22x	z
320		}
321		)
322	8x	list(
323	8x	groups = subjects,
324	8x	unique_values = subjects_vec,
325	8x	indexes = subjects_index
326		)
327		}
328
329		is_cmdstanr_available <- function() {
330	4x	requireNamespace("cmdstanr", quietly = TRUE)
331		}
332
333
334		is_connection <- function(obj) {
335	!	inherits(obj, "connection")
336		}

1		#' @include generics.R
2		NULL
3
4		# STAN_BLOCKS ----
5
6		#' List of Stan Blocks
7		#'
8		#' @description
9		#' A list with 1 element per standard Stan program blocks.
10		#' This object is mostly used internally as a reference for
11		#' what blocks are expected to exist within a given Stan program.
12		#'
13		#' @export
14		STAN_BLOCKS <- list(
15		functions = "functions",
16		data = "data",
17		transformed_data = "transformed data",
18		parameters = "parameters",
19		transformed_parameters = "transformed parameters",
20		model = "model",
21		generated_quantities = "generated quantities"
22		)
23
24		# add_missing_stan_blocks ----
25
26		#' Add Missing Stan Blocks
27		#'
28		#' @param x (`list`)\cr list of Stan code blocks
29		#' @param stan_blocks (`list`)\cr reference list of stan blocks.
30		#'
31		#' @return Amended list `x` such that all blocks in the global variable
32		#' `STAN_BLOCKS` are contained.
33		#'
34		#' @keywords internal
35		add_missing_stan_blocks <- function(x, stan_blocks = STAN_BLOCKS) {
36	384x	for (block in names(stan_blocks)) {
37	2683x	if (is.null(x[[block]])) {
38	146x	x[[block]] <- ""
39		}
40		}
41	384x	return(x)
42		}
43
44		# StanModule-class ----
45
46		#' `StanModule` Object and Constructor Function
47		#'
48		#' @param x (`string`)\cr file path to a Stan program or a character vector
49		#' of Stan code to be parsed.
50		#' @param ... additional arguments passed to the constructor.
51		#'
52		#' @slot functions (`character`)\cr the `functions` block.
53		#' @slot data (`character`)\cr the `data` block.
54		#' @slot transformed_data (`character`)\cr the `transformed_data` block.
55		#' @slot parameters (`character`)\cr the `parameters` block.
56		#' @slot transformed_parameters (`character`)\cr the `transformed_parameters` block.
57		#' @slot model (`character`)\cr the `model` block.
58		#' @slot generated_quantities (`character`)\cr the `generated_quantities` block.
59		#'
60		#' @exportClass StanModule
61		#' @export StanModule
62		#' @family StanModule
63		.StanModule <- setClass(
64		Class = "StanModule",
65		slots = list(
66		functions = "character",
67		data = "character",
68		transformed_data = "character",
69		parameters = "character",
70		transformed_parameters = "character",
71		model = "character",
72		generated_quantities = "character"
73		)
74		)
75
76		# StanModule-constructors ----
77
78		#' @rdname StanModule-class
79		StanModule <- function(
80		x = "",
81		...
82		) {
83	2223x	assert_that(
84	2223x	is.character(x),
85	2223x	length(x) == 1,
86	2223x	msg = "`x` must be a length 1 character vector"
87		)
88	2223x	code <- read_stan(x)
89	2223x	code_fragments <- as_stan_fragments(code)
90
91	2220x	if (paste0(unlist(code_fragments), collapse = "") == "" && paste0(x, collaspe = "") != "") {
92	1x	warning("Non-empty input resulted in an empty StanModule object. Is the input correct?")
93		}
94
95	2220x	.StanModule(
96	2220x	functions = code_fragments$functions,
97	2220x	data = code_fragments$data,
98	2220x	transformed_data = code_fragments$transformed_data,
99	2220x	parameters = code_fragments$parameters,
100	2220x	transformed_parameters = code_fragments$transformed_parameters,
101	2220x	model = code_fragments$model,
102	2220x	generated_quantities = code_fragments$generated_quantities,
103		...
104		)
105		}
106
107		# as.character-StanModule ----
108
109		#' `StanModule` -> `character`
110		#' @param x ([`StanModule`])\cr A stan program
111		#' @param ... Not Used.
112		#' @description
113		#' Converts a [`StanModule`] object into a valid Stan program file where each
114		#' line of the returned `character` vector represents a line of the program
115		#' @family StanModule
116		#' @export
117		as.character.StanModule <- function(x, ...) {
118	443x	as_stan_file(
119	443x	functions = x@functions,
120	443x	data = x@data,
121	443x	transformed_data = x@transformed_data,
122	443x	parameters = x@parameters,
123	443x	transformed_parameters = x@transformed_parameters,
124	443x	model = x@model,
125	443x	generated_quantities = x@generated_quantities
126		)
127		}
128
129
130
131		# merge-StanModule,StanModule ----
132
133		#' @param stan_blocks (`list`)\cr reference list of stan blocks.
134		#' @rdname merge
135		setMethod(
136		f = "merge",
137		signature = c("StanModule", "StanModule"),
138		definition = function(x, y, stan_blocks = STAN_BLOCKS, ...) {
139	867x	stan_fragments <- lapply(
140	867x	names(stan_blocks),
141	867x	\(par) {
142	6069x	if (all(slot(y, par) == "")) {
143	4081x	return(slot(x, par))
144		}
145	1988x	if (all(slot(x, par) == "")) {
146	561x	return(slot(y, par))
147		}
148	1427x	return(c(slot(x, par), slot(y, par)))
149		}
150		)
151	867x	names(stan_fragments) <- names(stan_blocks)
152	867x	stan_code <- do.call(as_stan_file, stan_fragments)
153	867x	StanModule(
154	867x	x = stan_code
155		)
156		}
157		)
158
159		# compileStanModel-StanModule,character ----
160
161		#' @rdname compileStanModel
162		compileStanModel.StanModule <- function(object) {
163	48x	exe_dir <- getOption("jmpost.cache_dir")
164	48x	if (!dir.exists(exe_dir)) {
165	3x	dir.create(exe_dir, recursive = TRUE)
166		}
167	48x	stan_code <- as.character(object)
168	48x	hash_name <- digest::digest(stan_code, "md5")
169	48x	exe_name <- paste0(
170	48x	"model_",
171	48x	hash_name,
172	48x	if (is_windows()) ".exe" else ""
173		)
174	48x	exe_file <- file.path(exe_dir, exe_name)
175	48x	source_file <- cmdstanr::write_stan_file(
176	48x	code = stan_code,
177	48x	dir = exe_dir,
178	48x	basename = sprintf("model_%s.stan", hash_name)
179		)
180
181		# Suppress "model executable is up to date" message as
182		# users are not in control of the cache so this message is meaningless
183	48x	withCallingHandlers(
184		{
185	48x	x <- cmdstanr::cmdstan_model(
186	48x	stan_file = source_file,
187	48x	exe_file = exe_file,
188	48x	quiet = TRUE
189		)
190		},
191	48x	message = function(m) {
192	!	if (m$message == "Model executable is up to date!\n") {
193	!	invokeRestart("muffleMessage")
194		}
195		}
196		)
197	48x	invisible(x)
198		}
199
200
201
202		# as.list-StanModule ----
203
204		#' `StanModule` -> `list`
205		#' @description
206		#' Returns a named list where each element of the list corresponds
207		#' to a Stan modelling block e.g. `data`, `model`, etc.
208		#' @param x ([`StanModule`])\cr A Stan Module
209		#' @param stan_blocks (`list`)\cr reference list of stan blocks.
210		#' @param ... Not Used.
211		#' @family StanModule
212		#' @export
213		as.list.StanModule <- function(x, stan_blocks = STAN_BLOCKS, ...) {
214	362x	string <- as.character(x)
215	362x	li <- as_stan_fragments(string)
216	362x	for (block in names(stan_blocks)) {
217	2534x	li[[block]] <- paste(li[[block]], collapse = "\n")
218		}
219	362x	return(li)
220		}
221
222		# is_file ----
223
224		#' Is String a Valid File?
225		#'
226		#' A utility function to check if a string is a valid file or not.
227		#' Used to help address short comings of [file.exists()] that will return `TRUE`
228		#' for a directory as well as a file.
229		#'
230		#' @param filename (`string`)\cr file name.
231		#'
232		#' @keywords internal
233		is_file <- function(filename = NULL) {
234	7455x	if (is.null(filename)) {
235	!	return(FALSE)
236		}
237	7455x	assert_that(
238	7455x	is.character(filename),
239	7455x	length(filename) == 1,
240	7455x	msg = "`filename` must be a length 1 character"
241		)
242	7455x	if (nchar(filename) > 1000) {
243	1446x	return(FALSE)
244		}
245	6009x	if (is.na(filename)) {
246	!	return(FALSE)
247		}
248	6009x	return(file.exists(filename) & !dir.exists(filename))
249		}
250
251		# read_stan ----
252
253		#' Stan Code as Character
254		#'
255		#' @param string Character, either the absolute path of a stan file, or the name of the stan
256		#' file in the package directory or the stan code as a string.
257		read_stan <- function(string) {
258	2493x	local_inst_file <- file.path("inst", "stan", string)
259	2493x	system_file <- system.file(file.path("stan", string), package = "jmpost")
260	2493x	local_file <- string
261	2493x	files <- c(local_file, local_inst_file, system_file)
262	2493x	for (fi in files) {
263	7455x	if (is_file(fi)) {
264	620x	string <- readLines(fi)
265	620x	break
266		}
267		}
268	2493x	string <- paste0(string, collapse = "\n")
269	2493x	return(string)
270		}
271
272		# as_stan_file ----
273
274		#' Merging Code Blocks into Stan Code Character Vector
275		#'
276		#' @param functions (`character`)\cr code block.
277		#' @param data (`character`)\cr code block.
278		#' @param transformed_data (`character`)\cr code block.
279		#' @param parameters (`character`)\cr code block.
280		#' @param transformed_parameters (`character`)\cr code block.
281		#' @param model (`character`)\cr code block.
282		#' @param generated_quantities (`character`)\cr code block.
283		#' @param stan_blocks (`list`)\cr reference list of stan blocks.
284		#'
285		#' @return Character vector of the complete Stan code.
286		#'
287		#' @keywords internal
288		as_stan_file <- function(
289		functions = "",
290		data = "",
291		transformed_data = "",
292		parameters = "",
293		transformed_parameters = "",
294		model = "",
295		generated_quantities = "",
296		stan_blocks = STAN_BLOCKS
297		) {
298	1310x	block_strings <- lapply(
299	1310x	names(stan_blocks),
300	1310x	function(id) {
301	9170x	char <- get(id)
302	9170x	if (any(nchar(char) >= 1)) {
303	4649x	return(sprintf("%s {\n%s\n}\n\n", stan_blocks[[id]], paste(char, collapse = "\n")))
304		} else {
305	4521x	return("")
306		}
307		}
308		)
309	1310x	return(paste0(block_strings, collapse = ""))
310		}
311
312		# as_stan_fragments ----
313
314		#' Conversion of Character Vector into Stan Code Block List
315		#'
316		#' @param x (`character`)\cr the single Stan code vector.
317		#' @param stan_blocks (`list`)\cr reference list of stan blocks.
318		#'
319		#' @return A list with the Stan code blocks.
320		#'
321		#' @details
322		#' Function only works if code is in format
323		#' ```
324		#' data {
325		#' <code>
326		#' }
327		#' model {
328		#' <code>
329		#' }
330		#' ```
331		#' That is to say we do not support code in inline format i.e.
332		#' ```
333		#' data { <code> }
334		#' model { <code> }
335		#' ```
336		#'
337		#' @keywords internal
338		as_stan_fragments <- function(x, stan_blocks = STAN_BLOCKS) {
339	2585x	code <- unlist(stringr::str_split(x, "\n"))
340
341	2585x	errmsg <- paste(
342	2585x	"There were problems parsing the `%s` block.",
343	2585x	"Please consult the `Formatting Stan Files` section of the",
344	2585x	"`Extending jmpost` vignette"
345		)
346
347		# Check to see if any block openings exist that have code on the same line
348		# e.g. `data { int i;}`. This is unsupported so we throw an error
349	2585x	for (block in stan_blocks) {
350	18090x	regex <- sprintf("^\\s%s\\s\\{\\s*[^\\s-]+", block)
351	18090x	if (any(grepl(regex, code, perl = TRUE))) {
352	2x	stop(sprintf(errmsg, block))
353		}
354		}
355
356		# We first look to identify the opening of a block e.g. `data {`
357		# We then regard all lines that follow as belonging to that block
358		# until we see another block being opened e.g. `model{`
359	2583x	results <- list()
360	2583x	target <- NULL
361	2583x	for (line in code) {
362	200954x	for (block in names(stan_blocks)) {
363	1383314x	regex <- sprintf("^\\s%s\\s\\{\\s*$", stan_blocks[[block]])
364	1383314x	if (stringr::str_detect(line, regex)) {
365	7489x	target <- block
366	7489x	line <- NULL
367	7489x	break
368		}
369		}
370	200954x	if (!is.null(target)) {
371		# This is memory inefficient but given the relatively small size of
372		# stan files its regarded as a acceptable simplification to ease the
373		# code burden
374	199992x	results[[target]] <- c(results[[target]], line)
375		}
376		}
377
378		# Loop over each block to remove trailing "}".
379	2583x	for (block in names(results)) {
380	7488x	block_length <- length(results[[block]])
381		# The following processing is only required if the block actually has content
382	7488x	if (block_length == 1 && results[[block]] == "") {
383	!	next
384		}
385	7488x	has_removed_char <- FALSE
386		# Walk backwards to find the closing `}` that corresponds to the `<block> {`
387	7488x	for (index in rev(seq_len(block_length))) {
388	16194x	line <- results[[block]][[index]]
389		# This code will exit the for loop as soon as it hits the closing `}`
390		# thus if we ever see a line that ends in text/numbers it means
391		# somethings gone wrong
392	16194x	if (stringr::str_detect(line, "[\\w\\d]+\\s*$")) {
393	1x	stop(sprintf(errmsg, block))
394		}
395	16193x	if (stringr::str_detect(line, "\\}\\s*$")) {
396	7487x	new_line <- stringr::str_replace(line, "\\s\\}\\s$", "")
397		# If the line is now blank after removing the closing `}` then drop the line
398	7487x	keep_offset <- if (nchar(new_line) == 0) -1 else 0
399		# Only keep lines from the start of the block to the closing `}`
400		# this is to ensure we drop blank lines that were between the end
401		# of the block and the start of the next
402	7487x	keep_range <- seq_len(index + keep_offset)
403	7487x	results[[block]][[index]] <- new_line
404	7487x	results[[block]] <- results[[block]][keep_range]
405	7487x	has_removed_char <- TRUE
406	7487x	break
407		}
408		}
409		# If we haven't actually removed a closing `}` then something has gone wrong...
410	7487x	if (!has_removed_char) {
411	!	stop(sprintf(errmsg, block))
412		}
413		}
414
415		# Add any missing blocks back in
416	2582x	for (block in names(stan_blocks)) {
417	18074x	if (is.null(results[[block]])) {
418	10587x	results[[block]] <- ""
419		}
420		}
421	2582x	results
422		}
423
424		#' `StanModule` -> Printable `Character`
425		#'
426		#' Converts [`StanModule`] object into a printable string.
427		#' @param object ([`StanModule`])\cr A stan program
428		#' @family StanModule
429		#' @param indent (`numeric`)\cr how much white space to prefix the print string with.
430		#' @keywords internal
431		#' @export
432		as_print_string.StanModule <- function(object, indent = 1, ...) {
433	1x	slots <- names(getSlots("StanModule"))
434	1x	components <- Filter(\(block) paste(slot(object, block), collapse = "") != "", slots)
435	1x	template <- c(
436	1x	"StanModule Object with components:",
437	1x	paste(" ", components)
438		)
439	1x	pad <- rep(" ", indent) \|> paste(collapse = "")
440	1x	template_padded <- paste(pad, template)
441	1x	sprintf(
442	1x	paste(template_padded, collapse = "\n")
443		)
444		}
445
446		#' @rdname show-object
447		#' @export
448		setMethod(
449		f = "show",
450		signature = "StanModule",
451		definition = function(object) {
452	1x	string <- as_print_string(object)
453	1x	cat("\n", string, "\n\n")
454		}
455		)

1		#' @include generics.R
2		#' @include utilities.R
3		NULL
4
5		#' Re-used documentation for `DataLongitudinal`
6		#'
7		#' @param object ([`DataLongitudinal`]) \cr Longitudinal Data.
8		#' @param x ([`DataLongitudinal`]) \cr Longitudinal Data.
9		#' @param ... Not Used.
10		#'
11		#' @name DataLongitudinal-Shared
12		#' @keywords internal
13		NULL
14
15
16		# DataLongitudinal-class ----
17
18		#' Longitudinal Data Object and Constructor Function
19		#'
20		#' The [`DataLongitudinal`] class handles the processing of the longitudinal data for fitting a Joint Model.
21		#'
22		#'
23		#' @slot data (`data.frame`)\cr See Arguments for details; Note that
24		#' observations that contain missing values in the required variables are removed.
25		#' @slot formula (`formula`)\cr See Arguments for details
26		#' @slot threshold (`numeric`)\cr See Arguments for details
27		#'
28		#' @family DataObjects
29		#' @family DataLongitudinal
30		#' @export DataLongitudinal
31		#' @exportClass DataLongitudinal
32		.DataLongitudinal <- setClass(
33		Class = "DataLongitudinal",
34		representation = list(
35		data = "data.frame",
36		formula = "formula",
37		threshold = "numeric_or_NULL"
38		)
39		)
40
41
42		#' @param data (`data.frame`)\cr containing the observed longitudinal data.
43		#' @param formula (`formula`)\cr of the form `outcome ~ time`, and cannot contain any additional covariates.
44		#' @param threshold (`numeric`)\cr cut-off value to be used to declare an observation as censored
45		#' (below detection limit).
46		#' @rdname DataLongitudinal-class
47		DataLongitudinal <- function(data, formula, threshold = NULL) {
48	45x	.DataLongitudinal(
49	45x	data = remove_missing_rows(data, formula),
50	45x	formula = formula,
51	45x	threshold = threshold
52		)
53		}
54
55		setValidity(
56		"DataLongitudinal",
57		function(object) {
58		if (!length(object@formula) == 3) {
59		return("`formula` should be a 2 sided formula")
60		}
61		if (!length(object@formula[[3]]) == 1) {
62		return("the RHS of `formula` should only have 1 value")
63		}
64		if (!length(object@threshold) <= 1) {
65		return("`threshold` must be of length 1 or `NULL`")
66		}
67		vars <- extractVariableNames(object)
68		vars$threshold <- NULL
69		vars$frm <- NULL
70		for (i in unlist(vars)) {
71		if (! i %in% names(object@data)) {
72		return(sprintf("Variable `%s` is not in data", i))
73		}
74		}
75		return(TRUE)
76		}
77		)
78
79
80
81
82		#' @rdname harmonise
83		harmonise.DataLongitudinal <- function(object, subject_var, subject_ord, ...) {
84	23x	data <- as.data.frame(object)
85	23x	vars <- extractVariableNames(object)
86	23x	assert_string(subject_var, na.ok = FALSE)
87	23x	assert_character(subject_ord, any.missing = FALSE)
88	23x	assert_that(
89	23x	subject_var %in% names(data),
90	23x	msg = sprintf("Subject variable `%s` not found in `longitudinal`", subject_var)
91		)
92	22x	assert_that(
93	22x	all(data[[subject_var]] %in% subject_ord),
94	22x	msg = "There are subjects in `longitudinal` that are not present in `subjects`"
95		)
96	21x	assert_that(
97	21x	all(subject_ord %in% data[[subject_var]]),
98	21x	msg = "There are subjects in `subjects` that are not present in `longitudinal`"
99		)
100	20x	data[[subject_var]] <- factor(
101	20x	as.character(data[[subject_var]]),
102	20x	levels = subject_ord
103		)
104	20x	data_re_ord <- order(
105	20x	data[[subject_var]],
106	20x	data[[vars$time]],
107	20x	data[[vars$outcome]]
108		)
109	20x	data_ord <- data[data_re_ord, ]
110	20x	DataLongitudinal(
111	20x	data = data_ord,
112	20x	formula = object@formula,
113	20x	threshold = object@threshold
114		)
115		}
116
117
118
119		#' `DataLongitudinal` -> `data.frame`
120		#'
121		#' @inheritParams DataLongitudinal-Shared
122		#'
123		#' @description
124		#' Converts a [`DataLongitudinal`] object into a `data.frame`.
125		#' The subject variable is cast to factor.
126		#' @family DataLongitudinal
127		#' @export
128		as.data.frame.DataLongitudinal <- function(x, ...) {
129	325x	x <- x@data
130	325x	rownames(x) <- NULL
131	325x	x
132		}
133
134
135
136
137		#' @inheritParams DataLongitudinal-Shared
138		#' @inherit extractVariableNames description title
139		#'
140		#' @returns
141		#' A list with the following named elements:
142		#' - `subject` (`character`)\cr The name of the variable containing the subject identifier
143		#' - `frm` (`formula`)\cr of the form `outcome ~ time`
144		#' - `time` (`character`)\cr The name of the variable containing the outcome time
145		#' - `outcome` (`character`)\cr The name of the variable containing the outcome values
146		#' - `threshold` (`numeric`)\cr cut-off value to be used to declare an observation as censored
147		#' (below detection limit).
148		#' @family DataLongitudinal
149		#' @family extractVariableNames
150		extractVariableNames.DataLongitudinal <- function(object) {
151	349x	list(
152	349x	frm = object@formula,
153	349x	time = as.character(object@formula[[3]]),
154	349x	outcome = as.character(object@formula[[2]]),
155	349x	threshold = object@threshold
156		)
157		}
158
159
160		#' @rdname as_stan_list.DataObject
161		#' @family DataLongitudinal
162		#' @export
163		as_stan_list.DataLongitudinal <- function(object, subject_var, ...) {
164
165	279x	df <- as.data.frame(object)
166	279x	vars <- extractVariableNames(object)
167
168	279x	assert_factor(df[[subject_var]])
169
170	279x	mat_sld_index <- stats::model.matrix(
171	279x	stats::as.formula(paste("~ -1 + ", subject_var)),
172	279x	data = df
173		) \|>
174	279x	t()
175
176	279x	adj_threshold <- if (is.null(vars$threshold)) {
177	56x	-999999
178		} else {
179	223x	vars$threshold
180		}
181
182	279x	index_obs <- which(df[[vars$outcome]] >= adj_threshold)
183	279x	index_cen <- which(df[[vars$outcome]] < adj_threshold)
184
185	279x	sparse_mat_inds_all_y <- rstan::extract_sparse_parts(mat_sld_index)
186	279x	sparse_mat_inds_obs_y <- rstan::extract_sparse_parts(mat_sld_index[, index_obs])
187	279x	sparse_mat_inds_cens_y <- rstan::extract_sparse_parts(mat_sld_index[, index_cen])
188
189	279x	model_data <- list(
190	279x	n_tumour_all = nrow(df),
191
192		# Number of individuals and tumour assessments.
193	279x	n_tumour_obs = length(index_obs),
194	279x	n_tumour_cens = length(index_cen),
195
196		# Index vectors
197	279x	subject_tumour_index = as.numeric(df[[subject_var]]),
198	279x	subject_tumour_index_obs = index_obs,
199	279x	subject_tumour_index_cens = index_cen,
200
201	279x	tumour_value = df[[vars$outcome]],
202	279x	tumour_time = df[[vars$time]],
203	279x	tumour_value_lloq = adj_threshold,
204
205		# Sparse matrix parameters
206		# Matrix of individuals x observed tumour assessments.
207	279x	n_mat_inds_obs_y = c(
208	279x	length(sparse_mat_inds_obs_y$w),
209	279x	length(sparse_mat_inds_obs_y$v),
210	279x	length(sparse_mat_inds_obs_y$u)
211		),
212	279x	w_mat_inds_obs_y = sparse_mat_inds_obs_y$w,
213	279x	v_mat_inds_obs_y = sparse_mat_inds_obs_y$v,
214	279x	u_mat_inds_obs_y = sparse_mat_inds_obs_y$u,
215
216		# Matrix of individuals x censored tumour assessments.
217	279x	n_mat_inds_cens_y = c(
218	279x	length(sparse_mat_inds_cens_y$w),
219	279x	length(sparse_mat_inds_cens_y$v),
220	279x	length(sparse_mat_inds_cens_y$u)
221		),
222	279x	w_mat_inds_cens_y = sparse_mat_inds_cens_y$w,
223	279x	v_mat_inds_cens_y = sparse_mat_inds_cens_y$v,
224	279x	u_mat_inds_cens_y = sparse_mat_inds_cens_y$u,
225
226		# Matrix of all individuals tumour assessments
227	279x	n_mat_inds_all_y = c(
228	279x	length(sparse_mat_inds_all_y$w),
229	279x	length(sparse_mat_inds_all_y$v),
230	279x	length(sparse_mat_inds_all_y$u)
231		),
232	279x	w_mat_inds_all_y = sparse_mat_inds_all_y$w,
233	279x	v_mat_inds_all_y = sparse_mat_inds_all_y$v,
234	279x	u_mat_inds_all_y = sparse_mat_inds_all_y$u
235
236		)
237
238	279x	return(model_data)
239		}
240
241		#' @rdname as_stan_list.DataObject
242		#' @export
243		as.list.DataLongitudinal <- function(x, ...) {
244	!	as_stan_list(x, ...)
245		}
246
247
248		#' `DataLongitudinal` -> Printable `Character`
249		#'
250		#' Converts [`DataLongitudinal`] object into a printable string.
251		#' @inheritParams DataLongitudinal-Shared
252		#' @family DataLongitudinal
253		#' @param indent (`numeric`)\cr how much white space to prefix the print string with.
254		#' @keywords internal
255		#' @export
256		as_print_string.DataLongitudinal <- function(object, indent = 1, ...) {
257	2x	template <- c(
258	2x	"Longitudinal-Data Object:",
259	2x	" # of Rows = %d",
260	2x	" # of Columns = %d",
261	2x	" # of Cen-Obvs = %d",
262	2x	" Formula = %s"
263		)
264	2x	pad <- rep(" ", indent) \|> paste(collapse = "")
265	2x	template_padded <- paste(pad, template)
266	2x	vars <- extractVariableNames(object)
267	2x	sprintf(
268	2x	paste(template_padded, collapse = "\n"),
269	2x	nrow(object@data),
270	2x	ncol(object@data),
271	2x	sum(object@data[[vars$outcome]] < vars$threshold),
272	2x	Reduce(paste, deparse(vars$frm))
273		)
274		}
275
276		#' @rdname show-object
277		#' @export
278		setMethod(
279		f = "show",
280		signature = "DataLongitudinal",
281		definition = function(object) {
282	1x	string <- as_print_string(object)
283	1x	cat("\n", string, "\n\n")
284		}
285		)

1		#' @include DataSurvival.R
2		#' @include DataLongitudinal.R
3		#' @include DataSubject.R
4		NULL
5
6
7		#' Re-used documentation for `DataJoint`
8		#'
9		#' @param object ([`DataJoint`]) \cr Survival and Longitudinal Data.
10		#' @param x ([`DataJoint`]) \cr Survival and Longitudinal Data.
11		#' @param ... Not Used.
12		#'
13		#' @name DataJoint-Shared
14		#' @keywords internal
15		NULL
16
17		setClassUnion("DataLongitudinal_or_NULL", c("DataLongitudinal", "NULL"))
18		setClassUnion("DataSurvival_or_NULL", c("DataSurvival", "NULL"))
19
20
21		# DataJoint-class ----
22
23		#' @title
24		#' Joint Data Object and Constructor Function
25		#'
26		#' @description
27		#' The `DataJoint` class handles combining data from a [`DataSurvival`] object and a
28		#' [`DataLongitudinal`] object.
29		#'
30		#' @slot subject (`DataSubject`)\cr See Argument for details.
31		#' @slot survival (`DataSurvival`)\cr See Argument for details.
32		#' @slot longitudinal (`DataLongitudinal`)\cr See Argument for details.
33		#'
34		#' @family DataObjects
35		#' @family DataJoint
36		#' @export DataJoint
37		#' @exportClass DataJoint
38		.DataJoint <- setClass(
39		Class = "DataJoint",
40		representation = list(
41		subject = "DataSubject",
42		survival = "DataSurvival_or_NULL",
43		longitudinal = "DataLongitudinal_or_NULL"
44		)
45		)
46
47		#' @param subject (`DataSubject`)\cr object created by [DataSubject()].
48		#' @param survival (`DataSurvival`)\cr object created by [DataSurvival()].
49		#' @param longitudinal (`DataLongitudinal`)\cr object created by [DataLongitudinal()].
50		#' @rdname DataJoint-class
51		DataJoint <- function(subject, survival = NULL, longitudinal = NULL) {
52
53	32x	subject_suited <- harmonise(subject)
54	31x	vars <- extractVariableNames(subject)
55	31x	subject_var <- vars$subject
56	31x	subject_ord <- levels(as.data.frame(subject_suited)[[vars$subject]])
57
58	31x	survival_suited <- harmonise(
59	31x	survival,
60	31x	subject_var = subject_var,
61	31x	subject_ord = subject_ord
62		)
63
64	28x	longitudinal_suited <- harmonise(
65	28x	longitudinal,
66	28x	subject_var = subject_var,
67	28x	subject_ord = subject_ord
68		)
69
70	25x	.DataJoint(
71	25x	subject = subject_suited,
72	25x	survival = survival_suited,
73	25x	longitudinal = longitudinal_suited
74		)
75		}
76
77
78
79		setValidity(
80		Class = "DataJoint",
81		method = function(object) {
82		vars <- extractVariableNames(object@subject)
83		subject_var <- vars$subject
84		subject_ord <- as.character(as.data.frame(object@subject)[[vars$subject]])
85		if (!is.null(object@survival)) {
86		survival_df <- as.data.frame(object@survival)
87		if (!subject_var %in% names(survival_df)) {
88		return(sprintf("Unable to find `%s` in `survival`", sujbect_var))
89		}
90		if (!all(survival_df[[subject_var]] %in% subject_ord)) {
91		return("There are subjects in `survival` that are not in `subject`")
92		}
93		if (!nrow(survival_df) == length(unique(survival_df[[subject_var]]))) {
94		return("There are duplicate subjects in `survival`")
95		}
96		}
97		if (!is.null(object@longitudinal)) {
98		long_df <- as.data.frame(object@longitudinal)
99		if (!subject_var %in% names(long_df)) {
100		return(sprintf("Unable to find `%s` in `longitudinal`", sujbect_var))
101		}
102		if (!all(long_df[[subject_var]] %in% subject_ord)) {
103		return("There are subjects in `longitudinal` that are not in `subject`")
104		}
105		}
106		subject_df <- as.data.frame(object@subject)
107		if (!subject_var %in% names(subject_df)) {
108		return(sprintf("Unable to find `%s` in `subject`", sujbect_var))
109		}
110		if (!nrow(subject_df) == length(unique(subject_df[[subject_var]]))) {
111		return("There are duplicate subjects in `subject`")
112		}
113		return(TRUE)
114		}
115		)
116
117
118		# DataJoint-as.list ----
119
120
121
122
123		#' Data Object -> `list`
124		#'
125		#' @param object (`DataSubject` or `DataLongitudinal` or `DataSurvival`) \cr
126		#' data object to convert to a `list`.
127		#' @param x (`DataSubject` or `DataLongitudinal` or `DataSurvival`) \cr
128		#' data object to convert to a `list`.
129		#' @param subject_var (`character`) \cr the name of the variable
130		#' containing the subject identifier.
131		#' @param ... not used.
132		#'
133		#' @description
134		#' Coerces a data object into a `list` of data components required
135		#' for fitting a [`JointModel`]. See the "Extending jmpost" vignette for more details.
136		#'
137		#' @name as_stan_list.DataObject
138		#' @family as_stan_list
139		#' @family DataJoint
140		#' @export
141		as_stan_list.DataJoint <- function(object, ...) {
142	292x	vars <- extractVariableNames(object@subject)
143	292x	subject_var <- vars$subject
144	292x	as_stan_list(object@subject) \|>
145	292x	append(as_stan_list(object@survival)) \|>
146	292x	append(as_stan_list(
147	292x	object@longitudinal,
148	292x	subject_var = subject_var
149		))
150		}
151
152		#' @rdname as_stan_list.DataObject
153		#' @export
154		as.list.DataJoint <- function(x, ...) {
155	241x	as_stan_list(x, ...)
156		}
157
158
159
160		#' Subsetting `DataJoint` as a `data.frame`
161		#'
162		#' @param x (`DataJoint`) \cr object created by [DataJoint()].
163		#' @param subjects (`character` or `list`)\cr subjects that you wish to subset the `data.frame`
164		#' to contain. See details.
165		#' @param ... Not used.
166		#'
167		#' @description
168		#'
169		#' Coerces the object into a `data.frame` containing just event times and status
170		#' filtering for specific subjects If `subjects` is a list then an additional variable `group` will be added
171		#' onto the dataset specifying which group the row belongs to.
172		#'
173		#' @examples
174		#' \dontrun{
175		#' subjects <- c("SUB1", "SUB2", "SUB3", "SUB4")
176		#' subset(x, subjects)
177		#'
178		#' groups <- list(
179		#' "g1" = c("SUB1", "SUB3", "SUB4"),
180		#' "g2" = c("SUB2", "SUB3")
181		#' )
182		#' subset(x, groups)
183		#' }
184		#' @family DataJoint
185		#' @export
186		subset.DataJoint <- function(x, subjects, ...) {
187	2x	data <- as.list(x)
188	2x	dat <- data.frame(
189	2x	time = data[["event_times"]],
190	2x	event = as.numeric(seq_along(data[["event_times"]]) %in% data[["subject_event_index"]]),
191	2x	subject = names(data[["subject_to_index"]])
192		)
193	2x	subset_and_add_grouping(dat, subjects)
194		}
195
196
197
198		#' `subset_and_add_grouping`
199		#'
200		#' @param dat (`data.frame`) \cr must have a column called `subject` which corresponds to the
201		#' values passed to `groupings`.
202		#' @param groupings (`character` or `list`)\cr subjects that you wish to subset the dataset
203		#' to contain. If `groupings` is a list then an additional variable `group` will be added
204		#' onto the dataset specifying which group the row belongs to.
205		#'
206		#' @details
207		#' Example of usage
208		#' ```
209		#' subjects <- c("SUB1", "SUB2", "SUB3", "SUB4")
210		#' subset_and_add_grouping(dat, subjects)
211		#'
212		#' groups <- list(
213		#' "g1" = c("SUB1", "SUB3", "SUB4"),
214		#' "g2" = c("SUB2", "SUB3")
215		#' )
216		#' subset_and_add_grouping(dat, groups)
217		#' ```
218		#'
219		#' @keywords internal
220		subset_and_add_grouping <- function(dat, groupings) {
221	8x	groupings <- decompose_subjects(groupings, dat$subject)$groups
222	5x	dat_subset_list <- lapply(
223	5x	seq_along(groupings),
224	5x	\(i) {
225	13x	dat_reduced <- dat[dat$subject %in% groupings[[i]], , drop = FALSE]
226	13x	dat_reduced[["group"]] <- names(groupings)[[i]]
227	13x	dat_reduced
228		}
229		)
230	5x	x <- Reduce(rbind, dat_subset_list)
231	5x	row.names(x) <- NULL
232	5x	x
233		}
234
235
236		#' Extract Observed Longitudinal Values
237		#'
238		#' Utility function to extract the observed longitudinal values from a [`DataJoint`] object
239		#' @param object ([`DataJoint`])\cr data used to fit a [`JointModel`].
240		#' @return A data.frame with the following columns
241		#' - `subject` (`character`)\cr The subject identifier
242		#' - `time` (`numeric`)\cr The time at which the observation occurred
243		#' - `Yob` (`numeric`)\cr The observed value
244		#' @keywords internal
245		extract_observed_values <- function(object) {
246	2x	assert_class(object, "DataJoint")
247	2x	data <- as.list(object)
248	2x	x <- data.frame(
249	2x	subject = names(data$subject_to_index)[data$subject_tumour_index],
250	2x	time = data$tumour_time,
251	2x	Yob = data$tumour_value
252		)
253	2x	row.names(x) <- NULL
254	2x	x
255		}
256
257		#' @rdname show-object
258		#' @export
259		setMethod(
260		f = "show",
261		signature = "DataJoint",
262		definition = function(object) {
263	1x	string_survival <- if (is.null(object@survival)) {
264	!	" Survival-Data Object:\n Not Specified"
265		} else {
266	1x	as_print_string(object@survival, indent = 5)
267		}
268
269	1x	string_longitudinal <- if (is.null(object@longitudinal)) {
270	!	" Longitudinal-Data Object:\n Not Specified"
271		} else {
272	1x	as_print_string(object@longitudinal, indent = 5)
273		}
274
275	1x	template <- c(
276	1x	"Joint-Data Object Containing:",
277		"",
278	1x	as_print_string(object@subject, indent = 5),
279		"",
280	1x	string_survival,
281		"",
282	1x	string_longitudinal
283		)
284	1x	cat("\n", paste(template, collapse = "\n"), "\n\n")
285		}
286		)

1
2		#' Re-used documentation for `Quantities`
3		#'
4		#' @param x ([`Quantities`]) \cr generated quantities.
5		#' @param object ([`Quantities`]) \cr generated quantities.
6		#' @param type (`character`)\cr sets the `type` variable.
7		#' @param conf.level (`numeric`) \cr confidence level of the interval.
8		#' @param ... not used.
9		#'
10		#' @keywords internal
11		#' @name Quantities-Shared
12		NULL
13
14
15		#' Generated Quantities Container
16		#'
17		#' A simple wrapper around a `matrix` to store required metadata
18		#'
19		#' @param quantities (`matrix`)\cr of generated quantities.
20		#' @param times (`numeric`)\cr labels specifying which time point the quantity was generated at.
21		#' @param groups (`character`)\cr labels for which group the quantity belongs to.
22		#'
23		#' @slot quantities (`matrix`)\cr See Arguments for details.
24		#' @slot times (`numeric`)\cr See Arguments for details.
25		#' @slot groups (`character`)\cr See Arguments for details.
26		#'
27		#' @details
28		#' Each row of the matrix represents a sample and each column represents a distinct quantity.
29		#' As such the number of columns in the matrix should equal the length of `times` and `groups`
30		#' which provide metadata for who the quantity belongs to and at what time point it was generated at.
31		#'
32		#' @keywords internal
33		#' @name Quantities-class
34		#' @family Quantities
35		.Quantities <- setClass(
36		"Quantities",
37		slots = list(
38		"quantities" = "matrix",
39		"times" = "numeric",
40		"groups" = "character"
41		)
42		)
43		#' @rdname Quantities-class
44		Quantities <- function(quantities, times, groups) {
45	34x	.Quantities(
46	34x	quantities = quantities,
47	34x	times = times,
48	34x	groups = groups
49		)
50		}
51
52		setValidity(
53		Class = "Quantities",
54		method = function(object) {
55		if (length(object@times) != length(object@groups)) {
56		return("Length of `times` must be equal to the length of `groups`")
57		}
58		if (length(object@times) != ncol(object@quantities)) {
59		return("Length of `times` must be equal to the number of columns in `quantities`")
60		}
61		return(TRUE)
62		}
63		)
64
65
66
67		#' @export
68		dim.Quantities <- function(x) {
69	9x	dim(x@quantities)
70		}
71
72
73		#' `Quantities` -> `data.frame`
74		#'
75		#' @inheritParams Quantities-Shared
76		#'
77		#' @keywords internal
78		#' @family Quantities
79		#' @export
80		as.data.frame.Quantities <- function(x, ...) {
81	2x	data.frame(
82	2x	group = rep(x@groups, each = nrow(x@quantities)),
83	2x	time = rep(x@times, each = nrow(x@quantities)),
84	2x	values = as.vector(x@quantities)
85		)
86		}
87
88
89		#' summary
90		#'
91		#' @description
92		#' This method returns a summary statistic `data.frame` of the quantities. Note that this
93		#' is just an internal utility method in order to share common code between
94		#' [LongitudinalQuantities] and [SurvivalQuantities]
95		#'
96		#' @inheritParams Quantities-Shared
97		#'
98		#' @returns
99		#' A `data.frame` with the following variables:
100		#' - `median` (`numeric`) \cr the median value of the quantity.
101		#' - `lower` (`numeric`) \cr the lower CI value of the quantity.
102		#' - `upper` (`numeric`) \cr the upper CI value of the quantity.
103		#' - `time` (`numeric`) \cr the time point which the quantity is for.
104		#' - `group` (`character`) \cr which group the quantity belongs to.
105		#' - `type` (`character`) \cr what type of quantity is it.
106		#'
107		#' @keywords internal
108		#' @family Quantities
109		#' @export
110		summary.Quantities <- function(object, conf.level = 0.95, ...) {
111	32x	quantities_summarised <- samples_median_ci(
112	32x	object@quantities,
113	32x	level = conf.level
114		)
115
116	32x	quantities_summarised$group <- object@groups
117	32x	quantities_summarised$time <- object@times
118	32x	quantities_summarised[, c("group", "time", "median", "lower", "upper")]
119		}
120
121
122
123		#' Create Grouped Quantities
124		#'
125		#' This function takes a matrix of quantity samples and aggregates them by calculating the
126		#' pointwise average.
127		#'
128		#' @param quantities_raw (`matrix`)\cr of samples with 1 row per sample and 1 column per
129		#' distinct quantity.
130		#'
131		#' @param collapser ([`QuantityCollapser`])\cr specifies which columns to combine together.
132		#'
133		#' @details
134		#' This function essentially implements the group wise average by collapsing multiple columns
135		#' together based on the specification provided by the `QuantityCollapser` object.
136		#' The [Grid-Dev] page provides an example of what this function implements
137		#'
138		#' @keywords internal
139		collapse_quantities <- function(quantities_raw, collapser) {
140	34x	assert_class(quantities_raw, "matrix")
141	34x	assert_class(collapser, "QuantityCollapser")
142
143	34x	quantities <- matrix(
144	34x	NA,
145	34x	nrow = nrow(quantities_raw),
146	34x	ncol = length(collapser)
147		)
148
149	34x	for (idx in seq_len(length(collapser))) {
150	12704x	quantities[, idx] <- quantities_raw[
151		,
152	12704x	collapser@indexes[[idx]],
153	12704x	drop = FALSE
154	12704x	] \|> rowMeans()
155		}
156
157	34x	return(quantities)
158		}
159
160		#' Extract Survival Quantities
161		#'
162		#' Utility function to extract generated quantities from a [cmdstanr::CmdStanGQ] object.
163		#' Multiple quantities are generated by default so this is a convenience function to extract
164		#' the desired ones and return them them as a user friendly [posterior::draws_matrix] object
165		#'
166		#' @param gq (`CmdStanGQ`) \cr a [cmdstanr::CmdStanGQ] object created by [generateQuantities()].
167		#' @param type (`character`)\cr quantity to be generated.
168		#' Must be one of `surv`, `haz`, `loghaz`, `cumhaz`, `lm_identity`.
169		#' @keywords internal
170		extract_quantities <- function(gq, type = c("surv", "haz", "loghaz", "cumhaz", "lm_identity")) {
171	35x	type <- match.arg(type)
172	35x	assert_class(gq, "CmdStanGQ")
173	35x	meta <- switch(type,
174	35x	surv = list("log_surv_fit_at_time_grid", exp),
175	35x	cumhaz = list("log_surv_fit_at_time_grid", \(x) -x),
176	35x	haz = list("log_haz_fit_at_time_grid", exp),
177	35x	loghaz = list("log_haz_fit_at_time_grid", identity),
178	35x	lm_identity = list("y_fit_at_time_grid", identity)
179		)
180	35x	result <- gq$draws(meta[[1]], format = "draws_matrix")
181	35x	result_transformed <- meta[[2]](result)
182	35x	cnames <- colnames(result_transformed)
183	35x	colnames(result_transformed) <- gsub(meta[[1]], "quantity", cnames)
184	35x	result_transformed
185		}
186
187
188		#' `Quantities` -> Printable `Character`
189		#'
190		#' Converts [`Quantities`] object into a printable string.
191		#' @inheritParams Quantities-Shared
192		#' @family Quantities
193		#' @keywords internal
194		#' @export
195		as_print_string.Quantities <- function(object, indent = 1, ...) {
196	1x	template <- c(
197	1x	"Quantities Object:",
198	1x	" # of samples = %d",
199	1x	" # of quantities = %d"
200		)
201	1x	pad <- rep(" ", indent) \|> paste(collapse = "")
202	1x	template_padded <- paste(pad, template)
203	1x	sprintf(
204	1x	paste(template_padded, collapse = "\n"),
205	1x	nrow(object),
206	1x	ncol(object)
207		)
208		}
209
210
211		#' @rdname show-object
212		#' @export
213		setMethod(
214		f = "show",
215		signature = "Quantities",
216		definition = function(object) {
217	1x	string <- as_print_string(object)
218	1x	cat("\n", string, "\n\n")
219		}
220		)

1
2		#' `SimSurvival` Function Arguments
3		#'
4		#' The documentation lists all the conventional arguments for [`SimSurvival`]
5		#' constructors.
6		#'
7		#' @param time_max (`number`)\cr the maximum time to simulate to.
8		#' @param time_step (`number`)\cr the time interval between evaluating the log-hazard function.
9		#' @param lambda_censor (`number`)\cr the censoring rate.
10		#' @param beta_cont (`number`)\cr the continuous covariate coefficient.
11		#' @param beta_cat (`numeric`)\cr the categorical covariate coefficients.
12		#' @param loghazard (`function`)\cr the log hazard function.
13		#' @param name (`character`)\cr the name of the object.
14		#' @param ... Not Used.
15		#'
16		#' @section Hazard Evaluation:
17		#'
18		#' Event times are simulated by sampling a cumulative hazard limit from a \eqn{U(0, 1)} distribution
19		#' for
20		#' each subject and then counting how much hazard they've been exposed to by evaluating the
21		#' log-hazard function at a set interval. The `time_max` argument sets the upper bound for the
22		#' number of time points to evaluate the log-hazard function at with subjects who have not had an
23		#' event being censored at `time_max`. The `time_step` argument sets the interval at which to
24		#' evaluate the log-hazard function. Setting smaller values for `time_step` will increase the
25		#' precision of the simulation at the cost of increased computation time. Likewise, setting large
26		#' values for `time_max` will minimize the number of censored subjects at the cost of
27		#' increased computation time.
28		#'
29		#' @name SimSurvival-Shared
30		#' @keywords internal
31		NULL
32
33
34		#' Abstract Simulation Class for Survival Data
35		#'
36		#' @inheritParams SimSurvival-Shared
37		#' @inheritSection SimSurvival-Shared Hazard Evaluation
38		#'
39		#' @slot time_max (`numeric`)\cr See arguments.
40		#' @slot time_step (`numeric`)\cr See arguments.
41		#' @slot lambda_censor (`numeric`)\cr See arguments.
42		#' @slot beta_cont (`numeric`)\cr See arguments.
43		#' @slot beta_cat (`numeric`)\cr See arguments.
44		#' @slot loghazard (`function`)\cr See arguments.
45		#' @slot name (`character`)\cr See arguments.
46		#'
47		#' @family SimSurvival
48		#' @exportClass SimSurvival
49		#' @name SimSurvival-class
50		.SimSurvival <- setClass(
51		"SimSurvival",
52		slots = c(
53		time_max = "numeric",
54		time_step = "numeric",
55		lambda_censor = "numeric",
56		beta_cont = "numeric",
57		beta_cat = "numeric",
58		loghazard = "function",
59		name = "character"
60		)
61		)
62
63		#' @rdname SimSurvival-class
64		#' @export
65		SimSurvival <- function(
66		time_max = 2000,
67		time_step = 1,
68		lambda_censor = 1 / 3000,
69		beta_cont = 0.2,
70		beta_cat = c("A" = 0, "B" = -0.4, "C" = 0.2),
71		loghazard,
72		name = "SimSurvival"
73		) {
74	24x	.SimSurvival(
75	24x	time_max = time_max,
76	24x	time_step = time_step,
77	24x	lambda_censor = lambda_censor,
78	24x	beta_cont = beta_cont,
79	24x	beta_cat = beta_cat,
80	24x	loghazard = loghazard,
81	24x	name = name
82		)
83		}
84
85
86		#' @rdname show-object
87		#' @export
88		setMethod(
89		f = "show",
90		signature = "SimSurvival",
91		definition = function(object) {
92	1x	x <- sprintf("\nA %s Object\n\n", as_print_string(object))
93	1x	cat(x)
94	1x	return(object)
95		}
96		)
97
98		#' @rdname as_print_string
99		as_print_string.SimSurvival <- function(object) {
100	1x	return(object@name)
101		}
102
103
104		#' Construct Time Intervals
105		#'
106		#' @param object (`SimSurvival`)\cr the survival simulation object to create evaluation points for.
107		#'
108		#' @return A `tibble` with `lower`, `upper`, `time`, `eval` and `width`.
109		#' @keywords internal
110		hazardWindows.SimSurvival <- function(object) {
111	22x	times <- seq(0, object@time_max, object@time_step)
112	22x	bound_lower <- times[-length(times)]
113	22x	bound_upper <- times[-1]
114	22x	bound_width <- bound_upper - bound_lower
115	22x	mid_point <- bound_upper - (bound_width / 2)
116	22x	tibble::tibble(
117	22x	lower = bound_lower,
118	22x	upper = bound_upper,
119	22x	midpoint = mid_point,
120	22x	width = bound_width
121		)
122		}
123
124		#' @rdname sampleSubjects
125		#' @export
126		sampleSubjects.SimSurvival <- function(object, subjects_df) {
127	22x	subjects_df \|>
128	22x	dplyr::mutate(cov_cont = stats::rnorm(dplyr::n())) \|>
129	22x	dplyr::mutate(cov_cat = factor(
130	22x	sample(names(object@beta_cat), replace = TRUE, size = dplyr::n()),
131	22x	levels = names(object@beta_cat)
132		)) \|>
133	22x	dplyr::mutate(log_haz_cov = .data$cov_cont * object@beta_cont + object@beta_cat[.data$cov_cat]) \|>
134	22x	dplyr::mutate(survival = stats::runif(dplyr::n())) \|>
135	22x	dplyr::mutate(chazard_limit = -log(.data$survival)) \|>
136	22x	dplyr::mutate(time_cen = stats::rexp(dplyr::n(), object@lambda_censor))
137		}
138
139
140		#' @rdname sampleObservations
141		#' @export
142		sampleObservations.SimSurvival <- function(object, times_df) {
143
144	22x	assert_that(
145	22x	all(times_df$time >= 0),
146	22x	msg = "All time points must be greater than or equal to 0"
147		)
148
149	22x	os_dat_chaz <- times_df \|>
150	22x	dplyr::mutate(log_bl_haz = object@loghazard(.data$midpoint)) \|>
151		# Fix to avoid issue with log(0) = NaN values
152	22x	dplyr::mutate(log_bl_haz = dplyr::if_else(.data$midpoint == 0, -999, .data$log_bl_haz)) \|>
153	22x	dplyr::mutate(hazard_instant = exp(.data$log_bl_haz + .data$log_haz_cov + .data$log_haz_link)) \|>
154		# Reset Inf values to large number to avoid NaN issues downstream
155		# This is suitable as Hazard limits tend to be in the range of -10 to 10 so large numbers
156		# are essentially equivalent to infinity for simulation purposes
157	22x	dplyr::mutate(hazard_instant = dplyr::if_else(.data$hazard_instant == Inf, 999, .data$hazard_instant)) \|>
158	22x	dplyr::mutate(hazard_instant = dplyr::if_else(.data$hazard_instant == -Inf, -999, .data$hazard_instant)) \|>
159	22x	dplyr::mutate(hazard_interval = .data$hazard_instant * .data$width) \|>
160	22x	dplyr::group_by(.data$subject) \|>
161	22x	dplyr::mutate(chazard = cumsum(.data$hazard_interval)) \|>
162	22x	dplyr::ungroup()
163
164	22x	os_had_event <- os_dat_chaz \|>
165	22x	dplyr::filter(.data$chazard >= .data$chazard_limit) \|>
166	22x	dplyr::group_by(.data$subject) \|>
167	22x	dplyr::slice(1) \|>
168	22x	dplyr::ungroup() \|>
169	22x	dplyr::mutate(event = 1)
170
171	22x	os_had_censor <- os_dat_chaz \|>
172	22x	dplyr::filter(!.data$subject %in% os_had_event$subject) \|>
173	22x	dplyr::group_by(.data$subject) \|>
174	22x	dplyr::slice(dplyr::n()) \|>
175	22x	dplyr::ungroup() \|>
176	22x	dplyr::mutate(event = 0)
177
178	22x	if (!(nrow(os_had_censor) == 0)) {
179	10x	message(sprintf("INFO: %i subjects did not die before max(times)", nrow(os_had_censor)))
180		}
181
182	22x	os_dat_complete <- os_had_event \|>
183	22x	dplyr::bind_rows(os_had_censor) \|>
184	22x	dplyr::mutate(real_time = .data$time) \|>
185	22x	dplyr::mutate(event = dplyr::if_else(.data$real_time <= .data$time_cen, .data$event, 0)) \|>
186	22x	dplyr::mutate(time = dplyr::if_else(.data$real_time <= .data$time_cen, .data$real_time, .data$time_cen)) \|>
187	22x	dplyr::arrange(.data$subject)
188
189	22x	os_dat_complete[, c("subject", "study", "arm", "time", "event", "cov_cont", "cov_cat")]
190		}
191
192
193		#' Simulate Survival Data from a Weibull Proportional Hazard Model
194		#'
195		#' @param lambda (`number`)\cr the scale parameter.
196		#' @param gamma (`number`)\cr the shape parameter.
197		#'
198		#' @inheritParams SimSurvival-Shared
199		#' @inheritSection SimSurvival-Shared Hazard Evaluation
200		#'
201		#' @family SimSurvival
202		#'
203		#' @export
204		SimSurvivalWeibullPH <- function(
205		lambda,
206		gamma,
207		time_max = 2000,
208		time_step = 1,
209		lambda_censor = 1 / 3000,
210		beta_cont = 0.2,
211		beta_cat = c("A" = 0, "B" = -0.4, "C" = 0.2)
212		) {
213	4x	SimSurvival(
214	4x	time_max = time_max,
215	4x	time_step = time_step,
216	4x	lambda_censor = lambda_censor,
217	4x	beta_cont = beta_cont,
218	4x	beta_cat = beta_cat,
219	4x	loghazard = function(time) {
220	4x	log(lambda) + log(gamma) + (gamma - 1) * log(time)
221		},
222	4x	name = "SimSurvivalWeibullPH"
223		)
224		}
225
226		#' Simulate Survival Data from a Log-Logistic Proportional Hazard Model
227		#'
228		#' @param a (`number`)\cr the scale parameter.
229		#' @param b (`number`)\cr the shape parameter.
230		#'
231		#' @inheritParams SimSurvival-Shared
232		#' @inheritSection SimSurvival-Shared Hazard Evaluation
233		#'
234		#' @family SimSurvival
235		#' @export
236		SimSurvivalLogLogistic <- function(
237		a,
238		b,
239		time_max = 2000,
240		time_step = 1,
241		lambda_censor = 1 / 3000,
242		beta_cont = 0.2,
243		beta_cat = c("A" = 0, "B" = -0.4, "C" = 0.2)
244		) {
245	2x	SimSurvival(
246	2x	time_max = time_max,
247	2x	time_step = time_step,
248	2x	lambda_censor = lambda_censor,
249	2x	beta_cont = beta_cont,
250	2x	beta_cat = beta_cat,
251	2x	loghazard = function(time) {
252	2x	c1 <- - log(a) + log(b) + (b - 1) * (- log(a) + log(time))
253	2x	c2 <- log(1 + (time / a)^b)
254	2x	return(c1 - c2)
255		},
256	2x	name = "SimSurvivalLogLogistic"
257		)
258		}
259
260
261
262		#' Simulate Survival Data from a Exponential Proportional Hazard Model
263		#'
264		#' @param lambda (`number`)\cr the rate parameter.
265		#'
266		#' @inheritParams SimSurvival-Shared
267		#' @inheritSection SimSurvival-Shared Hazard Evaluation
268		#'
269		#' @family SimSurvival
270		#'
271		#' @export
272		SimSurvivalExponential <- function(
273		lambda,
274		time_max = 2000,
275		time_step = 1,
276		lambda_censor = 1 / 3000,
277		beta_cont = 0.2,
278		beta_cat = c("A" = 0, "B" = -0.4, "C" = 0.2)
279		) {
280	17x	SimSurvival(
281	17x	time_max = time_max,
282	17x	time_step = time_step,
283	17x	lambda_censor = lambda_censor,
284	17x	beta_cont = beta_cont,
285	17x	beta_cat = beta_cat,
286	17x	loghazard = function(time) {
287	17x	log(lambda)
288		},
289	17x	name = "SimSurvivalExponential"
290		)
291		}

1
2		#' @include DataJoint.R
3		#' @include Quantities.R
4		NULL
5
6
7		#' Re-used documentation for `LongitudinalQuantities`
8		#'
9		#' @param x ([`LongitudinalQuantities`]) \cr longitudinal quantities.
10		#' @param object ([`LongitudinalQuantities`]) \cr longitudinal quantities.
11		#' @param ... not used.
12		#'
13		#' @keywords internal
14		#' @name LongitudinalQuantities-Shared
15		NULL
16
17
18		#' `LongitudinalQuantities` Object & Constructor Function
19		#'
20		#' Constructor function to generate a `LongitudinalQuantities` object.
21		#'
22		#' @details
23		#' Note that unlike [`SurvivalQuantities`], [`LongitudinalQuantities`] does not support
24		#' group aggregation.
25		#'
26		#' @slot quantities (`Quantities`)\cr The sampled quantities. Should contain 1 element per
27		#' element of `group`
28		#'
29		#' @slot data (`DataJoint`)\cr Survival and Longitudinal Data.
30		#'
31		#' @family LongitudinalQuantities
32		#' @name LongitudinalQuantities-class
33		#' @export LongitudinalQuantities
34		.LongitudinalQuantities <- setClass(
35		"LongitudinalQuantities",
36		slots = c(
37		"quantities" = "Quantities",
38		"data" = "DataJoint"
39		)
40		)
41
42		#' @param object ([`JointModelSamples`]) \cr samples as drawn from a Joint Model.
43		#'
44		#' @param grid (`Grid`) \cr object that specifies which subjects and time points to calculate the
45		#' quantities for. See [Grid-Functions].
46		#' @rdname LongitudinalQuantities-class
47		LongitudinalQuantities <- function(
48		object,
49		grid
50		) {
51	13x	assert_class(object, "JointModelSamples")
52	13x	assert_class(grid, "Grid")
53	13x	assert_that(
54	13x	!is(grid, "GridPrediction"),
55	13x	msg = "`GridPrediction` objects are not supported for `LongitudinalQuantities`"
56		)
57
58	13x	time_grid <- seq(
59	13x	from = 0,
60	13x	to = max(as.list(object@data)[["tumour_time"]]),
61	13x	length = 201
62		)
63
64	13x	grid <- coalesceGridTime(grid, time_grid)
65
66	13x	gq <- generateQuantities(
67	13x	object,
68	13x	generator = as.QuantityGenerator(grid, object@data),
69	13x	type = "longitudinal"
70		)
71
72	13x	quantities_raw <- extract_quantities(gq, type = "lm_identity")
73
74	13x	collapser <- as.QuantityCollapser(grid, object@data)
75	13x	quantities <- collapse_quantities(quantities_raw, collapser)
76
77	13x	.LongitudinalQuantities(
78	13x	quantities = Quantities(
79	13x	quantities,
80	13x	groups = collapser@groups,
81	13x	times = collapser@times
82		),
83	13x	data = object@data
84		)
85		}
86
87
88
89		#' `as.data.frame`
90		#'
91		#' @param x ([`LongitudinalQuantities`]) \cr longitudinal quantities.
92		#' @param ... not used.
93		#' @family LongitudinalQuantities
94		#' @export
95		as.data.frame.LongitudinalQuantities <- function(x, ...) {
96	1x	as.data.frame(x@quantities)
97		}
98
99
100
101		#' summary
102		#'
103		#' @description
104		#' This method returns a `data.frame` of the longitudinal quantities.
105		#'
106		#' @param conf.level (`numeric`) \cr confidence level of the interval.
107		#' @inheritParams LongitudinalQuantities-Shared
108		#'
109		#' @family LongitudinalQuantities
110		#' @family summary
111		#' @export
112		summary.LongitudinalQuantities <- function(
113		object,
114		conf.level = 0.95,
115		...
116		) {
117	12x	summary(object@quantities, conf.level = conf.level)
118		}
119
120
121		#' Longitudinal Plot
122		#'
123		#' Internal plotting function to create longitudinal plots
124		#' This function predominately exists to extract core logic into its own function
125		#' to enable easier unit testing.
126		#'
127		#' @param data (`data.frame`)\cr summary statistics for longitudinal
128		#' value estimates. See details.
129		#' @param data_obs (`data.frame`)\cr real observed values to be
130		#' overlaid for reference. See details.
131		#' @param add_ci (`logical`)\cr Should confidence intervals be added? Default = `TRUE`.
132		#' @details
133		#'
134		#' ## `data`
135		#' Should contain the following columns:
136		#' - `time` (`numeric`) \cr time point for the summary statistic.
137		#' - `group` (`character`) \cr the group in which the observation belongs to.
138		#' - `median` (`numeric`) \cr the median value for the summary statistic.
139		#' - `upper` (`numeric`) \cr the upper 95% CI for the summary statistic.
140		#' - `lower` (`numeric`) \cr the lower 95% CI for the summary statistic.
141		#'
142		#' ## `data_obs`
143		#' Should contain the following columns:
144		#' - `time` (`numeric`) \cr the time at which the observed value occurred.
145		#' - `Yob` (`numeric`) \cr the real observed value.
146		#' - `group` (`character`) \cr which group the event belongs to, should correspond to
147		#' values in `data$group`.
148		#' @keywords internal
149		longitudinal_plot <- function(
150		data,
151		data_obs = NULL,
152		add_ci = FALSE
153		) {
154	3x	p <- ggplot() +
155	3x	geom_line(aes(x = .data$time, y = .data$median), data = data) +
156	3x	xlab(expression(t)) +
157	3x	ylab(expression(y)) +
158	3x	facet_wrap(~group) +
159	3x	theme_bw()
160
161	3x	if (add_ci) {
162	2x	p <- p + geom_ribbon(
163	2x	aes(x = .data$time, ymin = .data$lower, ymax = .data$upper),
164	2x	data = data,
165	2x	alpha = 0.3
166		)
167		}
168
169	3x	if (!is.null(data_obs)) {
170	3x	p <- p + geom_point(aes(x = .data$time, y = .data$Yob), data = data_obs)
171		}
172	3x	return(p)
173		}
174
175
176		#' Automatic Plotting for `LongitudinalQuantities`
177		#'
178		#' @param conf.level (`numeric`) \cr confidence level of the interval. If values of `FALSE`,
179		#' `NULL` or `0` are provided then confidence regions will not be added to the plot.
180		#' @inheritParams LongitudinalQuantities-Shared
181		#'
182		#' @family LongitudinalQuantities
183		#' @family autoplot
184		#' @export
185		autoplot.LongitudinalQuantities <- function(object, conf.level = 0.95, ...) {
186	2x	include_ci <- !is.null(conf.level) && is.numeric(conf.level) && conf.level > 0
187		# If CI aren't needed supply a default 0.95 to summary function as it needs
188		# a value to be specified to work.
189	2x	conf.level <- if (include_ci) conf.level else 0.95
190	2x	data_sum <- summary(object, conf.level = conf.level)
191	2x	data_obs <- extract_observed_values(object@data)
192	2x	assert_that(
193	2x	"group" %in% names(data_sum),
194	2x	"subject" %in% names(data_obs)
195		)
196	2x	data_obs$group <- data_obs$subject
197	2x	data_obs <- data_obs[data_obs$group %in% data_sum$group, ]
198	2x	longitudinal_plot(
199	2x	data = data_sum,
200	2x	data_obs = data_obs,
201	2x	add_ci = include_ci
202		)
203		}
204
205
206
207		#' @rdname show-object
208		#' @export
209		setMethod(
210		f = "show",
211		signature = "LongitudinalQuantities",
212		definition = function(object) {
213	1x	template <- c(
214	1x	"LongitudinalQuantities Object:",
215	1x	" # of samples = %d",
216	1x	" # of quantities = %d"
217		)
218	1x	string <- sprintf(
219	1x	paste(template, collapse = "\n"),
220	1x	nrow(object@quantities),
221	1x	ncol(object@quantities)
222		)
223	1x	cat("\n", string, "\n\n")
224		}
225		)

1
2
3		#' Re-used documentation for `DataSubject`
4		#'
5		#' @param object ([`DataSubject`]) \cr subject-level data.
6		#' @param x ([`DataSubject`]) \cr subject-level data.
7		#' @param ... Not Used.
8		#'
9		#' @name DataSubject-Shared
10		#' @keywords internal
11		NULL
12
13
14
15		#' Subject Data Object and Constructor Function
16		#'
17		#' The [`DataSubject`] class handles the processing of the subject data for
18		#' fitting a [`JointModel`].
19		#'
20		#' @slot data (`data.frame`)\cr the subject-level data.
21		#' @slot subject (`character`)\cr the name of the variable containing the subject identifier.
22		#' @slot arm (`character`)\cr the name of the variable containing the arm identifier.
23		#' @slot study (`character`)\cr the name of the variable containing the study identifier.
24		#'
25		#' @family DataObjects
26		#' @family DataSubject
27		#' @exportClass DataSubject
28		#' @export DataSubject
29		.DataSubject <- setClass(
30		Class = "DataSubject",
31		representation = list(
32		data = "data.frame",
33		subject = "character",
34		arm = "character",
35		study = "character"
36		)
37		)
38
39
40		#' @param data (`data.frame`)\cr the subject-level data.
41		#' @param subject (`character`)\cr the name of the variable containing the subject identifier.
42		#' @param arm (`character`)\cr the name of the variable containing the arm identifier.
43		#' @param study (`character`)\cr the name of the variable containing the study identifier.
44		#' @rdname DataSubject-class
45		#' @export
46		DataSubject <- function(data, subject, arm, study) {
47	358x	vars <- c(subject, arm, study)
48	358x	vars_frm_chr <- paste0("~ ", paste(vars, collapse = " + "))
49	358x	.DataSubject(
50	358x	data = remove_missing_rows(data, stats::as.formula(vars_frm_chr)),
51	358x	subject = subject,
52	358x	arm = arm,
53	358x	study = study
54		)
55		}
56
57		setValidity(
58		"DataSubject",
59		method = function(object) {
60		if (length(object@subject) != 1) {
61		return("`subject` must be a length 1 character")
62		}
63		if (length(object@arm) != 1) {
64		return("`arm` must be a length 1 character")
65		}
66		if (length(object@study) != 1) {
67		return("`study` must be a length 1 character")
68		}
69		if (nrow(object@data) == 0) {
70		return("`data` should not have 0 rows")
71		}
72		sbj <- object@data[[object@subject]]
73		if (!(is(sbj, "character") \| is(sbj, "factor"))) {
74		return("`data[[subject]]` should be of type character or factor")
75		}
76		}
77		)
78
79
80
81		#' @inheritParams DataSubject-Shared
82		#' @inherit extractVariableNames description title
83		#'
84		#' @returns
85		#' A list with the following named elements:
86		#' - `subject` (`character`)\cr the name of the variable containing the subject identifier.
87		#' - `arm` (`character`)\cr the name of the variable containing the arm identifier.
88		#' - `study` (`character`) \cr the name of the variable containing the study identifier.
89		#' @family DataSubject
90		#' @family extractVariableNames
91		#' @export
92		#' @keywords internal
93		extractVariableNames.DataSubject <- function(object) {
94	970x	list(
95	970x	subject = object@subject,
96	970x	arm = object@arm,
97	970x	study = object@study
98		)
99		}
100
101
102		#' @rdname as_stan_list.DataObject
103		#' @family DataSubject
104		#' @export
105		as_stan_list.DataSubject <- function(object, ...) {
106	293x	df <- as.data.frame(harmonise(object))
107	293x	vars <- extractVariableNames(object)
108
109	293x	unique_arm_study_combos <- unique(
110	293x	data.frame(
111	293x	arm = as.numeric(df[[vars$arm]]),
112	293x	study = as.numeric(df[[vars$study]])
113		)
114		)
115
116	293x	list(
117	293x	n_subjects = nrow(df),
118	293x	n_studies = length(unique(df[[vars$study]])),
119	293x	n_arms = length(unique(df[[vars$arm]])),
120	293x	subject_study_index = as.numeric(df[[vars$study]]),
121	293x	subject_arm_index = as.numeric(df[[vars$arm]]),
122	293x	subject_to_index = stats::setNames(
123	293x	seq_len(nlevels(df[[vars$subject]])),
124	293x	levels(df[[vars$subject]])
125		),
126	293x	arm_to_index = stats::setNames(
127	293x	seq_len(nlevels(df[[vars$arm]])),
128	293x	levels(df[[vars$arm]])
129		),
130	293x	study_to_index = stats::setNames(
131	293x	seq_len(nlevels(df[[vars$study]])),
132	293x	levels(df[[vars$study]])
133		),
134	293x	pop_arm_index = unique_arm_study_combos$arm,
135	293x	pop_study_index = unique_arm_study_combos$study
136		)
137		}
138
139		#' @rdname as_stan_list.DataObject
140		#' @export
141		as.list.DataSubject <- function(x, ...) {
142	!	as_stan_list(x, ...)
143		}
144
145
146		#' `DataSubject` -> `data.frame`
147		#'
148		#' @inheritParams DataSubject-Shared
149		#'
150		#' @description
151		#' Converts a [`DataSubject`] object into a `data.frame`.
152		#' The subject variable is cast to factor.
153		#' @family DataSubject
154		#' @export
155		as.data.frame.DataSubject <- function(x, ...) {
156	703x	x <- x@data
157	703x	rownames(x) <- NULL
158	703x	x
159		}
160
161
162
163		#' @rdname harmonise
164		harmonise.DataSubject <- function(object, ...) {
165	325x	data <- as.data.frame(object)
166	325x	vars <- extractVariableNames(object)
167	325x	assert_that(
168	325x	vars$subject %in% names(data),
169	325x	vars$arm %in% names(data),
170	325x	vars$study %in% names(data)
171		)
172	325x	assert_character(
173	325x	as.character(data[[vars$subject]]),
174	325x	any.missing = FALSE,
175	325x	unique = TRUE
176		)
177	324x	data[[vars$subject]] <- factor(data[[vars$subject]])
178	324x	data[[vars$arm]] <- factor(data[[vars$arm]])
179	324x	data[[vars$study]] <- factor(data[[vars$study]])
180	324x	data <- data[order(data[[vars$subject]]), ]
181	324x	DataSubject(
182	324x	data = data,
183	324x	subject = object@subject,
184	324x	arm = object@arm,
185	324x	study = object@study
186		)
187		}
188
189		#' `DataSubject` -> Printable `Character`
190		#'
191		#' Converts [`DataSubject`] object into a printable string.
192		#' @inheritParams DataSubject-Shared
193		#' @family DataSubject
194		#' @keywords internal
195		#' @param indent (`numeric`)\cr how much white space to prefix the print string with.
196		#' @export
197		as_print_string.DataSubject <- function(object, indent = 1, ...) {
198	2x	template <- c(
199	2x	"Subject-Data Object:",
200	2x	" # of Subjects = %d",
201	2x	" # of Studies = %d",
202	2x	" # of Arms = %d"
203		)
204	2x	pad <- rep(" ", indent) \|> paste(collapse = "")
205	2x	template_padded <- paste(pad, template)
206	2x	vars <- extractVariableNames(object)
207	2x	sprintf(
208	2x	paste(template_padded, collapse = "\n"),
209	2x	nrow(object@data),
210	2x	length(unique(object@data[[vars$study]])),
211	2x	length(unique(object@data[[vars$arm]]))
212		)
213		}
214
215		#' @rdname show-object
216		#' @export
217		setMethod(
218		f = "show",
219		signature = "DataSubject",
220		definition = function(object) {
221	1x	string <- as_print_string(object)
222	1x	cat("\n", string, "\n\n")
223		}
224		)

1
2
3		#' Re-used documentation for Brier Score components
4		#'
5		#' @param t (`numeric`)\cr timepoints to calculate the desired quantity at.
6		#' @param times (`numeric`)\cr observed times.
7		#' @param events (`numeric`)\cr event indicator for `times`. Either 1 for an event or 0 for censor.
8		#' @param event_offset (`logical`)\cr If `TRUE` then \eqn{G(T_i)} is evaluated at \eqn{G(T_i-)}.
9		#' Setting this as `TRUE` mirrors the implementation of the `{pec}` package.
10		#' @param maintain_cen_order (`logical`)\cr If `TRUE` then, in the case of ties,
11		#' censor times are always considered
12		#' to have occurred after the event times when calculating the "reverse Kaplan-Meier" for the
13		#' IPCW estimates. Setting this to `TRUE` mirrors the implementation of the `{prodlim}`
14		#' package.
15		#' @param ... not used.
16		#'
17		#' @name Brier-Score-Shared
18		#' @keywords internal
19		NULL
20
21
22
23
24		#' Reverse Kaplan-Meier
25		#'
26		#' @inheritParams Brier-Score-Shared
27		#' @description
28		#' Calculates the survival estimates of the censoring distribution
29		#' using the Kaplan-Meier estimate.
30		#' This is primarily used in the calculation of the IPCW estimates.
31		#'
32		#' @details
33		#' With regards to ties between censor and event times; the standard
34		#' approach is to regard events as occurring before censors. However,
35		#' when modelling the censoring distribution we are regarding the
36		#' censors as "events" so which should come first in the case of ties?
37		#'
38		#' The `reverse_km_event_first()` function maintains the rule
39		#' that events always come first even if we are regarding the censors
40		#' as "events". This matches the implementation of
41		#' `prodlim::prodlim(..., reverse = TRUE)`.
42		#'
43		#' The `reverse_km_cen_first()` function provides the alternative
44		#' implementation assuming that in the case of ties the censor "events"
45		#' come before the event "censors". This is essentially a thin wrapper
46		#' around `survival::survfit(Surv(time, 1 - event), ...)`
47		#'
48		#' @name reverse_km
49		#' @keywords internal
50		reverse_km_event_first <- function(t, times, events) {
51	5x	assert_numeric(t, any.missing = FALSE, finite = TRUE)
52	5x	assert_numeric(times, any.missing = FALSE, finite = TRUE)
53	5x	assert_numeric(events, any.missing = FALSE, finite = TRUE)
54	5x	assert_that(
55	5x	length(times) == length(events),
56	5x	all(events == 1 \| events == 0)
57		)
58	5x	events_cen <- 1 - events
59	5x	ord <- order(times, events_cen)
60	5x	times <- times[ord]
61	5x	events <- events[ord]
62	5x	events_cen <- events_cen[ord]
63	5x	cs_events <- cumsum(events)
64	5x	cs_censor <- cumsum(events_cen)
65
66	5x	g_times <- unique(times)
67	5x	g_n_events_cen <- tapply(events_cen, times, sum)
68	5x	g_cs_events_cen <- tapply(cs_censor, times, max)
69	5x	g_cs_events <- tapply(cs_events, times, max)
70	5x	g_is_cen <- tapply(events_cen, times, max)
71
72	5x	nrisk <- length(times) - g_cs_events - g_cs_events_cen + g_n_events_cen
73	5x	surv_interval <- 1 - g_n_events_cen / nrisk
74	5x	surv_interval <- ifelse(nrisk == 0, 1, surv_interval)
75	5x	surv <- cumprod(surv_interval)
76
77	5x	ct <- g_times[which(g_is_cen == 1)]
78	5x	sv <- surv[which(g_is_cen == 1)]
79
80	5x	names(surv) <- NULL
81	5x	names(ct) <- NULL
82	5x	names(sv) <- NULL
83	5x	list(
84	5x	t = t,
85	5x	surv = c(1, sv)[findInterval(t, ct) + 1]
86		)
87		}
88
89
90		#' @rdname reverse_km
91		reverse_km_cen_first <- function(t, times, events) {
92	1x	assert_numeric(t, any.missing = FALSE, finite = TRUE)
93	1x	assert_numeric(times, any.missing = FALSE, finite = TRUE)
94	1x	assert_numeric(events, any.missing = FALSE, finite = TRUE)
95	1x	assert_that(
96	1x	length(times) == length(events),
97	1x	all(events == 1 \| events == 0)
98		)
99	1x	dat <- data.frame(
100	1x	times = times,
101	1x	events = events,
102	1x	cen_events = 1 - events
103		)
104	1x	mod <- survival::survfit(
105	1x	survival::Surv(times, cen_events) ~ 1,
106	1x	data = dat
107		)
108	1x	preds <- summary(mod, times = t[order(t)], extend = TRUE)$surv
109
110	1x	assert_that(
111	1x	length(preds) == length(t)
112		)
113	1x	list(
114	1x	t = t,
115	1x	surv = preds[match_order(t)]
116		)
117		}
118
119		#' Match Order
120		#'
121		#' @param x (`numeric`)\cr a vector for which we want to generate an index so other
122		#' vectors can be put into the same sort order
123		#'
124		#' Assuming we have a vector that is sorted then this function
125		#' will return the index vector to convert that sorted vector
126		#' into the same sort order as the input vector `x`.
127		#' For example let `x = 8, 7, 9 , 7`. If sorted we would get
128		#' `x_sort = 7, 7, 8, 9`. So in order to convert `x_sort`
129		#' back into `x` we'd need an index vector of `3, 1, 4, 2`.
130		#' This function is used to determine that corresponding
131		#' index vector for an arbitrarily sorted vector `x`.
132		#'
133		#' There is no specific handling of ties. It is assuming that in the case
134		#' of ties for `x` that the vector you are re-indexing also has tied values
135		#' thus the specific tied element selection does not matter.
136		#' @keywords internal
137		match_order <- function(x) {
138	2x	order(order(x))
139		}
140
141
142		#' Brier Score
143		#'
144		#' @inheritParams Brier-Score-Shared
145		#'
146		#' @description
147		#' Implements the Brier Score as detailed in \insertCite{blanche2015}{jmpost}
148		#'
149		#' @details
150		#' - `bs_get_squared_dist()` - implements the squared distance part
151		#' of the formula.
152		#' - `bs_get_weights()` - implements the IPCW weighting
153		#'
154		#' @references
155		#' \insertAllCited{}
156		#'
157		#' @keywords internal
158		brier_score <- function(
159		t,
160		times,
161		events,
162		pred_mat,
163		maintain_cen_order = TRUE,
164		event_offset = TRUE
165		) {
166
167	2x	square_diff_mat <- bs_get_squared_dist(
168	2x	t = t,
169	2x	times = times,
170	2x	events = events,
171	2x	pred_mat = pred_mat
172		)
173
174	2x	weight_mat <- bs_get_weights(
175	2x	t = t,
176	2x	times = times,
177	2x	events = events,
178	2x	event_offset = event_offset,
179	2x	maintain_cen_order = maintain_cen_order
180		)
181
182		# the following is a computational shortcut for diag(A %*% B)
183		# as we don't want to compute the off-diagonal entries of the
184		# matrix multiplication
185	2x	x <- colSums(weight_mat * square_diff_mat)
186	2x	names(x) <- t
187	2x	x / length(times)
188		}
189
190
191		#' @rdname brier_score
192		bs_get_squared_dist <- function(t, times, events, pred_mat) {
193
194	3x	assert_numeric(
195	3x	times,
196	3x	finite = TRUE,
197	3x	any.missing = FALSE
198		)
199	3x	assert_numeric(
200	3x	t,
201	3x	finite = TRUE,
202	3x	any.missing = FALSE
203		)
204	3x	assert_numeric(
205	3x	events,
206	3x	finite = TRUE,
207	3x	any.missing = FALSE,
208	3x	lower = 0,
209	3x	upper = 1
210		)
211	3x	assert_matrix(
212	3x	pred_mat,
213	3x	any.missing = FALSE,
214	3x	nrows = length(times),
215	3x	ncols = length(t)
216		)
217	3x	assert_that(
218	3x	length(events) == length(times)
219		)
220
221
222	3x	expected_mat <- mapply(
223	3x	\(ti, event) (ti <= t) * event * 1,
224	3x	ti = times,
225	3x	event = events,
226	3x	SIMPLIFY = FALSE
227		) \|>
228	3x	unlist() \|>
229	3x	matrix(ncol = length(t), byrow = TRUE)
230
231	3x	assert_that(
232	3x	nrow(pred_mat) == nrow(expected_mat),
233	3x	ncol(pred_mat) == ncol(expected_mat)
234		)
235
236	3x	(expected_mat - pred_mat)^2
237		}
238
239
240		#' @rdname brier_score
241		bs_get_weights <- function(
242		t,
243		times,
244		events,
245		event_offset = TRUE,
246		maintain_cen_order = TRUE
247		) {
248	3x	assert_numeric(
249	3x	times,
250	3x	finite = TRUE,
251	3x	any.missing = FALSE
252		)
253	3x	assert_numeric(
254	3x	t,
255	3x	finite = TRUE,
256	3x	any.missing = FALSE
257		)
258	3x	assert_numeric(
259	3x	events,
260	3x	finite = TRUE,
261	3x	any.missing = FALSE,
262	3x	lower = 0,
263	3x	upper = 1
264		)
265	3x	assert_flag(event_offset, na.ok = FALSE, null.ok = FALSE)
266	3x	assert_flag(maintain_cen_order, na.ok = FALSE, null.ok = FALSE)
267	3x	n_col <- length(t)
268	3x	n_row <- length(times)
269
270	3x	reverse_km <- if (maintain_cen_order) reverse_km_event_first else reverse_km_cen_first
271	3x	offset <- if (event_offset) -.Machine$double.eps^(1 / 2) else 0
272
273	3x	censor_dist_t <- reverse_km(t, times, events)
274	3x	weight_mat_t <- 1 / matrix(
275	3x	rep(censor_dist_t$surv, n_row),
276	3x	nrow = n_row,
277	3x	byrow = TRUE
278		)
279
280	3x	censor_dist_ti <- reverse_km(times + offset, times, events)
281	3x	weight_mat_ti <- 1 / matrix(
282	3x	rep(censor_dist_ti$surv, n_col),
283	3x	nrow = n_row
284		)
285
286	3x	indicator_mat_t <- mapply(
287	3x	\(ti) (ti > t) * 1,
288	3x	ti = times,
289	3x	SIMPLIFY = FALSE
290		) \|>
291	3x	unlist() \|>
292	3x	matrix(ncol = n_col, byrow = TRUE)
293
294	3x	indicator_mat_ti <- mapply(
295	3x	\(ti, event) (ti <= t) * event * 1,
296	3x	ti = times,
297	3x	event = events,
298	3x	SIMPLIFY = FALSE
299		) \|>
300	3x	unlist() \|>
301	3x	matrix(ncol = n_col, byrow = TRUE)
302
303	3x	assert_that(
304	3x	all(indicator_mat_t + indicator_mat_ti <= 1)
305		)
306
307	3x	weight_mat_t[indicator_mat_t == 0] <- 0
308	3x	weight_mat_ti[indicator_mat_ti == 0] <- 0
309
310	3x	(weight_mat_t + weight_mat_ti)
311		}

1		#' Simulating Joint Longitudinal and Time-to-Event Data
2		#'
3		#' @param design (`list`)\cr a list of [`SimGroup`] objects. See details.
4		#' @param longitudinal ([`SimLongitudinal`])\cr object specifying how to simulate the longitudinal data
5		#' @param survival ([`SimSurvival`])\cr object specifying how to simulate the survival data
6		#' @param .silent (`flag`)\cr whether to suppress info messages
7		#'
8		#' @slot longitudinal (`data.frame`)\cr the simulated longitudinal data.
9		#' @slot survival (`data.frame`)\cr the simulated survival data.
10		#'
11		#' @details
12		#'
13		#' The `design` argument is used to specify how many distinct groups should be simulated
14		#' including key information such as the number of subjects within the group as well as
15		#' which treatment arm and study the group belongs to. The `design` argument should be a
16		#' list of [`SimGroup`] objects e.g.
17		#' ```
18		#' design = list(
19		#' SimGroup(n = 50, study = "Study-1", arm = "Arm-A"),
20		#' SimGroup(n = 50, study = "Study-1", arm = "Arm-B")
21		#' )
22		#' ```
23		#'
24		#' @name SimJointData-class
25		#' @exportClass SimJointData
26		.SimJointData <- setClass(
27		"SimJointData",
28		slots = list(
29		longitudinal = "data.frame",
30		survival = "data.frame"
31		)
32		)
33
34
35
36		#' @rdname SimJointData-class
37		#' @export
38		SimJointData <- function(
39		design = list(
40		SimGroup(n = 50, study = "Study-1", arm = "Arm-A"),
41		SimGroup(n = 50, study = "Study-1", arm = "Arm-B")
42		),
43		longitudinal,
44		survival,
45		.silent = FALSE
46		) {
47
48	22x	assert(
49	22x	all(vapply(design, \(x) is(x, "SimGroup"), logical(1))),
50	22x	msg = "All elements of `design` must be of class `SimGroup`"
51		)
52
53	22x	hazard_evaluation_info <- hazardWindows(survival)
54
55	22x	n_group <- vapply(design, function(x) x@n, numeric(1))
56	22x	arms <- vapply(design, function(x) x@arm, character(1))
57	22x	studies <- vapply(design, function(x) x@study, character(1))
58	22x	n_subjects <- sum(n_group)
59	22x	n_times <- length(hazard_evaluation_info$midpoint)
60
61	22x	sprintf_string <- paste0("subject_%0", ceiling(log(n_subjects, 10)) + 1, "i")
62
63	22x	baseline <- dplyr::tibble(subject = sprintf(sprintf_string, seq_len(n_subjects))) \|>
64	22x	dplyr::mutate(arm = factor(rep(arms, times = n_group), levels = unique(arms))) \|>
65	22x	dplyr::mutate(study = factor(rep(studies, times = n_group), levels = unique(studies)))
66
67	22x	os_baseline <- sampleSubjects(survival, subjects_df = baseline)
68	22x	lm_baseline <- sampleSubjects(longitudinal, subjects_df = baseline)
69
70	22x	lm_dat_no_obvs <- lapply(
71	22x	longitudinal@times,
72	22x	\(time) {
73	4544x	baseline[["time"]] <- time
74	4544x	baseline
75		}
76		) \|>
77	22x	dplyr::bind_rows() \|>
78	22x	dplyr::left_join(lm_baseline, by = c("subject", "study", "arm"))
79
80	22x	lm_dat <- sampleObservations(longitudinal, lm_dat_no_obvs)
81
82
83	22x	hazard_eval_df <- dplyr::tibble(
84	22x	subject = rep(lm_baseline$subject, each = n_times),
85	22x	arm = rep(lm_baseline$arm, each = n_times),
86	22x	study = rep(lm_baseline$study, each = n_times),
87	22x	midpoint = rep(as.double(hazard_evaluation_info$midpoint), times = n_subjects),
88	22x	time = rep(as.double(hazard_evaluation_info$upper), times = n_subjects),
89	22x	width = rep(as.double(hazard_evaluation_info$width), times = n_subjects)
90		)
91
92	22x	lm_link_dat <- sampleObservations(
93	22x	longitudinal,
94	22x	dplyr::left_join(hazard_eval_df, lm_baseline, by = c("subject", "study", "arm"))
95	22x	)[, c("subject", "study", "arm", "log_haz_link", "time", "width", "midpoint")]
96
97	22x	os_eval_df <- lm_link_dat \|>
98	22x	dplyr::left_join(os_baseline, by = c("subject", "study", "arm"))
99
100	22x	withCallingHandlers(
101	22x	os_dat <- sampleObservations(survival, os_eval_df),
102	22x	message = function(e) {
103	!	if (!.silent) message(e)
104	10x	invokeRestart("muffleMessage")
105		}
106		)
107
108	22x	lm_dat2 <- lm_dat \|>
109	22x	dplyr::left_join(dplyr::select(os_dat, "subject", os_time = "time"), by = "subject") \|>
110	22x	dplyr::mutate(observed = (.data$time <= .data$os_time)) \|>
111	22x	dplyr::arrange(dplyr::pick(c("subject", "time")))
112
113	22x	assert_that(
114	22x	length(unique(os_dat$subject)) == length(os_dat$subject),
115	22x	length(os_dat$subject) == n_subjects,
116	22x	all(os_dat$time >= 0),
117	22x	all(os_dat$event %in% c(0, 1)),
118	22x	msg = "Assumptions for the Survival data are not met (please report this issue)"
119		)
120
121	22x	assert_that(
122	22x	nrow(lm_dat2) == n_subjects * length(longitudinal@times),
123	22x	length(unique(lm_dat2$subject)) == n_subjects,
124	22x	msg = "Assumptions for the Longitudinal data are not met (please report this issue)"
125		)
126
127	22x	return(
128	22x	.SimJointData(
129	22x	survival = os_dat,
130	22x	longitudinal = lm_dat2[, c("subject", "arm", "study", "time", "sld", "observed")]
131		)
132		)
133		}
134
135		#' @rdname show-object
136		#' @export
137		setMethod(
138		f = "show",
139		signature = "SimJointData",
140		definition = function(object) {
141	1x	x <- sprintf("\nA SimJointData Object\n\n")
142	1x	cat(x)
143	1x	return(object)
144		}
145		)

1
2		#' Quantity Grid Specification
3		#'
4		#' @param subjects (`character` or `NULL`)\cr vector of subjects to extract quantities for.
5		#' If `NULL` will default to all subjects within the dataset.
6		#'
7		#' @param times (`numeric` or `NULL`)\cr vector of time points to extract quantities at.
8		#' If `NULL` will default to 201 evenly spaced timepoints between 0 and either the max
9		#' observation time (for [`LongitudinalQuantities`]) or max event time (for [`SurvivalQuantities`]).
10		#'
11		#' @param groups (`list`)\cr named list of subjects to extract quantities for. See Group Specification.
12		#'
13		#' @param spec (`list`)\cr named list of subjects to extract quantities for. The names of each
14		#' element should be the required subjects with the element itself being a numeric vector of timepoints
15		#' to generate the quantity at.
16		#'
17		#' @param length.out (`numeric`)\cr number of evenly spaced timepoints to generate quantities at.
18		#'
19		#' @param newdata (`data.frame`) \cr new data to generate quantities for. Must contain the same columns
20		#' and factor levels of the original data used in the [`DataSurvival`] object.
21		#'
22		#' @param params (`list`)\cr named list of parameters to fix the longitudinal model parameters at when
23		#' predicting survival quantities. See [`getPredictionNames()`] for the required parameters.
24		#'
25		#' @description
26		#' These functions are used to specify which subjects and timepoints should be generated
27		#' when calculating quantities via [`SurvivalQuantities`] and [`LongitudinalQuantities`].
28		#'
29		#' @details
30		#'
31		#' - `GridFixed()` is used to specify a fixed set of timepoints to generate quantities at for
32		#' all the specified subjects.
33		#'
34		#' - `GridGrouped()` is similar to `GridFixed()` but allows for groupwise averaging
35		#' (see Group Specification).
36		#'
37		#' - `GridObserved()` generates quantities at the observed longitudinal timepoints for each
38		#' subject.
39		#'
40		#' - `GridManual()` allows for individual timepoint specification for each subject.
41		#'
42		#' - `GridEven()` generates quantities for each subject at N evenly spaced timepoints
43		#' between each subjects first and last longitudinal observations.
44		#'
45		#' - `GridEvent()` generates one quantity for each subject at their event/censor time
46		#' as indicated by the `time` variable in the survival dataset.
47		#'
48		#' - `GridPopulation()` generates longitudinal model quantities based on the population parameters at the
49		#' specified time points. Generates 1 set of quantities for each distinct combination of `arm`
50		#' and `study` within the [`DataSubject`] object provided to the [`JointModel`].
51		#'
52		#' - `GridPrediction()` generates survival quantities based on any user-defined values at the
53		#' specified time points. This is useful for generating quantities for a new dataset
54		#' on specific longitudinal model parameters. See [`getPredictionNames()`] to determine which
55		#' longitudinal model parameters need to be defined for a given longitudinal model.
56		#'
57		#' @section Group Specification:
58		#' For `GridGrouped()`, `groups` must be a named list of character vectors. Each element of the list
59		#' must be a character vector of the subjects that will form the group where the element name
60		#' is the corresponding name of the group. For example if the goal was to create two groups
61		#' named `Group-1` and `Group-2` which are composed of the subjects `sub-1`, `sub-2` and
62		#' `sub-3`, `sub-4` respectively then this would be specified as:
63		#' ```
64		#' GridGrouped(
65		#' groups = list(
66		#' "Group-1" = c("sub-1", "sub-2"),
67		#' "Group-2" = c("sub-3", "sub-4")
68		#' )
69		#' )
70		#' ```
71		#' @seealso [`SurvivalQuantities`], [`LongitudinalQuantities`]
72		#' @name Grid-Functions
73		NULL
74
75
76
77
78
79
80		#' Grid Developer Notes
81		#'
82		#' @description
83		#' Developer details for implementing / extending `Grid` objects for defining
84		#' generated quantities for [`SurvivalQuantities`] and [`LongitudinalQuantities`].
85		#'
86		#' @slot subjects (`character` or `NULL`)\cr vector of subjects to extract quantities for.
87		#' If `NULL` will default to all subjects within the dataset.
88		#' @slot times (`numeric` or `NULL`)\cr vector of time points to extract quantities at.
89		#' If `NULL` will default to 201 evenly spaced timepoints between 0 and either the max
90		#' @slot groups (`list`)\cr named list of subjects to extract quantities for. See details.
91		#'
92		#' @details
93		#' All grid classes must inherit from the abstract `Grid` class.
94		#' All grid classes must provide `as.QuantityGenerator(object, data)` and
95		#' `as.QuantityCollapser(object, data)` methods where `data` is a [`DataJoint`] object.
96		#' These methods must return a `QuantityGenerator` and `QuantityCollapser` object respectively.
97		#' The `QuantityGenerator` object specifies unique subject/timepoint combinations that samples
98		#' should be generated at.
99		#' The `QuantityCollapser` object specifies how to combine these generated samples
100		#' to form the desired quantities.
101		#' As an example say we want to generate grouped samples for the groups `Group-1` and `Group-2`
102		#' which consist of the subjects `sub-1`, `sub-2` and `sub-3`, `sub-4` respectively at two time points
103		#' `10` and `20`. We can achieve this as follows:
104		#' ```
105		#' QuantityGenerator(
106		#' times = c(10, 10, 10, 10, 20, 20, 20, 20),
107		#' subjects = c("sub-1" "sub-2", "sub-3", "sub-4", "sub-1" "sub-2", "sub-3", "sub-4")
108		#' )
109		#' QuantityCollapser(
110		#' times = c(10, 20, 10 , 20),
111		#' groups = c("Group-1", "Group-1", "Group-2", "Group-2"),
112		#' indexes = list(c(1, 2), c(5, 6), c(3, 4), c(7, 8))
113		#' )
114		#' ```
115		#' For population based quantities use the `arms` and `studies` arguments of `QuantityGenerator`
116		#' instead of `subjects`.
117		#'
118		#' @inheritSection Grid-Functions Group Specification
119		#'
120		#' @keywords internal
121		#' @seealso `Quant-Dev`
122		#' @name Grid-Dev
123		NULL
124
125
126
127		#' Quantity Developer Notes
128		#'
129		#' @description
130		#' Developer details for `QuantityX` objects/methods. This page just outlines the arguments
131		#' and slots of these objects/methods. For the full implementation details please see [Grid-Dev]
132		#'
133		#' @slot times (`numeric`)\cr See Arguments for details.
134		#' @slot subjects (`character`)\cr See Arguments for details.
135		#' @slot groups (`character`)\cr See Arguments for details.
136		#' @slot indexes (`list`)\cr See Arguments for details.
137		#'
138		#' @param times (`numeric`)\cr vector of time points to extract quantities at.
139		#' @param subjects (`character`)\cr vector of subjects to extract quantities for.
140		#' @param groups (`character`)\cr vector of labels to apply to the generated quantities.
141		#' @param indexes (`list`)\cr list of indexes that specify which observations from a
142		#' `QuantityGenerator` should be combined to form the desired quantities.
143		#'
144		#' @param object (`Grid`)\cr object to convert to a `QuantityGenerator` or `QuantityCollapser`.
145		#' @param data (`DataJoint`)\cr Survival and Longitudinal Data.
146		#' @param ... Not currently used.
147		#'
148		#' @details
149		#' The `as.QuantityGenerator` must return a `QuantityGenerator` object.
150		#' The `as.QuantityCollapser` must return a `QuantityCollapser` object.
151		#' @keywords internal
152		#' @name Quant-Dev
153		NULL
154
155
156		#' @rdname Grid-Dev
157		.Grid <- setClass("Grid")
158
159
160		#' @rdname Quant-Dev
161		.QuantityGenerator <- setClass(
162		"QuantityGenerator"
163		)
164		#' `QuantityGenerator` -> `list`
165		#' @description
166		#' Converts a `QuantityGenerator` object to a list containing the required input data for a stan
167		#' model.
168		#' @param object (`QuantityGenerator`)\cr object to convert to a list.
169		#' @param data (`DataJoint`)\cr Survival and Longitudinal Data.
170		#' @param ... Not currently used.
171		#' @keywords internal
172		as_stan_list.QuantityGenerator <- function(object, data, ...) {
173	!	stop("as_stan_list.QuantityGenerator not implemented")
174		}
175
176
177
178		#' @rdname Quant-Dev
179		.QuantityCollapser <- setClass(
180		"QuantityCollapser",
181		slots = c(
182		"times" = "numeric",
183		"groups" = "character",
184		"indexes" = "list"
185		)
186		)
187		#' @rdname Quant-Dev
188		QuantityCollapser <- function(times, groups, indexes) {
189	47x	.QuantityCollapser(
190	47x	times = times,
191	47x	groups = groups,
192	47x	indexes = indexes
193		)
194		}
195
196		setValidity(
197		"QuantityCollapser",
198		function(object) {
199		if (
200		length(object@times) != length(object@groups) \|\|
201		length(object@times) != length(object@indexes)
202		) {
203		return("Length of `times`, `groups`, and `indexes` must be equal")
204		}
205		}
206		)
207
208		#' @export
209		length.QuantityCollapser <- function(x) {
210	68x	length(x@indexes)
211		}
212
213
214		#' Expand and Validate Subjects
215		#'
216		#' @param subjects (`character`)\cr vector of subjects that should exist in `data`
217		#' @param data (`DataJoint`)\cr Survival and Longitudinal Data.
218		#'
219		#' @description
220		#' If `subjects` is `NULL` this will return a named list of all subjects in `data`.
221		#' Else it will return `subjects` as a named list ensuring that all subjects exist in `data`.
222		#'
223		#' @keywords internal
224		subjects_to_list <- function(subjects = NULL, data) {
225	66x	data_list <- as.list(data)
226	66x	subjects_exp <- if (is.null(subjects)) {
227	10x	subs <- as.list(names(data_list$subject_to_index))
228	10x	names(subs) <- names(data_list$subject_to_index)
229	10x	subs
230		} else {
231	56x	subs <- as.list(subjects)
232	56x	names(subs) <- subjects
233	56x	subs
234		}
235	66x	subjects_exp_vec <- unlist(subjects_exp, use.names = FALSE)
236	66x	assert_that(
237	66x	identical(subjects_exp_vec, unique(subjects_exp_vec)),
238	66x	msg = "All subject names must be unique"
239		)
240	66x	assert_that(
241	66x	all(subjects_exp_vec %in% names(data_list$subject_to_index)),
242	66x	msg = "Not all subjects exist within the data object"
243		)
244	65x	subjects_exp
245		}

1		# "missing" = no argument provided
2		# "NULL" = explicit NULL
3		setClassUnion("empty", c("missing", "NULL"))
4		setClassUnion("numeric_or_NULL", c("numeric", "NULL"))
5		setClassUnion("character_or_NULL", c("character", "NULL"))
6
7		# merge ----
8
9		#' `merge`
10		#'
11		#' Merge two `StanModule` or `ParameterList` objects.
12		#'
13		#' @param x first module.
14		#' @param y second module.
15		#' @param ... additional arguments.
16		#'
17		#' @export
18		# Needs to be S4 for multiple dispatch !
19		setGeneric(
20		name = "merge",
21		def = function(x, y, ...) standardGeneric("merge")
22		)
23
24
25		# show ----
26
27		#' Printing of Different Classes
28		#'
29		#' These methods print objects of different classes.
30		#'
31		#' @name show
32		#' @aliases show
33		#'
34		#' @param object what to print.
35		#'
36		#' @export
37		NULL
38
39
40		# write_stan ----
41
42		#' `write_stan`
43		#'
44		#' Write the Stan code for a Stan module.
45		#'
46		#' @param object the module.
47		#' @param destination (`character` or `connection`)\cr Where to write stan code to.
48		#' @param ... Additional arguments
49		#'
50		#' @export
51		write_stan <- function(object, destination, ...) {
52	!	UseMethod("write_stan")
53		}
54
55		# compileStanModel ----
56
57		#' `compileStanModel`
58		#'
59		#' Compile the Stan module.
60		#'
61		#' @param object the module.
62		#'
63		#' @export
64		compileStanModel <- function(object) {
65	65x	UseMethod("compileStanModel")
66		}
67
68
69		# sampleStanModel ----
70
71		#' `sampleStanModel`
72		#'
73		#' Sample from a Stan Module.
74		#'
75		#' @param object the module.
76		#' @param ... additional arguments.
77		#'
78		#' @export
79		sampleStanModel <- function(object, ...) {
80	17x	UseMethod("sampleStanModel")
81		}
82
83
84		# as.StanModule ----
85
86		#' `as.StanModule`
87		#'
88		#' Converts an object into a [`StanModule`].
89		#'
90		#' @param object what to convert.
91		#' @param ... additional options.
92		#' @family as.StanModule
93		#' @keywords internal
94		as.StanModule <- function(object, ...) {
95	1579x	UseMethod("as.StanModule")
96		}
97
98
99
100		#' `getParameters`
101		#'
102		#' Extract any modelling parameters as a [`ParameterList`] object
103		#' from a model.
104		#'
105		#' @param object where to obtain the parameters from.
106		#' @param ... additional options.
107		#'
108		#' @export
109		getParameters <- function(object, ...) {
110	134x	UseMethod("getParameters")
111		}
112
113
114		# extractVariableNames ----
115
116		#' Extract Mapping to Standardised Variable Names
117		#'
118		#' @description
119		#' Extract a `list` that maps the variable names in a user-defined
120		#' `data.frame` to standardised values.
121		#'
122		#' @param object the data object.
123		#' @family extractVariableNames
124		#' @keywords internal
125		extractVariableNames <- function(object) {
126	1672x	UseMethod("extractVariableNames")
127		}
128
129
130
131		# initialValues ----
132
133		#' `initialValues`
134		#'
135		#' Obtain the `list` of initial values to be passed to the Stan sampler.
136		#'
137		#' @param object where to get the initial values from.
138		#' @param n_chains the number of initial values to generate. See details.
139		#' @param ... Not currently used.
140		#'
141		#' @details
142		#' There are multiple ways of specifying initial values to Stan, see the `init` argument
143		#' in [cmdstanr::model-method-sample] for full details. Within this package we supply
144		#' initial values via a list of lists where each inner list contains the initial values
145		#' for a single chain. As such the `n_chains` argument specifies the number of inner lists
146		#' to generate.
147		#'
148		#' See the Vignette for further details of how to specify initial values.
149		#'
150		#' @export
151		initialValues <- function(object, ...) {
152	78501x	UseMethod("initialValues")
153		}
154
155
156		# size ----
157
158		#' `size`
159		#'
160		#' Obtain the `list` of parameter sizes.
161		#'
162		#' @param object where to get the parameter sizes from.
163		#'
164		#' @keywords internal
165		size <- function(object) {
166	136x	UseMethod("size")
167		}
168
169
170		# generateQuantities ----
171
172		#' `generateQuantities`
173		#'
174		#' Obtain the generated quantities from a Stan Model.
175		#'
176		#' @param object object to obtain generated quantities from
177		#' @param ... additional options.
178		#'
179		#' @export
180		generateQuantities <- function(object, ...) {
181	31x	UseMethod("generateQuantities")
182		}
183
184
185		#' Prepare Data Object
186		#'
187		#' @param object (`DataSubject` or `DataLongitudinal` or `DataSurvival`) \cr data object to "harmonise"
188		#' @param subject_var (`character`) \cr the name of the variable containing the subject identifier.
189		#' @param subject_ord (`character`) \cr the expected levels (in order) of the subject identifier.
190		#' @param ... not used.
191		#'
192		#' @details
193		#' This utility function prepares the datasets in the data objects in order to ensure they
194		#' are consistent and compatible with each other.
195		#'
196		#' In particular it ensures that the `subject` variable, as specified by `DataSubject`,
197		#' is available in `DataLongitudinal` and `DataSurvival` and that all levels are present
198		#' in all 3 data objects.
199		#'
200		#' It also sorts the datasets to ensure that indexes are consistent e.g. index 1 for
201		#' `DataSubject@data` corresponds to the same subject as index 1 for `DataSurvival@data`.
202		#' For `DataLongitudinal` the data is additionally sorted by time and outcome value.
203		#'
204		#' @seealso [`DataJoint`], [`DataSurvival`], [`DataSubject`], [`DataLongitudinal`]
205		#'
206		#' @keywords internal
207		#' @return Returns the original object but with the data standardised (see details)
208		harmonise <- function(object, ...) {
209	384x	UseMethod("harmonise")
210		}
211
212
213		#' @rdname harmonise
214		harmonise.default <- function(object, ...) {
215	13x	NULL
216		}
217
218
219		#' `as_stan_list`
220		#'
221		#' @description
222		#' Extracts a list of data elements from an object to be used as input
223		#' to a Stan Model
224		#'
225		#' @param object to be converted.
226		#' @param ... additional options.
227		#'
228		#' @family as_stan_list
229		#' @export
230		as_stan_list <- function(object, ...) {
231	2023x	UseMethod("as_stan_list")
232		}
233
234		#' @rdname as_stan_list
235		#' @export
236		as_stan_list.default <- function(object, ...) {
237	21x	NULL
238		}
239
240
241		#' `as_print_string`
242		#'
243		#' @description
244		#' Returns the character representation of an object which is suitable
245		#' for printing to the console
246		#'
247		#' @param object to be converted to string.
248		#' @param ... additional options.
249		#'
250		#' @family as_print_string
251		#' @keywords internal
252		as_print_string <- function(object, ...) {
253	57x	UseMethod("as_print_string")
254		}
255
256		#' Show an Object
257		#'
258		#' Prints an object to the console.
259		#'
260		#' @param object Object to be printed
261		#'
262		#' @name show-object
263		NULL
264
265
266
267		#' `brierScore`
268		#'
269		#' @description
270		#' Returns the Brier Score for a given model
271		#'
272		#' @param object to calculate Brier Score for.
273		#' @param ... additional options.
274		#'
275		#' @family brierScore
276		#' @export
277		brierScore <- function(object, ...) {
278	1x	UseMethod("brierScore")
279		}
280
281
282
283		#' Generate Simulated Observations
284		#'
285		#' @param object (`SimLongitudinal` or `SimSurvival`) \cr object to generate observations from.
286		#' @param times_df (`data.frame`) \cr the times at which to generate observations. See details.
287		#'
288		#' @details
289		#' The `times_df` argument should be a `data.frame` as created by `sampleSubjects` but
290		#' replicated for each time point at which observations are to be generated. That is if you want
291		#' to generate observations for times `c(0, 1, 2, 3)` then `times_df` should be created as:
292		#' ```
293		#' subject_dat <- sampleSubjects(object, ...)
294		#' times_df <- tidyr::expand_grid(
295		#' subject_dat,
296		#' time = c(0, 1, 2, 3)
297		#' )
298		#' ```
299		#'
300		#' @export
301		sampleObservations <- function(object, times_df) {
302	70x	UseMethod("sampleObservations")
303		}
304
305
306		#' Generate Simulated Subjects
307		#'
308		#' @param object (`SimLongitudinal` or `SimSurvival`) \cr object to generate subjects from.
309		#' @param subjects_df (`data.frame`) \cr the subjects to generate observations for. See details.
310		#'
311		#' @details
312		#' The `subjects_df` argument should be a `data.frame` with 1 row per desired subject to create
313		#' with the following columns:
314		#' - `study` (`factor`) the study identifier.
315		#' - `arm` (`factor`) the treatment arm identifier.
316		#' - `subject` (`character`) the subject identifier.
317		#'
318		#' This method takes care of generating all the individual subject data required for the
319		#' [`sampleObservations`] method to generate the observations.
320		#' @export
321		sampleSubjects <- function(object, subjects_df) {
322	48x	UseMethod("sampleSubjects")
323		}
324
325
326		#' Generate time windows for evaluating a hazard function
327		#'
328		#' @param object (`SurvivalModel`) \cr object to generate time windows for.
329		#' @param ... Not used.
330		#'
331		hazardWindows <- function(object, ...) {
332	22x	UseMethod("hazardWindows")
333		}
334
335
336		#' Coerce to `CmdStanMCMC`
337		#'
338		#' @param object to be converted
339		#' @param ... additional options
340		#'
341		#' @description
342		#' Coerces an object to a [`cmdstanr::CmdStanMCMC`] object
343		#'
344		#' @export
345		as.CmdStanMCMC <- function(object, ...) {
346	81x	UseMethod("as.CmdStanMCMC")
347		}
348
349		#' @rdname Quant-Dev
350		#' @export
351		as.QuantityGenerator <- function(object, ...) {
352	98x	UseMethod("as.QuantityGenerator")
353		}
354
355		#' @rdname Quant-Dev
356		#' @export
357		as.QuantityCollapser <- function(object, ...) {
358	37x	UseMethod("as.QuantityCollapser")
359		}
360
361
362		#' Coalesce Time
363		#'
364		#' @param object ([`Grid`]) \cr object to coalesce time for.
365		#' @param times (`numeric`) \cr the times to coalesce to.
366		#' @param ... Not used
367		#'
368		#' Method used to replace NULL times on grid objects (if appropriate)
369		#'
370		#' @keywords internal
371		coalesceGridTime <- function(object, times, ...) {
372	37x	UseMethod("coalesceGridTime")
373		}
374		#' @export
375		coalesceGridTime.default <- function(object, times, ...) {
376	2x	object
377		}
378
379
380		#' Resolve a Promise
381		#'
382		#' @param object (`ANY`)\cr an object to resolve.
383		#' @param ... (`ANY`)\cr additional arguments.
384		#'
385		#' If `object` is not a promise will just return itself else will resolve the promise
386		#' and return the promised object.
387		#'
388		#' @export
389		resolvePromise <- function(object, ...) {
390	62x	UseMethod("resolvePromise")
391		}
392
393		#' @rdname resolvePromise
394		#' @export
395		resolvePromise.default <- function(object, ...) {
396	!	object
397		}
398
399		#' Enable Link Generic
400		#'
401		#' @param object ([`LongitudinalModel`])\cr to enable link for.
402		#' @param ... Not used.
403		#'
404		#' Optional hook method that is called on a [`LongitudinalModel`] only if a link method
405		#' is provided to [`JointModel`]. This can be used to allow the model to include any
406		#' optional stan code that is only required if there are links present.
407		#'
408		#' @return [`LongitudinalModel`] object
409		#'
410		#' @export
411		enableLink <- function(object, ...) {
412	11x	UseMethod("enableLink")
413		}
414		#' @export
415		enableLink.default <- function(object, ...) {
416	!	object
417		}
418
419
420		#' Enable Generated Quantities Generic
421		#'
422		#' @param object ([`StanModel`])\cr to enable generated quantities for.
423		#' @param ... Not used.
424		#'
425		#' Optional hook method that is called on a [`StanModel`] if attempting to use
426		#' either [`LongitudinalQuantities`] or [`SurvivalQuantities`]
427		#'
428		#' @return [`StanModule`] object
429		#'
430		#' @export
431		enableGQ <- function(object, ...) {
432	114x	UseMethod("enableGQ")
433		}
434		#' @export
435		enableGQ.default <- function(object, ...) {
436	39x	StanModule()
437		}
438
439
440
441		#' Get Prediction Names
442		#'
443		#' Utility function that returns the names of the required parameters for predicting
444		#' survival quantities with [`GridPrediction`].
445		#'
446		#' @param object (`LongitudinalModel`) \cr A longitudinal model object
447		#' @param ... Not used.
448		#' @export
449		getPredictionNames <- function(object, ...) {
450	6x	UseMethod("getPredictionNames")
451		}
452
453		#' @rdname getPredictionNames
454		getPredictionNames.default <- function(object, ...) {
455	1x	NULL
456		}
457
458		#' As Formula
459		#'
460		#' Utility wrapper function to convert an object to a formula.
461		#' @param x (`ANY`) \cr object to convert to a formula.
462		#' @param ... Not used.
463		#' @export
464		as_formula <- function(x, ...) {
465	7x	UseMethod("as_formula")
466		}
467
468		#' @importFrom stats as.formula
469		#' @export
470		as_formula.default <- function(x, ...) {
471	!	as.formula(x, ...)
472		}
473
474
475		#' Set Constraints
476		#'
477		#' Applies constraints to a prior distribution to ensure any sampled numbers
478		#' from the distribution fall within the constraints
479		#'
480		#' @param object (`Prior`)\cr a prior distribution to apply constraints to
481		#' @param lower (`numeric`)\cr lower constraint boundary
482		#' @param upper (`numeric`)\cr upper constraint boundary
483		#'
484		#' @export
485		set_limits <- function(object, lower = -Inf, upper = Inf) {
486	285x	UseMethod("set_limits")
487		}

1
2
3		#' Promise
4		#'
5		#' Abstract class for promise objects to inherit off of
6		#' @aliases Promise
7		#' @exportClass Promise
8		.Promise <- setClass("Promise")
9
10
11		#' Promise of a `LongitudinalModel`
12		#'
13		#' An object that promises to resolve to a [`LongitudinalModel`] object.
14		#'
15		#' @exportClass PromiseLongitudinalModel
16		.PromiseLongitudinalModel <- setClass(
17		"PromiseLongitudinalModel",
18		contains = "Promise"
19		)
20
21		#' @rdname PromiseLongitudinalModel-class
22		#' @export
23		PromiseLongitudinalModel <- function() {
24	26x	.PromiseLongitudinalModel()
25		}
26
27
28
29
30		#' Promise of a `LinkComponent`
31		#'
32		#' An object that promises to resolve to a [`LinkComponent`] object.
33		#' Inheriting from [`Promise`] and [`LinkComponent`].
34		#'
35		#' @slot fun (`function`) \cr a function that returns a `LinkComponent`. See details.
36		#'
37		#' @param fun (`function`) \cr a function that returns a `LinkComponent`. See details.
38		#' @inheritParams LinkComponent
39		#'
40		#' @details
41		#'
42		#' The `fun` slot should be a function of signature `function(prior, model)` and should return
43		#' a [`LinkComponent`] object. An error will be thrown if the returned [`LinkComponent`] object
44		#' does not have the same `key` slot value as the original `PromiseLinkComponent`.
45		#'
46		#' @exportClass PromiseLinkComponent
47		.PromiseLinkComponent <- setClass(
48		"PromiseLinkComponent",
49		contains = c("Promise", "LinkComponent"),
50		slots = list(
51		fun = "function"
52		)
53		)
54
55		#' @rdname PromiseLinkComponent-class
56		#' @export
57		PromiseLinkComponent <- function(fun, prior, key) {
58	28x	.PromiseLinkComponent(
59	28x	fun = fun,
60	28x	stan = StanModule(),
61	28x	parameters = ParameterList(Parameter(name = key, prior = prior, size = 1)),
62	28x	key = key
63		)
64		}
65
66
67		#' @export
68		as.StanModule.PromiseLinkComponent <- function(object, model, ...) {
69	1x	resolved_object <- resolvePromise(object, model = model)
70	1x	as.StanModule(resolved_object, ...)
71		}
72
73		#' Resolve a `PromiseLinkComponent`
74		#'
75		#' Resolves a [`PromiseLinkComponent`] object to a [`LinkComponent`] object.
76		#' An error will be thrown if the returned [`LinkComponent`] object
77		#' does not have the same `key` slot value as the original [`PromiseLinkComponent`].
78		#'
79		#' @param object ([`PromiseLinkComponent`]) \cr the promise to resolve
80		#' @param model ([`LongitudinalModel`]) \cr the model to resolve the promise with
81		#' @param ... Not used.
82		#' @return ([`LinkComponent`]) \cr the resolved `LinkComponent` object
83		#'
84		#' @export
85		resolvePromise.PromiseLinkComponent <- function(object, model, ...) {
86	25x	x <- object@fun(
87	25x	prior = object@parameters@parameters[[1]]@prior,
88	25x	model = model
89		)
90	25x	assert_that(
91	25x	is(x, "LinkComponent"),
92	25x	msg = "Resolved `PromiseLinkComponent` did not produce a `LinkComponent` object"
93		)
94	25x	assert_that(
95	25x	names(object) == names(x),
96	25x	msg = paste(
97	25x	"Resolved `PromiseLinkComponent` did not produce a `LinkComponent` object",
98	25x	"with the same key as the promise"
99		)
100		)
101	24x	x
102		}

1		#' @include StanModel.R
2		NULL
3
4		# LongitudinalModel-class ----
5
6		#' `LongitudinalModel`
7		#'
8		#' This class extends the general [`StanModel`] class to comprise the longitudinal
9		#' model specification.
10		#'
11		#' @exportClass LongitudinalModel
12		.LongitudinalModel <- setClass(
13		Class = "LongitudinalModel",
14		contains = "StanModel"
15		)
16
17		# LongitudinalModel-constructors ----
18
19		#' @rdname LongitudinalModel-class
20		#'
21		#' @inheritParams stanmodel_arguments
22		#'
23		#' @export
24		LongitudinalModel <- function(
25		stan = StanModule(),
26		parameters = ParameterList(),
27		name = "<Unnamed>",
28		...
29		) {
30	74x	base_stan <- read_stan("base/longitudinal.stan")
31
32	74x	stan_full <- decorated_render(
33	74x	.x = base_stan,
34	74x	stan = add_missing_stan_blocks(as.list(stan))
35		)
36
37	74x	.LongitudinalModel(
38	74x	StanModel(
39	74x	stan = StanModule(stan_full),
40	74x	parameters = parameters,
41	74x	name = name,
42		...
43		)
44		)
45		}
46
47		#' @export
48		as_print_string.LongitudinalModel <- function(object, ...) {
49	5x	string <- sprintf(
50	5x	"\n%s Longitudinal Model with parameters:\n%s\n\n",
51	5x	object@name,
52	5x	paste(" ", as_print_string(object@parameters)) \|> paste(collapse = "\n")
53		)
54	5x	return(string)
55		}

1
2
3
4		#' jmpost settings
5		#'
6		#' @description
7		#' Define settings that modify the behaviour of the `jmpost` package
8		#'
9		#' Each of the following are the name of options that can be set via:
10		#' ```
11		#' options(<option_name> = <value>)
12		#' ```
13		#'
14		#' ## `jmpost.prior_shrinkage`
15		#'
16		#' Default = `0.5`
17		#'
18		#' By default all initial values are drawn as random sample from the respective prior
19		#' distribution with a shrinkage factor towards the mean. That is:
20		#' ```
21		#' initial_value = prior_mean * prior_shrinkage + (1 - prior_shrinkage) * prior_sample
22		#' ```
23		#' This setting controls the shrinkage factor. A value of 0 means no shrinkage (i.e.
24		#' pure random draw) whilst a value of 1 means the initial value is just the mean.
25		#'
26		#' ## `jmpost.cache_dir`
27		#'
28		#' Default = `tempfile()`
29		#'
30		#' Directory to store compiled stan models in. If not set, a temporary directory is used for
31		#' the given R session. Can also be set via the environment variable `JMPOST_CACHE_DIR`.
32		#'
33		#'
34		#'
35		#' ## `jmpost.gauss_quad_n`
36		#'
37		#' Default = 15
38		#'
39		#' In most cases the survival function of the joint model does not have a closed form
40		#' and as such it is calculated by integrating the hazard function. `jmpost` estimates this
41		#' via Gaussian Quadrature, in particular it uses [`statmod::gauss.quad`] with
42		#' `kind = "legendre"` to create the nodes and weights.
43		#'
44		#' This option specifies the `n` argument in the call to [`statmod::gauss.quad`]. In general
45		#' higher values of `n` lead to better accuracy of the approximation but at the cost of
46		#' increased computational time.
47		#'
48		#' @examples
49		#' \dontrun{
50		#' options(jmpost.prior_shrinkage = 0.5)
51		#' }
52		#' @name jmpost-settings
53		set_options <- function() {
54
55	!	cache_dir <- Sys.getenv("JMPOST_CACHE_DIR")
56
57	!	if (cache_dir == "" \|\| is.null(cache_dir)) {
58	!	cache_dir <- tempfile()
59		}
60
61	!	current_opts <- names(options())
62	!	jmpost_opts <- list(
63	!	jmpost.cache_dir = cache_dir,
64	!	jmpost.prior_shrinkage = 0.5,
65	!	jmpost.gauss_quad_n = 15
66		)
67	!	for (opt in names(jmpost_opts)) {
68	!	if (!opt %in% current_opts) {
69	!	options(jmpost_opts[opt])
70		}
71		}
72		}

1		#' @include SurvivalModel.R
2		NULL
3
4		# SurvivalExponential-class ----
5
6		#' `SurvivalExponential`
7		#'
8		#' This class extends the general [`SurvivalModel`] class for using the
9		#' exponential survival model.
10		#'
11		#' @exportClass SurvivalExponential
12		.SurvivalExponential <- setClass(
13		Class = "SurvivalExponential",
14		contains = "SurvivalModel"
15		)
16
17		# SurvivalExponential-constructors ----
18
19		#' @rdname SurvivalExponential-class
20		#'
21		#' @param lambda (`Prior`)\cr for the exponential rate `lambda`.
22		#' @param beta (`Prior`)\cr for covariates coefficients `beta`.
23		#'
24		#' @export
25		SurvivalExponential <- function(
26		lambda = prior_gamma(2, 5),
27		beta = prior_normal(0, 2)
28		) {
29	13x	.SurvivalExponential(
30	13x	SurvivalModel(
31	13x	name = "Exponential",
32	13x	stan = StanModule("sm-exponential/model.stan"),
33	13x	parameters = ParameterList(
34	13x	Parameter(name = "sm_exp_lambda", prior = lambda, size = 1),
35	13x	Parameter(name = "beta_os_cov", prior = beta, size = "p_os_cov_design")
36		)
37		)
38		)
39		}

1		#' @include generics.R
2		#' @include Parameter.R
3		#' @include Prior.R
4		NULL
5
6
7		#' ParameterList-Shared
8		#' @param object (`ParameterList`) \cr A List of [`Parameter`] Objects.
9		#' @param x (`ParameterList`) \cr A List of [`Parameter`] Objects.
10		#' @param ... Not Used.
11		#' @keywords internal
12		#' @name ParameterList-Shared
13		NULL
14
15
16		# ParameterList-class ----
17
18		#' `ParameterList`
19		#'
20		#' This class extends the general [`list`] type for containing [`Parameter`]
21		#' specifications.
22		#'
23		#'
24
25		#' @slot parameters (`list`) \cr a list of [`Parameter`] objects
26		#' @family ParameterList
27		#' @export ParameterList
28		#' @exportClass ParameterList
29		.ParameterList <- setClass(
30		Class = "ParameterList",
31		slots = c(
32		parameters = "list"
33		)
34		)
35		#' @param ... (`Parameter`)\cr which parameter specifications to include.
36		#' @rdname ParameterList-class
37		ParameterList <- function(...) {
38	211x	.ParameterList(parameters = list(...))
39		}
40
41		# ParameterList-validity ----
42
43		setValidity(
44		Class = "ParameterList",
45		method = function(object) {
46		is_parameters <- vapply(object@parameters, function(x) is(x, "Parameter"), logical(1))
47		if (!all(is_parameters)) {
48		return("all elements must be of class 'Parameter'")
49		}
50		return(TRUE)
51		}
52		)
53
54
55
56		# as.StanModule-ParameterList ----
57
58		#' `ParameterList` -> `StanModule`
59		#'
60		#' Converts a [`ParameterList`] object to a [`StanModule`] object
61		#'
62		#' @inheritParams ParameterList-Shared
63		#'
64		#' @family ParameterList
65		#' @family as.StanModule
66		#' @export
67		as.StanModule.ParameterList <- function(object, ...) {
68	71x	stan_modules <- lapply(
69	71x	object@parameters,
70	71x	as.StanModule
71		)
72	71x	assert_that(
73	71x	all(vapply(stan_modules, inherits, logical(1), "StanModule"))
74		)
75	71x	Reduce(merge, stan_modules)
76		}
77
78
79
80		#' `ParameterList` -> `list`
81		#'
82		#' Converts a ParameterList object to a list of parameter data values
83		#' for a Stan model.
84		#'
85		#' @inheritParams ParameterList-Shared
86		#'
87		#' @family as_stan_list
88		#' @family ParameterList
89		#' @export
90		as_stan_list.ParameterList <- function(object, ...) {
91	49x	stan_lists <- lapply(
92	49x	object@parameters,
93	49x	as_stan_list
94		)
95	49x	assert_that(
96	49x	all(vapply(stan_lists, is.list, logical(1)))
97		)
98	49x	Reduce(append, stan_lists)
99		}
100
101
102		# merge-ParameterList,ParameterList ----
103
104		#' @rdname merge
105		setMethod(
106		f = "merge",
107		signature = c(x = "ParameterList", y = "ParameterList"),
108		definition = function(x, y) {
109	43x	parameters <- append(x@parameters, y@parameters)
110	43x	do.call(ParameterList, parameters)
111		}
112		)
113
114		# as.list-ParameterList ----
115
116		#' `ParameterList` -> `list`
117		#' @description
118		#' Returns a named list where each element of the list corresponds
119		#' to a Stan modelling block e.g. `data`, `model`, etc.
120		#' @inheritParams ParameterList-Shared
121		#' @family ParameterList
122		#' @export
123		as.list.ParameterList <- function(x, ...) {
124	70x	as.list(as.StanModule(x))
125		}
126
127
128
129
130		#' Parameter-List Getter Functions
131		#' @description
132		#' Getter functions for the slots of a [`ParameterList`] object
133		#' @inheritParams ParameterList-Shared
134		#' @family ParameterList
135		#' @param n_chains (`integer`) \cr the number of chains.
136		#' @name ParameterList-Getter-Methods
137		NULL
138
139
140		#' @describeIn ParameterList-Getter-Methods The parameter-list's parameter names
141		#' @export
142		names.ParameterList <- function(x) {
143	89x	vapply(x@parameters, names, character(1))
144		}
145
146
147		#' @describeIn ParameterList-Getter-Methods The parameter-list's parameter initial values
148		#' @export
149		initialValues.ParameterList <- function(object, n_chains, ...) {
150		# Generate initial values as a list of lists. This is to ensure it is in the required
151		# format as specified by cmdstanr see the `init` argument of
152		# `help("model-method-sample", "cmdstanr")` for more details
153	45x	lapply(
154	45x	seq_len(n_chains),
155	45x	\(i) {
156	951x	vals <- lapply(object@parameters, initialValues)
157	936x	name <- vapply(object@parameters, names, character(1))
158	936x	names(vals) <- name
159	936x	vals
160		}
161		)
162		}
163
164
165		#' @describeIn ParameterList-Getter-Methods The parameter-list's parameter dimensionality
166		#' @export
167		size.ParameterList <- function(object) {
168	18x	x <- lapply(object@parameters, size)
169	18x	names(x) <- names(object)
170	18x	return(x)
171		}
172
173
174		#' `ParameterList` -> Printable `Character`
175		#'
176		#' Converts [`ParameterList`] object into a printable string.
177		#' @inheritParams ParameterList-Shared
178		#' @family ParameterList
179		#' @keywords internal
180		#' @export
181		as_print_string.ParameterList <- function(object, ...) {
182	13x	x <- vapply(object@parameters, as.character, character(1))
183	13x	if (length(x) == 0) {
184	!	x <- "<No Parameters>"
185		}
186	13x	return(x)
187		}
188
189
190
191		#' @rdname show-object
192		#' @export
193		setMethod(
194		f = "show",
195		signature = "ParameterList",
196		definition = function(object) {
197	1x	chrs <- as_print_string(object)
198	1x	string <- paste(" ", chrs) \|> paste(collapse = "\n")
199	1x	x <- sprintf("\nParameterList Object:\n%s\n\n", string)
200	1x	cat(x)
201	1x	return(object)
202		}
203		)

1
2
3		#' @include SimLongitudinal.R
4		#' @include generics.R
5		NULL
6
7		#' Simulate Longitudinal Data from a GSF Model
8		#'
9		#' @param times (`numeric`)\cr the times to generate observations at.
10		#' @param sigma (`number`)\cr the variance of the longitudinal values.
11		#' @param mu_s (`numeric`)\cr the mean shrinkage rates.
12		#' @param mu_g (`numeric`)\cr the mean growth rates.
13		#' @param mu_b (`numeric`)\cr the mean baseline values.
14		#' @param mu_phi (`numeric`)\cr the mean proportion of cells affected by the treatment
15		#' @param omega_b (`number`)\cr the baseline value standard deviation.
16		#' @param omega_s (`number`)\cr the shrinkage rate standard deviation.
17		#' @param omega_g (`number`)\cr the growth rate standard deviation.
18		#' @param omega_phi (`number`)\cr for the standard deviation of the proportion of cells
19		#' affected by the treatment `omega_phi`.
20		#' @param link_dsld (`number`)\cr the link coefficient for the derivative contribution.
21		#' @param link_ttg (`number`)\cr the link coefficient for the time-to-growth contribution.
22		#' @param link_identity (`number`)\cr the link coefficient for the SLD Identity contribution.
23		#' @param link_growth (`number`)\cr the link coefficient for the growth parameter contribution.
24		#'
25		#' @slot sigma (`numeric`)\cr See arguments.
26		#' @slot mu_s (`numeric`)\cr See arguments.
27		#' @slot mu_g (`numeric`)\cr See arguments.
28		#' @slot mu_b (`numeric`)\cr See arguments.
29		#' @slot mu_phi (`numeric`)\cr See arguments.
30		#' @slot omega_b (`numeric`)\cr See arguments.
31		#' @slot omega_s (`numeric`)\cr See arguments.
32		#' @slot omega_g (`numeric`)\cr See arguments.
33		#' @slot omega_phi (`numeric`)\cr See arguments.
34		#' @slot link_dsld (`numeric`)\cr See arguments.
35		#' @slot link_ttg (`numeric`)\cr See arguments.
36		#' @slot link_identity (`numeric`)\cr See arguments.
37		#' @slot link_growth (`numeric`)\cr See arguments.
38		#' @family SimLongitudinal
39		#' @name SimLongitudinalGSF-class
40		#' @exportClass SimLongitudinalGSF
41		.SimLongitudinalGSF <- setClass(
42		"SimLongitudinalGSF",
43		contains = "SimLongitudinal",
44		slots = c(
45		sigma = "numeric",
46		mu_s = "numeric",
47		mu_g = "numeric",
48		mu_b = "numeric",
49		mu_phi = "numeric",
50		omega_b = "numeric",
51		omega_s = "numeric",
52		omega_g = "numeric",
53		omega_phi = "numeric",
54		link_dsld = "numeric",
55		link_ttg = "numeric",
56		link_identity = "numeric",
57		link_growth = "numeric"
58		)
59		)
60
61		#' @rdname SimLongitudinalGSF-class
62		#' @export
63		SimLongitudinalGSF <- function(
64		times = c(-100, -50, 0, 50, 100, 150, 250, 350, 450, 550) / 365,
65		sigma = 0.01,
66		mu_s = log(c(0.6, 0.4)),
67		mu_g = log(c(0.25, 0.35)),
68		mu_b = log(60),
69		mu_phi = qlogis(c(0.4, 0.6)),
70		omega_b = 0.2,
71		omega_s = 0.2,
72		omega_g = 0.2,
73		omega_phi = 0.2,
74		link_dsld = 0,
75		link_ttg = 0,
76		link_identity = 0,
77		link_growth = 0
78		) {
79	8x	.SimLongitudinalGSF(
80	8x	times = times,
81	8x	sigma = sigma,
82	8x	mu_s = mu_s,
83	8x	mu_g = mu_g,
84	8x	mu_b = mu_b,
85	8x	mu_phi = mu_phi,
86	8x	omega_b = omega_b,
87	8x	omega_s = omega_s,
88	8x	omega_g = omega_g,
89	8x	omega_phi = omega_phi,
90	8x	link_dsld = link_dsld,
91	8x	link_ttg = link_ttg,
92	8x	link_identity = link_identity,
93	8x	link_growth = link_growth
94		)
95		}
96
97
98		setValidity(
99		"SimLongitudinalGSF",
100		function(object) {
101		par_lengths <- c(
102		length(object@mu_s),
103		length(object@mu_g),
104		length(object@mu_phi)
105		)
106		if (length(unique(par_lengths)) != 1) {
107		return("The parameters `mu_s`, `mu_g` and `mu_phi` must have the same length.")
108		}
109
110		len_1_pars <- c(
111		"sigma", "omega_b", "omega_s", "omega_g", "omega_phi",
112		"link_dsld", "link_ttg", "link_identity", "link_growth"
113		)
114		for (par in len_1_pars) {
115		if (length(slot(object, par)) != 1) {
116		return(sprintf("The `%s` parameter must be a length 1 numeric.", par))
117		}
118		}
119
120		return(TRUE)
121		}
122		)
123
124		#' @rdname as_print_string
125		as_print_string.SimLongitudinalGSF <- function(object) {
126	1x	return("SimLongitudinalGSF")
127		}
128
129		#' @rdname sampleObservations
130		#' @export
131		sampleObservations.SimLongitudinalGSF <- function(object, times_df) {
132	13x	times_df \|>
133	13x	dplyr::mutate(mu_sld = gsf_sld(.data$time, .data$psi_b, .data$psi_s, .data$psi_g, .data$psi_phi)) \|>
134	13x	dplyr::mutate(dsld = gsf_dsld(.data$time, .data$psi_b, .data$psi_s, .data$psi_g, .data$psi_phi)) \|>
135	13x	dplyr::mutate(ttg = gsf_ttg(.data$time, .data$psi_b, .data$psi_s, .data$psi_g, .data$psi_phi)) \|>
136	13x	dplyr::mutate(sld = stats::rnorm(dplyr::n(), .data$mu_sld, .data$mu_sld * object@sigma)) \|>
137	13x	dplyr::mutate(
138	13x	log_haz_link =
139	13x	(object@link_dsld * .data$dsld) +
140	13x	(object@link_ttg * .data$ttg) +
141	13x	(object@link_identity * .data$mu_sld) +
142	13x	(object@link_growth * log(.data$psi_g))
143		)
144		}
145
146
147		#' @rdname sampleSubjects
148		#' @export
149		sampleSubjects.SimLongitudinalGSF <- function(object, subjects_df) {
150	7x	assert_that(
151	7x	is.factor(subjects_df$study),
152	7x	is.factor(subjects_df$arm),
153	7x	length(levels(subjects_df$study)) == length(object@mu_b),
154	7x	length(levels(subjects_df$arm)) == length(object@mu_s),
155	7x	length(levels(subjects_df$arm)) == length(object@mu_g),
156	7x	length(levels(subjects_df$arm)) == length(object@mu_phi)
157		)
158
159	7x	res <- subjects_df \|>
160	7x	dplyr::distinct(.data$subject, .data$arm, .data$study) \|>
161	7x	dplyr::mutate(study_idx = as.numeric(.data$study)) \|>
162	7x	dplyr::mutate(arm_idx = as.numeric(.data$arm)) \|>
163	7x	dplyr::mutate(psi_b = stats::rlnorm(dplyr::n(), object@mu_b[.data$study_idx], object@omega_b)) \|>
164	7x	dplyr::mutate(psi_s = stats::rlnorm(dplyr::n(), object@mu_s[.data$arm_idx], object@omega_s)) \|>
165	7x	dplyr::mutate(psi_g = stats::rlnorm(dplyr::n(), object@mu_g[.data$arm_idx], object@omega_g)) \|>
166	7x	dplyr::mutate(psi_phi_logit = stats::rnorm(
167	7x	dplyr::n(),
168	7x	object@mu_phi[.data$arm_idx],
169	7x	object@omega_phi
170		)) \|>
171	7x	dplyr::mutate(psi_phi = stats::plogis(.data$psi_phi_logit))
172
173	7x	res[, c("subject", "arm", "study", "psi_b", "psi_s", "psi_g", "psi_phi")]
174		}
175
176
177
178
179		## sim_lm_gsf ----
180
181		#' Generalized Stein-Fojo Functionals
182		#'
183		#' @param time (`numeric`)\cr time grid.
184		#' @param b (`number`)\cr baseline.
185		#' @param s (`number`)\cr shrinkage.
186		#' @param g (`number`)\cr growth.
187		#' @param phi (`number`)\cr shrinkage proportion.
188		#'
189		#' @returns The function results.
190		#'
191		#' @keywords internal
192		gsf_sld <- function(time, b, s, g, phi) {
193	18x	phi <- dplyr::if_else(time >= 0, phi, 0)
194	18x	b * (phi * exp(-s * time) + (1 - phi) * exp(g * time))
195		}
196
197
198		#' @rdname gsf_sld
199		gsf_ttg <- function(time, b, s, g, phi) {
200	16x	t1 <- (log(s * phi / (g * (1 - phi))) / (g + s))
201	16x	t1[t1 <= 0] <- 0
202	16x	return(t1)
203		}
204
205
206		#' @rdname gsf_sld
207		gsf_dsld <- function(time, b, s, g, phi) {
208	16x	phi <- dplyr::if_else(time >= 0, phi, 0)
209	16x	t1 <- (1 - phi) * g * exp(g * time)
210	16x	t2 <- phi * s * exp(-s * time)
211	16x	return(b * (t1 - t2))
212		}

1		#' @include generics.R
2		NULL
3
4
5
6		#' Standard Links
7		#'
8		#' @param prior ([`Prior`]) \cr A [`Prior`] object.
9		#' @param model ([`LongitudinalModel`]) \cr A [`LongitudinalModel`] object.
10		#' @param ... Not used.
11		#'
12		#' @description
13		#' These functions are used to enable the use of the corresponding link function between
14		#' the survival and longitudinal models in a joint model. Note that the exact implementation
15		#' of the link function is model specific, see
16		#' \code{vignette("Statistical Specifications", package = "jmpost")} for more details.
17		#'
18		#' @name standard-link-user
19		NULL
20
21
22
23
24		#' @describeIn standard-link-user No link (fit the survival and longitudinal models independently)
25		#' @export
26		linkNone <- function() {
27	1x	Link()
28		}
29
30
31		#' @describeIn standard-link-user Time to growth link
32		#' @export
33		linkTTG <- function(prior, model = PromiseLongitudinalModel(), ...) {
34	12x	UseMethod("linkTTG", model)
35		}
36		#' @export
37		linkTTG.PromiseLongitudinalModel <- function(prior = prior_normal(0, 2), model, ...) {
38	5x	PromiseLinkComponent(fun = linkTTG, prior = prior, key = "link_ttg")
39		}
40		#' @export
41		linkTTG.default <- function(prior, model, ...) {
42	1x	stop(sprintf("Method `linkTTG` is not available for `%s`", class(model)[[1]]))
43		}
44
45
46
47
48		#' @describeIn standard-link-user Derivative of the SLD over time link
49		#' @export
50		linkDSLD <- function(prior, model = PromiseLongitudinalModel(), ...) {
51	29x	UseMethod("linkDSLD", model)
52		}
53		#' @export
54		linkDSLD.PromiseLongitudinalModel <- function(prior = prior_normal(0, 2), model, ...) {
55	15x	PromiseLinkComponent(fun = linkDSLD, prior = prior, key = "link_dsld")
56		}
57		#' @export
58		linkDSLD.default <- function(prior, model, ...) {
59	!	stop(sprintf("Method `linkDSLD` is not available for `%s`", class(model)[[1]]))
60		}
61
62
63
64
65		#' @describeIn standard-link-user Current SLD value link
66		#' @export
67		linkIdentity <- function(prior, model = PromiseLongitudinalModel(), ...) {
68	8x	UseMethod("linkIdentity", model)
69		}
70		#' @export
71		linkIdentity.PromiseLongitudinalModel <- function(prior = prior_normal(0, 2), model, ...) {
72	3x	PromiseLinkComponent(fun = linkIdentity, prior = prior, key = "link_identity")
73		}
74		#' @export
75		linkIdentity.default <- function(prior, model, ...) {
76	!	stop(sprintf("Method `linkIdentity` is not available for `%s`", class(model)[[1]]))
77		}
78
79
80		#' @describeIn standard-link-user Growth Parameter link
81		#' @export
82		linkGrowth <- function(prior, model = PromiseLongitudinalModel(), ...) {
83	6x	UseMethod("linkGrowth", model)
84		}
85		#' @export
86		linkGrowth.PromiseLongitudinalModel <- function(prior = prior_normal(0, 2), model, ...) {
87	3x	PromiseLinkComponent(fun = linkGrowth, prior = prior, key = "link_growth")
88		}
89		#' @export
90		linkGrowth.default <- function(prior, model, ...) {
91	!	stop(sprintf("Method `linkGrowth` is not available for `%s`", class(model)[[1]]))
92		}

1
2		.onLoad <- function(libname, pkgname) {
3	!	set_options()
4		}
5
6		.onAttach <- function(libname, pkgname) {
7	4x	if (!is_cmdstanr_available()) {
8	!	packageStartupMessage(
9	!	"jmpost uses cmdstanr for compiling and sampling from models, but it does not seem to be installed.\n",
10	!	"To install:\n",
11	!	"install.packages(\"cmdstanr\", repos = c(\"https://stan-dev.r-universe.dev/\", getOption(\"repos\")))"
12		)
13	4x	} else if (is.null(cmdstanr::cmdstan_version(error_on_NA = FALSE))) {
14	!	possible_paths <- unique(c(
15	!	cmdstanr::cmdstan_default_install_path(),
16	!	Sys.getenv("CMDSTAN"),
17	!	Sys.getenv("CMDSTAN_PATH"),
18	!	"/root/.cmdstan",
19	!	"~/.cmdstan"
20		))
21	!	possible_paths <- possible_paths[dir.exists(possible_paths)]
22
23	!	if (length(possible_paths)) {
24	!	for (try_path in possible_paths) {
25	!	new_path <- tryCatch(
26	!	suppressMessages(cmdstanr::set_cmdstan_path(try_path)),
27	!	warning = function(w) NULL,
28	!	error = function(e) NULL
29		)
30		}
31	!	if (!is.null(new_path)) {
32	!	packageStartupMessage("CmdStan path set to: ", new_path)
33		}
34		} else {
35	!	packageStartupMessage("jmpost could not identify CmdStan path. Please use cmdstanr::set_cmdstan_path()")
36		}
37		}
38	4x	return(invisible(NULL))
39		}
40
41
42		# This only exists to silence the false positive R CMD CHECK warning about
43		# importing but not using the posterior package. posterior is a dependency
44		# of rcmdstan that we use a lot implicitly. Also we link to their documentation
45		# pages in ours
46		.never_run <- function() {
47	!	posterior::as_draws()
48		}

1		#' @include StanModel.R
2		NULL
3
4		# SurvivalModel-class ----
5
6		#' `SurvivalModel`
7		#'
8		#' This class extends the general [`StanModel`] class to comprise the survival
9		#' model specification.
10		#'
11		#' @exportClass SurvivalModel
12		.SurvivalModel <- setClass(
13		Class = "SurvivalModel",
14		contains = "StanModel"
15		)
16
17		# SurvivalModel-constructors ----
18
19		#' @rdname SurvivalModel-class
20		#'
21		#' @inheritParams stanmodel_arguments
22		#'
23		#' @export
24		SurvivalModel <- function(
25		stan = StanModule(),
26		parameters = ParameterList(),
27		name = "<Unnamed>",
28		...
29		) {
30	28x	base_stan <- read_stan("base/survival.stan")
31	28x	stan_full <- decorated_render(
32	28x	.x = base_stan,
33	28x	stan = add_missing_stan_blocks(as.list(stan))
34		)
35	28x	.SurvivalModel(
36	28x	StanModel(
37	28x	name = name,
38	28x	stan = StanModule(stan_full),
39	28x	parameters = parameters,
40		...
41		)
42		)
43		}
44
45		#' @export
46		as_print_string.SurvivalModel <- function(object, ...) {
47	4x	string <- sprintf(
48	4x	"\n%s Survival Model with parameters:\n%s\n\n",
49	4x	object@name,
50	4x	paste(" ", as_print_string(object@parameters)) \|> paste(collapse = "\n")
51		)
52	4x	return(string)
53		}

1		#' @include SurvivalModel.R
2		NULL
3
4		# SurvivalLogLogistic-class ----
5
6		#' `SurvivalLogLogistic`
7		#'
8		#' This class extends the general [`SurvivalModel`] class for using the
9		#' log-logistic survival model.
10		#'
11		#' @exportClass SurvivalLogLogistic
12		.SurvivalLogLogistic <- setClass(
13		Class = "SurvivalLogLogistic",
14		contains = "SurvivalModel"
15		)
16
17		# SurvivalLogLogistic-constructors ----
18
19		#' @rdname SurvivalLogLogistic-class
20		#'
21		#' @param a (`Prior`)\cr Prior distribution for the scale parameter `a`.
22		#' @param b (`Prior`)\cr Prior distribution for the shape parameter `b`.
23		#' @param beta (`Prior`)\cr Prior distribution for covariates coefficients `beta`.
24		#'
25		#' @export
26		SurvivalLogLogistic <- function(
27		a = prior_lognormal(log(0.1), 5),
28		b = prior_gamma(2, 5),
29		beta = prior_normal(0, 2)
30		) {
31	4x	.SurvivalLogLogistic(
32	4x	SurvivalModel(
33	4x	name = "Log-Logistic",
34	4x	stan = StanModule("sm-loglogistic/model.stan"),
35	4x	parameters = ParameterList(
36	4x	Parameter(name = "sm_loglogis_a", prior = a, size = 1),
37	4x	Parameter(name = "sm_loglogis_b", prior = b, size = 1),
38	4x	Parameter(name = "beta_os_cov", prior = beta, size = "p_os_cov_design")
39		)
40		)
41		)
42		}

1
2		#' Define Simulation Group
3		#'
4		#' Specifies a simulation group to be used by [`SimJointData()`].
5		#'
6		#' @param n (`numeric`)\cr number of subjects in the group.
7		#' @param arm (`character`)\cr treatment arm.
8		#' @param study (`character`)\cr study name.
9		#'
10		#' @slot n (`numeric`)\cr See arguments.
11		#' @slot arm (`character`)\cr See arguments.
12		#' @slot study (`character`)\cr See arguments.
13		#'
14		#' @examples
15		#' SimGroup(n = 50, arm = "Arm-A", study = "Study-1")
16		#' @name SimGroup-class
17		#' @exportClass SimGroup
18		.SimGroup <- setClass(
19		"SimGroup",
20		slots = c(
21		n = "numeric",
22		arm = "character",
23		study = "character"
24		)
25		)
26
27		#' @export
28		#' @rdname SimGroup-class
29		SimGroup <- function(n, arm, study) {
30	49x	.SimGroup(
31	49x	n = n,
32	49x	arm = arm,
33	49x	study = study
34		)
35		}
36
37
38		setValidity(
39		"SimGroup",
40		function(object) {
41		if (length(object@n) != 1) {
42		return("`n` must be a length 1 integer")
43		}
44		if (length(object@arm) != 1) {
45		return("`arm` must be a length 1 string")
46		}
47		if (length(object@study) != 1) {
48		return("`study` must be a length 1 string")
49		}
50		if (any(object@n < 1) \| any(object@n %% 1 != 0)) {
51		return("`n` must be positive integer")
52		}
53		return(TRUE)
54		}
55		)

1
2		#' Abstract Simulation Class for Longitudinal Data
3		#'
4		#' @param times (`numeric`) the times to generate observations at.
5		#'
6		#' @description
7		#' This class exists to be extended by other classes that simulate longitudinal data.
8		#' It is not intended to be used directly.
9		#' @name SimLongitudinal-class
10		#' @family SimLongitudinal
11		#' @exportClass SimLongitudinal
12		.SimLongitudinal <- setClass(
13		"SimLongitudinal",
14		slots = list(
15		times = "numeric"
16		)
17		)
18
19		#' @rdname SimLongitudinal-class
20		#' @export
21		SimLongitudinal <- function(times = seq(0, 100, 50)) {
22	!	.SimLongitudinal(times = times)
23		}
24
25
26		#' @rdname show-object
27		#' @export
28		setMethod(
29		f = "show",
30		signature = "SimLongitudinal",
31		definition = function(object) {
32	4x	x <- sprintf("\nA %s Object\n\n", as_print_string(object))
33	4x	cat(x)
34	4x	return(object)
35		}
36		)
37
38		#' @rdname as_print_string
39		as_print_string.SimLongitudinal <- function(object) {
40	!	return("SimLongitudinal")
41		}

1		#' @include SurvivalModel.R
2		NULL
3
4		# SurvivalWeibullPH-class ----
5
6		#' `SurvivalWeibullPH`
7		#'
8		#' This class extends the general [`SurvivalModel`] class for using the
9		#' Weibull proportional hazards survival model.
10		#'
11		#' @exportClass SurvivalWeibullPH
12		.SurvivalWeibullPH <- setClass(
13		Class = "SurvivalWeibullPH",
14		contains = "SurvivalModel"
15		)
16
17		# SurvivalWeibullPH-constructors ----
18
19		#' @rdname SurvivalWeibullPH-class
20		#'
21		#' @param lambda (`Prior`)\cr for the scale `lambda`.
22		#' @param gamma (`Prior`)\cr for the shape `gamma`.
23		#' @param beta (`Prior`)\cr for covariates coefficients `beta`.
24		#'
25		#' @export
26		SurvivalWeibullPH <- function(
27		lambda = prior_gamma(2, 0.5),
28		gamma = prior_gamma(2, 0.5),
29		beta = prior_normal(0, 2)
30		) {
31	11x	.SurvivalWeibullPH(
32	11x	SurvivalModel(
33	11x	name = "Weibull-PH",
34	11x	stan = StanModule(x = "sm-weibull-ph/model.stan"),
35	11x	parameters = ParameterList(
36	11x	Parameter(name = "sm_weibull_ph_lambda", prior = lambda, size = 1),
37	11x	Parameter(name = "sm_weibull_ph_gamma", prior = gamma, size = 1),
38	11x	Parameter(name = "beta_os_cov", prior = beta, size = "p_os_cov_design")
39		)
40		)
41		)
42		}