jmpost coverage - 87.99%

Files
Source

#' @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.
#'
#' @typed stan: StanModule
#'   Stan code.
#' @typed x: LinkComponent
#'   a link component.
#' @typed object: LinkComponent
#'   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.
#'
#' @typed stan: StanModule
#'   Stan code. See Details.
#' @typed prior: Prior
#'   The prior for the scaling coeficient.
#' @typed key: character
#'   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
#'
#' @examples
#' LinkComponent(
#'   stan = StanModule(),
#'   prior = prior_normal(0, 1),
#'   key = "custom_link"
#' )
.LinkComponent <- setClass(
    Class = "LinkComponent",
    slots = list(
        "stan" = "StanModule",
        "parameters" = "ParameterList",
        "key" = "character"
    )
)


#' @rdname LinkComponent-class
#' @export
#'
#' @returns A `LinkComponent` object.
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
#'
#' @returns A `StanModule` object.
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
#'
#' @returns A named `list` representation of the object.
as.list.LinkComponent <- function(x, ...) {
    stan <- as.StanModule(x, ...)
    as.list(stan)
}

#' @family LinkComponent
#' @export
#'
#' @returns A character vector suitable for printing.
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
#'
#' @returns A character vector containing the component's parameter names.
names.LinkComponent <- function(x, ...) {
    names(x@parameters)
}

#' @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`
#'
#' @typed object: JointModel
#'   Joint model specification.
#' @typed x: JointModel
#'   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 declaration and prior 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"
    )
)

#' @typed longitudinal: "`LongitudinalModel` or `NULL`"
#'   the longitudinal model.
#' @typed survival: "`SurvivalModel` or `NULL`"
#'   the survival model.
#' @typed link: Link
#'   the link.
#' @rdname JointModel-class
#'
#' @returns A `JointModel` object.
#'
#' @examples
#' JointModel(LongitudinalGSF(), SurvivalWeibullPH(), linkDSLD())
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
#'
#' @returns A `StanModule` object containing the generated-quantities code.
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
#'
#' @returns A `StanModule` object.
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)),
        has_os_submodel = !is.null(object@survival),
        has_long_submodel = !is.null(object@longitudinal)
    )
    # 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)

    merge(
        StanModule("base/functions.stan"),
        StanModule(stan_full)
    )
}


#' `JointModel` -> `character`
#'
#' Renders a [`JointModel`] object to a stan program
#'
#' @inheritParams JointModel-Shared
#' @family JointModel
#' @export
#'
#' @returns A character vector.
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)
    invisible()
}


# compileStanModel-JointModel ----

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


# sampleStanModel-JointModel ----

#' @rdname sampleStanModel
#'
#' @typed data: "`DataJoint` or `list`"
#'   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
#'
#' @typed initial_values: list
#'   A list of lists containing the initial values
#'   must be 1 list per desired chain. All elements should have identical names
#' @typed data: list
#'   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.
#' @typed parameters: ParameterList
#'   the parameters object
#'
#' @details
#' This function is mostly a thin wrapper around `expand_initial_values` to
#' enable easier unit testing.
#'
#' @keywords internal
#'
#' @returns A list of initial values, one element per chain.
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
        ))
    }
)


#' @rdname getRandomEffectsNames
#' @export
getRandomEffectsNames.JointModel <- function(object, ...) {
    if (is.null(object@longitudinal)) {
        return(NULL)
    }
    getRandomEffectsNames(object@longitudinal)
}

#' @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
#'
#' @typed x: list
#'   list of Stan code blocks
#' @typed stan_blocks: list
#'   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
#'
#' @typed x: string
#'   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
#'
#' @returns A `StanModule` object.
#'
#' @examples
#' code <- "
#' parameters {
#'   real x;
#' }
#' model {
#'   x ~ normal(0, 1);
#' }
#' "
#' StanModule(code)
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`
#' @typed x: StanModule
#'   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
#'
#' @returns A character vector.
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 ----

#' @typed stan_blocks: list
#'   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
#' @export
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.
#' @typed x: StanModule
#'   A Stan Module
#' @typed stan_blocks: list
#'   reference list of stan blocks.
#' @param ... Not Used.
#' @family StanModule
#' @export
#'
#' @returns A named `list` representation of the object.
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.
#'
#' @typed filename: string
#'   file name.
#'
#' @keywords internal
#'
#' @returns `TRUE` if the input identifies an existing file; otherwise `FALSE`.
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.
#'
#' @returns A character vector containing Stan program code.
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
#'
#' @typed functions: character
#'   code block.
#' @typed data: character
#'   code block.
#' @typed transformed_data: character
#'   code block.
#' @typed parameters: character
#'   code block.
#' @typed transformed_parameters: character
#'   code block.
#' @typed model: character
#'   code block.
#' @typed generated_quantities: character
#'   code block.
#' @typed stan_blocks: list
#'   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
#'
#' @typed x: character
#'   the single Stan code vector.
#' @typed stan_blocks: list
#'   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.
#' @typed object: StanModule
#'   A stan program
#' @family StanModule
#' @typed indent: numeric
#'   how much white space to prefix the print string with.
#' @keywords internal
#' @export
#'
#' @returns A character vector suitable for printing.
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")
    }
)

#' Calculate Population Hazard Ratios
#'
#' Calculates hazard ratios marginalised over subject specific random effects using the
#' approach proposed by \insertCite{oudenhoven2020marginal}{jmpost}.
#'
#' @typed object: JointModelSamples
#'   samples as drawn from a Joint Model.
#' @typed hr_formula: formula
#'   defines the terms to include in the hazard ratio calculation.
#'   By default this uses the right side of the formula used in the survival model.
#'   Set to `NULL` not include any terms
#' @typed baseline: formula
#'   terms to model baseline hazard using variable `time`.
#'   Default is a B-spline from [splines]: `~bs(time, df = 10)`
#' @typed quantiles: numeric
#'   vector of two values in (0, 1) for calculating quantiles from log hazard ratio
#'   distributions.
#'
#' @references \insertAllCited{}
#'
#' @returns A list containing a summary of parameter distributions as a `data.frame` and a
#' matrix containing the parameter estimates for each sample.
#' @export
#' @importFrom splines bs
#'
#' @examples
#' \dontrun{
#' populationHR(fit, quantiles = c(0.01, 0.99))
#' }
populationHR <- function(
    object,
    hr_formula = object@data@survival@formula,
    baseline = ~ bs(time, df = 10),
    quantiles = c(0.025, 0.975)
) {
    assert_class(object, "JointModelSamples")
    assert_formula(hr_formula)
    assert_formula(baseline)
    assert_numeric(
        quantiles,
        lower = 0,
        upper = 1,
        any.missing = FALSE,
        unique = TRUE
    )
    if (!"time" %in% all.vars(baseline)) {
        stop("baseline formula should include a time term.")
    }

    # Extract the variable names used in the data
    subject_var <- object@data@subject@subject
    arm_var <- object@data@subject@arm
    long_time_var <- all.vars(delete.response(terms(
        object@data@longitudinal@formula
    )))
    surv_time_var <- all.vars(object@data@survival@formula[[2]][[2]])
    surv_covs <- all.vars(delete.response(terms(hr_formula)))
    if (!all(surv_covs %in% colnames(object@data@survival@data))) {
        stop("All variables in hr_formula must be in survival data")
    }

    marginal_formula <- stats::reformulate(c(
        attr(terms(baseline), "term.labels"),
        attr(terms(hr_formula), "term.labels")
    ))

    # Get the survival quantities at the observed longitudinal and survival times for each patient:
    times_df <- rbind(
        stats::setNames(
            object@data@longitudinal@data[, c(subject_var, long_time_var)],
            c("subject", "time")
        ),
        stats::setNames(
            object@data@survival@data[, c(subject_var, surv_time_var)],
            c("subject", "time")
        )
    )
    times_df <- times_df[order(times_df$subject, times_df$time), ]
    times_df <- times_df[times_df$time > 0, ]
    times_df <- times_df[!duplicated.data.frame(times_df), ]

    # Generate samples of the log hazard log h_i(t) for each patient i at these time points t.
    grid_spec <- split(times_df$time, times_df$subject)
    log_haz_samples <- SurvivalQuantities(
        object,
        grid = GridManual(grid_spec),
        type = "loghaz"
    )@quantities@quantities |>
        t()

    # Construct \tilde{X} in paper's notation with one row per patient's time
    W_df <- dplyr::left_join(
        times_df,
        stats::setNames(
            object@data@survival@data[, c(subject_var, surv_covs)],
            c("subject", surv_covs)
        ),
        by = "subject"
    )
    W_mat <- model.matrix(marginal_formula, W_df)

    # As model matrix contains baseline and covariates, we don't need the intercept term
    # but we want factor variables encoded relative to the intercept
    W_mat <- W_mat[, colnames(W_mat) != "(Intercept)", drop = FALSE]

    estimates <- stats::lm.fit(x = W_mat, y = log_haz_samples)$coefficients
    tidy_res <- apply(estimates, 1, function(x) {
        quantiles <- stats::quantile(x, probs = quantiles)
        c(mean = mean(x), median = median(x), quantiles)
    }) |>
        t() |>
        data.frame()

    list(summary = tidy_res, estimates)
}

#' Simulating Joint Longitudinal and Time-to-Event Data
#'
#' @typed design: list
#'   a list of [`SimGroup`] objects. See details.
#' @typed longitudinal: SimLongitudinal
#'   object specifying how to simulate the longitudinal data
#' @typed survival: SimSurvival
#'   object specifying how to simulate the survival data
#' @typed .silent: flag
#'   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
#'
#' @examples
#' set.seed(1)
#' SimJointData(
#'   design = list(SimGroup(3, "A", "Study 1")),
#'   longitudinal = SimLongitudinalRandomSlope(
#'     times = c(0, 10),
#'     slope_mu = 0.01,
#'     slope_sigma = 0.5
#'   ),
#'   survival = SimSurvivalExponential(lambda = 1 / 300)
#' )
.SimJointData <- setClass(
    "SimJointData",
    slots = list(
        longitudinal = "data.frame",
        survival = "data.frame"
    )
)


#' @rdname SimJointData-class
#' @export
#'
#' @returns A `SimJointData` object.
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)
    }
)


#' Add PFS events at Tumour Progression to Data
#'
#' Adds new columns `pfs_time` and `pfs_event` based on observed changes to SLD.
#'
#' @typed object: SimJointData
#'   where the PFS events should be added.
#' @typed relative_threshold: number
#'   a multiplicative threshold for the change in SLD compared to the `min(SLD)`.
#'   Default is 1.2 meaning a 20% increase.
#' @typed absolute_threshold: number
#'   an absolute threshold for the change in SLD compared to the minimum.
#'   Default is 5.
#' @typed from_time: number
#'   Ignore observations before this time for determining SLD minimum.
#' @typed observed_after: logical
#'   If `FALSE` set longitudinal observations after the progression time to
#'   `observed = FALSE`
#' @details
#' Both thresholds must be met for a progression to be declared.
#'
#' @export
#' @examples
#' data <- SimJointData(
#'   survival = SimSurvivalExponential(lambda = 1/10),
#'   longitudinal = SimLongitudinalSteinFojo()
#'   )
#' data <- add_pfs(data)
#' data@survival # now has pfs_time and pfs_event columns
#'
#' @returns The input `SimJointData` object with progression-free survival columns
#'   `pfs_time` and `pfs_event` added.
add_pfs <- function(
    object,
    relative_threshold = 1.2,
    absolute_threshold = 5,
    from_time = 0,
    observed_after = FALSE
) {
    assert_class(object, "SimJointData")

    pd_times <- object@longitudinal |>
        dplyr::filter(.data$time >= from_time) |>
        dplyr::mutate(
            min_sld = cummin(.data$sld),
            is_pd = .data$sld >=
                pmax(
                    .data$min_sld * relative_threshold,
                    .data$min_sld + absolute_threshold
                ) &
                .data$observed,
            pd_time = min(.data$time[.data$is_pd], Inf),
            .by = "subject"
        ) |>
        dplyr::select("subject", "pd_time") |>
        dplyr::slice_head(by = "subject")

    if (isFALSE(observed_after)) {
        object@longitudinal <- object@longitudinal |>
            dplyr::left_join(pd_times, by = "subject") |>
            dplyr::mutate(
                observed = dplyr::if_else(
                    .data$time > .data$pd_time,
                    FALSE,
                    .data$observed
                ),
                pd_time = NULL
            )
    }

    object@survival <-
        object@survival |>
        dplyr::left_join(pd_times, by = "subject") |>
        dplyr::mutate(
            pfs_time = pmin(.data$time, .data$pd_time, na.rm = TRUE),
            pfs_event = dplyr::if_else(
                .data$pfs_time < .data$time,
                1,
                .data$event
            ),
            pd_time = NULL
        )
    object
}


#' Cut Study Data
#' @typed object: SimJointData
#'   where the cut times should be applied.
#' @typed cut_time: numeric
#'   A vector of cut off times, either length 1 for all patients or
#'   `nrow(object@survival)` for a time per patient.
#' @details
#'   All observations after this time are remove. Survival is censored at this time and any longitudinal
#'   values are removed.
#' @export
#' @examples
#' data <- SimJointData(
#'   survival = SimSurvivalExponential(lambda = 1/10),
#'   longitudinal = SimLongitudinalSteinFojo()
#'   )
#' data <- cut_data(data, 5)
#' data@survival
#' # Now max time is 5
#' max(data@survival$time)
#'
#' @returns The input `SimJointData` object restricted to the requested follow-up time.
cut_data <- function(object, cut_time) {
    assert_class(object, "SimJointData")
    check_len <- if (length(cut_time) > 1) nrow(object@survival) else 1
    assert_numeric(cut_time, lower = 0, len = check_len)

    object@survival <- object@survival |>
        dplyr::mutate(
            cut_time = cut_time,
            event = dplyr::if_else(.data$time < .data$cut_time, .data$event, 0),
            time = pmin(.data$cut_time, .data$time)
        )
    if (all(c("pfs_event", "pfs_time") %in% colnames(object@survival))) {
        object@survival <- object@survival |>
            dplyr::mutate(
                pfs_event = dplyr::if_else(
                    .data$pfs_time < .data$cut_time,
                    .data$pfs_event,
                    0
                ),
                pfs_time = pmin(.data$cut_time, .data$pfs_time)
            )
    }

    object@longitudinal <- object@longitudinal |>
        dplyr::left_join(
            dplyr::select(object@survival, "subject", "cut_time"),
            by = "subject"
        ) |>
        dplyr::filter(.data$time <= .data$cut_time) |>
        dplyr::mutate(cut_time = NULL)

    object@survival$cut_time <- NULL

    object
}

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

#' `SimSurvival` Function Arguments
#'
#' The documentation lists all the conventional arguments for [`SimSurvival`]
#' constructors.
#'
#' @typed time_max: number
#'   the maximum time to simulate to.
#' @typed time_step: number
#'   the time interval between evaluating the log-hazard function.
#' @typed lambda_censor: number
#'   the censoring rate.
#' @typed beta_cont: number
#'   the continuous covariate coefficient.
#' @typed beta_cat: numeric
#'   the categorical covariate coefficients.
#' @typed loghazard: function
#'   the log hazard function.
#' @typed name: character
#'   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.
#' @slot beta_os_cov (`numeric`) Additional regression coefficients for survival models.
#'
#' @family SimSurvival
#' @exportClass SimSurvival
#' @name SimSurvival-class
#'
#' @examples
#' SimSurvival(loghazard = function(time, ...) rep(log(1 / 100), length(time)))
.SimSurvival <- setClass(
    "SimSurvival",
    slots = c(
        time_max = "numeric",
        time_step = "numeric",
        lambda_censor = "numeric",
        beta_cont = "numeric",
        beta_cat = "numeric",
        beta_os_cov = "numeric",
        loghazard = "function",
        name = "character"
    )
)

#' @rdname SimSurvival-class
#' @export
#'
#' @returns A `SimSurvival` object.
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
#' @exportS3Method NULL
as_print_string.SimSurvival <- function(object, ...) {
    return(object@name)
}


#' Construct Time Intervals
#'
#' @typed object: SimSurvival
#'   the survival simulation object to create evaluation points for.
#' @param ... Not Used.
#'
#' @return A `tibble` with `lower`, `upper`, `time`, `eval` and `width`.
#' @keywords internal
#' @exportS3Method NULL
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 subject(s) 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)

    keep_cols <- colnames(os_dat_complete) %in%
        c("subject", "study", "arm", "time", "event", "cov_cont", "cov_cat")
    os_dat_complete[, keep_cols]
}


#' Simulate Survival Data from a Weibull Proportional Hazard Model
#'
#' @typed lambda: number
#'   the scale parameter.
#' @typed gamma: number
#'   the shape parameter.
#'
#' @inheritParams SimSurvival-Shared
#' @inheritSection SimSurvival-Shared Hazard Evaluation
#'
#' @family SimSurvival
#'
#' @export
#'
#' @returns A `SimSurvival` object.
#'
#' @examples
#' SimSurvivalWeibullPH(lambda = 1 / 300, gamma = 1.2)
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
#'
#' @typed a: number
#'   the scale parameter.
#' @typed b: number
#'   the shape parameter.
#'
#' @inheritParams SimSurvival-Shared
#' @inheritSection SimSurvival-Shared Hazard Evaluation
#'
#' @family SimSurvival
#' @export
#'
#' @returns A `SimSurvival` object.
#'
#' @examples
#' SimSurvivalLogLogistic(a = 300, b = 2)
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
#'
#' @typed lambda: number
#'   the rate parameter.
#'
#' @inheritParams SimSurvival-Shared
#' @inheritSection SimSurvival-Shared Hazard Evaluation
#'
#' @family SimSurvival
#'
#' @export
#'
#' @returns A `SimSurvival` object.
#'
#' @examples
#' SimSurvivalExponential(lambda = 1 / 300)
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"
    )
}


#' Simulate Survival Data from a Gamma Proportional Hazard Model
#'
#' @typed k: number
#'   the shape parameter.
#' @typed theta: number
#'   the scale parameter.
#'
#' @inheritParams SimSurvival-Shared
#' @inheritSection SimSurvival-Shared Hazard Evaluation
#'
#' @family SimSurvival
#'
#' @importFrom stats dgamma pgamma
#'
#' @export
#'
#' @returns A `SimSurvival` object.
#'
#' @examples
#' SimSurvivalGamma(k = 2, theta = 300)
SimSurvivalGamma <- function(
    k,
    theta,
    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) {
            dgamma(time, k, scale = theta, log = TRUE) -
                pgamma(time, k, scale = theta, lower.tail = FALSE, log.p = TRUE)
        },
        name = "SimSurvivalGamma"
    )
}

#' Re-used documentation for `RandomEffectQuantities`
#'
#' @typed x: RandomEffectQuantities
#'   generated quantities.
#' @typed object: RandomEffectQuantities
#'   generated quantities.
#' @typed conf.level: numeric
#'   confidence level of the interval.
#' @param ... not used.
#'
#' @keywords internal
#' @name RandomEffectQuantities-Shared
NULL


#' Random Effects Quantities Container
#'
#' A simple wrapper around a `matrix` to store required metadata for patient level
#' random effects data
#'
#' @typed quantities: matrix
#'   of random effects values.
#' @typed subject: character
#'   labels specifying which subjects the values belong to.
#' @typed parameter: character
#'   labels specifying which parameter the value is.
#'
#' @slot quantities (`matrix`)\cr See Arguments for details.
#' @slot subject (`numeric`)\cr See Arguments for details.
#' @slot parameter (`character`)\cr See Arguments for details.
#'
#' @details
#' Each row of the matrix represents a sample and each column represents a distinct subject
#' specific parameter.
#' As such the number of columns in the matrix should equal the length of `subject` and `parameter`
#' which provide metadata for who the parameter corresponds to as well as which parameter it is.
#'
#' @keywords internal
#' @name RandomEffectQuantities-class
#' @family RandomEffectQuantities
.RandomEffectQuantities <- setClass(
    "RandomEffectQuantities",
    slots = list(
        "quantities" = "matrix",
        "subject" = "character",
        "parameter" = "character"
    )
)
#' @rdname RandomEffectQuantities-class
#'
#' @returns A `RandomEffectQuantities` object.
RandomEffectQuantities <- function(quantities, subject, parameter) {
    .RandomEffectQuantities(
        quantities = quantities,
        subject = subject,
        parameter = parameter
    )
}

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


#' `RandomEffectQuantities` -> Printable `Character`
#'
#' Converts [`RandomEffectQuantities`] object into a printable string.
#' @inheritParams RandomEffectQuantities-Shared
#' @typed indent: numeric
#'   the number of spaces to indent the string by.
#' @family RandomEffectQuantities
#' @keywords internal
#' @export
#'
#' @returns A character vector suitable for printing.
as_print_string.RandomEffectQuantities <- function(object, indent = 1, ...) {
    parameter_string <- paste0("        ", unique(object@parameter))
    template <- c(
        "RandomEffectQuantities Object:",
        "    # of samples         = %d",
        "    # of unique subjects = %d",
        "    For parameters:",
        parameter_string
    )
    pad <- rep(" ", indent) |> paste(collapse = "")
    template_padded <- paste(pad, template)
    sprintf(
        paste(template_padded, collapse = "\n"),
        nrow(object@quantities),
        length(unique(object@subject))
    )
}


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


#' `RandomEffectQuantities` -> `data.frame`
#'
#' Returns a `data.frame` of the subject-level random effect parameter samples.
#'
#' @inheritParams RandomEffectQuantities-Shared
#'
#' @keywords internal
#' @family RandomEffectQuantities
#' @export
#'
#' @returns A `data.frame` representation of the object.
as.data.frame.RandomEffectQuantities <- function(x, ...) {
    data.frame(
        subject = rep(x@subject, each = nrow(x@quantities)),
        parameter = rep(x@parameter, each = nrow(x@quantities)),
        values = as.vector(x@quantities)
    )
}


#' summary
#'
#' @description
#' This method returns a summary statistic `data.frame` of the random effect parameters
#'
#' @inheritParams RandomEffectQuantities-Shared
#'
#' @returns
#' A `data.frame` with the following variables:
#' - `subject` (`character`) \cr the subject identifier.
#' - `parameter` (`character`) \cr the parameter identifier.
#' - `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.
#'
#' @keywords internal
#' @family RandomEffectQuantities
#' @export
summary.RandomEffectQuantities <- function(object, conf.level = 0.95, ...) {
    quantities_summarised <- samples_median_ci(
        object@quantities,
        level = conf.level
    )

    quantities_summarised$subject <- object@subject
    quantities_summarised$parameter <- object@parameter
    quantities_summarised[, c(
        "subject",
        "parameter",
        "median",
        "lower",
        "upper"
    )]
}


#' Extract Random Effects Samples from a Longitudinal Model
#'
#' Helper function to extract subject-level random effects samples from the longitudinal
#' sub-model of a joint model samples object.
#'
#' @typed object: JointModelSamples
#'   samples as drawn from a Joint Model.
#' @family RandomEffectQuantities
#' @family JointModelSamples
#' @export
#'
#' @returns A `RandomEffectQuantities` object.
#'
#' @examples
#' \dontrun{
#' LongitudinalRandomEffects(fit)
#' }
LongitudinalRandomEffects <- function(object) {
    assert_class(object, "JointModelSamples")

    subject_indexes <- as.list(object@data)$subject_to_index
    random_effects_names <- getRandomEffectsNames(object@model)
    expanded <- expand.grid(
        parameter_long = random_effects_names,
        subject_indexes = subject_indexes
    )
    expanded$subject <- names(subject_indexes)[expanded$subject_indexes]
    expanded$parameter <- names(random_effects_names)[expanded$parameter_long]
    stan_parameter_names <- sprintf(
        "%s[%i]",
        expanded$parameter_long,
        expanded$subject_indexes
    )
    draw_matrix <- posterior::as_draws_matrix(object@results$draws(
        stan_parameter_names
    ))
    class(draw_matrix) <- "matrix"
    rownames(draw_matrix) <- NULL
    colnames(draw_matrix) <- NULL

    RandomEffectQuantities(
        draw_matrix,
        subject = expanded$subject,
        parameter = expanded$parameter
    )
}

#' @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
#'
#' @returns A named `list` representation of the object.
as.list.GridObserved <- function(x, data, ...) {
    subjects_to_list(x@subjects, data)
}

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

#' Re-used documentation for `DataLongitudinal`
#'
#' @typed object: DataLongitudinal
#'   Longitudinal Data.
#' @typed x: DataLongitudinal
#'   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"
    )
)


#' @typed data: data.frame
#'   containing the observed longitudinal data.
#' @typed formula: formula
#'   of the form `outcome ~ time`, and cannot contain any additional covariates.
#' @typed threshold: numeric
#'   cut-off value to be used to declare an observation as censored
#'   (below detection limit).
#' @rdname DataLongitudinal-class
#'
#' @returns A `DataLongitudinal` object.
#'
#' @examples
#' DataLongitudinal(tumor_data, sld ~ year, threshold = 5)
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
#' @exportS3Method NULL
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
#'
#' @returns A `data.frame` representation of the object.
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
#' @exportS3Method NULL
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
#' @typed indent: numeric
#'   how much white space to prefix the print string with.
#' @keywords internal
#' @export
#'
#' @returns A character vector suitable for printing.
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 StanModel.R
NULL

# LongitudinalModel-class ----

#' `LongitudinalModel`
#'
#' This class extends the general [`StanModel`] class to comprise the longitudinal
#' model specification.
#'
#' @slot scaled_variance logical scalar, whether the variance should be scaled
#'   by the expected value, corresponding to a multiplicative model. If not,
#'   then an additive error model is used.
#'
#' @exportClass LongitudinalModel
.LongitudinalModel <- setClass(
    Class = "LongitudinalModel",
    contains = "StanModel",
    slots = c(
        scaled_variance = "logical"
    )
)

# LongitudinalModel-constructors ----

#' @rdname LongitudinalModel-class
#'
#' @inheritParams stanmodel_arguments
#' @typed scaled_variance: flag
#'   whether the variance should be scaled
#'   by the expected value, corresponding to a multiplicative model. If not,
#'   then an additive error model is used.
#'
#' @export
#'
#' @returns A `LongitudinalModel` object.
#'
#' @examples
#' LongitudinalModel(scaled_variance = FALSE)
LongitudinalModel <- function(
    stan = StanModule(),
    parameters = ParameterList(),
    name = "<Unnamed>",
    scaled_variance = NA,
    ...
) {
    assert_flag(scaled_variance)
    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,
            ...
        ),
        scaled_variance = scaled_variance
    )
}

#' @export
#'
#' @returns A character vector suitable for printing.
as_print_string.LongitudinalModel <- function(object, ...) {
    string <- sprintf(
        "\n%s Longitudinal Model (%s error) with parameters:\n%s\n\n",
        object@name,
        if (is.na(object@scaled_variance)) {
            # Should not happen but just to be robust against this case.
            "unknown"
        } else if (object@scaled_variance) {
            "multiplicative"
        } else {
            "additive"
        },
        paste("   ", as_print_string(object@parameters)) |>
            paste(collapse = "\n")
    )
    return(string)
}

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

#' Simulate Longitudinal Data from a Random Slope Model
#'
#' @typed times: numeric
#'   the times to generate observations at.
#' @typed intercept: numeric
#'   the mean baseline value for each study.
#' @typed slope_mu: numeric
#'   the population slope for each treatment arm.
#' @typed slope_sigma: numeric
#'   the random slope standard deviation for each treatment arm.
#' @typed sigma: number
#'   the variance of the longitudinal values.
#' @typed link_dsld: number
#'   the link coefficient for the DSLD contribution.
#' @typed link_identity: number
#'   the link coefficient for the identity contribution.
#' @typed scaled_variance: logical
#'   whether the variance should be scaled by the expected value
#'   (see the "Statistical Specifications" vignette for more details)
#'
#' @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.
#' @slot scaled_variance (`logical`)\cr See arguments.
#'
#' @family SimLongitudinal
#' @name SimLongitudinalRandomSlope-class
#' @exportClass SimLongitudinalRandomSlope
#'
#' @examples
#' SimLongitudinalRandomSlope(times = c(0, 30, 60))
.SimLongitudinalRandomSlope <- setClass(
    "SimLongitudinalRandomSlope",
    contains = "SimLongitudinal",
    slots = c(
        intercept = "numeric",
        slope_mu = "numeric",
        slope_sigma = "numeric",
        sigma = "numeric",
        link_dsld = "numeric",
        link_identity = "numeric",
        scaled_variance = "logical"
    )
)

#' @rdname SimLongitudinalRandomSlope-class
#' @export
#'
#' @returns A `SimLongitudinalRandomSlope` object.
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 = c(0.5, 0.6),
    sigma = 2,
    link_dsld = 0,
    link_identity = 0,
    scaled_variance = FALSE
) {
    if (length(slope_sigma) == 1) {
        slope_sigma <- rep(slope_sigma, length(slope_mu))
    }

    .SimLongitudinalRandomSlope(
        times = times,
        intercept = intercept,
        slope_mu = slope_mu,
        slope_sigma = slope_sigma,
        sigma = sigma,
        link_dsld = link_dsld,
        link_identity = link_identity,
        scaled_variance = scaled_variance
    )
}

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

#' @rdname sampleObservations
#' @export
sampleObservations.SimLongitudinalRandomSlope <- function(object, times_df) {
    times_df |>
        dplyr::mutate(
            sld_mu = .data$intercept + .data$slope_ind * .data$time,
            sld_sd = ifelse(
                object@scaled_variance,
                pmax(.data$sld_mu * object@sigma, .Machine$double.eps),
                object@sigma
            ),
            sld = stats::rnorm(dplyr::n(), .data$sld_mu, .data$sld_sd),
            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) == nlevels(subjects_df[["arm"]]),
        msg = "`length(slope_mu)` should be equal to the number of arms"
    )

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

    assert_that(
        length(object@intercept) == nlevels(subjects_df[["study"]]),
        msg = "`length(intercept)` should be equal to the number of 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[as.numeric(.data$arm)]
            )
        )
}

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


#' Re-used documentation for `SurvivalQuantities`
#'
#' @typed object: SurvivalQuantities
#'   survival quantities.
#' @typed x: SurvivalQuantities
#'   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
#'
#' @examples
#' \dontrun{
#' SurvivalQuantities(fit, GridFixed(times = seq(0, 100, by = 10)))
#' }
.SurvivalQuantities <- setClass(
    Class = "SurvivalQuantities",
    slots = c(
        "quantities" = "Quantities",
        "grid" = "Grid",
        "type" = "character",
        "data" = "DataJoint"
    )
)

#' @typed object: JointModelSamples
#'   samples as drawn from a Joint Model.
#'
#' @typed grid: Grid
#'   object that specifies which subjects and time points to calculate the
#'   quantities for. See [Grid-Functions].
#'
#' @typed type: character
#'   quantity to be generated.
#'   Must be one of `surv`, `haz`, `loghaz`, `cumhaz`.
#'
#' @rdname SurvivalQuantities-class
#'
#' @returns A `SurvivalQuantities` object.
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`
#'
#' @typed x: SurvivalQuantities
#'   longitudinal quantities.
#' @param ... not used.
#' @family SurvivalQuantities
#' @export
#'
#' @returns A `data.frame` representation of the object.
as.data.frame.SurvivalQuantities <- function(x, ...) {
    as.data.frame(x@quantities)
}


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


#' Automatic Plotting for `SurvivalQuantities``
#'
#' @inheritParams SurvivalQuantities-Shared
#' @typed add_km: logical
#'   if `TRUE` Kaplan-Meier curves will be added to the plot for
#'   each group/subject.
#' @typed add_wrap: logical
#'   if `TRUE` will apply a [ggplot2::facet_wrap()] to the plot
#'   by each group/subject.
#' @typed conf.level: numeric
#'   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
#'
#' @returns A `ggplot` object.
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.
#'
#' @typed data: data.frame
#'   summary statistics for a survival
#'   curve to be plotted. See details.
#' @typed add_ci: logical
#'   should confidence intervals be added? Default = `TRUE`.
#' @typed add_wrap: logical
#'   should the plots be wrapped by `data$group`? Default = `TRUE`.
#' @typed kmdf: "`data.frame` or `NULL`"
#'   event times and status used to plot
#'   overlaying KM curves. If `NULL` no KM curve will be plotted. See details.
#' @typed y_label: "`character` or `expression`"
#'   label to display on the y-axis.
#'   Default = `expression(S(t))`.
#' @typed x_label: "`character` or `expression`"
#'   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
#'
#' @returns A `ggplot` object.
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
#'
#' @returns A `data.frame` containing Brier scores by time and group.
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 <- as.numeric(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
    )
}

#' Re-used documentation for `DataSubject`
#'
#' @typed object: DataSubject
#'   subject-level data.
#' @typed x: DataSubject
#'   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"
    )
)


#' @typed data: data.frame
#'   the subject-level data.
#' @typed subject: character
#'   the name of the variable containing the subject identifier.
#' @typed arm: character
#'   the name of the variable containing the arm identifier.
#' @typed study: character
#'   the name of the variable containing the study identifier.
#' @rdname DataSubject-class
#' @export
#'
#' @returns A `DataSubject` object.
#'
#' @examples
#' dat <- transform(os_data, study = "Study 1")
#' DataSubject(dat, subject = "id", arm = "arm", study = "study")
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
#'
#' @returns A `data.frame` representation of the object.
as.data.frame.DataSubject <- function(x, ...) {
    x <- x@data
    rownames(x) <- NULL
    x
}


#' @rdname harmonise
#' @exportS3Method NULL
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
#' @typed indent: numeric
#'   how much white space to prefix the print string with.
#' @export
#'
#' @returns A character vector suitable for printing.
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")
    }
)

#' Row Numbers of Data with Missing Variables
#'
#' @typed df: data.frame
#'   input data.
#' @typed formula: "`formula` or `NULL`"
#'   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.
#'
#' @typed data: data.frame
#'   input data.
#' @typed formula: "`formula` or `NULL`"
#'   which variables to inspect for missingness.
#' @typed extra_vars: character
#'   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
#'
#' @typed initial_values: list
#'   initial values with names.
#' @typed sizes: list
#'   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
#'
#' @typed sizes: list
#'   may include character elements that correspond to
#'   names in the data list.
#' @typed data: list
#'   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
#'
#' @typed samples: matrix
#'   with samples in rows and parameters in columns.
#' @typed level: number
#'   credibility level to use for the credible intervals.
#'
#' @returns A `data.frame` with columns `median`, `lower` and `upper`.
#' @keywords internal
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 = getOption("jmpost.double_eps"),
        machine_double_neg_eps = getOption("jmpost.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
#'
#' @typed time_grid: numeric
#'   A vector of times which quantities will be
#'   evaluated at.
#'
#' @keywords internal
#'
#' @returns Invisibly returns `time_grid`; invalid grids raise an error.
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`
#'
#' @typed subjects: "`character` or `NULL`"
#'   Character vector representing the subjects.
#'   If NULL, it will be set to the value of `all_subjects`.
#' @typed all_subjects: character
#'   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()])
#'
#' @typed subjects: "`character` or `list`"
#'   subject identifiers. If `NULL` will be set to `all_subjects`.
#'
#' @typed all_subjects: character
#'   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 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
#'
#' @returns This is just a placeholder which will error when used.
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 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
#'
#' @typed mu_bsld: Prior
#'   for the mean baseline value `mu_bsld`.
#' @typed mu_ks: Prior
#'   for the mean shrinkage rate `mu_ks`.
#' @typed mu_kg: Prior
#'   for the mean growth rate `mu_kg`.
#' @typed mu_phi: Prior
#'   for the mean proportion of cells affected by the treatment `mu_phi`.
#'
#' @typed omega_bsld: Prior
#'   for the baseline value standard deviation `omega_bsld`.
#' @typed omega_ks: Prior
#'   for the shrinkage rate standard deviation `omega_ks`.
#' @typed omega_kg: Prior
#'   for the growth rate standard deviation `omega_kg`.
#' @typed omega_phi: Prior
#'   for the standard deviation of the proportion of cells
#'   affected by the treatment `omega_phi`.
#'
#' @typed sigma: Prior
#'   for the variance of the longitudinal values `sigma`.
#'
#' @typed centred: logical
#'   whether to use the centred parameterization.
#' @typed scaled_variance: logical
#'   whether the variance should be scaled by the expected value, corresponding to a multiplicative model.
#'   As a default, the variance is not scaled by the expected value, corresponding to an additive model.
#'   (See the "Statistical Specifications" vignette for more details.)
#'
#' @importFrom stats qlogis
#' @export
#'
#' @returns A `LongitudinalGSF` object.
#'
#' @examples
#' LongitudinalGSF()
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),

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

    # Apply constraints
    omega_bsld <- set_limits(omega_bsld, lower = getOption("jmpost.double_eps"))
    omega_ks <- set_limits(omega_ks, lower = getOption("jmpost.double_eps"))
    omega_kg <- set_limits(omega_kg, lower = getOption("jmpost.double_eps"))
    omega_phi <- set_limits(omega_phi, lower = getOption("jmpost.double_eps"))
    sigma <- set_limits(sigma, lower = getOption("jmpost.double_eps"))

    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 = "n_studies"
        ),
        Parameter(name = "lm_gsf_omega_ks", prior = omega_ks, size = "n_arms"),
        Parameter(name = "lm_gsf_omega_kg", prior = omega_kg, size = "n_arms"),
        Parameter(
            name = "lm_gsf_omega_phi",
            prior = omega_phi,
            size = "n_arms"
        ),

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

    assert_flag(centred)
    parameters_extra <- if (centred) {
        list(
            Parameter(
                name = "lm_gsf_psi_bsld",
                prior = set_limits(prior_init_only(prior_lognormal(
                    median(mu_bsld),
                    median(omega_bsld)
                )), lower = getOption("jmpost.double_eps")),
                size = "n_subjects"
            ),
            Parameter(
                name = "lm_gsf_psi_ks",
                prior = set_limits(prior_init_only(prior_lognormal(
                    median(mu_ks),
                    median(omega_ks)
                )), lower = getOption("jmpost.double_eps")),
                size = "n_subjects"
            ),
            Parameter(
                name = "lm_gsf_psi_kg",
                prior = set_limits(prior_init_only(prior_lognormal(
                    median(mu_kg),
                    median(omega_kg)
                )), lower = getOption("jmpost.double_eps")),
                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),
        scaled_variance = scaled_variance
    )
    .LongitudinalGSF(x)
}


#' @export
#'
#' @returns A `StanModule` object containing the generated-quantities code.
enableGQ.LongitudinalGSF <- function(object, ...) {
    StanModule("lm-gsf/quantities.stan")
}


#' @export
#'
#' @returns The longitudinal model with its link-related Stan code enabled.
enableLink.LongitudinalGSF <- function(object, ...) {
    object@stan <- merge(
        object@stan,
        StanModule("lm-gsf/link.stan")
    )
    object
}

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

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

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

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


#' @export
#'
#' @returns A `LinkComponent` object.
linkShrinkage.LongitudinalGSF <- function(
    prior = prior_normal(0, 2),
    model,
    ...
) {
    LinkComponent(
        key = "link_shrinkage",
        stan = StanModule("lm-gsf/link_shrinkage.stan"),
        prior = prior
    )
}

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

#' @rdname getRandomEffectsNames
#' @export
getRandomEffectsNames.LongitudinalGSF <- function(object, ...) {
    c(
        "b" = "lm_gsf_psi_bsld",
        "s" = "lm_gsf_psi_ks",
        "g" = "lm_gsf_psi_kg",
        "phi" = "lm_gsf_psi_phi"
    )
}

#' Re-used documentation for Brier Score components
#'
#' @typed t: numeric
#'   timepoints to calculate the desired quantity at.
#' @typed times: numeric
#'   observed times.
#' @typed events: numeric
#'   event indicator for `times`. Either 1 for an event or 0 for censor.
#' @typed event_offset: logical
#'   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.
#' @typed maintain_cen_order: logical
#'   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
#'
#' @returns The reverse Kaplan-Meier estimate at each requested time.
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
#'
#' @typed x: numeric
#'   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
#'
#' @returns An integer vector giving the order that groups identical values.
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
#'
#' @returns A numeric vector or matrix of Brier-score components.
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 generics.R
#' @include utilities.R
NULL

#' Re-used documentation for `DataSurvival`
#'
#' @typed object: DataSurvival
#'   Survival Data.
#' @typed x: DataSurvival
#'   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"
    )
)

#' @typed data: data.frame
#'   the observed time-to-event data.
#' @typed formula: formula
#'   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
#'
#' @returns A `DataSurvival` object.
#'
#' @examples
#' DataSurvival(os_data, Surv(os_time, os_event) ~ age + sex)
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]])
        lhs_surv <- identical(LHS, "Surv")
        lhs_colon_surv <- identical(LHS, c("::", "survival", "Surv"))
        if (!(lhs_surv || lhs_colon_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
#'
#' @returns A `data.frame` representation of the object.
as.data.frame.DataSurvival <- function(x, ...) {
    x <- x@data
    rownames(x) <- NULL
    x
}


#' @rdname as_stan_list.DataObject
#' @family DataSurvival
#' @param df (`data.frame`)\cr The data frame to be used for generating the design matrix.
#' @param vars (`list`)\cr The list of variable names to be used for generating the design matrix.
#' @export
model.matrix.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)
    if (length(remove_index) > 0) {
        design_mat <- design_mat[, -remove_index, drop = FALSE]
    }
    rownames(design_mat) <- NULL
    design_mat
}

#' Extract Covariate Names for Survival Data
#'
#' @typed object: DataSurvival
#'  Survival data object to extract covariate names from.
#' @param ... not used.
#'
#' @family covariates
#' @export
#'
#' @returns A character vector containing the covariate names.
covariates.DataSurvival <- function(object, ...) {
    design_mat <- model.matrix(object, ...)
    colnames(design_mat)
}

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

    design_mat <- model.matrix(object, df = df, vars = vars)

    # 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
#' @exportS3Method NULL
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
#' @typed indent: numeric
#'   how much white space to prefix the print string with.
#' @keywords internal
#' @export
#'
#' @returns A character vector suitable for printing.
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
#'
#' @typed olddata: DataSurvival
#'   The original data to be used as a template for the new data
#' @typed newdata: data.frame
#'   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
#'
#' @returns A model matrix whose columns match the original design matrix.
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
#'
#' @returns A `formula` object.
as_formula.DataSurvival <- function(x, ...) {
    vars <- extractVariableNames(x)
    vars$frm
}

#' Simulate Patients from Posterior Predictive Distribution
#' @typed object: JointModelSamples
#'   the object containing the posterior samples.
#' @typed newdata: data.frame
#'   A data frame containing data in the same format as
#'   the `object@data@survival@data`. Importantly, it should contain the same covariates and factor levels
#'   as the variables used in the survival formula `object@data@survival@formula` and
#'   the same columns used for `study`, `id`, and `arm`.
#' @param ... unused
#'   Unused.
#' @typed times: numeric
#'   times to simulate SLD for all patients.
#' @typed jitter_var: numeric
#'   variances to add noise to the observed SLD `times`. The first value is for any times
#'   less than 0 and the second for any times after 0. All positive (negative) times will remain positive (negative).
#'    Jitter values are generated from a normal distribution with mean 0 and the given variances.
#' @typed time_max: number
#'   the maximum time to simulate to.
#' @typed time_step: number
#'   the time interval between evaluating the log-hazard function.
#' @typed lambda_censor: number
#'   the censoring rate, as the parameter of an exponential distribution.
#' @typed scaled_variance: flag
#'   Should variance be scaled by the expected value. Must be set the same as was used for
#' model fitting.
#' @param seed ignored.
#' @param nsim ignored.
#' @details Simulates a set of patients based on the covariates of those
#' used in the model fit or from `newdata`, which must contain the
#' same column names and factor levels.
#'
#' @importFrom stats simulate
#' @export
#'
#' @returns A `SimJointData` object containing posterior-predictive simulations.
simulate.JointModelSamples <- function(
    object,
    nsim = NULL,
    seed = NULL,
    newdata = NULL,
    times = c(-2, 0, 10, 50, 100),
    jitter_var = c(0, 0),
    time_max = 2000,
    time_step = 1,
    lambda_censor = 1 / 3000,
    scaled_variance = FALSE,
    ...
) {
    subj_data <- if (is.null(newdata)) {
        dplyr::left_join(
            object@data@subject@data[, c(
                object@data@subject@subject,
                object@data@subject@arm,
                object@data@subject@study
            )],
            object@data@survival@data[, c(
                object@data@subject@subject,
                all.vars(delete.response(terms(object@data@survival@formula)))
            )],
            by = object@data@subject@subject
        )
    } else {
        newdata
    }

    subject_data <- DataSubject(
        subj_data,
        subject = object@data@subject@subject,
        arm = object@data@subject@arm,
        study = object@data@subject@study
    )

    n_patients <- nrow(subject_data@data)

    requiredVars <- list(
        longitudinal = lapply(
            object@model@survival@parameters@parameters,
            function(x) x@name
        ),
        survival = lapply(
            object@model@longitudinal@parameters@parameters,
            function(x) x@name
        ),
        link = lapply(object@model@link@components, function(x) x@key)
    )

    draws <- object@results$draws(
        variables = unlist(requiredVars),
        format = "draws_matrix"
    )
    draw_id <- sample.int(nrow(draws), n_patients, replace = FALSE)

    long_models <- list()
    surv_models <- list()
    for (i in seq.int(n_patients)) {
        draw <- draws[draw_id[i], ]
        long_models[[i]] <- createLongitudinalSimObject(
            object@model@longitudinal,
            draw,
            times = add_jitter(times, jitter_var = jitter_var),
            scaled_variance = scaled_variance
        )
        surv_models[[i]] <- createSurvivalSimObject(
            object@model@survival,
            draw,
            lambda_censor = lambda_censor,
            time_step = time_step,
            time_max = time_max
        )
    }

    SimJointDataResults(
        subject = subject_data,
        surv_formula = object@data@survival@formula,
        longitudinal = long_models,
        survival = surv_models
    )
}


add_jitter <- function(times, jitter_var = c(0, 0)) {
    stopifnot(all(jitter_var >= 0))
    if (jitter_var[1] > 0) {
        neg_times <- times < 0
        times[neg_times] <- pmin(
            times[neg_times] + rnorm(sum(neg_times), sd = sqrt(jitter_var[1])),
            -.Machine$double.eps / 2
        )
    }
    if (jitter_var[2] > 0) {
        pos_times <- times > 0
        times[pos_times] <- pmax(
            times[pos_times] + rnorm(sum(pos_times), sd = sqrt(jitter_var[2])),
            .Machine$double.eps / 2
        )
    }
    times
}


# Longitudinal Sim Object constructors --------
#' @noRd
#'
#' @returns A `SimLongitudinal` object configured from the posterior draw.
createLongitudinalSimObject <- function(object, draw, ...) {
    UseMethod("createLongitudinalSimObject")
}

#' Get draws value by name
#' @typed draws: matrix
#'   the draws input.
#' @typed name: character
#'   to match column names of `draws` with [startsWith].
#' @return A vector of matching values or `NULL` if no match.
#' @noRd
get_vars <- function(draws, name) {
    result <- as.numeric(draws[, startsWith(colnames(draws), name)])
    if (length(result) == 0) {
        NULL
    } else {
        result
    }
}

#' @exportS3Method
#'
#' @returns A `SimLongitudinal` object configured from the posterior draw.
createLongitudinalSimObject.LongitudinalRandomSlope <- function(
    object,
    draw,
    ...
) {
    args <- list(...)

    args$intercept <- get_vars(draw, "lm_rs_intercept")
    args$slope_mu <- get_vars(draw, "lm_rs_slope_mu")
    args$slope_sigma <- get_vars(draw, "lm_rs_slope_sigma")
    args$sigma <- get_vars(draw, "lm_rs_sigma")
    args$link_dsld <- get_vars(draw, "link_dsld")
    args$link_identity <- get_vars(draw, "link_identity")
    args$scaled_variance <- NULL # not defined for random slope
    do.call(SimLongitudinalRandomSlope, args)
}

#' @exportS3Method
#'
#' @returns A `SimLongitudinal` object configured from the posterior draw.
createLongitudinalSimObject.LongitudinalSteinFojo <- function(
    object,
    draw,
    ...
) {
    args <- list(...)
    args$sigma <- get_vars(draw, "lm_sf_sigma")
    args$mu_s <- get_vars(draw, "lm_sf_mu_ks")
    args$mu_g <- get_vars(draw, "lm_sf_mu_kg")
    args$mu_b <- get_vars(draw, "lm_sf_mu_bsld")
    args$omega_s <- get_vars(draw, "lm_sf_omega_ks")
    args$omega_g <- get_vars(draw, "lm_sf_omega_kg")
    args$omega_b <- get_vars(draw, "lm_sf_omega_bsld")
    args$link_dsld <- get_vars(draw, "link_dsld")
    args$link_ttg <- get_vars(draw, "link_ttg")
    args$link_identity <- get_vars(draw, "link_identity")
    args$link_growth <- get_vars(draw, "link_growth")
    args$link_shrinkage <- get_vars(draw, "link_shrinkage")
    do.call(SimLongitudinalSteinFojo, args)
}

#' @exportS3Method
#'
#' @returns A `SimLongitudinal` object configured from the posterior draw.
createLongitudinalSimObject.LongitudinalGSF <- function(object, draw, ...) {
    args <- list(...)
    args$sigma <- get_vars(draw, "lm_gsf_sigma")
    args$mu_s <- get_vars(draw, "lm_gsf_mu_ks")
    args$mu_g <- get_vars(draw, "lm_gsf_mu_kg")
    args$mu_b <- get_vars(draw, "lm_gsf_mu_bsld")
    args$mu_phi <- get_vars(draw, "lm_gsf_mu_phi")
    args$omega_s <- get_vars(draw, "lm_gsf_omega_ks")
    args$omega_g <- get_vars(draw, "lm_gsf_omega_kg")
    args$omega_b <- get_vars(draw, "lm_gsf_omega_bsld")
    args$omega_phi <- get_vars(draw, "lm_gsf_omega_phi")
    args$link_dsld <- get_vars(draw, "link_dsld")
    args$link_ttg <- get_vars(draw, "link_ttg")
    args$link_identity <- get_vars(draw, "link_identity")
    args$link_growth <- get_vars(draw, "link_growth")
    args$link_shrinkage <- get_vars(draw, "link_shrinkage")

    do.call(SimLongitudinalGSF, args)
}

#' @exportS3Method
#'
#' @returns A `SimLongitudinal` object configured from the posterior draw.
createLongitudinalSimObject.LongitudinalClaretBruno <- function(
    object,
    draw,
    ...
) {
    args <- list(...)
    args$sigma <- get_vars(draw, "lm_clbr_sigma")
    args$mu_b <- get_vars(draw, "lm_clbr_mu_b")
    args$mu_g <- get_vars(draw, "lm_clbr_mu_g")
    args$mu_c <- get_vars(draw, "lm_clbr_mu_c")
    args$mu_p <- get_vars(draw, "lm_clbr_mu_p")
    args$omega_b <- get_vars(draw, "lm_clbr_omega_b")
    args$omega_g <- get_vars(draw, "lm_clbr_omega_g")
    args$omega_c <- get_vars(draw, "lm_clbr_omega_c")
    args$omega_p <- get_vars(draw, "lm_clbr_omega_p")
    args$link_dsld <- get_vars(draw, "link_dsld")
    args$link_ttg <- get_vars(draw, "link_ttg")
    args$link_identity <- get_vars(draw, "link_identity")
    args$link_growth <- get_vars(draw, "link_growth")
    args$link_shrinkage <- get_vars(draw, "link_shrinkage")
    do.call(SimLongitudinalClaretBruno, args)
}


# Survival Sim Object constructors --------

#' @noRd
#'
#' @returns A `SimSurvival` object configured from the posterior draw.
createSurvivalSimObject <- function(object, draw, ...) {
    UseMethod("createSurvivalSimObject")
}

#' @exportS3Method
#'
#' @returns A `SimSurvival` object configured from the posterior draw.
createSurvivalSimObject.SurvivalWeibullPH <- function(object, draw, ...) {
    args <- list(...)
    args$lambda <- get_vars(draw, "sm_weibull_ph_lambda")
    args$gamma <- get_vars(draw, "sm_weibull_ph_gamma")

    result <- do.call(SimSurvivalWeibullPH, args)
    result@beta_os_cov <- get_vars(draw, "beta_os_cov")
    result
}

#' @exportS3Method
#'
#' @returns A `SimSurvival` object configured from the posterior draw.
createSurvivalSimObject.SurvivalExponential <- function(object, draw, ...) {
    args <- list(...)
    args$lambda <- get_vars(draw, "sm_exp_lambda")

    result <- do.call(SimSurvivalExponential, args)
    result@beta_os_cov <- get_vars(draw, "beta_os_cov")
    result
}

#' @exportS3Method
#'
#' @returns A `SimSurvival` object configured from the posterior draw.
createSurvivalSimObject.SurvivalGamma <- function(object, draw, ...) {
    lapply(object@model@survival@parameters@parameters, function(x) x@name)
    lapply(object@model@link@components, function(x) x@key)
    args <- list(...)
    args$k <- get_vars(draw, "sm_gamma_k")
    args$theta <- get_vars(draw, "sm_gamma_theta")

    result <- do.call(SimSurvivalGamma, args)
    result@beta_os_cov <- get_vars(draw, "beta_os_cov")
    result
}

#' @exportS3Method
#'
#' @returns A `SimSurvival` object configured from the posterior draw.
createSurvivalSimObject.SurvivalLogLogistic <- function(object, draw, ...) {
    lapply(object@model@survival@parameters@parameters, function(x) x@name)
    lapply(object@model@link@components, function(x) x@key)
    args <- list(...)
    args$a <- get_vars(draw, "sm_loglogis_a")
    args$b <- get_vars(draw, "sm_loglogis_b")

    result <- do.call(SimSurvivalLogLogistic, args)
    result@beta_os_cov <- get_vars(draw, "beta_os_cov")
    result
}


sampleSubjectsFromObs <- function(object, subjects_df, covs_matrix) {
    subjects_df$log_haz_cov <- rowSums(
        covs_matrix[, -1] * t(sapply(object, function(x) x@beta_os_cov))
    )
    subjects_df |>
        dplyr::mutate(survival = stats::runif(dplyr::n())) |>
        dplyr::mutate(chazard_limit = -log(.data$survival)) |>
        dplyr::mutate(
            time_cen = stats::rexp(
                dplyr::n(),
                sapply(object, function(x) x@lambda_censor)
            )
        )
}

# Simulate patient function ------
SimJointDataResults <- function(
    subject,
    surv_formula,
    longitudinal,
    survival,
    .silent = FALSE
) {
    # take hazard windows from first element. All have the same
    # @time_step and @time_max parameters
    hazard_evaluation_info <- hazardWindows(survival[[1]])

    group_counts <- as.data.frame(table(subject@data[, c(
        subject@study,
        subject@arm
    )]))

    n_group <- group_counts$Freq
    arms <- group_counts[, 2]
    studies <- group_counts[, 1]
    n_subjects <- sum(n_group)
    n_times <- length(hazard_evaluation_info$midpoint)

    baseline <- data.frame(
        subject = subject@data[[subject@subject]],
        study = subject@data[[subject@study]],
        arm = subject@data[[subject@arm]]
    )

    cov_cols <- cbind(
        subject = subject@data[[subject@subject]],
        subject@data[, all.vars(delete.response(terms(surv_formula)))]
    )

    os_baseline <- sampleSubjectsFromObs(
        survival,
        subjects_df = baseline,
        model.matrix(delete.response(terms(surv_formula)), subject@data)
    )

    lm_baseline <- dplyr::bind_rows(
        lapply(seq.int(longitudinal), function(i) {
            sampleSubjects(longitudinal[[i]], subjects_df = baseline[i, ])
        })
    )

    lm_times <- lapply(seq.int(longitudinal), function(i) {
        longitudinal[[i]]@times
    })

    lm_dat_no_obvs <-
        cbind(
            baseline[rep(seq.int(nrow(baseline)), times = lengths(lm_times)), ],
            time = unlist(lm_times)
        ) |>
        dplyr::left_join(lm_baseline, by = c("subject", "study", "arm"))

    lm_dat <- dplyr::bind_rows(
        .mapply(
            sampleObservations,
            dots = list(
                times_df = split(lm_dat_no_obvs, ~subject),
                object = longitudinal
            ),
            MoreArgs = NULL
        )
    )

    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 <- dplyr::bind_rows(
        .mapply(
            sampleObservations,
            dots = list(
                times_df = split(
                    dplyr::left_join(
                        hazard_eval_df,
                        lm_baseline,
                        by = c("subject", "study", "arm")
                    ),
                    ~subject
                ),
                object = longitudinal
            ),
            MoreArgs = NULL
        )
    )[, 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"))

    times_df <- split(os_eval_df, ~subject)
    os_data_list <- list()
    for (i in seq_along(times_df)) {
        times_df[[i]]$log_bl_haz <- survival[[i]]@loghazard(
            times_df[[i]]$midpoint
        )
    }
    os_dat <- sampleObservations2(dplyr::bind_rows(times_df))

    os_dat <- dplyr::left_join(os_dat, cov_cols, by = "subject")

    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[[1]]@times),
        length(unique(lm_dat2$subject)) == n_subjects,
        msg = "Assumptions for the Longitudinal data are not met (please report this issue)"
    )

    original_names <- c("subject", "arm", "study")
    names(original_names) <- c(subject@subject, subject@arm, subject@study)

    return(
        .SimJointData(
            survival = os_dat |> dplyr::rename(dplyr::all_of(original_names)),
            longitudinal = lm_dat2[, c(
                "subject",
                "arm",
                "study",
                "time",
                "sld",
                "observed"
            )] |>
                dplyr::rename(dplyr::all_of(original_names))
        )
    )
}


# adapted from sampleObservations so that $midpoint and $log_bl_haz are already in times_df.
sampleObservations2 <- function(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 |>
        # 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
            ),
            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
            hazard_instant = dplyr::if_else(
                .data$hazard_instant == Inf,
                999,
                .data$hazard_instant
            ),
            hazard_instant = dplyr::if_else(
                .data$hazard_instant == -Inf,
                -999,
                .data$hazard_instant
            ),
            hazard_interval = .data$hazard_instant * .data$width
        ) |>
        dplyr::mutate(chazard = cumsum(.data$hazard_interval), .by = "subject")

    os_had_event <- os_dat_chaz |>
        dplyr::filter(.data$chazard >= .data$chazard_limit) |>
        dplyr::slice_head(by = "subject") |>
        dplyr::mutate(event = 1)

    os_had_censor <- os_dat_chaz |>
        dplyr::filter(!.data$subject %in% os_had_event$subject) |>
        dplyr::slice_tail(by = "subject") |>

        dplyr::mutate(event = 0)

    if (!(nrow(os_had_censor) == 0)) {
        message(sprintf(
            "INFO: %i subject(s) 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,
            event = dplyr::if_else(
                .data$real_time <= .data$time_cen,
                .data$event,
                0
            ),
            time = dplyr::if_else(
                .data$real_time <= .data$time_cen,
                .data$real_time,
                .data$time_cen
            )
        ) |>
        dplyr::arrange(.data$subject)

    keep_cols <- colnames(os_dat_complete) %in%
        c("subject", "study", "arm", "time", "event", "cov_cont", "cov_cat")
    os_dat_complete[, keep_cols]
}

#' @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
#'
#' @returns This is just a placeholder which will error when used.
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 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)
}

# "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
#'
#' @examples
#' show(prior_normal(0, 1))
NULL


# write_stan ----

#' `write_stan`
#'
#' Write the Stan code for a Stan module.
#'
#' @param object the module.
#' @typed destination: "`character` or  `connection`"
#'   Where to write stan code to.
#' @param ... Additional arguments
#'
#' @export
#'
#' @returns Invisibly returns `NULL` after writing the Stan program.
#'
#' @examples
#' path <- tempfile(fileext = ".stan")
#' write_stan(JointModel(LongitudinalGSF()), path)
#' unlink(path)
write_stan <- function(object, destination, ...) {
    UseMethod("write_stan")
}

# compileStanModel ----

#' `compileStanModel`
#'
#' Compile the Stan module.
#'
#' @param object the module.
#'
#' @export
#'
#' @returns A compiled `cmdstanr::CmdStanModel` object.
#'
#' @examples
#' \dontrun{
#' compileStanModel(JointModel(LongitudinalGSF()))
#' }
compileStanModel <- function(object) {
    UseMethod("compileStanModel")
}


# sampleStanModel ----

#' `sampleStanModel`
#'
#' Sample from a Stan Module.
#'
#' @param object the module.
#' @param ... additional arguments.
#'
#' @export
#'
#' @returns A `JointModelSamples` object.
#'
#' @examples
#' \dontrun{
#' sampleStanModel(JointModel(LongitudinalGSF()), data = joint_data)
#' }
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
#'
#' @returns A `StanModule` object.
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
#'
#' @returns A named list mapping source variables to standard names.
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
#'
#' @returns Initial values for the object, usually as a numeric vector or list.
#'
#' @examples
#' initialValues(prior_normal(0, 1))
initialValues <- function(object, ...) {
    UseMethod("initialValues")
}


# size ----

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


# auxiliaryInitialValues ----

#' `auxiliaryInitialValues`
#'
#' Obtain initial values for any additional Stan parameters introduced by an object.
#'
#' @param object where to get the auxiliary initial values from.
#' @param ... additional options.
#'
#' @details
#' Some objects, such as shrinkage priors, declare additional Stan parameters
#' beyond the main model parameter. This helper returns named initial values for
#' those additional parameters.
#'
#' @keywords internal
#'
#' @returns A named `list` of initial values for auxiliary parameters.
auxiliaryInitialValues <- function(object, ...) {
    UseMethod("auxiliaryInitialValues")
}


# auxiliarySize ----

#' `auxiliarySize`
#'
#' Obtain sizes for any additional Stan parameters introduced by an object.
#'
#' @param object where to get the auxiliary parameter sizes from.
#' @param ... additional options.
#'
#' @details
#' Some objects, such as shrinkage priors, declare additional Stan parameters
#' beyond the main model parameter. This helper returns named sizes for those
#' additional parameters so their initial values can be expanded consistently.
#'
#' @keywords internal
#'
#' @returns A named `list` of auxiliary parameter sizes.
auxiliarySize <- function(object, ...) {
    UseMethod("auxiliarySize")
}


# generateQuantities ----

#' `generateQuantities`
#'
#' Obtain the generated quantities from a Stan Model.
#'
#' @param object object to obtain generated quantities from
#' @param ... additional options.
#'
#' @export
#'
#' @returns A `cmdstanr::CmdStanGQ` object containing generated quantities.
#'
#' @examples
#' \dontrun{
#' generateQuantities(fit, generator, "survival")
#' }
generateQuantities <- function(object, ...) {
    UseMethod("generateQuantities")
}


#' Prepare Data Object
#'
#' @typed object: "`DataSubject` or `DataLongitudinal` or `DataSurvival`"
#'   data object to "harmonise"
#' @typed subject_var: character
#'   the name of the variable containing the subject identifier.
#' @typed subject_ord: character
#'   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
#'
#' @returns A named `list` suitable for use as Stan data.
#'
#' @examples
#' as_stan_list(Parameter(prior_normal(0, 1), "beta"))
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
#'
#' @returns A character vector suitable for printing.
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
#'
#' @returns A numeric Brier score or a `data.frame` of scores over time.
#'
#' @examples
#' \dontrun{
#' brierScore(survival_quantities)
#' }
brierScore <- function(object, ...) {
    UseMethod("brierScore")
}


#' Generate Simulated Observations
#'
#' @typed object: "`SimLongitudinal` or `SimSurvival`"
#'   object to generate observations from.
#' @typed times_df: data.frame
#'   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
#'
#' @returns A `data.frame` containing simulated observations.
#'
#' @examples
#' sim <- SimLongitudinalRandomSlope(slope_mu = 0.01, slope_sigma = 0.5)
#' subjects <- data.frame(study = factor("S"), arm = factor("A"), subject = "1")
#' subject_parameters <- sampleSubjects(sim, subjects)
#' sampleObservations(sim, transform(subject_parameters, time = 0))
sampleObservations <- function(object, times_df) {
    UseMethod("sampleObservations")
}


#' Generate Simulated Subjects
#'
#' @typed object: "`SimLongitudinal` or `SimSurvival`"
#'   object to generate subjects from.
#' @typed subjects_df: data.frame
#'   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
#'
#' @returns A `data.frame` containing simulated subject-level parameters.
#'
#' @examples
#' subjects <- data.frame(study = factor("S"), arm = factor("A"), subject = "1")
#' sim <- SimLongitudinalRandomSlope(slope_mu = 0.01, slope_sigma = 0.5)
#' sampleSubjects(sim, subjects)
sampleSubjects <- function(object, subjects_df) {
    UseMethod("sampleSubjects")
}


#' Generate time windows for evaluating a hazard function
#'
#' @typed object: SurvivalModel
#'   object to generate time windows for.
#' @param ... Not used.
#'
#'
#' @returns A `data.frame` describing the hazard-evaluation intervals.
hazardWindows <- function(object, ...) {
    UseMethod("hazardWindows")
}

#' @rdname Quant-Dev
#' @export
#'
#' @returns A `QuantityGenerator` or `QuantityCollapser` object, according to the function called.
#'
#' @examples
#' \dontrun{
#' grid <- GridFixed(times = c(0, 10))
#' as.QuantityGenerator(grid, joint_data)
#' as.QuantityCollapser(grid, joint_data)
#' }
as.QuantityGenerator <- function(object, ...) {
    UseMethod("as.QuantityGenerator")
}

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


#' Coalesce Time
#'
#' @typed object: Grid
#'   object to coalesce time for.
#' @typed times: numeric
#'   the times to coalesce to.
#' @param ... Not used
#'
#' Method used to replace NULL times on grid objects (if appropriate)
#'
#' @keywords internal
#'
#' @returns A `Grid` object with missing time points replaced.
coalesceGridTime <- function(object, times, ...) {
    UseMethod("coalesceGridTime")
}
#' @export
#'
#' @returns The input object, unchanged.
coalesceGridTime.default <- function(object, times, ...) {
    object
}


#' Resolve a Promise
#'
#' @typed object: ANY
#'   an object to resolve.
#' @typed ...: ANY
#'   additional arguments.
#'
#' If `object` is not a promise will just return itself else will resolve the promise
#' and return the promised object.
#'
#' @export
#'
#' @returns The input object, or the value obtained by resolving a promise.
#'
#' @examples
#' resolvePromise(1)
#' resolvePromise(linkDSLD(), LongitudinalGSF())
resolvePromise <- function(object, ...) {
    UseMethod("resolvePromise")
}

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

#' Enable Link Generic
#'
#' @typed object: LongitudinalModel
#'   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
#'
#' @examples
#' enableLink(LongitudinalGSF())
enableLink <- function(object, ...) {
    UseMethod("enableLink")
}
#' @export
#'
#' @returns The longitudinal model with its link-related Stan code enabled.
enableLink.default <- function(object, ...) {
    object
}


#' Enable Generated Quantities Generic
#'
#' @typed object: StanModel
#'   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
#'
#' @examples
#' enableGQ(LongitudinalGSF())
enableGQ <- function(object, ...) {
    UseMethod("enableGQ")
}
#' @export
#'
#' @returns A `StanModule` object containing the generated-quantities code.
enableGQ.default <- function(object, ...) {
    StanModule()
}


#' Get Prediction Names
#'
#' Utility function that returns the names of the required parameters for predicting
#' survival quantities with [`GridPrediction`].
#'
#' @typed object: LongitudinalModel
#'   A longitudinal model object
#' @param ... Not used.
#' @export
#'
#' @returns A character vector of required prediction parameter names, or `NULL`.
#'
#' @examples
#' getPredictionNames(LongitudinalGSF())
getPredictionNames <- function(object, ...) {
    UseMethod("getPredictionNames")
}

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


#' Get Random Effects Names
#'
#' Utility function that returns the names of the random effects parameters.
#' The main use for this is to allow the [`LongitudinalRandomEffects`] function
#' to know which parameters it needs to extract and to what common names
#' it should map the parameters to.
#'
#' @typed object: LongitudinalModel
#'   A longitudinal model object
#' @param ... Not used.
#' @export
#'
#' @returns A named character vector of random-effect parameter names, or `NULL`.
#'
#' @examples
#' getRandomEffectsNames(LongitudinalGSF())
getRandomEffectsNames <- function(object, ...) {
    UseMethod("getRandomEffectsNames")
}

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


#' As Formula
#'
#' Utility wrapper function to convert an object to a formula.
#' @typed x: ANY
#'   object to convert to a formula.
#' @param ... Not used.
#' @export
#'
#' @returns A `formula` object.
#'
#' @examples
#' as_formula("response ~ time")
as_formula <- function(x, ...) {
    UseMethod("as_formula")
}

#' @importFrom stats as.formula
#' @export
#'
#' @returns A `formula` object.
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
#'
#' @typed object: Prior
#'   a prior distribution to apply constraints to
#' @typed lower: numeric
#'   lower constraint boundary
#' @typed upper: numeric
#'   upper constraint boundary
#'
#' @export
#'
#' @returns A `Prior` object with the requested bounds.
#'
#' @examples
#' set_limits(prior_normal(0, 1), lower = 0)
set_limits <- function(object, lower = -Inf, upper = Inf) {
    UseMethod("set_limits")
}


#' Save Object to File
#'
#' @typed object: ANY
#'   object to save.
#' @typed file: character
#'   file to save object to.
#' @typed ...: ANY
#'   additional arguments.
#'
#' @family saveObject
#' @export
#'
#' @returns Invisibly returns `NULL` after saving the object.
#'
#' @examples
#' \dontrun{
#' saveObject(fit, "joint-model-fit.rds")
#' }
saveObject <- function(object, file, ...) {
    UseMethod("saveObject")
}


#' Extract Covariate Names
#'
#' @typed object: ANY
#'   the object to extract covariate names from.
#' @param ... additional arguments added by methods.
#'
#' @family covariates
#' @export
#'
#' @returns A character vector containing the covariate names.
#'
#' @examples
#' surv_data <- DataSurvival(os_data, Surv(os_time, os_event) ~ age + sex)
#' covariates(surv_data)
covariates <- function(object, ...) {
    UseMethod("covariates")
}

#' Extract Covariate Shrinkage Factors from the Posterior Samples
#'
#' @typed object: ANY
#'   the object to extract shrinkage factors from.
#' @param ... additional arguments added by methods.
#'
#' @family shrinkage
#' @export
#'
#' @returns A named numeric vector or matrix of shrinkage factors.
#'
#' @examples
#' \dontrun{
#' shrinkage(fit)
#' }
shrinkage <- function(object, ...) {
    UseMethod("shrinkage")
}

.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()"
            )
        }
    }
    invisible()
}

.onLoad <- function(...) {
    set_options()
    s3_register("cmdstanr::as.CmdStanMCMC", "JointModelSamples")
}

# 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 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
#'
#' @typed lambda: Prior
#'   for the scale `lambda`.
#' @typed gamma: Prior
#'   for the shape `gamma`.
#' @typed beta: Prior
#'   for covariates coefficients `beta`.
#'
#' @export
#'
#' @returns A `SurvivalWeibullPH` object.
#'
#' @examples
#' SurvivalWeibullPH()
SurvivalWeibullPH <- function(
    lambda = prior_gamma(2, 0.5),
    gamma = prior_gamma(2, 0.5),
    beta = prior_normal(0, 2)
) {
    lambda <- set_limits(lambda, lower = getOption("jmpost.double_eps"))
    gamma <- set_limits(gamma, lower = getOption("jmpost.double_eps"))

    .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"
                )
            )
        )
    )
}

#' @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.
#'
#' @typed x: Link
#'   a link object.
#' @typed object: Link
#'   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.
#'
#' @typed ...: "[`LinkComponent`] or [`PromiseLinkComponent`]"
#'   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
#'
#' @returns A `Link` object.
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.
#'
#' @typed object: Link
#'   a link object.
#' @typed model: LongitudinalModel
#'   the model object.
#' @param ... Not Used.
#'
#' @export
#'
#' @returns A `Link` object with all promised components resolved.
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
#'
#' @returns A `StanModule` object.
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
#'
#' @returns A named `list` representation of the object.
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
#'
#' @returns A single integer giving the number of link components.
length.Link <- function(x) {
    length(x@components)
}


#' @export
#'
#' @returns A character vector suitable for printing.
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 generics.R
#' @include Parameter.R
#' @include Prior.R
NULL


#' ParameterList-Shared
#' @typed object: ParameterList
#'   A List of [`Parameter`] Objects.
#' @typed x: ParameterList
#'   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. When converted to a [`StanModule`], the list supplies the
#' Stan declarations and prior statements for those parameters.
#'
#'

#' @slot parameters (`list`) \cr a list of [`Parameter`] objects
#' @family ParameterList
#' @export ParameterList
#' @exportClass ParameterList
.ParameterList <- setClass(
    Class = "ParameterList",
    slots = c(
        parameters = "list"
    )
)
#' @typed ...: Parameter
#'   which parameter specifications to include.
#' @rdname ParameterList-class
#'
#' @returns A `ParameterList` object.
#'
#' @examples
#' ParameterList(
#'   Parameter(prior_normal(0, 1), "intercept"),
#'   Parameter(prior_normal(0, 2), "slope")
#' )
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
#'
#' @returns A `StanModule` object.
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
#'
#' @returns A named `list` suitable for use as Stan data.
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
#'
#' @returns A named `list` representation of the object.
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
#' @typed n_chains: integer
#'   the number of chains.
#' @name ParameterList-Getter-Methods
NULL


#' @describeIn ParameterList-Getter-Methods The parameter-list's parameter names
#' @export
#'
#' @returns The requested parameter-list property: names, initial values, or total size.
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, ...) {
    parameters <- Filter(\(x) !x@prior@.is_const, object@parameters)
    # 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(parameters, function(parameter) {
                c(
                    stats::setNames(
                        list(initialValues(parameter)),
                        names(parameter)
                    ),
                    auxiliaryInitialValues(
                        parameter@prior,
                        name = names(parameter),
                        size = size(parameter)
                    )
                )
            })
            unlist(vals, recursive = FALSE)
        }
    )
}


#' @describeIn ParameterList-Getter-Methods The parameter-list's parameter dimensionality
#' @export
size.ParameterList <- function(object) {
    parameters <- Filter(\(x) !x@prior@.is_const, object@parameters)
    sizes <- lapply(parameters, function(parameter) {
        c(
            stats::setNames(list(size(parameter)), names(parameter)),
            auxiliarySize(
                parameter@prior,
                name = names(parameter),
                size = size(parameter)
            )
        )
    })
    unlist(sizes, recursive = FALSE)
}


#' `ParameterList` -> Printable `Character`
#'
#' Converts [`ParameterList`] object into a printable string.
#' @inheritParams ParameterList-Shared
#' @family ParameterList
#' @keywords internal
#' @export
#'
#' @returns A character vector suitable for printing.
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 generics.R
NULL


#' Standard Links
#'
#' @typed prior: Prior
#'   A [`Prior`] object.
#' @typed model: LongitudinalModel
#'   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
#'
#' @examples
#' linkNone()
#' linkDSLD()
#' linkTTG(prior_normal(0, 1), model = LongitudinalGSF())
NULL


#' @describeIn standard-link-user No link (fit the survival and longitudinal models independently)
#' @export
#'
#' @returns A `Link` or `LinkComponent` object representing the requested association.
linkNone <- function() {
    Link()
}


#' @describeIn standard-link-user Time to growth link
#' @export
linkTTG <- function(prior, model = PromiseLongitudinalModel(), ...) {
    UseMethod("linkTTG", model)
}
#' @export
#'
#' @returns A `PromiseLinkComponent` object.
linkTTG.PromiseLongitudinalModel <- function(
    prior = prior_normal(0, 2),
    model,
    ...
) {
    PromiseLinkComponent(fun = linkTTG, prior = prior, key = "link_ttg")
}
#' @export
#'
#' @returns This method always raises an error for unsupported models.
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
#'
#' @returns A `PromiseLinkComponent` object.
linkDSLD.PromiseLongitudinalModel <- function(
    prior = prior_normal(0, 2),
    model,
    ...
) {
    PromiseLinkComponent(fun = linkDSLD, prior = prior, key = "link_dsld")
}
#' @export
#'
#' @returns This method always raises an error for unsupported models.
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
#'
#' @returns A `PromiseLinkComponent` object.
linkIdentity.PromiseLongitudinalModel <- function(
    prior = prior_normal(0, 2),
    model,
    ...
) {
    PromiseLinkComponent(
        fun = linkIdentity,
        prior = prior,
        key = "link_identity"
    )
}
#' @export
#'
#' @returns This method always raises an error for unsupported models.
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
#'
#' @returns A `PromiseLinkComponent` object.
linkGrowth.PromiseLongitudinalModel <- function(
    prior = prior_normal(0, 2),
    model,
    ...
) {
    PromiseLinkComponent(fun = linkGrowth, prior = prior, key = "link_growth")
}
#' @export
#'
#' @returns This method always raises an error for unsupported models.
linkGrowth.default <- function(prior, model, ...) {
    stop(sprintf(
        "Method `linkGrowth` is not available for `%s`",
        class(model)[[1]]
    ))
}


#' @describeIn standard-link-user Shrinkage Parameter link
#' @export
linkShrinkage <- function(prior, model = PromiseLongitudinalModel(), ...) {
    UseMethod("linkShrinkage", model)
}
#' @export
#'
#' @returns A `PromiseLinkComponent` object.
linkShrinkage.PromiseLongitudinalModel <- function(
    prior = prior_normal(0, 2),
    model,
    ...
) {
    PromiseLinkComponent(
        fun = linkShrinkage,
        prior = prior,
        key = "link_shrinkage"
    )
}
#' @export
#'
#' @returns This method always raises an error for unsupported models.
linkShrinkage.default <- function(prior, model, ...) {
    stop(sprintf(
        "Method `linkShrinkage` is not available for `%s`",
        class(model)[[1]]
    ))
}

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

#' Simulate Longitudinal Data from a GSF Model
#'
#' @typed times: numeric
#'   the times to generate observations at.
#' @typed sigma: number
#'   the variance of the longitudinal values.
#' @typed mu_s: numeric
#'   the mean shrinkage rates.
#' @typed mu_g: numeric
#'   the mean growth rates.
#' @typed mu_b: numeric
#'   the mean baseline values.
#' @typed mu_phi: numeric
#'   the mean proportion of cells affected by the treatment
#' @typed omega_b: number
#'   the baseline value standard deviation.
#' @typed omega_s: number
#'   the shrinkage rate standard deviation.
#' @typed omega_g: number
#'   the growth rate standard deviation.
#' @typed omega_phi: number
#'   for the standard deviation of the proportion of cells
#'   affected by the treatment `omega_phi`.
#' @typed link_dsld: number
#'   the link coefficient for the derivative contribution.
#' @typed link_ttg: number
#'   the link coefficient for the time-to-growth contribution.
#' @typed link_identity: number
#'   the link coefficient for the SLD Identity contribution.
#' @typed link_growth: number
#'   the link coefficient for the log-growth parameter contribution.
#' @typed link_shrinkage: number
#'   the link coefficient for the log-shrinkage parameter contribution.
#' @typed scaled_variance: logical
#'   whether the variance should be scaled by the expected value
#'   (see the "Statistical Specifications" vignette for more details)
#'
#' @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.
#' @slot link_shrinkage (`numeric`)\cr See arguments.
#' @slot scaled_variance (`logical`)\cr See arguments.
#' @family SimLongitudinal
#' @name SimLongitudinalGSF-class
#' @exportClass SimLongitudinalGSF
#'
#' @examples
#' SimLongitudinalGSF(times = c(0, 30, 60))
.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",
        link_shrinkage = "numeric",
        scaled_variance = "logical"
    )
)

#' @rdname SimLongitudinalGSF-class
#' @export
#'
#' @returns A `SimLongitudinalGSF` object.
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,
    link_shrinkage = 0,
    scaled_variance = FALSE
) {
    if (length(omega_b) == 1) {
        omega_b <- rep(omega_b, length(mu_b))
    }
    if (length(omega_s) == 1) {
        omega_s <- rep(omega_s, length(mu_s))
    }
    if (length(omega_g) == 1) {
        omega_g <- rep(omega_g, length(mu_g))
    }
    if (length(omega_phi) == 1) {
        omega_phi <- rep(omega_phi, length(mu_phi))
    }

    .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,
        link_shrinkage = link_shrinkage,
        scaled_variance = scaled_variance
    )
}


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."
            )
        }

        pairs <- list(
            "omega_b" = "mu_b",
            "omega_s" = "mu_s",
            "omega_g" = "mu_g",
            "omega_phi" = "mu_phi"
        )
        for (i in seq_along(pairs)) {
            omega <- slot(object, names(pairs)[[i]])
            mu <- slot(object, pairs[[i]])
            if (!(length(omega) == length(mu))) {
                return(
                    sprintf(
                        "`%s` must be length 1 or the same length as `%s`",
                        omega,
                        mu
                    )
                )
            }
        }

        len_1_pars <- c(
            "sigma",
            "link_dsld",
            "link_ttg",
            "link_identity",
            "link_growth",
            "link_shrinkage"
        )
        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
#' @exportS3Method NULL
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
            ),
            dsld = gsf_dsld(
                .data$time,
                .data$psi_b,
                .data$psi_s,
                .data$psi_g,
                .data$psi_phi
            ),
            ttg = gsf_ttg(
                .data$time,
                .data$psi_b,
                .data$psi_s,
                .data$psi_g,
                .data$psi_phi
            ),
            sld_sd = ifelse(
                object@scaled_variance,
                .data$mu_sld * object@sigma,
                object@sigma
            ),
            sld = stats::rnorm(dplyr::n(), .data$mu_sld, .data$sld_sd),
            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)) +
                (object@link_shrinkage * log(.data$psi_s))
        )
}


#' @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),
            arm_idx = as.numeric(.data$arm),
            psi_b = stats::rlnorm(
                dplyr::n(),
                object@mu_b[.data$study_idx],
                object@omega_b[.data$study_idx]
            ),
            psi_s = stats::rlnorm(
                dplyr::n(),
                object@mu_s[.data$arm_idx],
                object@omega_s[.data$arm_idx]
            ),
            psi_g = stats::rlnorm(
                dplyr::n(),
                object@mu_g[.data$arm_idx],
                object@omega_g[.data$arm_idx]
            ),
            psi_phi_logit = stats::rnorm(
                dplyr::n(),
                object@mu_phi[.data$arm_idx],
                object@omega_phi[.data$arm_idx]
            ),
            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
#'
#' @typed time: numeric
#'   time grid.
#' @typed b: number
#'   baseline.
#' @typed s: number
#'   shrinkage.
#' @typed g: number
#'   growth.
#' @typed phi: number
#'   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 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
#'
#' @typed intercept: Prior
#'   for the `intercept`.
#' @typed slope_mu: Prior
#'   for the population slope `slope_mu` (one per arm).
#' @typed slope_sigma: Prior
#'   for the random slope standard deviation `slope_sigma` (one per arm).
#' @typed sigma: Prior
#'   for the variance of the longitudinal values `sigma`.
#' @typed scaled_variance: logical
#'   whether the variance should be scaled by the expected value, corresponding to a multiplicative model.
#'   As a default, the variance is not scaled by the expected value, corresponding to an additive model.
#'   (See the "Statistical Specifications" vignette for more details.)
#'
#' @export
#'
#' @returns A `LongitudinalRandomSlope` object.
#'
#' @examples
#' LongitudinalRandomSlope()
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),
    scaled_variance = FALSE
) {
    stan <- StanModule(decorated_render(
        .x = read_stan("lm-random-slope/model.stan"),
        scaled_variance = scaled_variance
    ))

    # Apply constriants
    sigma <- set_limits(sigma, lower = getOption("jmpost.double_eps"))
    slope_sigma <- set_limits(slope_sigma, lower = getOption("jmpost.double_eps"))

    .LongitudinalRandomSlope(
        LongitudinalModel(
            name = "Random Slope",
            stan = stan,
            scaled_variance = scaled_variance,
            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 = "n_arms"
                ),
                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
#'
#' @returns A `StanModule` object containing the generated-quantities code.
enableGQ.LongitudinalRandomSlope <- function(object, ...) {
    StanModule("lm-random-slope/quantities.stan")
}

#' @export
#'
#' @returns The longitudinal model with its link-related Stan code enabled.
enableLink.LongitudinalRandomSlope <- function(object, ...) {
    object@stan <- merge(
        object@stan,
        StanModule("lm-random-slope/link.stan")
    )
    object
}


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

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

#' @export
#'
#' @returns A `LinkComponent` object.
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")
}


#' @rdname getRandomEffectsNames
#' @export
getRandomEffectsNames.LongitudinalRandomSlope <- function(object, ...) {
    c("slope" = "lm_rs_ind_rnd_slope")
}

#' @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
#'
#' @returns A named `list` representation of the object.
as.list.GridManual <- function(x, data, ...) {
    subs <- as.list(names(x@spec))
    names(subs) <- names(x@spec)
    subs
}

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


#' Re-used documentation for `DataJoint`
#'
#' @typed object: DataJoint
#'   Survival and Longitudinal Data.
#' @typed x: DataJoint
#'   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"
    )
)

#' @typed subject: DataSubject
#'   object created by [DataSubject()].
#' @typed survival: DataSurvival
#'   object created by [DataSurvival()].
#' @typed longitudinal: DataLongitudinal
#'   object created by [DataLongitudinal()].
#' @rdname DataJoint-class
#'
#' @returns A `DataJoint` object.
#'
#' @examples
#' subjects <- data.frame(id = c("1", "2"), arm = "A", study = "S")
#' survival <- data.frame(id = c("1", "2"), time = c(5, 8), event = c(1, 0))
#' longitudinal <- data.frame(
#'   id = rep(c("1", "2"), each = 2),
#'   time = rep(c(0, 1), 2),
#'   response = c(10, 9, 12, 11)
#' )
#' DataJoint(
#'   DataSubject(subjects, "id", "arm", "study"),
#'   DataSurvival(survival, Surv(time, event) ~ 1),
#'   DataLongitudinal(longitudinal, response ~ time)
#' )
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`
#'
#' @typed object: "`DataSubject` or `DataLongitudinal` or `DataSurvival`"
#'   data object to convert to a `list`.
#' @typed x: "`DataSubject` or `DataLongitudinal` or `DataSurvival`"
#'   data object to convert to a `list`.
#' @typed subject_var: character
#'   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
#'
#' @returns A named `list` suitable for use as Stan data.
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`
#'
#' @typed x: DataJoint
#'   object created by [DataJoint()].
#' @typed subjects: "`character` or `list`"
#'   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
#'
#' @returns A `DataJoint` object restricted to the requested subjects.
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`
#'
#' @typed dat: data.frame
#'   must have a column called `subject` which corresponds to the
#'   values passed to `groupings`.
#' @typed groupings: "`character` or `list`"
#'   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
#'
#' @returns A subsetted `data.frame` with a grouping column.
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
#' @typed object: DataJoint
#'   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
#'
#' @returns Invisibly returns `NULL` after printing the object.
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")
    }
)

#' @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
#'
#' @returns A named `list` representation of the object.
as.list.GridEvent <- function(x, data, ...) {
    subjects_to_list(x@subjects, data)
}

#' @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
#'
#' @returns A named `list` representation of the object.
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 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
#'
#' @typed a: Prior
#'   Prior distribution for the scale parameter `a`.
#' @typed b: Prior
#'   Prior distribution for the shape parameter `b`.
#' @typed beta: Prior
#'   Prior distribution for covariates coefficients `beta`.
#'
#' @export
#'
#' @returns A `SurvivalLogLogistic` object.
#'
#' @examples
#' SurvivalLogLogistic()
SurvivalLogLogistic <- function(
    a = prior_lognormal(log(0.1), 5),
    b = prior_gamma(2, 5),
    beta = prior_normal(0, 2)
) {
    a <- set_limits(a, lower = getOption("jmpost.double_eps"))
    b <- set_limits(b, lower = getOption("jmpost.double_eps"))

    .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"
                )
            )
        )
    )
}

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


#' Re-used documentation for `LongitudinalQuantities`
#'
#' @typed x: LongitudinalQuantities
#'   longitudinal quantities.
#' @typed object: LongitudinalQuantities
#'   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
#'
#' @examples
#' \dontrun{
#' LongitudinalQuantities(fit, GridEven(length.out = 25))
#' }
.LongitudinalQuantities <- setClass(
    "LongitudinalQuantities",
    slots = c(
        "quantities" = "Quantities",
        "data" = "DataJoint"
    )
)

#' @typed object: JointModelSamples
#'   samples as drawn from a Joint Model.
#'
#' @typed grid: Grid
#'   object that specifies which subjects and time points to calculate the
#'   quantities for. See [Grid-Functions].
#' @rdname LongitudinalQuantities-class
#'
#' @returns A `LongitudinalQuantities` object.
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`
#'
#' @typed x: LongitudinalQuantities
#'   longitudinal quantities.
#' @param ... not used.
#' @family LongitudinalQuantities
#' @export
#'
#' @returns A `data.frame` representation of the object.
as.data.frame.LongitudinalQuantities <- function(x, ...) {
    as.data.frame(x@quantities)
}


#' summary
#'
#' @description
#' This method returns a `data.frame` of the longitudinal quantities.
#'
#' @typed conf.level: numeric
#'   confidence level of the interval.
#' @inheritParams LongitudinalQuantities-Shared
#'
#' @family LongitudinalQuantities
#' @family summary
#' @export
#'
#' @returns A `data.frame` containing posterior summary statistics.
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.
#'
#' @typed data: data.frame
#'   summary statistics for longitudinal
#'   value estimates. See details.
#' @typed data_obs: data.frame
#'   real observed values to be
#'   overlaid for reference.  See details.
#' @typed add_ci: logical
#'   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
#'
#' @returns A `ggplot` object.
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`
#'
#' @typed conf.level: numeric
#'   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
#'
#' @returns A `ggplot` object.
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 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
#'
#' @returns A `JointModelSamples` object.
#'
#' @examples
#' \dontrun{
#' # JointModelSamples objects are returned by sampleStanModel().
#' fit <- sampleStanModel(JointModel(LongitudinalGSF()), data = joint_data)
#' }
.JointModelSamples <- setClass(
    "JointModelSamples",
    slots = c(
        model = "JointModel",
        data = "DataJoint",
        results = "CmdStanMCMC"
    )
)


#' @rdname generateQuantities
#' @typed generator: QuantityGenerator
#'   object that specifies which subjects and time points
#'   to calculate the quantities at
#' @typed type: character
#'   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
#'
#' @returns A `StanModule` object.
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.
#'
#' @typed object: JointModelSamples
#'   samples as drawn from a [`JointModel`].
#' @family JointModelSamples
#' @typed indent: numeric
#'   how much white space to prefix the print string with.
#' @keywords internal
#' @export
#'
#' @returns A character vector suitable for printing.
as_print_string.JointModelSamples <- function(object, indent = 1, ...) {
    sizes <- vapply(
        cmdstanr::as.CmdStanMCMC(object)$metadata()[["stan_variable_sizes"]],
        \(x) {
            if (length(x) == 1 && x == 1) {
                return("")
            }
            paste0("[", paste(x, collapse = ", "), "]")
        },
        character(1)
    )
    variable_string <- paste0(
        "        ",
        cmdstanr::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"),
        cmdstanr::as.CmdStanMCMC(object)$metadata()$iter_sampling,
        cmdstanr::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
#'
#' @returns The underlying `cmdstanr::CmdStanMCMC` object.
as.CmdStanMCMC.JointModelSamples <- function(object, ...) {
    return(object@results)
}


#' Save a `JointModelSamples` object to a file.
#'
#' This function is just a wrapper around `saveRDS` that saves the object to a file
#' ensuring that all of the Stan samples are correctly stored. Note that as
#' `cmdstanr` objects store their samples as a csv file the samples may be lost
#' if you call `saveRDS` directly on the object.
#'
#' @typed object: JointModelSamples
#'   the object to save.
#' @typed file: character
#'   the file to save the object to.
#' @typed ...: ANY
#'   additional arguments to [`saveRDS`].
#'
#' @family saveObject
#'
#' @export
#'
#' @returns Invisibly returns `NULL` after saving the object.
saveObject.JointModelSamples <- function(object, file, ...) {
    object@results$draws()
    try(object@results$sampler_diagnostics(), silent = TRUE)
    try(object@results$init(), silent = TRUE)
    try(object@results$profiles(), silent = TRUE)
    saveRDS(object, file, ...)
}

#' Extract shrinkage factors from a `JointModelSamples` object
#'
#' @typed object: JointModelSamples
#'   the object to extract shrinkage factors from.
#' @param ... not used.
#' @return the shrinkage factors correctly named after the covariate names.
#'
#' @family shrinkage
#' @export
shrinkage.JointModelSamples <- function(object, ...) {
    draws <- cmdstanr::as.CmdStanMCMC(object)$draws()
    assert_class(object@data@survival, "DataSurvival")
    covariate_names <- covariates(object@data@survival)
    assert_subset(
        "prior_shrinkage_factors_beta_os_cov[1]",
        posterior::variables(draws)
    )
    shrinkage_draws <- subset(
        draws,
        variable = "prior_shrinkage_factors_beta_os_cov"
    )
    shrinkage_vars <- posterior::variables(shrinkage_draws)
    assert_true(identical(length(covariate_names), length(shrinkage_vars)))
    posterior::variables(shrinkage_draws) <- covariate_names
    shrinkage_draws
}

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

#' Simulate Longitudinal Data from a Stein-Fojo Model
#'
#' @typed times: numeric
#'   the times to generate observations at.
#' @typed sigma: number
#'   the variance of the longitudinal values.
#' @typed mu_s: numeric
#'   the mean shrinkage rates for the two treatment arms.
#' @typed mu_g: numeric
#'   the mean growth rates for the two treatment arms.
#' @typed mu_b: numeric
#'   the mean baseline values for the two treatment arms.
#' @typed omega_b: number
#'   the baseline value standard deviation.
#' @typed omega_s: number
#'   the shrinkage rate standard deviation.
#' @typed omega_g: number
#'   the growth rate standard deviation.
#' @typed link_dsld: number
#'   the link coefficient for the derivative contribution.
#' @typed link_ttg: number
#'   the link coefficient for the time-to-growth contribution.
#' @typed link_identity: number
#'   the link coefficient for the SLD Identity contribution.
#' @typed link_growth: number
#'   the link coefficient for the log-growth parameter contribution.
#' @typed link_shrinkage: number
#'   the link coefficient for the log-shrinkage parameter contribution.
#' @typed scaled_variance: logical
#'   whether the variance should be scaled by the expected value
#'   (see the "Statistical Specifications" vignette for more details)
#'
#' @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.
#' @slot link_shrinkage (`numeric`)\cr See arguments.
#' @slot scaled_variance (`logical`)\cr See arguments.
#'
#' @family SimLongitudinal
#' @name SimLongitudinalSteinFojo-class
#' @exportClass SimLongitudinalSteinFojo
#'
#' @examples
#' SimLongitudinalSteinFojo(times = c(0, 30, 60))
.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",
        link_shrinkage = "numeric",
        scaled_variance = "logical"
    )
)

#' @rdname SimLongitudinalSteinFojo-class
#' @export
#'
#' @returns A `SimLongitudinalSteinFojo` object.
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,
    link_shrinkage = 0,
    scaled_variance = FALSE
) {
    if (length(omega_b) == 1) {
        omega_b <- rep(omega_b, length(mu_b))
    }
    if (length(omega_s) == 1) {
        omega_s <- rep(omega_s, length(mu_s))
    }
    if (length(omega_g) == 1) {
        omega_g <- rep(omega_g, length(mu_g))
    }

    .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,
        link_shrinkage = link_shrinkage,
        scaled_variance = scaled_variance
    )
}


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."
            )
        }
        pairs <- list(
            "omega_b" = "mu_b",
            "omega_s" = "mu_s",
            "omega_g" = "mu_g"
        )
        for (i in seq_along(pairs)) {
            omega <- slot(object, names(pairs)[[i]])
            mu <- slot(object, pairs[[i]])
            if (!(length(omega) == length(mu))) {
                return(
                    sprintf(
                        "`%s` must be length 1 or the same length as `%s`",
                        omega,
                        mu
                    )
                )
            }
        }
        len_1_pars <- c(
            "sigma",
            "link_dsld",
            "link_ttg",
            "link_identity",
            "link_growth",
            "link_shrinkage"
        )
        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
#' @exportS3Method NULL
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),
            dsld = sf_dsld(.data$time, .data$psi_b, .data$psi_s, .data$psi_g),
            ttg = sf_ttg(.data$time, .data$psi_b, .data$psi_s, .data$psi_g),
            sld_sd = ifelse(
                object@scaled_variance,
                .data$mu_sld * object@sigma,
                object@sigma
            ),
            sld = stats::rnorm(dplyr::n(), .data$mu_sld, .data$sld_sd),
            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)) +
                (object@link_shrinkage * log(.data$psi_s))
        )
}


#' @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),
            arm_idx = as.numeric(.data$arm),
            psi_b = stats::rlnorm(
                dplyr::n(),
                object@mu_b[.data$study_idx],
                object@omega_b[.data$study_idx]
            ),
            psi_s = stats::rlnorm(
                dplyr::n(),
                object@mu_s[.data$arm_idx],
                object@omega_s[.data$arm_idx]
            ),
            psi_g = stats::rlnorm(
                dplyr::n(),
                object@mu_g[.data$arm_idx],
                object@omega_g[.data$arm_idx]
            )
        )

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


#' Stein-Fojo Functionals
#'
#' @typed time: numeric
#'   time grid.
#' @typed b: number
#'   baseline.
#' @typed s: number
#'   shrinkage.
#' @typed g: number
#'   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 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
#'
#' @returns A `Grid` object of the requested type.
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
#'
#' @returns A named `list` representation of the object.
as.list.GridEven <- function(x, data, ...) {
    subjects_to_list(x@subjects, data)
}

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

#' @include SurvivalModel.R
NULL


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

# SurvivalGamma-constructors ----

#' @rdname SurvivalGamma-class
#'
#' @typed k: Prior
#'   for the shape `k`.
#' @typed theta: Prior
#'   for the scale `theta`.
#' @typed beta: Prior
#'   for covariates coefficients `beta`.
#'
#' @export
#'
#' @returns A `SurvivalGamma` object.
#'
#' @examples
#' SurvivalGamma()
SurvivalGamma <- function(
    k = prior_gamma(2, 0.5),
    theta = prior_gamma(2, 0.5),
    beta = prior_normal(0, 2)
) {
    k <- set_limits(k, lower = getOption("jmpost.double_eps"))
    theta <- set_limits(theta, lower = getOption("jmpost.double_eps"))

    .SurvivalGamma(
        SurvivalModel(
            name = "Gamma",
            stan = StanModule(x = "sm-gamma/model.stan"),
            parameters = ParameterList(
                Parameter(name = "sm_gamma_k", prior = k, size = 1),
                Parameter(name = "sm_gamma_theta", prior = theta, size = 1),
                Parameter(
                    name = "beta_os_cov",
                    prior = beta,
                    size = "p_os_cov_design"
                )
            )
        )
    )
}

#' Re-used documentation for `Quantities`
#'
#' @typed x: Quantities
#'   generated quantities.
#' @typed object: Quantities
#'   generated quantities.
#' @typed type: character
#'   sets the `type` variable.
#' @typed conf.level: numeric
#'   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
#'
#' @typed quantities: matrix
#'   of generated quantities.
#' @typed times: numeric
#'   labels specifying which time point the quantity was generated at.
#' @typed groups: character
#'   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
#'
#' @returns A `Quantities` object.
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
#'
#' @returns An integer vector containing the dimensions of the `Quantities` object.
dim.Quantities <- function(x) {
    dim(x@quantities)
}


#' `Quantities` -> `data.frame`
#'
#' @inheritParams Quantities-Shared
#'
#' @keywords internal
#' @family Quantities
#' @export
#'
#' @returns A `data.frame` representation of the object.
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.
#'
#' @typed quantities_raw: matrix
#'   of samples with 1 row per sample and 1 column per
#'   distinct quantity.
#'
#' @typed collapser: QuantityCollapser
#'   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
#'
#' @returns An array of collapsed posterior quantities.
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_along(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
#'
#' @typed gq: CmdStanGQ
#'   a [cmdstanr::CmdStanGQ] object created by [generateQuantities()].
#' @typed type: character
#'   quantity to be generated.
#'   Must be one of `surv`, `haz`, `loghaz`, `cumhaz`, `lm_identity`.
#' @keywords internal
#'
#' @returns A [posterior::draws_matrix] containing the requested generated quantities.
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
#'
#' @returns A character vector suitable for printing.
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")
    }
)

#' Quantity Grid Specification
#'
#' @typed subjects: "`character` or `NULL`"
#'   vector of subjects to extract quantities for.
#'   If `NULL` will default to all subjects within the dataset.
#'
#' @typed times: "`numeric` or `NULL`"
#'   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`]).
#'
#' @typed groups: list
#'   named list of subjects to extract quantities for. See Group Specification.
#'
#' @typed spec: list
#'   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.
#'
#' @typed length.out: numeric
#'   number of evenly spaced timepoints to generate quantities at.
#'
#' @typed newdata: data.frame
#'   new data to generate quantities for. Must contain the same columns
#'   and factor levels of the original data used in the [`DataSurvival`] object.
#'
#' @typed params: list
#'   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
#'
#' @examples
#' GridFixed(subjects = c("subject 1", "subject 2"), times = c(0, 10))
#' GridGrouped(groups = list(control = c("subject 1", "subject 2")))
#' GridEven(subjects = c("subject 1", "subject 2"), length.out = 20)
#' GridObserved()
#' GridEvent()
#' GridPopulation(times = seq(0, 100, by = 10))
#' GridManual(list("subject 1" = c(0, 5), "subject 2" = c(0, 10)))
#' GridPrediction(times = c(0, 10), newdata = data.frame(age = 60))
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.
#'
#' @typed times: numeric
#'   vector of time points to extract quantities at.
#' @typed subjects: character
#'   vector of subjects to extract quantities for.
#' @typed groups: character
#'   vector of labels to apply to the generated quantities.
#' @typed indexes: list
#'   list of indexes that specify which observations from a
#'   `QuantityGenerator` should be combined to form the desired quantities.
#'
#' @typed object: Grid
#'   object to convert to a `QuantityGenerator` or `QuantityCollapser`.
#' @typed data: DataJoint
#'   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
#' Just a placeholder which will error when used.
#'
#' @typed object: QuantityGenerator
#'   object to convert to a list.
#' @typed data: DataJoint
#'   Survival and Longitudinal Data.
#' @param ... Not currently used.
#' @keywords internal
#' @export
#'
#' @returns Will always throw an error.
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
#'
#' @returns A single integer giving the number of collapsed quantities.
length.QuantityCollapser <- function(x) {
    length(x@indexes)
}


#' Expand and Validate Subjects
#'
#' @typed subjects: character
#'   vector of subjects that should exist in `data`
#' @typed data: DataJoint
#'   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
#'
#' @returns A named list grouping the requested subject identifiers.
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 generics.R
#' @include Link.R
#' @include LongitudinalModel.R
#' @include SurvivalModel.R
#' @include StanModule.R
#' @include StanModel.R
#' @include Parameter.R
NULL


#' @rdname getParameters
#' @export
#'
#' @returns A `ParameterList` object.
#'
#' @examples
#' getParameters(LongitudinalGSF())
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
#'
#' @returns A merged object of the same class as `x` and `y`.
#'
#' @examples
#' x <- ParameterList(Parameter(prior_normal(0, 1), "x"))
#' y <- ParameterList(Parameter(prior_normal(0, 1), "y"))
#' merge(x, y)
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 SimLongitudinal.R
#' @include generics.R
NULL

#' Simulate Longitudinal Data from a Claret-Bruno Model
#'
#' @typed times: numeric
#'   the times to generate observations at.
#' @typed sigma: number
#'   the variance of the longitudinal values.
#'
#' @typed mu_b: numeric
#'   the mean population baseline sld value.
#' @typed mu_g: numeric
#'   the mean population growth rate.
#' @typed mu_c: numeric
#'   the mean population resistance rate.
#' @typed mu_p: numeric
#'   the mean population growth inhibition.
#'
#' @typed omega_b: number
#'   the population standard deviation for the baseline sld value.
#' @typed omega_g: number
#'   the population standard deviation for the growth rate.
#' @typed omega_c: number
#'   the population standard deviation for the resistance rate.
#' @typed omega_p: number
#'   the population standard deviation for the growth inhibition.
#'
#' @typed link_dsld: number
#'   the link coefficient for the derivative contribution.
#' @typed link_ttg: number
#'   the link coefficient for the time-to-growth contribution.
#' @typed link_identity: number
#'   the link coefficient for the SLD Identity contribution.
#' @typed link_growth: number
#'   the link coefficient for the growth parameter contribution.
#'
#' @typed scaled_variance: logical
#'   whether the variance should be scaled by the expected value
#'   (see the "Statistical Specifications" vignette for more details)
#'
#' @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.
#'
#' @slot scaled_variance (`logical`)\cr See arguments.
#'
#' @family SimLongitudinal
#' @name SimLongitudinalClaretBruno-class
#' @exportClass SimLongitudinalClaretBruno
#'
#' @examples
#' SimLongitudinalClaretBruno(times = c(0, 30, 60))
.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",
        scaled_variance = "logical"
    )
)

#' @rdname SimLongitudinalClaretBruno-class
#' @export
#'
#' @returns A `SimLongitudinalClaretBruno` object.
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,
    scaled_variance = FALSE
) {
    if (length(omega_b) == 1) {
        omega_b <- rep(omega_b, length(mu_b))
    }
    if (length(omega_g) == 1) {
        omega_g <- rep(omega_g, length(mu_g))
    }
    if (length(omega_c) == 1) {
        omega_c <- rep(omega_c, length(mu_c))
    }
    if (length(omega_p) == 1) {
        omega_p <- rep(omega_p, length(mu_p))
    }

    .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,
        scaled_variance = scaled_variance
    )
}


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."
            )
        }
        pairs <- list(
            "omega_b" = "mu_b",
            "omega_g" = "mu_g",
            "omega_c" = "mu_c",
            "omega_p" = "mu_p"
        )
        for (i in seq_along(pairs)) {
            omega <- slot(object, names(pairs)[[i]])
            mu <- slot(object, pairs[[i]])
            if (!(length(omega) == length(mu))) {
                return(
                    sprintf(
                        "`%s` must be length 1 or the same length as `%s`",
                        omega,
                        mu
                    )
                )
            }
        }
        len_1_pars <- c(
            "sigma",
            "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
#' @exportS3Method NULL
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
            ),
            dsld = clbr_dsld(
                .data$time,
                .data$ind_b,
                .data$ind_g,
                .data$ind_c,
                .data$ind_p
            ),
            ttg = clbr_ttg(
                .data$time,
                .data$ind_b,
                .data$ind_g,
                .data$ind_c,
                .data$ind_p
            ),
            sld_sd = ifelse(
                object@scaled_variance,
                .data$mu_sld * object@sigma,
                object@sigma
            ),
            sld = stats::rnorm(dplyr::n(), .data$mu_sld, .data$sld_sd),

            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[.data$study_idx]
            )
        ) |>
        dplyr::mutate(
            ind_g = stats::rlnorm(
                dplyr::n(),
                object@mu_g[.data$arm_idx],
                object@omega_g[.data$arm_idx]
            )
        ) |>
        dplyr::mutate(
            ind_c = stats::rlnorm(
                dplyr::n(),
                object@mu_c[.data$arm_idx],
                object@omega_c[.data$arm_idx]
            )
        ) |>
        dplyr::mutate(
            ind_p = stats::rlnorm(
                dplyr::n(),
                object@mu_p[.data$arm_idx],
                object@omega_p[.data$arm_idx]
            )
        )

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


#' Claret-Bruno Functionals
#'
#' @typed t: numeric
#'   time grid.
#' @typed b: number
#'   baseline sld.
#' @typed g: number
#'   growth rate.
#' @typed c: number
#'   resistance rate.
#' @typed p: number
#'   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 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
#'
#' @returns A `SurvivalModel` object.
#'
#' @examples
#' SurvivalModel()
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
#'
#' @returns A character vector suitable for printing.
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 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
#'
#' @typed mu_b: Prior
#'   for the mean population baseline sld value.
#' @typed mu_g: Prior
#'   for the mean population growth rate.
#' @typed mu_c: Prior
#'   for the mean population resistance rate.
#' @typed mu_p: Prior
#'   for the mean population growth inhibition
#'
#' @typed omega_b: Prior
#'   for the population standard deviation for the baseline sld value.
#' @typed omega_g: Prior
#'   for the population standard deviation for the growth rate.
#' @typed omega_c: Prior
#'   for the population standard deviation for the resistance rate.
#' @typed omega_p: Prior
#'   for the population standard deviation for the growth inhibition.
#'
#' @typed sigma: Prior
#'   for the variance of the longitudinal values.
#'
#' @typed centred: logical
#'   whether to use the centred parameterization.
#' @typed scaled_variance: logical
#'   whether the variance should be scaled by the expected value, corresponding to a multiplicative model.
#'   As a default, the variance is not scaled by the expected value, corresponding to an additive model.
#'   (See the "Statistical Specifications" vignette for more details.)
#'
#' @export
#'
#' @returns A `LongitudinalClaretBruno` object.
#'
#' @examples
#' LongitudinalClaretBruno()
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),

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

    # Apply constraints
    omega_b <- set_limits(omega_b, lower = getOption("jmpost.double_eps"))
    omega_g <- set_limits(omega_g, lower = getOption("jmpost.double_eps"))
    omega_c <- set_limits(omega_c, lower = getOption("jmpost.double_eps"))
    omega_p <- set_limits(omega_p, lower = getOption("jmpost.double_eps"))
    sigma <- set_limits(sigma, lower = getOption("jmpost.double_eps"))

    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 = "n_studies"
        ),
        Parameter(name = "lm_clbr_omega_g", prior = omega_g, size = "n_arms"),
        Parameter(name = "lm_clbr_omega_c", prior = omega_c, size = "n_arms"),
        Parameter(name = "lm_clbr_omega_p", prior = omega_p, size = "n_arms"),

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

    assert_flag(centred)
    parameters_extra <- if (centred) {
        list(
            Parameter(
                name = "lm_clbr_ind_b",
                prior = set_limits(prior_init_only(prior_lognormal(
                    median(mu_b),
                    median(omega_b)
                )), lower = getOption("jmpost.double_eps")),
                size = "n_subjects"
            ),
            Parameter(
                name = "lm_clbr_ind_g",
                prior = set_limits(prior_init_only(prior_lognormal(
                    median(mu_g),
                    median(omega_g)
                )), lower = getOption("jmpost.double_eps")),
                size = "n_subjects"
            ),
            Parameter(
                name = "lm_clbr_ind_c",
                prior = set_limits(prior_init_only(prior_lognormal(
                    median(mu_c),
                    median(omega_c)
                )), lower = getOption("jmpost.double_eps")),
                size = "n_subjects"
            ),
            Parameter(
                name = "lm_clbr_ind_p",
                prior = set_limits(prior_init_only(prior_lognormal(
                    median(mu_p),
                    median(omega_p)
                )), lower = getOption("jmpost.double_eps")),
                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),
        scaled_variance = scaled_variance
    )
    .LongitudinalClaretBruno(x)
}


#' @export
#'
#' @returns The longitudinal model with its link-related Stan code enabled.
enableLink.LongitudinalClaretBruno <- function(object, ...) {
    object@stan <- merge(
        object@stan,
        StanModule("lm-claret-bruno/link.stan")
    )
    object
}

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

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

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

#' @export
#'
#' @returns A `LinkComponent` object.
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
#'
#' @returns A `StanModule` object containing the generated-quantities code.
enableGQ.LongitudinalClaretBruno <- function(object, ...) {
    StanModule("lm-claret-bruno/quantities.stan")
}


#' @rdname getRandomEffectsNames
#' @export
getRandomEffectsNames.LongitudinalGSF <- function(object, ...) {
    c(
        "b" = "lm_clbr_ind_b",
        "g" = "lm_clbr_ind_g",
        "c" = "lm_clbr_ind_c",
        "p" = "lm_clbr_ind_p"
    )
}

#' @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
#'
#' @returns A named `list` representation of the object.
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 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
#'
#' @typed mu_bsld: Prior
#'   for the mean baseline value `mu_bsld`.
#' @typed mu_ks: Prior
#'   for the mean shrinkage rate `mu_ks`.
#' @typed mu_kg: Prior
#'   for the mean growth rate `mu_kg`.
#'
#' @typed omega_bsld: Prior
#'   for the baseline value standard deviation `omega_bsld`.
#' @typed omega_ks: Prior
#'   for the shrinkage rate standard deviation `omega_ks`.
#' @typed omega_kg: Prior
#'   for the growth rate standard deviation `omega_kg`.
#'
#' @typed sigma: Prior
#'   for the variance of the longitudinal values `sigma`.
#'
#' @typed scaled_variance: logical
#'   whether the variance should be scaled by the expected value, corresponding to a multiplicative model.
#'   As a default, the variance is not scaled by the expected value, corresponding to an additive model.
#'   (See the "Statistical Specifications" vignette for more details.)
#' @typed centred: logical
#'   whether to use the centred parameterization.
#'
#' @export
#'
#' @returns A `LongitudinalSteinFojo` object.
#'
#' @examples
#' LongitudinalSteinFojo()
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),

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

    # Apply constriants
    omega_bsld <- set_limits(omega_bsld, lower = getOption("jmpost.double_eps"))
    omega_ks <- set_limits(omega_ks, lower = getOption("jmpost.double_eps"))
    omega_kg <- set_limits(omega_kg, lower = getOption("jmpost.double_eps"))
    sigma <- set_limits(sigma, lower = getOption("jmpost.double_eps"))

    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 = "n_studies"
        ),
        Parameter(name = "lm_sf_omega_ks", prior = omega_ks, size = "n_arms"),
        Parameter(name = "lm_sf_omega_kg", prior = omega_kg, size = "n_arms"),

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

    assert_flag(centred)
    parameters_extra <- if (centred) {
        list(
            Parameter(
                name = "lm_sf_psi_bsld",
                prior = set_limits(prior_init_only(prior_lognormal(
                    median(mu_bsld),
                    median(omega_bsld)
                )), lower = getOption("jmpost.double_eps")),
                size = "n_subjects"
            ),
            Parameter(
                name = "lm_sf_psi_ks",
                prior = set_limits(prior_init_only(prior_lognormal(
                    median(mu_ks),
                    median(omega_ks)
                )), lower = getOption("jmpost.double_eps")),
                size = "n_subjects"
            ),
            Parameter(
                name = "lm_sf_psi_kg",
                prior = set_limits(prior_init_only(prior_lognormal(
                    median(mu_kg),
                    median(omega_kg)
                )), lower = getOption("jmpost.double_eps")),
                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),
        scaled_variance = scaled_variance
    )
    .LongitudinalSteinFojo(x)
}


#' @export
#'
#' @returns A `StanModule` object containing the generated-quantities code.
enableGQ.LongitudinalSteinFojo <- function(object, ...) {
    StanModule("lm-stein-fojo/quantities.stan")
}

#' @export
#'
#' @returns The longitudinal model with its link-related Stan code enabled.
enableLink.LongitudinalSteinFojo <- function(object, ...) {
    object@stan <- merge(
        object@stan,
        StanModule("lm-stein-fojo/link.stan")
    )
    object
}

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

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

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

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

#' @export
#'
#' @returns A `LinkComponent` object.
linkShrinkage.LongitudinalSteinFojo <- function(
    prior = prior_normal(0, 2),
    model,
    ...
) {
    LinkComponent(
        key = "link_shrinkage",
        stan = StanModule("lm-stein-fojo/link_shrinkage.stan"),
        prior = prior
    )
}

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


#' @rdname getRandomEffectsNames
#' @export
getRandomEffectsNames.LongitudinalSteinFojo <- function(object, ...) {
    c(
        "b" = "lm_sf_psi_bsld",
        "s" = "lm_sf_psi_ks",
        "g" = "lm_sf_psi_kg"
    )
}

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

#' `StanModel` Function Arguments
#'
#' The documentation lists all the conventional arguments for wrappers around
#' [StanModel()].
#'
#' @typed stan: StanModule
#'   code containing the model-specific Stan code specification.
#' @typed parameters: ParameterList
#'   the parameter declaration and prior specification.
#' @typed parameter: ParameterList
#'   the (single) parameter declaration and prior specification.
#' @typed name: character
#'   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
#'
#' @returns A `StanModel` object.
#'
#' @examples
#' StanModel(StanModule(), ParameterList(), name = "Example model")
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.
#' @typed x: StanModel
#'   A Stan Model
#' @param ... Not Used.
#' @family StanModel
#' @export
#'
#' @returns A named `list` representation of the object.
as.list.StanModel <- function(x, ...) {
    as.list(x@stan)
}

# getParameters-StanModel ----

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


#' @export
#'
#' @returns A character vector suitable for printing.
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)
}

#' 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.
#'
#' ## `jmpost.double_eps` and `jmpost.double_neg_eps`
#'
#' Default = `0` and `sqrt(.Machine$double.neg.eps)`
#'
#' These are small floating-point numbers `x` which are used in parameter
#' definitions as lower boundaries to prevent numerical issues with the MCMC
#' sampler in the beginning of the sampling process.
#'
#' @examples
#' \dontrun{
#' options(jmpost.prior_shrinkage = 0.5)
#' }
#' @name jmpost-settings
#'
#' @returns Invisibly returns `NULL` after initializing the package options.
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,
        jmpost.double_eps = 0,
        jmpost.double_neg_eps = sqrt(.Machine$double.neg.eps)
    )
    for (opt in names(jmpost_opts)) {
        if (!opt %in% current_opts) {
            options(jmpost_opts[opt])
        }
    }
    invisible()
}

#' 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
#'
#' @returns A `PromiseLongitudinalModel` object.
#'
#' @examples
#' PromiseLongitudinalModel()
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.
#'
#' @typed fun: function
#'   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
#'
#' @returns A `PromiseLinkComponent` object.
#'
#' @examples
#' promise <- PromiseLinkComponent(linkDSLD, prior_normal(0, 2), "link_dsld")
#' resolvePromise(promise, LongitudinalGSF())
PromiseLinkComponent <- function(fun, prior, key) {
    .PromiseLinkComponent(
        fun = fun,
        stan = StanModule(),
        parameters = ParameterList(Parameter(
            name = key,
            prior = prior,
            size = 1
        )),
        key = key
    )
}


#' @export
#'
#' @returns A `StanModule` object.
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`].
#'
#' @typed object: PromiseLinkComponent
#'   the promise to resolve
#' @typed model: LongitudinalModel
#'   the model to resolve the promise with
#' @param ... Not used.
#' @typedreturn LinkComponent
#'   the resolved `LinkComponent` object
#'
#' @export
#'
#' @returns The resolved `LinkComponent` object.
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 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
#'
#' @typed lambda: Prior
#'   for the exponential rate `lambda`.
#' @typed beta: Prior
#'   for covariates coefficients `beta`.
#'
#' @export
#'
#' @returns A `SurvivalExponential` object.
#'
#' @examples
#' SurvivalExponential()
SurvivalExponential <- function(
    lambda = prior_gamma(2, 5),
    beta = prior_normal(0, 2)
) {
    lambda <- set_limits(lambda, lower = getOption("jmpost.double_eps"))
    .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"
                )
            )
        )
    )
}

#' Abstract Simulation Class for Longitudinal Data
#'
#' @typed 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
#'
#' @examples
#' SimLongitudinal(times = c(0, 30, 60))
.SimLongitudinal <- setClass(
    "SimLongitudinal",
    slots = list(
        times = "numeric"
    )
)

#' @rdname SimLongitudinal-class
#' @export
#'
#' @returns A `SimLongitudinal` object.
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
#' @exportS3Method NULL
as_print_string.SimLongitudinal <- function(object, ...) {
    return("SimLongitudinal")
}

#' Define Simulation Group
#'
#' Specifies a simulation group to be used by [`SimJointData()`].
#'
#' @typed n: numeric
#'   number of subjects in the group.
#' @typed arm: character
#'   treatment arm.
#' @typed study: character
#'   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
#'
#' @returns A `SimGroup` object.
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)
    }
)

1		#' Calculate Population Hazard Ratios
2		#'
3		#' Calculates hazard ratios marginalised over subject specific random effects using the
4		#' approach proposed by \insertCite{oudenhoven2020marginal}{jmpost}.
5		#'
6		#' @typed object: JointModelSamples
7		#' samples as drawn from a Joint Model.
8		#' @typed hr_formula: formula
9		#' defines the terms to include in the hazard ratio calculation.
10		#' By default this uses the right side of the formula used in the survival model.
11		#' Set to `NULL` not include any terms
12		#' @typed baseline: formula
13		#' terms to model baseline hazard using variable `time`.
14		#' Default is a B-spline from [splines]: `~bs(time, df = 10)`
15		#' @typed quantiles: numeric
16		#' vector of two values in (0, 1) for calculating quantiles from log hazard ratio
17		#' distributions.
18		#'
19		#' @references \insertAllCited{}
20		#'
21		#' @returns A list containing a summary of parameter distributions as a `data.frame` and a
22		#' matrix containing the parameter estimates for each sample.
23		#' @export
24		#' @importFrom splines bs
25		#'
26		#' @examples
27		#' \dontrun{
28		#' populationHR(fit, quantiles = c(0.01, 0.99))
29		#' }
30		populationHR <- function(
31		object,
32		hr_formula = object@data@survival@formula,
33		baseline = ~ bs(time, df = 10),
34		quantiles = c(0.025, 0.975)
35		) {
36	8x	assert_class(object, "JointModelSamples")
37	7x	assert_formula(hr_formula)
38	6x	assert_formula(baseline)
39	6x	assert_numeric(
40	6x	quantiles,
41	6x	lower = 0,
42	6x	upper = 1,
43	6x	any.missing = FALSE,
44	6x	unique = TRUE
45		)
46	5x	if (!"time" %in% all.vars(baseline)) {
47	1x	stop("baseline formula should include a time term.")
48		}
49
50		# Extract the variable names used in the data
51	4x	subject_var <- object@data@subject@subject
52	4x	arm_var <- object@data@subject@arm
53	4x	long_time_var <- all.vars(delete.response(terms(
54	4x	object@data@longitudinal@formula
55		)))
56	4x	surv_time_var <- all.vars(object@data@survival@formula[[2]][[2]])
57	4x	surv_covs <- all.vars(delete.response(terms(hr_formula)))
58	4x	if (!all(surv_covs %in% colnames(object@data@survival@data))) {
59	1x	stop("All variables in hr_formula must be in survival data")
60		}
61
62	3x	marginal_formula <- stats::reformulate(c(
63	3x	attr(terms(baseline), "term.labels"),
64	3x	attr(terms(hr_formula), "term.labels")
65		))
66
67		# Get the survival quantities at the observed longitudinal and survival times for each patient:
68	3x	times_df <- rbind(
69	3x	stats::setNames(
70	3x	object@data@longitudinal@data[, c(subject_var, long_time_var)],
71	3x	c("subject", "time")
72		),
73	3x	stats::setNames(
74	3x	object@data@survival@data[, c(subject_var, surv_time_var)],
75	3x	c("subject", "time")
76		)
77		)
78	3x	times_df <- times_df[order(times_df$subject, times_df$time), ]
79	3x	times_df <- times_df[times_df$time > 0, ]
80	3x	times_df <- times_df[!duplicated.data.frame(times_df), ]
81
82		# Generate samples of the log hazard log h_i(t) for each patient i at these time points t.
83	3x	grid_spec <- split(times_df$time, times_df$subject)
84	3x	log_haz_samples <- SurvivalQuantities(
85	3x	object,
86	3x	grid = GridManual(grid_spec),
87	3x	type = "loghaz"
88	3x	)@quantities@quantities \|>
89	3x	t()
90
91		# Construct \tilde{X} in paper's notation with one row per patient's time
92	3x	W_df <- dplyr::left_join(
93	3x	times_df,
94	3x	stats::setNames(
95	3x	object@data@survival@data[, c(subject_var, surv_covs)],
96	3x	c("subject", surv_covs)
97		),
98	3x	by = "subject"
99		)
100	3x	W_mat <- model.matrix(marginal_formula, W_df)
101
102		# As model matrix contains baseline and covariates, we don't need the intercept term
103		# but we want factor variables encoded relative to the intercept
104	3x	W_mat <- W_mat[, colnames(W_mat) != "(Intercept)", drop = FALSE]
105
106	3x	estimates <- stats::lm.fit(x = W_mat, y = log_haz_samples)$coefficients
107	3x	tidy_res <- apply(estimates, 1, function(x) {
108	26x	quantiles <- stats::quantile(x, probs = quantiles)
109	26x	c(mean = mean(x), median = median(x), quantiles)
110		}) \|>
111	3x	t() \|>
112	3x	data.frame()
113
114	3x	list(summary = tidy_res, estimates)
115		}

1		#' @include generics.R
2		#' @include Grid.R
3		NULL
4
5
6		#' @rdname Quant-Dev
7		.QuantityGeneratorSubject <- setClass(
8		"QuantityGeneratorSubject",
9		contains = "QuantityGenerator",
10		slots = c(
11		"times" = "numeric",
12		"subjects" = "character_or_NULL"
13		)
14		)
15
16
17		#' @rdname Quant-Dev
18		QuantityGeneratorSubject <- function(times, subjects = NULL) {
19	92x	.QuantityGeneratorSubject(
20	92x	times = times,
21	92x	subjects = subjects
22		)
23		}
24
25
26		setValidity(
27		"QuantityGeneratorSubject",
28		function(object) {
29		if (length(object@times) != length(object@subjects)) {
30		return("Length of `times` and `subjects` must be equal")
31		}
32		return(TRUE)
33		}
34		)
35
36
37		#' @rdname as_stan_list.QuantityGenerator
38		#' @export
39		as_stan_list.QuantityGeneratorSubject <- function(object, data, ...) {
40	32x	assert_that(
41	32x	is(data, "DataJoint")
42		)
43	32x	ret <- list()
44	32x	data_list <- as.list(data)
45	32x	ret[["gq_subject_index"]] <- data_list$subject_to_index[as.character(
46	32x	object@subjects
47		)]
48	32x	ret[["gq_n_quant"]] <- length(object@subjects)
49	32x	ret[["gq_times"]] <- object@times
50
51		# dummy pop indexes in order for stan code to actualy compile. In this setting
52		# this matrix isn't actually used so doesn't matter what these values are
53		# but don't want to have to burden individual longitudinal models with the
54		# conditional logic to check if they are generating population quantities or not
55	32x	ret[["gq_long_pop_arm_index"]] <- rep(1, ret[["gq_n_quant"]])
56	32x	ret[["gq_long_pop_study_index"]] <- rep(1, ret[["gq_n_quant"]])
57
58		# Sanity checks
59	32x	assert_that(
60	32x	length(ret[["gq_times"]]) == ret[["gq_n_quant"]],
61	32x	all(object@subjects %in% names(data_list$subject_to_index))
62		)
63	32x	return(ret)
64		}

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(
30	4x	data_list$tumour_time,
31	4x	data_list$subject_tumour_index,
32	4x	unique
33		)
34	4x	subject_visits <- unique_visits[data_list$subject_to_index[subjects]]
35	4x	visit_lengths <- vapply(subject_visits, length, numeric(1))
36	4x	QuantityGeneratorSubject(
37	4x	times = unlist(subject_visits, use.names = FALSE),
38	4x	subjects = rep(subjects, visit_lengths)
39		)
40		}
41
42		#' @rdname Quant-Dev
43		#' @export
44		as.QuantityCollapser.GridObserved <- function(object, data, ...) {
45	2x	generator <- as.QuantityGenerator(object, data)
46
47	2x	QuantityCollapser(
48	2x	times = generator@times,
49	2x	groups = generator@subjects,
50	2x	indexes = as.list(seq_along(generator@times))
51		)
52		}
53
54		#' @export
55		#'
56		#' @returns A named `list` representation of the object.
57		as.list.GridObserved <- function(x, data, ...) {
58	4x	subjects_to_list(x@subjects, data)
59		}

1		#' @include SimLongitudinal.R
2		#' @include generics.R
3		NULL
4
5		#' Simulate Longitudinal Data from a Random Slope Model
6		#'
7		#' @typed times: numeric
8		#' the times to generate observations at.
9		#' @typed intercept: numeric
10		#' the mean baseline value for each study.
11		#' @typed slope_mu: numeric
12		#' the population slope for each treatment arm.
13		#' @typed slope_sigma: numeric
14		#' the random slope standard deviation for each treatment arm.
15		#' @typed sigma: number
16		#' the variance of the longitudinal values.
17		#' @typed link_dsld: number
18		#' the link coefficient for the DSLD contribution.
19		#' @typed link_identity: number
20		#' the link coefficient for the identity contribution.
21		#' @typed scaled_variance: logical
22		#' whether the variance should be scaled by the expected value
23		#' (see the "Statistical Specifications" vignette for more details)
24		#'
25		#' @slot intercept (`numeric`)\cr See arguments.
26		#' @slot slope_mu (`numeric`)\cr See arguments.
27		#' @slot slope_sigma (`numeric`)\cr See arguments.
28		#' @slot sigma (`numeric`)\cr See arguments.
29		#' @slot link_dsld (`numeric`)\cr See arguments.
30		#' @slot link_identity (`numeric`)\cr See arguments.
31		#' @slot scaled_variance (`logical`)\cr See arguments.
32		#'
33		#' @family SimLongitudinal
34		#' @name SimLongitudinalRandomSlope-class
35		#' @exportClass SimLongitudinalRandomSlope
36		#'
37		#' @examples
38		#' SimLongitudinalRandomSlope(times = c(0, 30, 60))
39		.SimLongitudinalRandomSlope <- setClass(
40		"SimLongitudinalRandomSlope",
41		contains = "SimLongitudinal",
42		slots = c(
43		intercept = "numeric",
44		slope_mu = "numeric",
45		slope_sigma = "numeric",
46		sigma = "numeric",
47		link_dsld = "numeric",
48		link_identity = "numeric",
49		scaled_variance = "logical"
50		)
51		)
52
53		#' @rdname SimLongitudinalRandomSlope-class
54		#' @export
55		#'
56		#' @returns A `SimLongitudinalRandomSlope` object.
57		SimLongitudinalRandomSlope <- function(
58		times = c(-100, -50, 0, 50, 100, 150, 250, 350, 450, 550),
59		intercept = 50,
60		slope_mu = c(0.01, 0.03),
61		slope_sigma = c(0.5, 0.6),
62		sigma = 2,
63		link_dsld = 0,
64		link_identity = 0,
65		scaled_variance = FALSE
66		) {
67	17x	if (length(slope_sigma) == 1) {
68	4x	slope_sigma <- rep(slope_sigma, length(slope_mu))
69		}
70
71	17x	.SimLongitudinalRandomSlope(
72	17x	times = times,
73	17x	intercept = intercept,
74	17x	slope_mu = slope_mu,
75	17x	slope_sigma = slope_sigma,
76	17x	sigma = sigma,
77	17x	link_dsld = link_dsld,
78	17x	link_identity = link_identity,
79	17x	scaled_variance = scaled_variance
80		)
81		}
82
83		#' @rdname as_print_string
84		#' @exportS3Method NULL
85		as_print_string.SimLongitudinalRandomSlope <- function(object, ...) {
86	1x	return("SimLongitudinalRandomSlope")
87		}
88
89		#' @rdname sampleObservations
90		#' @export
91		sampleObservations.SimLongitudinalRandomSlope <- function(object, times_df) {
92	31x	times_df \|>
93	31x	dplyr::mutate(
94	31x	sld_mu = .data$intercept + .data$slope_ind * .data$time,
95	31x	sld_sd = ifelse(
96	31x	object@scaled_variance,
97	31x	pmax(.data$sld_mu * object@sigma, .Machine$double.eps),
98	31x	object@sigma
99		),
100	31x	sld = stats::rnorm(dplyr::n(), .data$sld_mu, .data$sld_sd),
101	31x	log_haz_link = object@link_dsld *
102	31x	.data$slope_ind +
103	31x	object@link_identity * .data$sld_mu
104		)
105		}
106
107		#' @rdname sampleSubjects
108		#' @export
109		sampleSubjects.SimLongitudinalRandomSlope <- function(object, subjects_df) {
110	16x	assert_that(
111	16x	is.factor(subjects_df[["study"]]),
112	16x	is.factor(subjects_df[["arm"]])
113		)
114
115	16x	assert_that(
116	16x	length(object@slope_mu) == nlevels(subjects_df[["arm"]]),
117	16x	msg = "`length(slope_mu)` should be equal to the number of arms"
118		)
119
120	16x	assert_that(
121	16x	length(object@slope_sigma) == nlevels(subjects_df[["arm"]]),
122	16x	msg = "`length(slope_sigma)` should be equal to the number of arms"
123		)
124
125	16x	assert_that(
126	16x	length(object@intercept) == nlevels(subjects_df[["study"]]),
127	16x	msg = "`length(intercept)` should be equal to the number of studies"
128		)
129
130	16x	assert_that(
131	16x	nrow(subjects_df) == length(unique(subjects_df[["subject"]])),
132	16x	msg = "The number of rows in `subjects_df` should be equal to the number of unique subjects"
133		)
134
135	16x	subjects_df \|>
136	16x	dplyr::mutate(intercept = object@intercept[as.numeric(.data$study)]) \|>
137	16x	dplyr::mutate(
138	16x	slope_ind = stats::rnorm(
139	16x	n = dplyr::n(),
140	16x	mean = object@slope_mu[as.numeric(.data$arm)],
141	16x	sd = object@slope_sigma[as.numeric(.data$arm)]
142		)
143		)
144		}

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(
38		"`newdata` must not contain a column named '..new_subject..'"
39		)
40		}
41		return(TRUE)
42		}
43		)
44
45
46		#' @rdname Quant-Dev
47		#' @export
48		as.QuantityGenerator.GridPrediction <- function(object, data, ...) {
49	6x	assert_class(data, "DataJoint")
50	6x	data_list <- as.list(data)
51	6x	validate_time_grid(object@times)
52
53	6x	n_times <- length(object@times)
54	6x	n_obs <- nrow(object@newdata)
55	6x	newdata <- object@newdata
56	6x	newdata[["..new_subject.."]] <- sprintf(
57	6x	"new_subject_%i",
58	6x	seq_len(nrow(newdata))
59		)
60
61	6x	QuantityGeneratorPrediction(
62	6x	times = rep(object@times, each = n_obs),
63	6x	newdata = replicate(newdata, n = n_times, simplify = FALSE) \|>
64	6x	dplyr::bind_rows(),
65	6x	params = object@params
66		)
67		}
68
69		#' @rdname Quant-Dev
70		#' @export
71		as.QuantityCollapser.GridPrediction <- function(object, data, ...) {
72	2x	assert_class(data, "DataJoint")
73	2x	data_list <- as.list(data)
74	2x	generator <- as.QuantityGenerator(object, data)
75	2x	QuantityCollapser(
76	2x	times = generator@times,
77	2x	groups = generator@newdata[["..new_subject.."]],
78	2x	indexes = as.list(seq_along(generator@times))
79		)
80		}
81
82
83		#' @export
84		#'
85		#' @returns This is just a placeholder which will error when used.
86		as.list.GridPrediction <- function(x, data, ...) {
87	!	stop("`as.list()` is not implemented for `GridPrediction` objects")
88		}
89
90
91		#' @rdname coalesceGridTime
92		#' @export
93		coalesceGridTime.GridPrediction <- function(object, times, ...) {
94	4x	if (is.null(object@times)) {
95	!	object <- GridPrediction(
96	!	times = times,
97	!	newdata = object@newdata,
98	!	params = object@params
99		)
100		}
101	4x	object
102		}

1		#' @include StanModule.R
2		#' @include LongitudinalModel.R
3		#' @include ParameterList.R
4		#' @include generics.R
5		#' @include Prior.R
6		NULL
7
8
9		#' `LinkComponent` Function Arguments
10		#'
11		#' This exists to contain all the common arguments for [`LinkComponent`] methods.
12		#'
13		#' @typed stan: StanModule
14		#' Stan code.
15		#' @typed x: LinkComponent
16		#' a link component.
17		#' @typed object: LinkComponent
18		#' a link component.
19		#' @param ... Not Used.
20		#'
21		#' @name LinkComponent-Shared
22		#' @keywords internal
23		NULL
24
25
26		#' `LinkComponent`
27		#'
28		#' @slot stan (`StanModule`)\cr See Arguments.
29		#' @slot name (`character`)\cr See Arguments.
30		#' @slot parameters (`ParameterList`)\cr The parameter specification.
31		#'
32		#' @typed stan: StanModule
33		#' Stan code. See Details.
34		#' @typed prior: Prior
35		#' The prior for the scaling coeficient.
36		#' @typed key: character
37		#' Link identifier. See Details.
38		#'
39		#' @details
40		#'
41		#' This object provides key information needed to construct a link contribution in the
42		#' survival model based on the parameters of the longitudinal model.
43		#'
44		#' Each link component defines a stan function of the longitudinal model parameters which is
45		#' multiplied by a model coefficient and added to the survival models hazard function.
46		#'
47		#' For full details about the specification of a `LinkComponent` please see
48		#' \code{vignette("extending-jmpost", package = "jmpost")}.
49		#'
50		#' @inheritParams stanmodel_arguments
51		#' @family LinkComponent
52		#' @name LinkComponent-class
53		#' @exportClass Link
54		#'
55		#' @examples
56		#' LinkComponent(
57		#' stan = StanModule(),
58		#' prior = prior_normal(0, 1),
59		#' key = "custom_link"
60		#' )
61		.LinkComponent <- setClass(
62		Class = "LinkComponent",
63		slots = list(
64		"stan" = "StanModule",
65		"parameters" = "ParameterList",
66		"key" = "character"
67		)
68		)
69
70
71		#' @rdname LinkComponent-class
72		#' @export
73		#'
74		#' @returns A `LinkComponent` object.
75		LinkComponent <- function(stan, prior, key, ...) {
76	58x	.LinkComponent(
77	58x	stan = stan,
78	58x	key = key,
79	58x	parameters = ParameterList(Parameter(
80	58x	name = key,
81	58x	prior = prior,
82	58x	size = 1
83		)),
84		...
85		)
86		}
87
88
89		#' @family LinkComponent
90		#' @rdname getParameters
91		#' @export
92		getParameters.LinkComponent <- function(object, ...) {
93	49x	object@parameters
94		}
95
96
97		#' @family LinkComponent
98		#' @rdname initialValues
99		#' @export
100		initialValues.LinkComponent <- function(object, n_chains, ...) {
101	!	initialValues(object@parameters, n_chains)
102		}
103
104
105		#' `LinkComponent` -> `StanModule`
106		#'
107		#' Converts a [`LinkComponent`] object to a [`StanModule`] object
108		#'
109		#' @inheritParams LinkComponent-Shared
110		#'
111		#' @family LinkComponent
112		#' @family as.StanModule
113		#' @export
114		#'
115		#' @returns A `StanModule` object.
116		as.StanModule.LinkComponent <- function(object, ...) {
117	57x	object@stan
118		}
119
120
121		#' `LinkComponent` -> `list`
122		#'
123		#' @inheritParams LinkComponent-Shared
124		#'
125		#' @description
126		#' Returns a named list where each element of the list corresponds
127		#' to a Stan modelling block e.g. `data`, `model`, etc.
128		#'
129		#' @family LinkComponent
130		#' @export
131		#'
132		#' @returns A named `list` representation of the object.
133		as.list.LinkComponent <- function(x, ...) {
134	!	stan <- as.StanModule(x, ...)
135	!	as.list(stan)
136		}
137
138		#' @family LinkComponent
139		#' @export
140		#'
141		#' @returns A character vector suitable for printing.
142		as_print_string.LinkComponent <- function(object, ...) {
143	13x	as_print_string(object@parameters)
144		}
145
146
147		#' @rdname show-object
148		#' @export
149		setMethod(
150		f = "show",
151		signature = "LinkComponent",
152		definition = function(object) {
153	3x	cat(
154	3x	paste0(
155	3x	"\nLinkComponent with parameter:\n ",
156	3x	as_print_string(object),
157	3x	"\n\n"
158		)
159		)
160		}
161		)
162
163		#' @family LinkComponent
164		#' @export
165		#'
166		#' @returns A character vector containing the component's parameter names.
167		names.LinkComponent <- function(x, ...) {
168	102x	names(x@parameters)
169		}

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

1		#' Simulating Joint Longitudinal and Time-to-Event Data
2		#'
3		#' @typed design: list
4		#' a list of [`SimGroup`] objects. See details.
5		#' @typed longitudinal: SimLongitudinal
6		#' object specifying how to simulate the longitudinal data
7		#' @typed survival: SimSurvival
8		#' object specifying how to simulate the survival data
9		#' @typed .silent: flag
10		#' whether to suppress info messages
11		#'
12		#' @slot longitudinal (`data.frame`)\cr the simulated longitudinal data.
13		#' @slot survival (`data.frame`)\cr the simulated survival data.
14		#'
15		#' @details
16		#'
17		#' The `design` argument is used to specify how many distinct groups should be simulated
18		#' including key information such as the number of subjects within the group as well as
19		#' which treatment arm and study the group belongs to. The `design` argument should be a
20		#' list of [`SimGroup`] objects e.g.
21		#' ```
22		#' design = list(
23		#' SimGroup(n = 50, study = "Study-1", arm = "Arm-A"),
24		#' SimGroup(n = 50, study = "Study-1", arm = "Arm-B")
25		#' )
26		#' ```
27		#'
28		#' @name SimJointData-class
29		#' @exportClass SimJointData
30		#'
31		#' @examples
32		#' set.seed(1)
33		#' SimJointData(
34		#' design = list(SimGroup(3, "A", "Study 1")),
35		#' longitudinal = SimLongitudinalRandomSlope(
36		#' times = c(0, 10),
37		#' slope_mu = 0.01,
38		#' slope_sigma = 0.5
39		#' ),
40		#' survival = SimSurvivalExponential(lambda = 1 / 300)
41		#' )
42		.SimJointData <- setClass(
43		"SimJointData",
44		slots = list(
45		longitudinal = "data.frame",
46		survival = "data.frame"
47		)
48		)
49
50
51		#' @rdname SimJointData-class
52		#' @export
53		#'
54		#' @returns A `SimJointData` object.
55		SimJointData <- function(
56		design = list(
57		SimGroup(n = 50, study = "Study-1", arm = "Arm-A"),
58		SimGroup(n = 50, study = "Study-1", arm = "Arm-B")
59		),
60		longitudinal,
61		survival,
62		.silent = FALSE
63		) {
64	24x	assert(
65	24x	all(vapply(design, \(x) is(x, "SimGroup"), logical(1))),
66	24x	msg = "All elements of `design` must be of class `SimGroup`"
67		)
68
69	24x	hazard_evaluation_info <- hazardWindows(survival)
70
71	24x	n_group <- vapply(design, function(x) x@n, numeric(1))
72	24x	arms <- vapply(design, function(x) x@arm, character(1))
73	24x	studies <- vapply(design, function(x) x@study, character(1))
74	24x	n_subjects <- sum(n_group)
75	24x	n_times <- length(hazard_evaluation_info$midpoint)
76
77	24x	sprintf_string <- paste0(
78	24x	"subject_%0",
79	24x	ceiling(log(n_subjects, 10)) + 1,
80	24x	"i"
81		)
82
83	24x	baseline <- dplyr::tibble(
84	24x	subject = sprintf(sprintf_string, seq_len(n_subjects))
85		) \|>
86	24x	dplyr::mutate(
87	24x	arm = factor(rep(arms, times = n_group), levels = unique(arms))
88		) \|>
89	24x	dplyr::mutate(
90	24x	study = factor(
91	24x	rep(studies, times = n_group),
92	24x	levels = unique(studies)
93		)
94		)
95
96	24x	os_baseline <- sampleSubjects(survival, subjects_df = baseline)
97	24x	lm_baseline <- sampleSubjects(longitudinal, subjects_df = baseline)
98
99	24x	lm_dat_no_obvs <- lapply(
100	24x	longitudinal@times,
101	24x	\(time) {
102	4558x	baseline[["time"]] <- time
103	4558x	baseline
104		}
105		) \|>
106	24x	dplyr::bind_rows() \|>
107	24x	dplyr::left_join(lm_baseline, by = c("subject", "study", "arm"))
108
109	24x	lm_dat <- sampleObservations(longitudinal, lm_dat_no_obvs)
110
111	24x	hazard_eval_df <- dplyr::tibble(
112	24x	subject = rep(lm_baseline$subject, each = n_times),
113	24x	arm = rep(lm_baseline$arm, each = n_times),
114	24x	study = rep(lm_baseline$study, each = n_times),
115	24x	midpoint = rep(
116	24x	as.double(hazard_evaluation_info$midpoint),
117	24x	times = n_subjects
118		),
119	24x	time = rep(as.double(hazard_evaluation_info$upper), times = n_subjects),
120	24x	width = rep(as.double(hazard_evaluation_info$width), times = n_subjects)
121		)
122
123	24x	lm_link_dat <- sampleObservations(
124	24x	longitudinal,
125	24x	dplyr::left_join(
126	24x	hazard_eval_df,
127	24x	lm_baseline,
128	24x	by = c("subject", "study", "arm")
129		)
130	24x	)[, c(
131	24x	"subject",
132	24x	"study",
133	24x	"arm",
134	24x	"log_haz_link",
135	24x	"time",
136	24x	"width",
137	24x	"midpoint"
138		)]
139
140	24x	os_eval_df <- lm_link_dat \|>
141	24x	dplyr::left_join(os_baseline, by = c("subject", "study", "arm"))
142
143	24x	withCallingHandlers(
144	24x	os_dat <- sampleObservations(survival, os_eval_df),
145	24x	message = function(e) {
146	8x	if (!.silent) {
147	!	message(e)
148		}
149	8x	invokeRestart("muffleMessage")
150		}
151		)
152
153	24x	lm_dat2 <- lm_dat \|>
154	24x	dplyr::left_join(
155	24x	dplyr::select(os_dat, "subject", os_time = "time"),
156	24x	by = "subject"
157		) \|>
158	24x	dplyr::mutate(observed = (.data$time <= .data$os_time)) \|>
159	24x	dplyr::arrange(dplyr::pick(c("subject", "time")))
160
161	24x	assert_that(
162	24x	length(unique(os_dat$subject)) == length(os_dat$subject),
163	24x	length(os_dat$subject) == n_subjects,
164	24x	all(os_dat$time >= 0),
165	24x	all(os_dat$event %in% c(0, 1)),
166	24x	msg = "Assumptions for the Survival data are not met (please report this issue)"
167		)
168
169	24x	assert_that(
170	24x	nrow(lm_dat2) == n_subjects * length(longitudinal@times),
171	24x	length(unique(lm_dat2$subject)) == n_subjects,
172	24x	msg = "Assumptions for the Longitudinal data are not met (please report this issue)"
173		)
174
175	24x	return(
176	24x	.SimJointData(
177	24x	survival = os_dat,
178	24x	longitudinal = lm_dat2[, c(
179	24x	"subject",
180	24x	"arm",
181	24x	"study",
182	24x	"time",
183	24x	"sld",
184	24x	"observed"
185		)]
186		)
187		)
188		}
189
190		#' @rdname show-object
191		#' @export
192		setMethod(
193		f = "show",
194		signature = "SimJointData",
195		definition = function(object) {
196	1x	x <- sprintf("\nA SimJointData Object\n\n")
197	1x	cat(x)
198	1x	return(object)
199		}
200		)
201
202
203		#' Add PFS events at Tumour Progression to Data
204		#'
205		#' Adds new columns `pfs_time` and `pfs_event` based on observed changes to SLD.
206		#'
207		#' @typed object: SimJointData
208		#' where the PFS events should be added.
209		#' @typed relative_threshold: number
210		#' a multiplicative threshold for the change in SLD compared to the `min(SLD)`.
211		#' Default is 1.2 meaning a 20% increase.
212		#' @typed absolute_threshold: number
213		#' an absolute threshold for the change in SLD compared to the minimum.
214		#' Default is 5.
215		#' @typed from_time: number
216		#' Ignore observations before this time for determining SLD minimum.
217		#' @typed observed_after: logical
218		#' If `FALSE` set longitudinal observations after the progression time to
219		#' `observed = FALSE`
220		#' @details
221		#' Both thresholds must be met for a progression to be declared.
222		#'
223		#' @export
224		#' @examples
225		#' data <- SimJointData(
226		#' survival = SimSurvivalExponential(lambda = 1/10),
227		#' longitudinal = SimLongitudinalSteinFojo()
228		#' )
229		#' data <- add_pfs(data)
230		#' data@survival # now has pfs_time and pfs_event columns
231		#'
232		#' @returns The input `SimJointData` object with progression-free survival columns
233		#' `pfs_time` and `pfs_event` added.
234		add_pfs <- function(
235		object,
236		relative_threshold = 1.2,
237		absolute_threshold = 5,
238		from_time = 0,
239		observed_after = FALSE
240		) {
241	3x	assert_class(object, "SimJointData")
242
243	3x	pd_times <- object@longitudinal \|>
244	3x	dplyr::filter(.data$time >= from_time) \|>
245	3x	dplyr::mutate(
246	3x	min_sld = cummin(.data$sld),
247	3x	is_pd = .data$sld >=
248	3x	pmax(
249	3x	.data$min_sld * relative_threshold,
250	3x	.data$min_sld + absolute_threshold
251		) &
252	3x	.data$observed,
253	3x	pd_time = min(.data$time[.data$is_pd], Inf),
254	3x	.by = "subject"
255		) \|>
256	3x	dplyr::select("subject", "pd_time") \|>
257	3x	dplyr::slice_head(by = "subject")
258
259	3x	if (isFALSE(observed_after)) {
260	2x	object@longitudinal <- object@longitudinal \|>
261	2x	dplyr::left_join(pd_times, by = "subject") \|>
262	2x	dplyr::mutate(
263	2x	observed = dplyr::if_else(
264	2x	.data$time > .data$pd_time,
265	2x	FALSE,
266	2x	.data$observed
267		),
268	2x	pd_time = NULL
269		)
270		}
271
272	3x	object@survival <-
273	3x	object@survival \|>
274	3x	dplyr::left_join(pd_times, by = "subject") \|>
275	3x	dplyr::mutate(
276	3x	pfs_time = pmin(.data$time, .data$pd_time, na.rm = TRUE),
277	3x	pfs_event = dplyr::if_else(
278	3x	.data$pfs_time < .data$time,
279	3x	1,
280	3x	.data$event
281		),
282	3x	pd_time = NULL
283		)
284	3x	object
285		}
286
287
288		#' Cut Study Data
289		#' @typed object: SimJointData
290		#' where the cut times should be applied.
291		#' @typed cut_time: numeric
292		#' A vector of cut off times, either length 1 for all patients or
293		#' `nrow(object@survival)` for a time per patient.
294		#' @details
295		#' All observations after this time are remove. Survival is censored at this time and any longitudinal
296		#' values are removed.
297		#' @export
298		#' @examples
299		#' data <- SimJointData(
300		#' survival = SimSurvivalExponential(lambda = 1/10),
301		#' longitudinal = SimLongitudinalSteinFojo()
302		#' )
303		#' data <- cut_data(data, 5)
304		#' data@survival
305		#' # Now max time is 5
306		#' max(data@survival$time)
307		#'
308		#' @returns The input `SimJointData` object restricted to the requested follow-up time.
309		cut_data <- function(object, cut_time) {
310	3x	assert_class(object, "SimJointData")
311	3x	check_len <- if (length(cut_time) > 1) nrow(object@survival) else 1
312	3x	assert_numeric(cut_time, lower = 0, len = check_len)
313
314	2x	object@survival <- object@survival \|>
315	2x	dplyr::mutate(
316	2x	cut_time = cut_time,
317	2x	event = dplyr::if_else(.data$time < .data$cut_time, .data$event, 0),
318	2x	time = pmin(.data$cut_time, .data$time)
319		)
320	2x	if (all(c("pfs_event", "pfs_time") %in% colnames(object@survival))) {
321	!	object@survival <- object@survival \|>
322	!	dplyr::mutate(
323	!	pfs_event = dplyr::if_else(
324	!	.data$pfs_time < .data$cut_time,
325	!	.data$pfs_event,
326	!	0
327		),
328	!	pfs_time = pmin(.data$cut_time, .data$pfs_time)
329		)
330		}
331
332	2x	object@longitudinal <- object@longitudinal \|>
333	2x	dplyr::left_join(
334	2x	dplyr::select(object@survival, "subject", "cut_time"),
335	2x	by = "subject"
336		) \|>
337	2x	dplyr::filter(.data$time <= .data$cut_time) \|>
338	2x	dplyr::mutate(cut_time = NULL)
339
340	2x	object@survival$cut_time <- NULL
341
342	2x	object
343		}

1		#' `SimSurvival` Function Arguments
2		#'
3		#' The documentation lists all the conventional arguments for [`SimSurvival`]
4		#' constructors.
5		#'
6		#' @typed time_max: number
7		#' the maximum time to simulate to.
8		#' @typed time_step: number
9		#' the time interval between evaluating the log-hazard function.
10		#' @typed lambda_censor: number
11		#' the censoring rate.
12		#' @typed beta_cont: number
13		#' the continuous covariate coefficient.
14		#' @typed beta_cat: numeric
15		#' the categorical covariate coefficients.
16		#' @typed loghazard: function
17		#' the log hazard function.
18		#' @typed name: character
19		#' the name of the object.
20		#' @param ... Not Used.
21		#'
22		#' @section Hazard Evaluation:
23		#'
24		#' Event times are simulated by sampling a cumulative hazard limit from a \eqn{U(0, 1)} distribution
25		#' for
26		#' each subject and then counting how much hazard they've been exposed to by evaluating the
27		#' log-hazard function at a set interval. The `time_max` argument sets the upper bound for the
28		#' number of time points to evaluate the log-hazard function at with subjects who have not had an
29		#' event being censored at `time_max`. The `time_step` argument sets the interval at which to
30		#' evaluate the log-hazard function. Setting smaller values for `time_step` will increase the
31		#' precision of the simulation at the cost of increased computation time. Likewise, setting large
32		#' values for `time_max` will minimize the number of censored subjects at the cost of
33		#' increased computation time.
34		#'
35		#' @name SimSurvival-Shared
36		#' @keywords internal
37		NULL
38
39
40		#' Abstract Simulation Class for Survival Data
41		#'
42		#' @inheritParams SimSurvival-Shared
43		#' @inheritSection SimSurvival-Shared Hazard Evaluation
44		#'
45		#' @slot time_max (`numeric`)\cr See arguments.
46		#' @slot time_step (`numeric`)\cr See arguments.
47		#' @slot lambda_censor (`numeric`)\cr See arguments.
48		#' @slot beta_cont (`numeric`)\cr See arguments.
49		#' @slot beta_cat (`numeric`)\cr See arguments.
50		#' @slot loghazard (`function`)\cr See arguments.
51		#' @slot name (`character`)\cr See arguments.
52		#' @slot beta_os_cov (`numeric`) Additional regression coefficients for survival models.
53		#'
54		#' @family SimSurvival
55		#' @exportClass SimSurvival
56		#' @name SimSurvival-class
57		#'
58		#' @examples
59		#' SimSurvival(loghazard = function(time, ...) rep(log(1 / 100), length(time)))
60		.SimSurvival <- setClass(
61		"SimSurvival",
62		slots = c(
63		time_max = "numeric",
64		time_step = "numeric",
65		lambda_censor = "numeric",
66		beta_cont = "numeric",
67		beta_cat = "numeric",
68		beta_os_cov = "numeric",
69		loghazard = "function",
70		name = "character"
71		)
72		)
73
74		#' @rdname SimSurvival-class
75		#' @export
76		#'
77		#' @returns A `SimSurvival` object.
78		SimSurvival <- function(
79		time_max = 2000,
80		time_step = 1,
81		lambda_censor = 1 / 3000,
82		beta_cont = 0.2,
83		beta_cat = c("A" = 0, "B" = -0.4, "C" = 0.2),
84		loghazard,
85		name = "SimSurvival"
86		) {
87	30x	.SimSurvival(
88	30x	time_max = time_max,
89	30x	time_step = time_step,
90	30x	lambda_censor = lambda_censor,
91	30x	beta_cont = beta_cont,
92	30x	beta_cat = beta_cat,
93	30x	loghazard = loghazard,
94	30x	name = name
95		)
96		}
97
98
99		#' @rdname show-object
100		#' @export
101		setMethod(
102		f = "show",
103		signature = "SimSurvival",
104		definition = function(object) {
105	4x	x <- sprintf("\nA %s Object\n\n", as_print_string(object))
106	4x	cat(x)
107	4x	return(object)
108		}
109		)
110
111		#' @rdname as_print_string
112		#' @exportS3Method NULL
113		as_print_string.SimSurvival <- function(object, ...) {
114	4x	return(object@name)
115		}
116
117
118		#' Construct Time Intervals
119		#'
120		#' @typed object: SimSurvival
121		#' the survival simulation object to create evaluation points for.
122		#' @param ... Not Used.
123		#'
124		#' @return A `tibble` with `lower`, `upper`, `time`, `eval` and `width`.
125		#' @keywords internal
126		#' @exportS3Method NULL
127		hazardWindows.SimSurvival <- function(object, ...) {
128	24x	times <- seq(0, object@time_max, object@time_step)
129	24x	bound_lower <- times[-length(times)]
130	24x	bound_upper <- times[-1]
131	24x	bound_width <- bound_upper - bound_lower
132	24x	mid_point <- bound_upper - (bound_width / 2)
133	24x	tibble::tibble(
134	24x	lower = bound_lower,
135	24x	upper = bound_upper,
136	24x	midpoint = mid_point,
137	24x	width = bound_width
138		)
139		}
140
141		#' @rdname sampleSubjects
142		#' @export
143		sampleSubjects.SimSurvival <- function(object, subjects_df) {
144	24x	subjects_df \|>
145	24x	dplyr::mutate(cov_cont = stats::rnorm(dplyr::n())) \|>
146	24x	dplyr::mutate(
147	24x	cov_cat = factor(
148	24x	sample(
149	24x	names(object@beta_cat),
150	24x	replace = TRUE,
151	24x	size = dplyr::n()
152		),
153	24x	levels = names(object@beta_cat)
154		)
155		) \|>
156	24x	dplyr::mutate(
157	24x	log_haz_cov = .data$cov_cont *
158	24x	object@beta_cont +
159	24x	object@beta_cat[.data$cov_cat]
160		) \|>
161	24x	dplyr::mutate(survival = stats::runif(dplyr::n())) \|>
162	24x	dplyr::mutate(chazard_limit = -log(.data$survival)) \|>
163	24x	dplyr::mutate(time_cen = stats::rexp(dplyr::n(), object@lambda_censor))
164		}
165
166
167		#' @rdname sampleObservations
168		#' @export
169		sampleObservations.SimSurvival <- function(object, times_df) {
170	24x	assert_that(
171	24x	all(times_df$time >= 0),
172	24x	msg = "All time points must be greater than or equal to 0"
173		)
174
175	24x	os_dat_chaz <- times_df \|>
176	24x	dplyr::mutate(log_bl_haz = object@loghazard(.data$midpoint)) \|>
177		# Fix to avoid issue with log(0) = NaN values
178	24x	dplyr::mutate(
179	24x	log_bl_haz = dplyr::if_else(
180	24x	.data$midpoint == 0,
181	24x	-999,
182	24x	.data$log_bl_haz
183		)
184		) \|>
185	24x	dplyr::mutate(
186	24x	hazard_instant = exp(
187	24x	.data$log_bl_haz + .data$log_haz_cov + .data$log_haz_link
188		)
189		) \|>
190		# Reset Inf values to large number to avoid NaN issues downstream
191		# This is suitable as Hazard limits tend to be in the range of -10 to 10 so large numbers
192		# are essentially equivalent to infinity for simulation purposes
193	24x	dplyr::mutate(
194	24x	hazard_instant = dplyr::if_else(
195	24x	.data$hazard_instant == Inf,
196	24x	999,
197	24x	.data$hazard_instant
198		)
199		) \|>
200	24x	dplyr::mutate(
201	24x	hazard_instant = dplyr::if_else(
202	24x	.data$hazard_instant == -Inf,
203	24x	-999,
204	24x	.data$hazard_instant
205		)
206		) \|>
207	24x	dplyr::mutate(hazard_interval = .data$hazard_instant * .data$width) \|>
208	24x	dplyr::group_by(.data$subject) \|>
209	24x	dplyr::mutate(chazard = cumsum(.data$hazard_interval)) \|>
210	24x	dplyr::ungroup()
211
212	24x	os_had_event <- os_dat_chaz \|>
213	24x	dplyr::filter(.data$chazard >= .data$chazard_limit) \|>
214	24x	dplyr::group_by(.data$subject) \|>
215	24x	dplyr::slice(1) \|>
216	24x	dplyr::ungroup() \|>
217	24x	dplyr::mutate(event = 1)
218
219	24x	os_had_censor <- os_dat_chaz \|>
220	24x	dplyr::filter(!.data$subject %in% os_had_event$subject) \|>
221	24x	dplyr::group_by(.data$subject) \|>
222	24x	dplyr::slice(dplyr::n()) \|>
223	24x	dplyr::ungroup() \|>
224	24x	dplyr::mutate(event = 0)
225
226	24x	if (!(nrow(os_had_censor) == 0)) {
227	8x	message(sprintf(
228	8x	"INFO: %i subject(s) did not die before max(times)",
229	8x	nrow(os_had_censor)
230		))
231		}
232
233	24x	os_dat_complete <- os_had_event \|>
234	24x	dplyr::bind_rows(os_had_censor) \|>
235	24x	dplyr::mutate(real_time = .data$time) \|>
236	24x	dplyr::mutate(
237	24x	event = dplyr::if_else(
238	24x	.data$real_time <= .data$time_cen,
239	24x	.data$event,
240	24x	0
241		)
242		) \|>
243	24x	dplyr::mutate(
244	24x	time = dplyr::if_else(
245	24x	.data$real_time <= .data$time_cen,
246	24x	.data$real_time,
247	24x	.data$time_cen
248		)
249		) \|>
250	24x	dplyr::arrange(.data$subject)
251
252	24x	keep_cols <- colnames(os_dat_complete) %in%
253	24x	c("subject", "study", "arm", "time", "event", "cov_cont", "cov_cat")
254	24x	os_dat_complete[, keep_cols]
255		}
256
257
258		#' Simulate Survival Data from a Weibull Proportional Hazard Model
259		#'
260		#' @typed lambda: number
261		#' the scale parameter.
262		#' @typed gamma: number
263		#' the shape parameter.
264		#'
265		#' @inheritParams SimSurvival-Shared
266		#' @inheritSection SimSurvival-Shared Hazard Evaluation
267		#'
268		#' @family SimSurvival
269		#'
270		#' @export
271		#'
272		#' @returns A `SimSurvival` object.
273		#'
274		#' @examples
275		#' SimSurvivalWeibullPH(lambda = 1 / 300, gamma = 1.2)
276		SimSurvivalWeibullPH <- function(
277		lambda,
278		gamma,
279		time_max = 2000,
280		time_step = 1,
281		lambda_censor = 1 / 3000,
282		beta_cont = 0.2,
283		beta_cat = c("A" = 0, "B" = -0.4, "C" = 0.2)
284		) {
285	4x	SimSurvival(
286	4x	time_max = time_max,
287	4x	time_step = time_step,
288	4x	lambda_censor = lambda_censor,
289	4x	beta_cont = beta_cont,
290	4x	beta_cat = beta_cat,
291	4x	loghazard = function(time) {
292	3x	log(lambda) + log(gamma) + (gamma - 1) * log(time)
293		},
294	4x	name = "SimSurvivalWeibullPH"
295		)
296		}
297
298		#' Simulate Survival Data from a Log-Logistic Proportional Hazard Model
299		#'
300		#' @typed a: number
301		#' the scale parameter.
302		#' @typed b: number
303		#' the shape parameter.
304		#'
305		#' @inheritParams SimSurvival-Shared
306		#' @inheritSection SimSurvival-Shared Hazard Evaluation
307		#'
308		#' @family SimSurvival
309		#' @export
310		#'
311		#' @returns A `SimSurvival` object.
312		#'
313		#' @examples
314		#' SimSurvivalLogLogistic(a = 300, b = 2)
315		SimSurvivalLogLogistic <- function(
316		a,
317		b,
318		time_max = 2000,
319		time_step = 1,
320		lambda_censor = 1 / 3000,
321		beta_cont = 0.2,
322		beta_cat = c("A" = 0, "B" = -0.4, "C" = 0.2)
323		) {
324	2x	SimSurvival(
325	2x	time_max = time_max,
326	2x	time_step = time_step,
327	2x	lambda_censor = lambda_censor,
328	2x	beta_cont = beta_cont,
329	2x	beta_cat = beta_cat,
330	2x	loghazard = function(time) {
331	1x	c1 <- -log(a) + log(b) + (b - 1) * (-log(a) + log(time))
332	1x	c2 <- log(1 + (time / a)^b)
333	1x	return(c1 - c2)
334		},
335	2x	name = "SimSurvivalLogLogistic"
336		)
337		}
338
339
340		#' Simulate Survival Data from a Exponential Proportional Hazard Model
341		#'
342		#' @typed lambda: number
343		#' the rate parameter.
344		#'
345		#' @inheritParams SimSurvival-Shared
346		#' @inheritSection SimSurvival-Shared Hazard Evaluation
347		#'
348		#' @family SimSurvival
349		#'
350		#' @export
351		#'
352		#' @returns A `SimSurvival` object.
353		#'
354		#' @examples
355		#' SimSurvivalExponential(lambda = 1 / 300)
356		SimSurvivalExponential <- function(
357		lambda,
358		time_max = 2000,
359		time_step = 1,
360		lambda_censor = 1 / 3000,
361		beta_cont = 0.2,
362		beta_cat = c("A" = 0, "B" = -0.4, "C" = 0.2)
363		) {
364	22x	SimSurvival(
365	22x	time_max = time_max,
366	22x	time_step = time_step,
367	22x	lambda_censor = lambda_censor,
368	22x	beta_cont = beta_cont,
369	22x	beta_cat = beta_cat,
370	22x	loghazard = function(time) {
371	21x	log(lambda)
372		},
373	22x	name = "SimSurvivalExponential"
374		)
375		}
376
377
378		#' Simulate Survival Data from a Gamma Proportional Hazard Model
379		#'
380		#' @typed k: number
381		#' the shape parameter.
382		#' @typed theta: number
383		#' the scale parameter.
384		#'
385		#' @inheritParams SimSurvival-Shared
386		#' @inheritSection SimSurvival-Shared Hazard Evaluation
387		#'
388		#' @family SimSurvival
389		#'
390		#' @importFrom stats dgamma pgamma
391		#'
392		#' @export
393		#'
394		#' @returns A `SimSurvival` object.
395		#'
396		#' @examples
397		#' SimSurvivalGamma(k = 2, theta = 300)
398		SimSurvivalGamma <- function(
399		k,
400		theta,
401		time_max = 2000,
402		time_step = 1,
403		lambda_censor = 1 / 3000,
404		beta_cont = 0.2,
405		beta_cat = c("A" = 0, "B" = -0.4, "C" = 0.2)
406		) {
407	1x	SimSurvival(
408	1x	time_max = time_max,
409	1x	time_step = time_step,
410	1x	lambda_censor = lambda_censor,
411	1x	beta_cont = beta_cont,
412	1x	beta_cat = beta_cat,
413	1x	loghazard = function(time) {
414	1x	dgamma(time, k, scale = theta, log = TRUE) -
415	1x	pgamma(time, k, scale = theta, lower.tail = FALSE, log.p = TRUE)
416		},
417	1x	name = "SimSurvivalGamma"
418		)
419		}

1		#' Re-used documentation for `RandomEffectQuantities`
2		#'
3		#' @typed x: RandomEffectQuantities
4		#' generated quantities.
5		#' @typed object: RandomEffectQuantities
6		#' generated quantities.
7		#' @typed conf.level: numeric
8		#' confidence level of the interval.
9		#' @param ... not used.
10		#'
11		#' @keywords internal
12		#' @name RandomEffectQuantities-Shared
13		NULL
14
15
16		#' Random Effects Quantities Container
17		#'
18		#' A simple wrapper around a `matrix` to store required metadata for patient level
19		#' random effects data
20		#'
21		#' @typed quantities: matrix
22		#' of random effects values.
23		#' @typed subject: character
24		#' labels specifying which subjects the values belong to.
25		#' @typed parameter: character
26		#' labels specifying which parameter the value is.
27		#'
28		#' @slot quantities (`matrix`)\cr See Arguments for details.
29		#' @slot subject (`numeric`)\cr See Arguments for details.
30		#' @slot parameter (`character`)\cr See Arguments for details.
31		#'
32		#' @details
33		#' Each row of the matrix represents a sample and each column represents a distinct subject
34		#' specific parameter.
35		#' As such the number of columns in the matrix should equal the length of `subject` and `parameter`
36		#' which provide metadata for who the parameter corresponds to as well as which parameter it is.
37		#'
38		#' @keywords internal
39		#' @name RandomEffectQuantities-class
40		#' @family RandomEffectQuantities
41		.RandomEffectQuantities <- setClass(
42		"RandomEffectQuantities",
43		slots = list(
44		"quantities" = "matrix",
45		"subject" = "character",
46		"parameter" = "character"
47		)
48		)
49		#' @rdname RandomEffectQuantities-class
50		#'
51		#' @returns A `RandomEffectQuantities` object.
52		RandomEffectQuantities <- function(quantities, subject, parameter) {
53	1x	.RandomEffectQuantities(
54	1x	quantities = quantities,
55	1x	subject = subject,
56	1x	parameter = parameter
57		)
58		}
59
60		setValidity(
61		Class = "RandomEffectQuantities",
62		method = function(object) {
63		if (length(object@subject) != ncol(object@quantities)) {
64		return(
65		"Length of `subject` must be equal to the number of columns in `quantities`"
66		)
67		}
68		if (length(object@parameter) != ncol(object@quantities)) {
69		return(
70		"Length of `parameter` must be equal to the number of columns in `quantities`"
71		)
72		}
73		TRUE
74		}
75		)
76
77
78		#' `RandomEffectQuantities` -> Printable `Character`
79		#'
80		#' Converts [`RandomEffectQuantities`] object into a printable string.
81		#' @inheritParams RandomEffectQuantities-Shared
82		#' @typed indent: numeric
83		#' the number of spaces to indent the string by.
84		#' @family RandomEffectQuantities
85		#' @keywords internal
86		#' @export
87		#'
88		#' @returns A character vector suitable for printing.
89		as_print_string.RandomEffectQuantities <- function(object, indent = 1, ...) {
90	!	parameter_string <- paste0(" ", unique(object@parameter))
91	!	template <- c(
92	!	"RandomEffectQuantities Object:",
93	!	" # of samples = %d",
94	!	" # of unique subjects = %d",
95	!	" For parameters:",
96	!	parameter_string
97		)
98	!	pad <- rep(" ", indent) \|> paste(collapse = "")
99	!	template_padded <- paste(pad, template)
100	!	sprintf(
101	!	paste(template_padded, collapse = "\n"),
102	!	nrow(object@quantities),
103	!	length(unique(object@subject))
104		)
105		}
106
107
108		#' @rdname show-object
109		#' @export
110		setMethod(
111		f = "show",
112		signature = "RandomEffectQuantities",
113		definition = function(object) {
114	!	string <- as_print_string(object)
115	!	cat("\n", string, "\n\n")
116		}
117		)
118
119
120		#' `RandomEffectQuantities` -> `data.frame`
121		#'
122		#' Returns a `data.frame` of the subject-level random effect parameter samples.
123		#'
124		#' @inheritParams RandomEffectQuantities-Shared
125		#'
126		#' @keywords internal
127		#' @family RandomEffectQuantities
128		#' @export
129		#'
130		#' @returns A `data.frame` representation of the object.
131		as.data.frame.RandomEffectQuantities <- function(x, ...) {
132	1x	data.frame(
133	1x	subject = rep(x@subject, each = nrow(x@quantities)),
134	1x	parameter = rep(x@parameter, each = nrow(x@quantities)),
135	1x	values = as.vector(x@quantities)
136		)
137		}
138
139
140		#' summary
141		#'
142		#' @description
143		#' This method returns a summary statistic `data.frame` of the random effect parameters
144		#'
145		#' @inheritParams RandomEffectQuantities-Shared
146		#'
147		#' @returns
148		#' A `data.frame` with the following variables:
149		#' - `subject` (`character`) \cr the subject identifier.
150		#' - `parameter` (`character`) \cr the parameter identifier.
151		#' - `median` (`numeric`) \cr the median value of the quantity.
152		#' - `lower` (`numeric`) \cr the lower CI value of the quantity.
153		#' - `upper` (`numeric`) \cr the upper CI value of the quantity.
154		#'
155		#' @keywords internal
156		#' @family RandomEffectQuantities
157		#' @export
158		summary.RandomEffectQuantities <- function(object, conf.level = 0.95, ...) {
159	1x	quantities_summarised <- samples_median_ci(
160	1x	object@quantities,
161	1x	level = conf.level
162		)
163
164	1x	quantities_summarised$subject <- object@subject
165	1x	quantities_summarised$parameter <- object@parameter
166	1x	quantities_summarised[, c(
167	1x	"subject",
168	1x	"parameter",
169	1x	"median",
170	1x	"lower",
171	1x	"upper"
172		)]
173		}
174
175
176		#' Extract Random Effects Samples from a Longitudinal Model
177		#'
178		#' Helper function to extract subject-level random effects samples from the longitudinal
179		#' sub-model of a joint model samples object.
180		#'
181		#' @typed object: JointModelSamples
182		#' samples as drawn from a Joint Model.
183		#' @family RandomEffectQuantities
184		#' @family JointModelSamples
185		#' @export
186		#'
187		#' @returns A `RandomEffectQuantities` object.
188		#'
189		#' @examples
190		#' \dontrun{
191		#' LongitudinalRandomEffects(fit)
192		#' }
193		LongitudinalRandomEffects <- function(object) {
194	1x	assert_class(object, "JointModelSamples")
195
196	1x	subject_indexes <- as.list(object@data)$subject_to_index
197	1x	random_effects_names <- getRandomEffectsNames(object@model)
198	1x	expanded <- expand.grid(
199	1x	parameter_long = random_effects_names,
200	1x	subject_indexes = subject_indexes
201		)
202	1x	expanded$subject <- names(subject_indexes)[expanded$subject_indexes]
203	1x	expanded$parameter <- names(random_effects_names)[expanded$parameter_long]
204	1x	stan_parameter_names <- sprintf(
205	1x	"%s[%i]",
206	1x	expanded$parameter_long,
207	1x	expanded$subject_indexes
208		)
209	1x	draw_matrix <- posterior::as_draws_matrix(object@results$draws(
210	1x	stan_parameter_names
211		))
212	1x	class(draw_matrix) <- "matrix"
213	1x	rownames(draw_matrix) <- NULL
214	1x	colnames(draw_matrix) <- NULL
215
216	1x	RandomEffectQuantities(
217	1x	draw_matrix,
218	1x	subject = expanded$subject,
219	1x	parameter = expanded$parameter
220		)
221		}

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

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		#' @slot scaled_variance logical scalar, whether the variance should be scaled
12		#' by the expected value, corresponding to a multiplicative model. If not,
13		#' then an additive error model is used.
14		#'
15		#' @exportClass LongitudinalModel
16		.LongitudinalModel <- setClass(
17		Class = "LongitudinalModel",
18		contains = "StanModel",
19		slots = c(
20		scaled_variance = "logical"
21		)
22		)
23
24		# LongitudinalModel-constructors ----
25
26		#' @rdname LongitudinalModel-class
27		#'
28		#' @inheritParams stanmodel_arguments
29		#' @typed scaled_variance: flag
30		#' whether the variance should be scaled
31		#' by the expected value, corresponding to a multiplicative model. If not,
32		#' then an additive error model is used.
33		#'
34		#' @export
35		#'
36		#' @returns A `LongitudinalModel` object.
37		#'
38		#' @examples
39		#' LongitudinalModel(scaled_variance = FALSE)
40		LongitudinalModel <- function(
41		stan = StanModule(),
42		parameters = ParameterList(),
43		name = "<Unnamed>",
44		scaled_variance = NA,
45		...
46		) {
47	99x	assert_flag(scaled_variance)
48	99x	base_stan <- read_stan("base/longitudinal.stan")
49
50	99x	stan_full <- decorated_render(
51	99x	.x = base_stan,
52	99x	stan = add_missing_stan_blocks(as.list(stan))
53		)
54
55	99x	.LongitudinalModel(
56	99x	StanModel(
57	99x	stan = StanModule(stan_full),
58	99x	parameters = parameters,
59	99x	name = name,
60		...
61		),
62	99x	scaled_variance = scaled_variance
63		)
64		}
65
66		#' @export
67		#'
68		#' @returns A character vector suitable for printing.
69		as_print_string.LongitudinalModel <- function(object, ...) {
70	12x	string <- sprintf(
71	12x	"\n%s Longitudinal Model (%s error) with parameters:\n%s\n\n",
72	12x	object@name,
73	12x	if (is.na(object@scaled_variance)) {
74		# Should not happen but just to be robust against this case.
75	!	"unknown"
76	12x	} else if (object@scaled_variance) {
77	!	"multiplicative"
78		} else {
79	12x	"additive"
80		},
81	12x	paste(" ", as_print_string(object@parameters)) \|>
82	12x	paste(collapse = "\n")
83		)
84	12x	return(string)
85		}

1		#' @include DataJoint.R
2		#' @include Quantities.R
3		NULL
4
5
6		#' Re-used documentation for `SurvivalQuantities`
7		#'
8		#' @typed object: SurvivalQuantities
9		#' survival quantities.
10		#' @typed x: SurvivalQuantities
11		#' survival quantities.
12		#' @param ... not used.
13		#'
14		#' @keywords internal
15		#' @name SurvivalQuantities-Shared
16		NULL
17
18
19		#' `SurvivalQuantities` Object & Constructor Function
20		#'
21		#' Constructor function to generate a `SurvivalQuantities` object.
22		#'
23		#' @slot quantities (`Quantities`)\cr The sampled quantities. Should contain 1 element per
24		#' element of `group`
25		#' @slot groups (`list`)\cr See argument section for details
26		#' @slot type (`character`)\cr See See argument section for details
27		#' @slot time_grid (`numeric`)\cr See argument section for details
28		#' @slot data ([`DataJoint`])\cr The data that the Joint Model was fitted to to produce
29		#' the samples/quantities
30		#'
31		#' @section Group Specification:
32		#' If `groups` is a character vector of subject IDs then the survival quantities will
33		#' only be calculated for those specific subjects.
34		#'
35		#' If `groups` is a list then any elements with more than 1 subject ID will be grouped together
36		#' and their quantities will be calculated by taking a point-wise average.
37		#' For example: `groups = list("g1" = c("sub1", "sub2"), "g2" = c("sub3", "sub4"))` would result
38		#' in 2 groups being created whose values are the pointwise average
39		#' of `c("sub1", "sub2")` and `c("sub3", "sub4")` respectively.
40		#'
41		#' If `groups=NULL` then all subjects from original dataset will be selected
42		#'
43		#' @family SurvivalQuantities
44		#' @name SurvivalQuantities-class
45		#' @export SurvivalQuantities
46		#'
47		#' @examples
48		#' \dontrun{
49		#' SurvivalQuantities(fit, GridFixed(times = seq(0, 100, by = 10)))
50		#' }
51		.SurvivalQuantities <- setClass(
52		Class = "SurvivalQuantities",
53		slots = c(
54		"quantities" = "Quantities",
55		"grid" = "Grid",
56		"type" = "character",
57		"data" = "DataJoint"
58		)
59		)
60
61		#' @typed object: JointModelSamples
62		#' samples as drawn from a Joint Model.
63		#'
64		#' @typed grid: Grid
65		#' object that specifies which subjects and time points to calculate the
66		#' quantities for. See [Grid-Functions].
67		#'
68		#' @typed type: character
69		#' quantity to be generated.
70		#' Must be one of `surv`, `haz`, `loghaz`, `cumhaz`.
71		#'
72		#' @rdname SurvivalQuantities-class
73		#'
74		#' @returns A `SurvivalQuantities` object.
75		SurvivalQuantities <- function(
76		object,
77		grid,
78		type = c("surv", "haz", "loghaz", "cumhaz")
79		) {
80	28x	type <- match.arg(type)
81	28x	assert_class(object, "JointModelSamples")
82	28x	assert_class(grid, "Grid")
83	26x	assert_that(
84	26x	!is(grid, "GridPopulation"),
85	26x	msg = "GridPopulation objects are not supported for `SurvivalQuantities`"
86		)
87
88	25x	time_grid <- seq(
89	25x	from = 0,
90	25x	to = max(as.list(object@data)[["event_times"]]),
91	25x	length = 201
92		)
93
94	25x	grid <- coalesceGridTime(grid, time_grid)
95
96	25x	generator <- as.QuantityGenerator(grid, data = object@data)
97
98	25x	assert_that(
99	25x	all(generator@times >= 0),
100	25x	msg = "Time points must be greater than or equal to 0"
101		)
102
103	23x	gq <- generateQuantities(
104	23x	object,
105	23x	generator = generator,
106	23x	type = "survival"
107		)
108
109	22x	quantities_raw <- extract_quantities(gq, type)
110	22x	collapser <- as.QuantityCollapser(grid, object@data)
111	22x	quantities <- collapse_quantities(quantities_raw, collapser)
112
113	22x	.SurvivalQuantities(
114	22x	quantities = Quantities(
115	22x	quantities,
116	22x	groups = collapser@groups,
117	22x	times = collapser@times
118		),
119	22x	grid = grid,
120	22x	data = object@data,
121	22x	type = type
122		)
123		}
124
125
126		#' `as.data.frame`
127		#'
128		#' @typed x: SurvivalQuantities
129		#' longitudinal quantities.
130		#' @param ... not used.
131		#' @family SurvivalQuantities
132		#' @export
133		#'
134		#' @returns A `data.frame` representation of the object.
135		as.data.frame.SurvivalQuantities <- function(x, ...) {
136	!	as.data.frame(x@quantities)
137		}
138
139
140		#' summary
141		#'
142		#' @description
143		#' This method returns a `data.frame` of the longitudinal quantities.
144		#'
145		#' @typed conf.level: numeric
146		#' confidence level of the interval.
147		#' @inheritParams SurvivalQuantities-Shared
148		#'
149		#' @family SurvivalQuantities
150		#' @family summary
151		#' @export
152		#'
153		#' @returns A `data.frame` containing posterior summary statistics.
154		summary.SurvivalQuantities <- function(
155		object,
156		conf.level = 0.95,
157		...
158		) {
159	21x	summary(object@quantities, conf.level = conf.level)
160		}
161
162
163		#' Automatic Plotting for `SurvivalQuantities``
164		#'
165		#' @inheritParams SurvivalQuantities-Shared
166		#' @typed add_km: logical
167		#' if `TRUE` Kaplan-Meier curves will be added to the plot for
168		#' each group/subject.
169		#' @typed add_wrap: logical
170		#' if `TRUE` will apply a [ggplot2::facet_wrap()] to the plot
171		#' by each group/subject.
172		#' @typed conf.level: numeric
173		#' confidence level of the interval. If values of `FALSE`,
174		#' `NULL` or `0` are provided then confidence regions will not be added to the plot
175		#' @param ... not used.
176		#'
177		#' @family SurvivalQuantities
178		#' @family autoplot
179		#' @export
180		#'
181		#' @returns A `ggplot` object.
182		autoplot.SurvivalQuantities <- function(
183		object,
184		conf.level = 0.95,
185		add_km = FALSE,
186		add_wrap = TRUE,
187		...
188		) {
189	3x	include_ci <- !is.null(conf.level) &&
190	3x	is.numeric(conf.level) &&
191	3x	conf.level > 0
192		# If CI aren't needed supply a default 0.95 to summary function as it needs
193		# a value to be specified to work
194	3x	conf.level <- if (include_ci) conf.level else 0.95
195	3x	assert_that(
196	3x	is.flag(add_km),
197	3x	length(conf.level) == 1,
198	3x	conf.level < 1,
199	3x	is.flag(add_wrap)
200		)
201	3x	kmdf <- if (add_km) {
202	1x	subset(
203	1x	object@data,
204	1x	as.list(object@grid, data = object@data)
205		)
206		} else {
207	2x	NULL
208		}
209	3x	all_fit_df <- summary(object, conf.level = conf.level)
210	3x	label <- switch(
211	3x	object@type,
212	3x	"surv" = expression(S(t)),
213	3x	"cumhaz" = expression(H(t)),
214	3x	"haz" = expression(h(t)),
215	3x	"loghaz" = expression(log(h(t)))
216		)
217	3x	survival_plot(
218	3x	data = all_fit_df,
219	3x	add_ci = include_ci,
220	3x	add_wrap = add_wrap,
221	3x	kmdf = kmdf,
222	3x	y_label = label
223		)
224		}
225
226
227		#' Survival Plot
228		#'
229		#' Internal plotting function to create survival plots with KM curve overlays
230		#' This function predominately exists to extract core logic into its own function
231		#' to enable easier unit testing.
232		#'
233		#' @typed data: data.frame
234		#' summary statistics for a survival
235		#' curve to be plotted. See details.
236		#' @typed add_ci: logical
237		#' should confidence intervals be added? Default = `TRUE`.
238		#' @typed add_wrap: logical
239		#' should the plots be wrapped by `data$group`? Default = `TRUE`.
240		#' @typed kmdf: "`data.frame` or `NULL`"
241		#' event times and status used to plot
242		#' overlaying KM curves. If `NULL` no KM curve will be plotted. See details.
243		#' @typed y_label: "`character` or `expression`"
244		#' label to display on the y-axis.
245		#' Default = `expression(S(t))`.
246		#' @typed x_label: "`character` or `expression`"
247		#' label to display on the x-axis.
248		#'
249		#' @details
250		#'
251		#' ## `data`
252		#' Should contain the following columns:
253		#' - `time` (`numeric`) \cr time point for the summary statistic.
254		#' - `group` (`character`) \cr the group in which the observation belongs to.
255		#' - `median` (`numeric`) \cr the median value for the summary statistic.
256		#' - `upper` (`numeric`) \cr the upper 95% CI for the summary statistic.
257		#' - `lower` (`numeric`) \cr the lower 95% CI for the summary statistic.
258		#'
259		#' ## `kmdf`
260		#' Should contain the following columns:
261		#' - `time` (`numeric`) \cr the time at which an event occurred.
262		#' - `event` (`numeric`) \cr 1/0 status indicator for the event.
263		#' - `group` (`character`) \cr which group the event belongs to, should correspond to values in `data$group`.
264		#' @keywords internal
265		#'
266		#' @returns A `ggplot` object.
267		survival_plot <- function(
268		data,
269		add_ci = TRUE,
270		add_wrap = TRUE,
271		kmdf = NULL,
272		y_label = expression(S(t)),
273		x_label = expression(t)
274		) {
275	3x	assert_that(
276	3x	is.flag(add_ci),
277	3x	is.flag(add_wrap),
278	3x	is.expression(y_label) \|\| is.character(y_label),
279	3x	is.expression(x_label) \|\| is.character(x_label),
280	3x	is.null(kmdf) \| is.data.frame(kmdf)
281		)
282
283	3x	p <- ggplot() +
284	3x	xlab(x_label) +
285	3x	ylab(y_label) +
286	3x	theme_bw()
287
288	3x	if (add_wrap) {
289	1x	p <- p + facet_wrap(~group)
290	1x	aes_ci <- aes(x = .data$time, ymin = .data$lower, ymax = .data$upper)
291	1x	aes_line <- aes(x = .data$time, y = .data$median)
292	1x	aes_km <- aes(time = .data$time, status = .data$event)
293		} else {
294	2x	aes_ci <- aes(
295	2x	x = .data$time,
296	2x	ymin = .data$lower,
297	2x	ymax = .data$upper,
298	2x	fill = .data$group,
299	2x	group = .data$group
300		)
301	2x	aes_line <- aes(
302	2x	x = .data$time,
303	2x	y = .data$median,
304	2x	colour = .data$group,
305	2x	group = .data$group
306		)
307	2x	aes_km <- aes(
308	2x	time = .data$time,
309	2x	status = .data$event,
310	2x	group = .data$group,
311	2x	colour = .data$group
312		)
313		}
314	3x	p <- p + geom_line(aes_line, data = data)
315	3x	if (add_ci) {
316	1x	p <- p + geom_ribbon(aes_ci, data = data, alpha = 0.3)
317		}
318	3x	if (!is.null(kmdf)) {
319	1x	p <- p +
320	1x	ggplot2.utils::geom_km(aes_km, data = kmdf) +
321	1x	ggplot2.utils::geom_km_ticks(aes_km, data = kmdf)
322		}
323	3x	p
324		}
325
326
327		#' @rdname show-object
328		#' @export
329		setMethod(
330		f = "show",
331		signature = "SurvivalQuantities",
332		definition = function(object) {
333	2x	template <- c(
334	2x	"SurvivalQuantities Object:",
335	2x	" # of samples = %d",
336	2x	" # of quantities = %d",
337	2x	" Type = %s"
338		)
339	2x	string <- sprintf(
340	2x	paste(template, collapse = "\n"),
341	2x	nrow(object@quantities),
342	2x	ncol(object@quantities),
343	2x	object@type
344		)
345	2x	cat("\n", string, "\n\n")
346		}
347		)
348
349
350		#' `brierScore`
351		#'
352		#' @description
353		#' Derives the Brier Scores (using Inverse Probability of Censoring Weighting)
354		#' for the Survival estimates as detailed in \insertCite{blanche2015}{jmpost}.
355		#'
356		#' @inheritParams SurvivalQuantities-Shared
357		#' @inheritParams Brier-Score-Shared
358		#'
359		#' @family brierScore
360		#' @family SurvivalQuantities
361		#' @references
362		#' \insertAllCited{}
363		#' @export
364		#'
365		#' @returns A `data.frame` containing Brier scores by time and group.
366		brierScore.SurvivalQuantities <- function(
367		object,
368		maintain_cen_order = TRUE,
369		event_offset = TRUE,
370		...
371		) {
372	3x	assert_that(
373	3x	object@type == "surv",
374	3x	msg = paste(
375	3x	"Brier Score can only be calculated when the survival quantities were",
376	3x	"generated with `type = 'surv'`",
377	3x	collapse = " "
378		)
379		)
380	3x	assert_that(
381	3x	is(object@grid, "GridFixed"),
382	3x	msg = paste(
383	3x	"Brier Score can only be calculated when the survival quantities were",
384	3x	"generated with `grid = GridFixed()`",
385	3x	collapse = " "
386		)
387		)
388
389	3x	sdat <- summary(object)
390	3x	times <- unique(as.QuantityGenerator(object@grid, object@data)@times)
391	3x	times <- times[order(times)]
392	3x	assert_that(
393	3x	nrow(sdat) == length(times) * length(unique(sdat$group))
394		)
395
396	3x	subject_col <- extractVariableNames(object@data@subject)$subject
397	3x	time_col <- extractVariableNames(object@data@survival)$time
398	3x	event_col <- extractVariableNames(object@data@survival)$event
399	3x	groups <- as.character(object@data@survival@data[[subject_col]])
400	3x	orig_times <- object@data@survival@data[[time_col]]
401	3x	events <- as.numeric(object@data@survival@data[[event_col]])
402
403	3x	pred_mat <- matrix(
404	3x	ncol = length(times),
405	3x	nrow = length(unique(sdat$group))
406		)
407	3x	for (i in seq_along(times)) {
408	18x	pred_mat[, i] <- sdat[sdat["time"] == times[i], "median"]
409	18x	assert_that(
410	18x	all(groups == sdat[sdat["time"] == times[i], "group"])
411		)
412		}
413	3x	brier_score(
414	3x	t = times,
415	3x	times = orig_times,
416	3x	events = events,
417	3x	pred_mat = 1 - pred_mat,
418	3x	maintain_cen_order = maintain_cen_order,
419	3x	event_offset = event_offset
420		)
421		}

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

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

1		#' @include generics.R
2		#' @include StanModule.R
3		NULL
4
5		#' `Prior` Function Arguments
6		#'
7		#' The documentation lists all the conventional arguments for [`Prior`]
8		#' constructors.
9		#'
10		#' @typed centre: number
11		#' the central point of distribution to shrink sampled values towards
12		#' (for most distributions this is the mean or median if the mean is undefined)
13		#' @typed x: Prior
14		#' a prior Distribution
15		#' @typed object: Prior
16		#' a prior Distribution
17		#' @typed name: character
18		#' the name of the parameter the prior distribution is for
19		#' @typed size: "`numeric` or `character`"
20		#' the parameter size.
21		#' @param ... Not Used.
22		#'
23		#' @name Prior-Shared
24		#' @keywords internal
25		NULL
26
27		# Prior-class ----
28
29		#' Prior Object and Constructor Function
30		#'
31		#' Specifies the prior distribution in a Stan Model
32		#'
33		#' @slot parameters (`list`)\cr See arguments.
34		#' @slot repr_model (`string`)\cr See arguments.
35		#' @slot repr_data (`string`)\cr See arguments.
36		#' @slot repr_parameters (`string`)\cr See arguments.
37		#' @slot repr_transformed_parameters (`string`)\cr See arguments.
38		#' @slot repr_generated_quantities (`string`)\cr See arguments.
39		#' @slot auxiliary_initial_values (`function`)\cr See arguments.
40		#' @slot auxiliary_size (`function`)\cr See arguments.

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

1		#' Simulate Patients from Posterior Predictive Distribution
2		#' @typed object: JointModelSamples
3		#' the object containing the posterior samples.
4		#' @typed newdata: data.frame
5		#' A data frame containing data in the same format as
6		#' the `object@data@survival@data`. Importantly, it should contain the same covariates and factor levels
7		#' as the variables used in the survival formula `object@data@survival@formula` and
8		#' the same columns used for `study`, `id`, and `arm`.
9		#' @param ... unused
10		#' Unused.
11		#' @typed times: numeric
12		#' times to simulate SLD for all patients.
13		#' @typed jitter_var: numeric
14		#' variances to add noise to the observed SLD `times`. The first value is for any times
15		#' less than 0 and the second for any times after 0. All positive (negative) times will remain positive (negative).
16		#' Jitter values are generated from a normal distribution with mean 0 and the given variances.
17		#' @typed time_max: number
18		#' the maximum time to simulate to.
19		#' @typed time_step: number
20		#' the time interval between evaluating the log-hazard function.
21		#' @typed lambda_censor: number
22		#' the censoring rate, as the parameter of an exponential distribution.
23		#' @typed scaled_variance: flag
24		#' Should variance be scaled by the expected value. Must be set the same as was used for
25		#' model fitting.
26		#' @param seed ignored.
27		#' @param nsim ignored.
28		#' @details Simulates a set of patients based on the covariates of those
29		#' used in the model fit or from `newdata`, which must contain the
30		#' same column names and factor levels.
31		#'
32		#' @importFrom stats simulate
33		#' @export
34		#'
35		#' @returns A `SimJointData` object containing posterior-predictive simulations.
36		simulate.JointModelSamples <- function(
37		object,
38		nsim = NULL,
39		seed = NULL,
40		newdata = NULL,
41		times = c(-2, 0, 10, 50, 100),
42		jitter_var = c(0, 0),
43		time_max = 2000,
44		time_step = 1,
45		lambda_censor = 1 / 3000,
46		scaled_variance = FALSE,
47		...
48		) {
49	!	subj_data <- if (is.null(newdata)) {
50	!	dplyr::left_join(
51	!	object@data@subject@data[, c(
52	!	object@data@subject@subject,
53	!	object@data@subject@arm,
54	!	object@data@subject@study
55		)],
56	!	object@data@survival@data[, c(
57	!	object@data@subject@subject,
58	!	all.vars(delete.response(terms(object@data@survival@formula)))
59		)],
60	!	by = object@data@subject@subject
61		)
62		} else {
63	!	newdata
64		}
65
66	!	subject_data <- DataSubject(
67	!	subj_data,
68	!	subject = object@data@subject@subject,
69	!	arm = object@data@subject@arm,
70	!	study = object@data@subject@study
71		)
72
73	!	n_patients <- nrow(subject_data@data)
74
75	!	requiredVars <- list(
76	!	longitudinal = lapply(
77	!	object@model@survival@parameters@parameters,
78	!	function(x) x@name
79		),
80	!	survival = lapply(
81	!	object@model@longitudinal@parameters@parameters,
82	!	function(x) x@name
83		),
84	!	link = lapply(object@model@link@components, function(x) x@key)
85		)
86
87	!	draws <- object@results$draws(
88	!	variables = unlist(requiredVars),
89	!	format = "draws_matrix"
90		)
91	!	draw_id <- sample.int(nrow(draws), n_patients, replace = FALSE)
92
93	!	long_models <- list()
94	!	surv_models <- list()
95	!	for (i in seq.int(n_patients)) {
96	!	draw <- draws[draw_id[i], ]
97	!	long_models[[i]] <- createLongitudinalSimObject(
98	!	object@model@longitudinal,
99	!	draw,
100	!	times = add_jitter(times, jitter_var = jitter_var),
101	!	scaled_variance = scaled_variance
102		)
103	!	surv_models[[i]] <- createSurvivalSimObject(
104	!	object@model@survival,
105	!	draw,
106	!	lambda_censor = lambda_censor,
107	!	time_step = time_step,
108	!	time_max = time_max
109		)
110		}
111
112	!	SimJointDataResults(
113	!	subject = subject_data,
114	!	surv_formula = object@data@survival@formula,
115	!	longitudinal = long_models,
116	!	survival = surv_models
117		)
118		}
119
120
121		add_jitter <- function(times, jitter_var = c(0, 0)) {
122	!	stopifnot(all(jitter_var >= 0))
123	!	if (jitter_var[1] > 0) {
124	!	neg_times <- times < 0
125	!	times[neg_times] <- pmin(
126	!	times[neg_times] + rnorm(sum(neg_times), sd = sqrt(jitter_var[1])),
127	!	-.Machine$double.eps / 2
128		)
129		}
130	!	if (jitter_var[2] > 0) {
131	!	pos_times <- times > 0
132	!	times[pos_times] <- pmax(
133	!	times[pos_times] + rnorm(sum(pos_times), sd = sqrt(jitter_var[2])),
134	!	.Machine$double.eps / 2
135		)
136		}
137	!	times
138		}
139
140
141		# Longitudinal Sim Object constructors --------
142		#' @noRd
143		#'
144		#' @returns A `SimLongitudinal` object configured from the posterior draw.
145		createLongitudinalSimObject <- function(object, draw, ...) {
146	!	UseMethod("createLongitudinalSimObject")
147		}
148
149		#' Get draws value by name
150		#' @typed draws: matrix
151		#' the draws input.
152		#' @typed name: character
153		#' to match column names of `draws` with [startsWith].
154		#' @return A vector of matching values or `NULL` if no match.
155		#' @noRd
156		get_vars <- function(draws, name) {
157	!	result <- as.numeric(draws[, startsWith(colnames(draws), name)])
158	!	if (length(result) == 0) {
159	!	NULL
160		} else {
161	!	result
162		}
163		}
164
165		#' @exportS3Method
166		#'
167		#' @returns A `SimLongitudinal` object configured from the posterior draw.
168		createLongitudinalSimObject.LongitudinalRandomSlope <- function(
169		object,
170		draw,
171		...
172		) {
173	!	args <- list(...)
174
175	!	args$intercept <- get_vars(draw, "lm_rs_intercept")
176	!	args$slope_mu <- get_vars(draw, "lm_rs_slope_mu")
177	!	args$slope_sigma <- get_vars(draw, "lm_rs_slope_sigma")
178	!	args$sigma <- get_vars(draw, "lm_rs_sigma")
179	!	args$link_dsld <- get_vars(draw, "link_dsld")
180	!	args$link_identity <- get_vars(draw, "link_identity")
181	!	args$scaled_variance <- NULL # not defined for random slope
182	!	do.call(SimLongitudinalRandomSlope, args)
183		}
184
185		#' @exportS3Method
186		#'
187		#' @returns A `SimLongitudinal` object configured from the posterior draw.
188		createLongitudinalSimObject.LongitudinalSteinFojo <- function(
189		object,
190		draw,
191		...
192		) {
193	!	args <- list(...)
194	!	args$sigma <- get_vars(draw, "lm_sf_sigma")
195	!	args$mu_s <- get_vars(draw, "lm_sf_mu_ks")
196	!	args$mu_g <- get_vars(draw, "lm_sf_mu_kg")
197	!	args$mu_b <- get_vars(draw, "lm_sf_mu_bsld")
198	!	args$omega_s <- get_vars(draw, "lm_sf_omega_ks")
199	!	args$omega_g <- get_vars(draw, "lm_sf_omega_kg")
200	!	args$omega_b <- get_vars(draw, "lm_sf_omega_bsld")
201	!	args$link_dsld <- get_vars(draw, "link_dsld")
202	!	args$link_ttg <- get_vars(draw, "link_ttg")
203	!	args$link_identity <- get_vars(draw, "link_identity")
204	!	args$link_growth <- get_vars(draw, "link_growth")
205	!	args$link_shrinkage <- get_vars(draw, "link_shrinkage")
206	!	do.call(SimLongitudinalSteinFojo, args)
207		}
208
209		#' @exportS3Method
210		#'
211		#' @returns A `SimLongitudinal` object configured from the posterior draw.
212		createLongitudinalSimObject.LongitudinalGSF <- function(object, draw, ...) {
213	!	args <- list(...)
214	!	args$sigma <- get_vars(draw, "lm_gsf_sigma")
215	!	args$mu_s <- get_vars(draw, "lm_gsf_mu_ks")
216	!	args$mu_g <- get_vars(draw, "lm_gsf_mu_kg")
217	!	args$mu_b <- get_vars(draw, "lm_gsf_mu_bsld")
218	!	args$mu_phi <- get_vars(draw, "lm_gsf_mu_phi")
219	!	args$omega_s <- get_vars(draw, "lm_gsf_omega_ks")
220	!	args$omega_g <- get_vars(draw, "lm_gsf_omega_kg")
221	!	args$omega_b <- get_vars(draw, "lm_gsf_omega_bsld")
222	!	args$omega_phi <- get_vars(draw, "lm_gsf_omega_phi")
223	!	args$link_dsld <- get_vars(draw, "link_dsld")
224	!	args$link_ttg <- get_vars(draw, "link_ttg")
225	!	args$link_identity <- get_vars(draw, "link_identity")
226	!	args$link_growth <- get_vars(draw, "link_growth")
227	!	args$link_shrinkage <- get_vars(draw, "link_shrinkage")
228
229	!	do.call(SimLongitudinalGSF, args)
230		}
231
232		#' @exportS3Method
233		#'
234		#' @returns A `SimLongitudinal` object configured from the posterior draw.
235		createLongitudinalSimObject.LongitudinalClaretBruno <- function(
236		object,
237		draw,
238		...
239		) {
240	!	args <- list(...)
241	!	args$sigma <- get_vars(draw, "lm_clbr_sigma")
242	!	args$mu_b <- get_vars(draw, "lm_clbr_mu_b")
243	!	args$mu_g <- get_vars(draw, "lm_clbr_mu_g")
244	!	args$mu_c <- get_vars(draw, "lm_clbr_mu_c")
245	!	args$mu_p <- get_vars(draw, "lm_clbr_mu_p")
246	!	args$omega_b <- get_vars(draw, "lm_clbr_omega_b")
247	!	args$omega_g <- get_vars(draw, "lm_clbr_omega_g")
248	!	args$omega_c <- get_vars(draw, "lm_clbr_omega_c")
249	!	args$omega_p <- get_vars(draw, "lm_clbr_omega_p")
250	!	args$link_dsld <- get_vars(draw, "link_dsld")
251	!	args$link_ttg <- get_vars(draw, "link_ttg")
252	!	args$link_identity <- get_vars(draw, "link_identity")
253	!	args$link_growth <- get_vars(draw, "link_growth")
254	!	args$link_shrinkage <- get_vars(draw, "link_shrinkage")
255	!	do.call(SimLongitudinalClaretBruno, args)
256		}
257
258
259		# Survival Sim Object constructors --------
260
261		#' @noRd
262		#'
263		#' @returns A `SimSurvival` object configured from the posterior draw.
264		createSurvivalSimObject <- function(object, draw, ...) {
265	!	UseMethod("createSurvivalSimObject")
266		}
267
268		#' @exportS3Method
269		#'
270		#' @returns A `SimSurvival` object configured from the posterior draw.
271		createSurvivalSimObject.SurvivalWeibullPH <- function(object, draw, ...) {
272	!	args <- list(...)
273	!	args$lambda <- get_vars(draw, "sm_weibull_ph_lambda")
274	!	args$gamma <- get_vars(draw, "sm_weibull_ph_gamma")
275
276	!	result <- do.call(SimSurvivalWeibullPH, args)
277	!	result@beta_os_cov <- get_vars(draw, "beta_os_cov")
278	!	result
279		}
280
281		#' @exportS3Method
282		#'
283		#' @returns A `SimSurvival` object configured from the posterior draw.
284		createSurvivalSimObject.SurvivalExponential <- function(object, draw, ...) {
285	!	args <- list(...)
286	!	args$lambda <- get_vars(draw, "sm_exp_lambda")
287
288	!	result <- do.call(SimSurvivalExponential, args)
289	!	result@beta_os_cov <- get_vars(draw, "beta_os_cov")
290	!	result
291		}
292
293		#' @exportS3Method
294		#'
295		#' @returns A `SimSurvival` object configured from the posterior draw.
296		createSurvivalSimObject.SurvivalGamma <- function(object, draw, ...) {
297	!	lapply(object@model@survival@parameters@parameters, function(x) x@name)
298	!	lapply(object@model@link@components, function(x) x@key)
299	!	args <- list(...)
300	!	args$k <- get_vars(draw, "sm_gamma_k")
301	!	args$theta <- get_vars(draw, "sm_gamma_theta")
302
303	!	result <- do.call(SimSurvivalGamma, args)
304	!	result@beta_os_cov <- get_vars(draw, "beta_os_cov")
305	!	result
306		}
307
308		#' @exportS3Method
309		#'
310		#' @returns A `SimSurvival` object configured from the posterior draw.
311		createSurvivalSimObject.SurvivalLogLogistic <- function(object, draw, ...) {
312	!	lapply(object@model@survival@parameters@parameters, function(x) x@name)
313	!	lapply(object@model@link@components, function(x) x@key)
314	!	args <- list(...)
315	!	args$a <- get_vars(draw, "sm_loglogis_a")
316	!	args$b <- get_vars(draw, "sm_loglogis_b")
317
318	!	result <- do.call(SimSurvivalLogLogistic, args)
319	!	result@beta_os_cov <- get_vars(draw, "beta_os_cov")
320	!	result
321		}
322
323
324		sampleSubjectsFromObs <- function(object, subjects_df, covs_matrix) {
325	!	subjects_df$log_haz_cov <- rowSums(
326	!	covs_matrix[, -1] * t(sapply(object, function(x) x@beta_os_cov))
327		)
328	!	subjects_df \|>
329	!	dplyr::mutate(survival = stats::runif(dplyr::n())) \|>
330	!	dplyr::mutate(chazard_limit = -log(.data$survival)) \|>
331	!	dplyr::mutate(
332	!	time_cen = stats::rexp(
333	!	dplyr::n(),
334	!	sapply(object, function(x) x@lambda_censor)
335		)
336		)
337		}
338
339		# Simulate patient function ------
340		SimJointDataResults <- function(
341		subject,
342		surv_formula,
343		longitudinal,
344		survival,
345		.silent = FALSE
346		) {
347		# take hazard windows from first element. All have the same
348		# @time_step and @time_max parameters
349	!	hazard_evaluation_info <- hazardWindows(survival[[1]])
350
351	!	group_counts <- as.data.frame(table(subject@data[, c(
352	!	subject@study,
353	!	subject@arm
354		)]))
355
356	!	n_group <- group_counts$Freq
357	!	arms <- group_counts[, 2]
358	!	studies <- group_counts[, 1]
359	!	n_subjects <- sum(n_group)
360	!	n_times <- length(hazard_evaluation_info$midpoint)
361
362	!	baseline <- data.frame(
363	!	subject = subject@data[[subject@subject]],
364	!	study = subject@data[[subject@study]],
365	!	arm = subject@data[[subject@arm]]
366		)
367
368	!	cov_cols <- cbind(
369	!	subject = subject@data[[subject@subject]],
370	!	subject@data[, all.vars(delete.response(terms(surv_formula)))]
371		)
372
373	!	os_baseline <- sampleSubjectsFromObs(
374	!	survival,
375	!	subjects_df = baseline,
376	!	model.matrix(delete.response(terms(surv_formula)), subject@data)
377		)
378
379	!	lm_baseline <- dplyr::bind_rows(
380	!	lapply(seq.int(longitudinal), function(i) {
381	!	sampleSubjects(longitudinal[[i]], subjects_df = baseline[i, ])
382		})
383		)
384
385	!	lm_times <- lapply(seq.int(longitudinal), function(i) {
386	!	longitudinal[[i]]@times
387		})
388
389	!	lm_dat_no_obvs <-
390	!	cbind(
391	!	baseline[rep(seq.int(nrow(baseline)), times = lengths(lm_times)), ],
392	!	time = unlist(lm_times)
393		) \|>
394	!	dplyr::left_join(lm_baseline, by = c("subject", "study", "arm"))
395
396	!	lm_dat <- dplyr::bind_rows(
397	!	.mapply(
398	!	sampleObservations,
399	!	dots = list(
400	!	times_df = split(lm_dat_no_obvs, ~subject),
401	!	object = longitudinal
402		),
403	!	MoreArgs = NULL
404		)
405		)
406
407	!	hazard_eval_df <- dplyr::tibble(
408	!	subject = rep(lm_baseline$subject, each = n_times),
409	!	arm = rep(lm_baseline$arm, each = n_times),
410	!	study = rep(lm_baseline$study, each = n_times),
411	!	midpoint = rep(
412	!	as.double(hazard_evaluation_info$midpoint),
413	!	times = n_subjects
414		),
415	!	time = rep(as.double(hazard_evaluation_info$upper), times = n_subjects),
416	!	width = rep(as.double(hazard_evaluation_info$width), times = n_subjects)
417		)
418
419	!	lm_link_dat <- dplyr::bind_rows(
420	!	.mapply(
421	!	sampleObservations,
422	!	dots = list(
423	!	times_df = split(
424	!	dplyr::left_join(
425	!	hazard_eval_df,
426	!	lm_baseline,
427	!	by = c("subject", "study", "arm")
428		),
429	!	~subject
430		),
431	!	object = longitudinal
432		),
433	!	MoreArgs = NULL
434		)
435	!	)[, c(
436	!	"subject",
437	!	"study",
438	!	"arm",
439	!	"log_haz_link",
440	!	"time",
441	!	"width",
442	!	"midpoint"
443		)]
444
445	!	os_eval_df <- lm_link_dat \|>
446	!	dplyr::left_join(os_baseline, by = c("subject", "study", "arm"))
447
448	!	times_df <- split(os_eval_df, ~subject)
449	!	os_data_list <- list()
450	!	for (i in seq_along(times_df)) {
451	!	times_df[[i]]$log_bl_haz <- survival[[i]]@loghazard(
452	!	times_df[[i]]$midpoint
453		)
454		}
455	!	os_dat <- sampleObservations2(dplyr::bind_rows(times_df))
456
457	!	os_dat <- dplyr::left_join(os_dat, cov_cols, by = "subject")
458
459	!	lm_dat2 <- lm_dat \|>
460	!	dplyr::left_join(
461	!	dplyr::select(os_dat, "subject", os_time = "time"),
462	!	by = "subject"
463		) \|>
464	!	dplyr::mutate(observed = (.data$time <= .data$os_time)) \|>
465	!	dplyr::arrange(dplyr::pick(c("subject", "time")))
466
467	!	assert_that(
468	!	length(unique(os_dat$subject)) == length(os_dat$subject),
469	!	length(os_dat$subject) == n_subjects,
470	!	all(os_dat$time >= 0),
471	!	all(os_dat$event %in% c(0, 1)),
472	!	msg = "Assumptions for the Survival data are not met (please report this issue)"
473		)
474
475	!	assert_that(
476	!	nrow(lm_dat2) == n_subjects * length(longitudinal[[1]]@times),
477	!	length(unique(lm_dat2$subject)) == n_subjects,
478	!	msg = "Assumptions for the Longitudinal data are not met (please report this issue)"
479		)
480
481	!	original_names <- c("subject", "arm", "study")
482	!	names(original_names) <- c(subject@subject, subject@arm, subject@study)
483
484	!	return(
485	!	.SimJointData(
486	!	survival = os_dat \|> dplyr::rename(dplyr::all_of(original_names)),
487	!	longitudinal = lm_dat2[, c(
488	!	"subject",
489	!	"arm",
490	!	"study",
491	!	"time",
492	!	"sld",
493	!	"observed"
494		)] \|>
495	!	dplyr::rename(dplyr::all_of(original_names))
496		)
497		)
498		}
499
500
501		# adapted from sampleObservations so that $midpoint and $log_bl_haz are already in times_df.
502		sampleObservations2 <- function(times_df) {
503	!	assert_that(
504	!	all(times_df$time >= 0),
505	!	msg = "All time points must be greater than or equal to 0"
506		)
507
508	!	os_dat_chaz <- times_df \|>
509		# Fix to avoid issue with log(0) = NaN values
510	!	dplyr::mutate(
511	!	log_bl_haz = dplyr::if_else(
512	!	.data$midpoint == 0,
513	!	-999,
514	!	.data$log_bl_haz
515		),
516	!	hazard_instant = exp(
517	!	.data$log_bl_haz + .data$log_haz_cov + .data$log_haz_link
518		),
519		# Reset Inf values to large number to avoid NaN issues downstream
520		# This is suitable as Hazard limits tend to be in the range of -10 to 10 so large numbers
521		# are essentially equivalent to infinity for simulation purposes
522	!	hazard_instant = dplyr::if_else(
523	!	.data$hazard_instant == Inf,
524	!	999,
525	!	.data$hazard_instant
526		),
527	!	hazard_instant = dplyr::if_else(
528	!	.data$hazard_instant == -Inf,
529	!	-999,
530	!	.data$hazard_instant
531		),
532	!	hazard_interval = .data$hazard_instant * .data$width
533		) \|>
534	!	dplyr::mutate(chazard = cumsum(.data$hazard_interval), .by = "subject")
535
536	!	os_had_event <- os_dat_chaz \|>
537	!	dplyr::filter(.data$chazard >= .data$chazard_limit) \|>
538	!	dplyr::slice_head(by = "subject") \|>
539	!	dplyr::mutate(event = 1)
540
541	!	os_had_censor <- os_dat_chaz \|>
542	!	dplyr::filter(!.data$subject %in% os_had_event$subject) \|>
543	!	dplyr::slice_tail(by = "subject") \|>
544
545	!	dplyr::mutate(event = 0)
546
547	!	if (!(nrow(os_had_censor) == 0)) {
548	!	message(sprintf(
549	!	"INFO: %i subject(s) did not die before max(times)",
550	!	nrow(os_had_censor)
551		))
552		}
553
554	!	os_dat_complete <- os_had_event \|>
555	!	dplyr::bind_rows(os_had_censor) \|>
556	!	dplyr::mutate(
557	!	real_time = .data$time,
558	!	event = dplyr::if_else(
559	!	.data$real_time <= .data$time_cen,
560	!	.data$event,
561	!	0
562		),
563	!	time = dplyr::if_else(
564	!	.data$real_time <= .data$time_cen,
565	!	.data$real_time,
566	!	.data$time_cen
567		)
568		) \|>
569	!	dplyr::arrange(.data$subject)
570
571	!	keep_cols <- colnames(os_dat_complete) %in%
572	!	c("subject", "study", "arm", "time", "event", "cov_cont", "cov_cat")
573	!	os_dat_complete[, keep_cols]
574		}

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		#'
38		#' @examples
39		#' show(prior_normal(0, 1))
40		NULL
41
42
43		# write_stan ----
44
45		#' `write_stan`
46		#'
47		#' Write the Stan code for a Stan module.
48		#'
49		#' @param object the module.
50		#' @typed destination: "`character` or `connection`"
51		#' Where to write stan code to.
52		#' @param ... Additional arguments
53		#'
54		#' @export
55		#'
56		#' @returns Invisibly returns `NULL` after writing the Stan program.
57		#'
58		#' @examples
59		#' path <- tempfile(fileext = ".stan")
60		#' write_stan(JointModel(LongitudinalGSF()), path)
61		#' unlink(path)
62		write_stan <- function(object, destination, ...) {
63	!	UseMethod("write_stan")
64		}
65
66		# compileStanModel ----
67
68		#' `compileStanModel`
69		#'
70		#' Compile the Stan module.
71		#'
72		#' @param object the module.
73		#'
74		#' @export
75		#'
76		#' @returns A compiled `cmdstanr::CmdStanModel` object.
77		#'
78		#' @examples
79		#' \dontrun{
80		#' compileStanModel(JointModel(LongitudinalGSF()))
81		#' }
82		compileStanModel <- function(object) {
83	73x	UseMethod("compileStanModel")
84		}
85
86
87		# sampleStanModel ----
88
89		#' `sampleStanModel`
90		#'
91		#' Sample from a Stan Module.
92		#'
93		#' @param object the module.
94		#' @param ... additional arguments.
95		#'
96		#' @export
97		#'
98		#' @returns A `JointModelSamples` object.
99		#'
100		#' @examples
101		#' \dontrun{
102		#' sampleStanModel(JointModel(LongitudinalGSF()), data = joint_data)
103		#' }
104		sampleStanModel <- function(object, ...) {
105	18x	UseMethod("sampleStanModel")
106		}
107
108
109		# as.StanModule ----
110
111		#' `as.StanModule`
112		#'
113		#' Converts an object into a [`StanModule`].
114		#'
115		#' @param object what to convert.
116		#' @param ... additional options.
117		#' @family as.StanModule
118		#' @keywords internal
119		#'
120		#' @returns A `StanModule` object.
121		as.StanModule <- function(object, ...) {
122	2419x	UseMethod("as.StanModule")
123		}
124
125
126		#' `getParameters`
127		#'
128		#' Extract any modelling parameters as a [`ParameterList`] object
129		#' from a model.
130		#'
131		#' @param object where to obtain the parameters from.
132		#' @param ... additional options.
133		#'
134		#' @export
135		getParameters <- function(object, ...) {
136	246x	UseMethod("getParameters")
137		}
138
139
140		# extractVariableNames ----
141
142		#' Extract Mapping to Standardised Variable Names
143		#'
144		#' @description
145		#' Extract a `list` that maps the variable names in a user-defined
146		#' `data.frame` to standardised values.
147		#'
148		#' @param object the data object.
149		#' @family extractVariableNames
150		#' @keywords internal
151		#'
152		#' @returns A named list mapping source variables to standard names.
153		extractVariableNames <- function(object) {
154	1904x	UseMethod("extractVariableNames")
155		}
156
157
158		# initialValues ----
159
160		#' `initialValues`
161		#'
162		#' Obtain the `list` of initial values to be passed to the Stan sampler.
163		#'
164		#' @param object where to get the initial values from.
165		#' @param n_chains the number of initial values to generate. See details.
166		#' @param ... Not currently used.
167		#'
168		#' @details
169		#' There are multiple ways of specifying initial values to Stan, see the `init` argument
170		#' in [cmdstanr::model-method-sample] for full details. Within this package we supply
171		#' initial values via a list of lists where each inner list contains the initial values
172		#' for a single chain. As such the `n_chains` argument specifies the number of inner lists
173		#' to generate.
174		#'
175		#' See the Vignette for further details of how to specify initial values.
176		#'
177		#' @export
178		#'
179		#' @returns Initial values for the object, usually as a numeric vector or list.
180		#'
181		#' @examples
182		#' initialValues(prior_normal(0, 1))
183		initialValues <- function(object, ...) {
184	90691x	UseMethod("initialValues")
185		}
186
187
188		# size ----
189
190		#' `size`
191		#'
192		#' Obtain the `list` of parameter sizes.
193		#'
194		#' @param object where to get the parameter sizes from.
195		#'
196		#' @keywords internal
197		#'
198		#' @returns A numeric value or named list describing parameter sizes.
199		size <- function(object) {
200	164x	UseMethod("size")
201		}
202
203
204		# auxiliaryInitialValues ----
205
206		#' `auxiliaryInitialValues`
207		#'
208		#' Obtain initial values for any additional Stan parameters introduced by an object.
209		#'
210		#' @param object where to get the auxiliary initial values from.
211		#' @param ... additional options.
212		#'
213		#' @details
214		#' Some objects, such as shrinkage priors, declare additional Stan parameters
215		#' beyond the main model parameter. This helper returns named initial values for
216		#' those additional parameters.
217		#'
218		#' @keywords internal
219		#'
220		#' @returns A named `list` of initial values for auxiliary parameters.
221		auxiliaryInitialValues <- function(object, ...) {
222	9927x	UseMethod("auxiliaryInitialValues")
223		}
224
225
226		# auxiliarySize ----
227
228		#' `auxiliarySize`
229		#'
230		#' Obtain sizes for any additional Stan parameters introduced by an object.
231		#'
232		#' @param object where to get the auxiliary parameter sizes from.
233		#' @param ... additional options.
234		#'
235		#' @details
236		#' Some objects, such as shrinkage priors, declare additional Stan parameters
237		#' beyond the main model parameter. This helper returns named sizes for those
238		#' additional parameters so their initial values can be expanded consistently.
239		#'
240		#' @keywords internal
241		#'
242		#' @returns A named `list` of auxiliary parameter sizes.
243		auxiliarySize <- function(object, ...) {
244	138x	UseMethod("auxiliarySize")
245		}
246
247
248		# generateQuantities ----
249
250		#' `generateQuantities`
251		#'
252		#' Obtain the generated quantities from a Stan Model.
253		#'
254		#' @param object object to obtain generated quantities from
255		#' @param ... additional options.
256		#'
257		#' @export
258		#'
259		#' @returns A `cmdstanr::CmdStanGQ` object containing generated quantities.
260		#'
261		#' @examples
262		#' \dontrun{
263		#' generateQuantities(fit, generator, "survival")
264		#' }
265		generateQuantities <- function(object, ...) {
266	37x	UseMethod("generateQuantities")
267		}
268
269
270		#' Prepare Data Object
271		#'
272		#' @typed object: "`DataSubject` or `DataLongitudinal` or `DataSurvival`"
273		#' data object to "harmonise"
274		#' @typed subject_var: character
275		#' the name of the variable containing the subject identifier.
276		#' @typed subject_ord: character
277		#' the expected levels (in order) of the subject identifier.
278		#' @param ... not used.
279		#'
280		#' @details
281		#' This utility function prepares the datasets in the data objects in order to ensure they
282		#' are consistent and compatible with each other.
283		#'
284		#' In particular it ensures that the `subject` variable, as specified by `DataSubject`,
285		#' is available in `DataLongitudinal` and `DataSurvival` and that all levels are present
286		#' in all 3 data objects.
287		#'
288		#' It also sorts the datasets to ensure that indexes are consistent e.g. index 1 for
289		#' `DataSubject@data` corresponds to the same subject as index 1 for `DataSurvival@data`.
290		#' For `DataLongitudinal` the data is additionally sorted by time and outcome value.
291		#'
292		#' @seealso [`DataJoint`], [`DataSurvival`], [`DataSubject`], [`DataLongitudinal`]
293		#'
294		#' @keywords internal
295		#' @return Returns the original object but with the data standardised (see details)
296		harmonise <- function(object, ...) {
297	432x	UseMethod("harmonise")
298		}
299
300
301		#' @rdname harmonise
302		harmonise.default <- function(object, ...) {
303	14x	NULL
304		}
305
306
307		#' `as_stan_list`
308		#'
309		#' @description
310		#' Extracts a list of data elements from an object to be used as input
311		#' to a Stan Model
312		#'
313		#' @param object to be converted.
314		#' @param ... additional options.
315		#'
316		#' @family as_stan_list
317		#' @export
318		#'
319		#' @returns A named `list` suitable for use as Stan data.
320		#'
321		#' @examples
322		#' as_stan_list(Parameter(prior_normal(0, 1), "beta"))
323		as_stan_list <- function(object, ...) {
324	2332x	UseMethod("as_stan_list")
325		}
326
327		#' @rdname as_stan_list
328		#' @export
329		as_stan_list.default <- function(object, ...) {
330	32x	NULL
331		}
332
333
334		#' `as_print_string`
335		#'
336		#' @description
337		#' Returns the character representation of an object which is suitable
338		#' for printing to the console
339		#'
340		#' @param object to be converted to string.
341		#' @param ... additional options.
342		#'
343		#' @family as_print_string
344		#' @keywords internal
345		#'
346		#' @returns A character vector suitable for printing.
347		as_print_string <- function(object, ...) {
348	135x	UseMethod("as_print_string")
349		}
350
351		#' Show an Object
352		#'
353		#' Prints an object to the console.
354		#'
355		#' @param object Object to be printed
356		#'
357		#' @name show-object
358		NULL
359
360
361		#' `brierScore`
362		#'
363		#' @description
364		#' Returns the Brier Score for a given model
365		#'
366		#' @param object to calculate Brier Score for.
367		#' @param ... additional options.
368		#'
369		#' @family brierScore
370		#' @export
371		#'
372		#' @returns A numeric Brier score or a `data.frame` of scores over time.
373		#'
374		#' @examples
375		#' \dontrun{
376		#' brierScore(survival_quantities)
377		#' }
378		brierScore <- function(object, ...) {
379	3x	UseMethod("brierScore")
380		}
381
382
383		#' Generate Simulated Observations
384		#'
385		#' @typed object: "`SimLongitudinal` or `SimSurvival`"
386		#' object to generate observations from.
387		#' @typed times_df: data.frame
388		#' the times at which to generate observations. See details.
389		#'
390		#' @details
391		#' The `times_df` argument should be a `data.frame` as created by `sampleSubjects` but
392		#' replicated for each time point at which observations are to be generated. That is if you want
393		#' to generate observations for times `c(0, 1, 2, 3)` then `times_df` should be created as:
394		#' ```
395		#' subject_dat <- sampleSubjects(object, ...)
396		#' times_df <- tidyr::expand_grid(
397		#' subject_dat,
398		#' time = c(0, 1, 2, 3)
399		#' )
400		#' ```
401		#'
402		#' @export
403		#'
404		#' @returns A `data.frame` containing simulated observations.
405		#'
406		#' @examples
407		#' sim <- SimLongitudinalRandomSlope(slope_mu = 0.01, slope_sigma = 0.5)
408		#' subjects <- data.frame(study = factor("S"), arm = factor("A"), subject = "1")
409		#' subject_parameters <- sampleSubjects(sim, subjects)
410		#' sampleObservations(sim, transform(subject_parameters, time = 0))
411		sampleObservations <- function(object, times_df) {
412	76x	UseMethod("sampleObservations")
413		}
414
415
416		#' Generate Simulated Subjects
417		#'
418		#' @typed object: "`SimLongitudinal` or `SimSurvival`"
419		#' object to generate subjects from.
420		#' @typed subjects_df: data.frame
421		#' the subjects to generate observations for. See details.
422		#'
423		#' @details
424		#' The `subjects_df` argument should be a `data.frame` with 1 row per desired subject to create
425		#' with the following columns:
426		#' - `study` (`factor`) the study identifier.
427		#' - `arm` (`factor`) the treatment arm identifier.
428		#' - `subject` (`character`) the subject identifier.
429		#'
430		#' This method takes care of generating all the individual subject data required for the
431		#' [`sampleObservations`] method to generate the observations.
432		#' @export
433		#'
434		#' @returns A `data.frame` containing simulated subject-level parameters.
435		#'
436		#' @examples
437		#' subjects <- data.frame(study = factor("S"), arm = factor("A"), subject = "1")
438		#' sim <- SimLongitudinalRandomSlope(slope_mu = 0.01, slope_sigma = 0.5)
439		#' sampleSubjects(sim, subjects)
440		sampleSubjects <- function(object, subjects_df) {
441	52x	UseMethod("sampleSubjects")
442		}
443
444
445		#' Generate time windows for evaluating a hazard function
446		#'
447		#' @typed object: SurvivalModel
448		#' object to generate time windows for.
449		#' @param ... Not used.
450		#'
451		#'
452		#' @returns A `data.frame` describing the hazard-evaluation intervals.
453		hazardWindows <- function(object, ...) {
454	24x	UseMethod("hazardWindows")
455		}
456
457		#' @rdname Quant-Dev
458		#' @export
459		#'
460		#' @returns A `QuantityGenerator` or `QuantityCollapser` object, according to the function called.
461		#'
462		#' @examples
463		#' \dontrun{
464		#' grid <- GridFixed(times = c(0, 10))
465		#' as.QuantityGenerator(grid, joint_data)
466		#' as.QuantityCollapser(grid, joint_data)
467		#' }
468		as.QuantityGenerator <- function(object, ...) {
469	112x	UseMethod("as.QuantityGenerator")
470		}
471
472		#' @rdname Quant-Dev
473		#' @export
474		as.QuantityCollapser <- function(object, ...) {
475	43x	UseMethod("as.QuantityCollapser")
476		}
477
478
479		#' Coalesce Time
480		#'
481		#' @typed object: Grid
482		#' object to coalesce time for.
483		#' @typed times: numeric
484		#' the times to coalesce to.
485		#' @param ... Not used
486		#'
487		#' Method used to replace NULL times on grid objects (if appropriate)
488		#'
489		#' @keywords internal
490		#'
491		#' @returns A `Grid` object with missing time points replaced.
492		coalesceGridTime <- function(object, times, ...) {
493	43x	UseMethod("coalesceGridTime")
494		}
495		#' @export
496		#'
497		#' @returns The input object, unchanged.
498		coalesceGridTime.default <- function(object, times, ...) {
499	5x	object
500		}
501
502
503		#' Resolve a Promise
504		#'
505		#' @typed object: ANY
506		#' an object to resolve.
507		#' @typed ...: ANY
508		#' additional arguments.
509		#'
510		#' If `object` is not a promise will just return itself else will resolve the promise
511		#' and return the promised object.
512		#'
513		#' @export
514		#'
515		#' @returns The input object, or the value obtained by resolving a promise.
516		#'
517		#' @examples
518		#' resolvePromise(1)
519		#' resolvePromise(linkDSLD(), LongitudinalGSF())
520		resolvePromise <- function(object, ...) {
521	119x	UseMethod("resolvePromise")
522		}
523
524		#' @rdname resolvePromise
525		#' @export
526		resolvePromise.default <- function(object, ...) {
527	!	object
528		}
529
530		#' Enable Link Generic
531		#'
532		#' @typed object: LongitudinalModel
533		#' to enable link for.
534		#' @param ... Not used.
535		#'
536		#' Optional hook method that is called on a [`LongitudinalModel`] only if a link method
537		#' is provided to [`JointModel`]. This can be used to allow the model to include any
538		#' optional stan code that is only required if there are links present.
539		#'
540		#' @return [`LongitudinalModel`] object
541		#'
542		#' @export
543		#'
544		#' @examples
545		#' enableLink(LongitudinalGSF())
546		enableLink <- function(object, ...) {
547	24x	UseMethod("enableLink")
548		}
549		#' @export
550		#'
551		#' @returns The longitudinal model with its link-related Stan code enabled.
552		enableLink.default <- function(object, ...) {
553	!	object
554		}
555
556
557		#' Enable Generated Quantities Generic
558		#'
559		#' @typed object: StanModel
560		#' to enable generated quantities for.
561		#' @param ... Not used.
562		#'
563		#' Optional hook method that is called on a [`StanModel`] if attempting to use
564		#' either [`LongitudinalQuantities`] or [`SurvivalQuantities`]
565		#'
566		#' @return [`StanModule`] object
567		#'
568		#' @export
569		#'
570		#' @examples
571		#' enableGQ(LongitudinalGSF())
572		enableGQ <- function(object, ...) {
573	132x	UseMethod("enableGQ")
574		}
575		#' @export
576		#'
577		#' @returns A `StanModule` object containing the generated-quantities code.
578		enableGQ.default <- function(object, ...) {
579	46x	StanModule()
580		}
581
582
583		#' Get Prediction Names
584		#'
585		#' Utility function that returns the names of the required parameters for predicting
586		#' survival quantities with [`GridPrediction`].
587		#'
588		#' @typed object: LongitudinalModel
589		#' A longitudinal model object
590		#' @param ... Not used.
591		#' @export
592		#'
593		#' @returns A character vector of required prediction parameter names, or `NULL`.
594		#'
595		#' @examples
596		#' getPredictionNames(LongitudinalGSF())
597		getPredictionNames <- function(object, ...) {
598	6x	UseMethod("getPredictionNames")
599		}
600
601		#' @rdname getPredictionNames
602		#' @export
603		getPredictionNames.default <- function(object, ...) {
604	1x	NULL
605		}
606
607
608		#' Get Random Effects Names
609		#'
610		#' Utility function that returns the names of the random effects parameters.
611		#' The main use for this is to allow the [`LongitudinalRandomEffects`] function
612		#' to know which parameters it needs to extract and to what common names
613		#' it should map the parameters to.
614		#'
615		#' @typed object: LongitudinalModel
616		#' A longitudinal model object
617		#' @param ... Not used.
618		#' @export
619		#'
620		#' @returns A named character vector of random-effect parameter names, or `NULL`.
621		#'
622		#' @examples
623		#' getRandomEffectsNames(LongitudinalGSF())
624		getRandomEffectsNames <- function(object, ...) {
625	2x	UseMethod("getRandomEffectsNames")
626		}
627
628		#' @rdname getRandomEffectsNames
629		#' @export
630		getRandomEffectsNames.default <- function(object, ...) {
631	!	NULL
632		}
633
634
635		#' As Formula
636		#'
637		#' Utility wrapper function to convert an object to a formula.
638		#' @typed x: ANY
639		#' object to convert to a formula.
640		#' @param ... Not used.
641		#' @export
642		#'
643		#' @returns A `formula` object.
644		#'
645		#' @examples
646		#' as_formula("response ~ time")
647		as_formula <- function(x, ...) {
648	7x	UseMethod("as_formula")
649		}
650
651		#' @importFrom stats as.formula
652		#' @export
653		#'
654		#' @returns A `formula` object.
655		as_formula.default <- function(x, ...) {
656	!	as.formula(x, ...)
657		}
658
659
660		#' Set Constraints
661		#'
662		#' Applies constraints to a prior distribution to ensure any sampled numbers
663		#' from the distribution fall within the constraints
664		#'
665		#' @typed object: Prior
666		#' a prior distribution to apply constraints to
667		#' @typed lower: numeric
668		#' lower constraint boundary
669		#' @typed upper: numeric
670		#' upper constraint boundary
671		#'
672		#' @export
673		#'
674		#' @returns A `Prior` object with the requested bounds.
675		#'
676		#' @examples
677		#' set_limits(prior_normal(0, 1), lower = 0)
678		set_limits <- function(object, lower = -Inf, upper = Inf) {
679	534x	UseMethod("set_limits")
680		}
681
682
683		#' Save Object to File
684		#'
685		#' @typed object: ANY
686		#' object to save.
687		#' @typed file: character
688		#' file to save object to.
689		#' @typed ...: ANY
690		#' additional arguments.
691		#'
692		#' @family saveObject
693		#' @export
694		#'
695		#' @returns Invisibly returns `NULL` after saving the object.
696		#'
697		#' @examples
698		#' \dontrun{
699		#' saveObject(fit, "joint-model-fit.rds")
700		#' }
701		saveObject <- function(object, file, ...) {
702	1x	UseMethod("saveObject")
703		}
704
705
706		#' Extract Covariate Names
707		#'
708		#' @typed object: ANY
709		#' the object to extract covariate names from.
710		#' @param ... additional arguments added by methods.
711		#'
712		#' @family covariates
713		#' @export
714		#'
715		#' @returns A character vector containing the covariate names.
716		#'
717		#' @examples
718		#' surv_data <- DataSurvival(os_data, Surv(os_time, os_event) ~ age + sex)
719		#' covariates(surv_data)
720		covariates <- function(object, ...) {
721	2x	UseMethod("covariates")
722		}
723
724		#' Extract Covariate Shrinkage Factors from the Posterior Samples
725		#'
726		#' @typed object: ANY
727		#' the object to extract shrinkage factors from.
728		#' @param ... additional arguments added by methods.
729		#'
730		#' @family shrinkage
731		#' @export
732		#'
733		#' @returns A named numeric vector or matrix of shrinkage factors.
734		#'
735		#' @examples
736		#' \dontrun{
737		#' shrinkage(fit)
738		#' }
739		shrinkage <- function(object, ...) {
740	1x	UseMethod("shrinkage")
741		}

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

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		#' @typed lambda: Prior
22		#' for the scale `lambda`.
23		#' @typed gamma: Prior
24		#' for the shape `gamma`.
25		#' @typed beta: Prior
26		#' for covariates coefficients `beta`.
27		#'
28		#' @export
29		#'
30		#' @returns A `SurvivalWeibullPH` object.
31		#'
32		#' @examples
33		#' SurvivalWeibullPH()
34		SurvivalWeibullPH <- function(
35		lambda = prior_gamma(2, 0.5),
36		gamma = prior_gamma(2, 0.5),
37		beta = prior_normal(0, 2)
38		) {
39	24x	lambda <- set_limits(lambda, lower = getOption("jmpost.double_eps"))
40	24x	gamma <- set_limits(gamma, lower = getOption("jmpost.double_eps"))
41
42	24x	.SurvivalWeibullPH(
43	24x	SurvivalModel(
44	24x	name = "Weibull-PH",
45	24x	stan = StanModule(x = "sm-weibull-ph/model.stan"),
46	24x	parameters = ParameterList(
47	24x	Parameter(
48	24x	name = "sm_weibull_ph_lambda",
49	24x	prior = lambda,
50	24x	size = 1
51		),
52	24x	Parameter(
53	24x	name = "sm_weibull_ph_gamma",
54	24x	prior = gamma,
55	24x	size = 1
56		),
57	24x	Parameter(
58	24x	name = "beta_os_cov",
59	24x	prior = beta,
60	24x	size = "p_os_cov_design"
61		)
62		)
63		)
64		)
65		}

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

1		#' @include generics.R
2		NULL
3
4
5		#' Standard Links
6		#'
7		#' @typed prior: Prior
8		#' A [`Prior`] object.
9		#' @typed model: LongitudinalModel
10		#' A [`LongitudinalModel`] object.
11		#' @param ... Not used.
12		#'
13		#' @description
14		#' These functions are used to enable the use of the corresponding link function between
15		#' the survival and longitudinal models in a joint model. Note that the exact implementation
16		#' of the link function is model specific, see
17		#' \code{vignette("Statistical Specifications", package = "jmpost")} for more details.
18		#'
19		#' @name standard-link-user
20		#'
21		#' @examples
22		#' linkNone()
23		#' linkDSLD()
24		#' linkTTG(prior_normal(0, 1), model = LongitudinalGSF())
25		NULL
26
27
28		#' @describeIn standard-link-user No link (fit the survival and longitudinal models independently)
29		#' @export
30		#'
31		#' @returns A `Link` or `LinkComponent` object representing the requested association.
32		linkNone <- function() {
33	1x	Link()
34		}
35
36
37		#' @describeIn standard-link-user Time to growth link
38		#' @export
39		linkTTG <- function(prior, model = PromiseLongitudinalModel(), ...) {
40	15x	UseMethod("linkTTG", model)
41		}
42		#' @export
43		#'
44		#' @returns A `PromiseLinkComponent` object.
45		linkTTG.PromiseLongitudinalModel <- function(
46		prior = prior_normal(0, 2),
47		model,
48		...
49		) {
50	7x	PromiseLinkComponent(fun = linkTTG, prior = prior, key = "link_ttg")
51		}
52		#' @export
53		#'
54		#' @returns This method always raises an error for unsupported models.
55		linkTTG.default <- function(prior, model, ...) {
56	1x	stop(sprintf(
57	1x	"Method `linkTTG` is not available for `%s`",
58	1x	class(model)[[1]]
59		))
60		}
61
62
63		#' @describeIn standard-link-user Derivative of the SLD over time link
64		#' @export
65		linkDSLD <- function(prior, model = PromiseLongitudinalModel(), ...) {
66	49x	UseMethod("linkDSLD", model)
67		}
68		#' @export
69		#'
70		#' @returns A `PromiseLinkComponent` object.
71		linkDSLD.PromiseLongitudinalModel <- function(
72		prior = prior_normal(0, 2),
73		model,
74		...
75		) {
76	26x	PromiseLinkComponent(fun = linkDSLD, prior = prior, key = "link_dsld")
77		}
78		#' @export
79		#'
80		#' @returns This method always raises an error for unsupported models.
81		linkDSLD.default <- function(prior, model, ...) {
82	!	stop(sprintf(
83	!	"Method `linkDSLD` is not available for `%s`",
84	!	class(model)[[1]]
85		))
86		}
87
88
89		#' @describeIn standard-link-user Current SLD value link
90		#' @export
91		linkIdentity <- function(prior, model = PromiseLongitudinalModel(), ...) {
92	15x	UseMethod("linkIdentity", model)
93		}
94		#' @export
95		#'
96		#' @returns A `PromiseLinkComponent` object.
97		linkIdentity.PromiseLongitudinalModel <- function(
98		prior = prior_normal(0, 2),
99		model,
100		...
101		) {
102	7x	PromiseLinkComponent(
103	7x	fun = linkIdentity,
104	7x	prior = prior,
105	7x	key = "link_identity"
106		)
107		}
108		#' @export
109		#'
110		#' @returns This method always raises an error for unsupported models.
111		linkIdentity.default <- function(prior, model, ...) {
112	!	stop(sprintf(
113	!	"Method `linkIdentity` is not available for `%s`",
114	!	class(model)[[1]]
115		))
116		}
117
118
119		#' @describeIn standard-link-user Growth Parameter link
120		#' @export
121		linkGrowth <- function(prior, model = PromiseLongitudinalModel(), ...) {
122	20x	UseMethod("linkGrowth", model)
123		}
124		#' @export
125		#'
126		#' @returns A `PromiseLinkComponent` object.
127		linkGrowth.PromiseLongitudinalModel <- function(
128		prior = prior_normal(0, 2),
129		model,
130		...
131		) {
132	10x	PromiseLinkComponent(fun = linkGrowth, prior = prior, key = "link_growth")
133		}
134		#' @export
135		#'
136		#' @returns This method always raises an error for unsupported models.
137		linkGrowth.default <- function(prior, model, ...) {
138	!	stop(sprintf(
139	!	"Method `linkGrowth` is not available for `%s`",
140	!	class(model)[[1]]
141		))
142		}
143
144
145		#' @describeIn standard-link-user Shrinkage Parameter link
146		#' @export
147		linkShrinkage <- function(prior, model = PromiseLongitudinalModel(), ...) {
148	12x	UseMethod("linkShrinkage", model)
149		}
150		#' @export
151		#'
152		#' @returns A `PromiseLinkComponent` object.
153		linkShrinkage.PromiseLongitudinalModel <- function(
154		prior = prior_normal(0, 2),
155		model,
156		...
157		) {
158	6x	PromiseLinkComponent(
159	6x	fun = linkShrinkage,
160	6x	prior = prior,
161	6x	key = "link_shrinkage"
162		)
163		}
164		#' @export
165		#'
166		#' @returns This method always raises an error for unsupported models.
167		linkShrinkage.default <- function(prior, model, ...) {
168	!	stop(sprintf(
169	!	"Method `linkShrinkage` is not available for `%s`",
170	!	class(model)[[1]]
171		))
172		}

1		#' @include LongitudinalModel.R
2		#' @include Link.R
3		NULL
4
5		# LongitudinalRandomSlope-class ----
6
7		#' `LongitudinalRandomSlope`
8		#'
9		#' This class extends the general [`LongitudinalModel`] class for using the
10		#' random slope linear model for the longitudinal outcome.
11		#'
12		#' @section Available Links:
13		#' - [`linkDSLD()`]
14		#' - [`linkIdentity()`]
15		#' @exportClass LongitudinalRandomSlope
16		.LongitudinalRandomSlope <- setClass(
17		Class = "LongitudinalRandomSlope",
18		contains = "LongitudinalModel"
19		)
20
21
22		# LongitudinalRandomSlope-constructors ----
23
24		#' @rdname LongitudinalRandomSlope-class
25		#'
26		#' @typed intercept: Prior
27		#' for the `intercept`.
28		#' @typed slope_mu: Prior
29		#' for the population slope `slope_mu` (one per arm).
30		#' @typed slope_sigma: Prior
31		#' for the random slope standard deviation `slope_sigma` (one per arm).
32		#' @typed sigma: Prior
33		#' for the variance of the longitudinal values `sigma`.
34		#' @typed scaled_variance: logical
35		#' whether the variance should be scaled by the expected value, corresponding to a multiplicative model.
36		#' As a default, the variance is not scaled by the expected value, corresponding to an additive model.
37		#' (See the "Statistical Specifications" vignette for more details.)
38		#'
39		#' @export
40		#'
41		#' @returns A `LongitudinalRandomSlope` object.
42		#'
43		#' @examples
44		#' LongitudinalRandomSlope()
45		LongitudinalRandomSlope <- function(
46		intercept = prior_normal(30, 10),
47		slope_mu = prior_normal(1, 3),
48		slope_sigma = prior_lognormal(0, 1.5),
49		sigma = prior_lognormal(0, 1.5),
50		scaled_variance = FALSE
51		) {
52	33x	stan <- StanModule(decorated_render(
53	33x	.x = read_stan("lm-random-slope/model.stan"),
54	33x	scaled_variance = scaled_variance
55		))
56
57		# Apply constriants
58	33x	sigma <- set_limits(sigma, lower = getOption("jmpost.double_eps"))
59	33x	slope_sigma <- set_limits(slope_sigma, lower = getOption("jmpost.double_eps"))
60
61	33x	.LongitudinalRandomSlope(
62	33x	LongitudinalModel(
63	33x	name = "Random Slope",
64	33x	stan = stan,
65	33x	scaled_variance = scaled_variance,
66	33x	parameters = ParameterList(
67	33x	Parameter(
68	33x	name = "lm_rs_intercept",
69	33x	prior = intercept,
70	33x	size = "n_studies"
71		),
72	33x	Parameter(
73	33x	name = "lm_rs_slope_mu",
74	33x	prior = slope_mu,
75	33x	size = "n_arms"
76		),
77	33x	Parameter(
78	33x	name = "lm_rs_slope_sigma",
79	33x	prior = slope_sigma,
80	33x	size = "n_arms"
81		),
82	33x	Parameter(name = "lm_rs_sigma", prior = sigma, size = 1),
83	33x	Parameter(
84	33x	name = "lm_rs_ind_rnd_slope",
85	33x	prior = prior_init_only(prior_normal(
86	33x	median(slope_mu),
87	33x	median(slope_sigma)
88		)),
89	33x	size = "n_subjects"
90		)
91		)
92		)
93		)
94		}
95
96
97		#' @export
98		#'
99		#' @returns A `StanModule` object containing the generated-quantities code.
100		enableGQ.LongitudinalRandomSlope <- function(object, ...) {
101	31x	StanModule("lm-random-slope/quantities.stan")
102		}
103
104		#' @export
105		#'
106		#' @returns The longitudinal model with its link-related Stan code enabled.
107		enableLink.LongitudinalRandomSlope <- function(object, ...) {
108	8x	object@stan <- merge(
109	8x	object@stan,
110	8x	StanModule("lm-random-slope/link.stan")
111		)
112	8x	object
113		}
114
115
116		#' @export
117		#'
118		#' @returns A `LinkComponent` object.
119		linkDSLD.LongitudinalRandomSlope <- function(
120		prior = prior_normal(0, 2),
121		model,
122		...
123		) {
124	9x	LinkComponent(
125	9x	key = "link_dsld",
126	9x	stan = StanModule("lm-random-slope/link_dsld.stan"),
127	9x	prior = prior
128		)
129		}
130
131		#' @export
132		#'
133		#' @returns A `LinkComponent` object.
134		linkIdentity.LongitudinalRandomSlope <- function(
135		prior = prior_normal(0, 2),
136		model,
137		...
138		) {
139	3x	LinkComponent(
140	3x	key = "link_identity",
141	3x	stan = StanModule("lm-random-slope/link_identity.stan"),
142	3x	prior = prior
143		)
144		}
145
146		#' @export
147		#'
148		#' @returns A `LinkComponent` object.
149		linkGrowth.LongitudinalRandomSlope <- function(
150		prior = prior_normal(0, 2),
151		model,
152		...
153		) {
154	!	LinkComponent(
155	!	key = "link_growth",
156	!	stan = StanModule("lm-random-slope/link_growth.stan"),
157	!	prior = prior
158		)
159		}
160
161		#' @rdname getPredictionNames
162		#' @export
163		getPredictionNames.LongitudinalRandomSlope <- function(object, ...) {
164	1x	c("intercept", "slope")
165		}
166
167
168		#' @rdname getRandomEffectsNames
169		#' @export
170		getRandomEffectsNames.LongitudinalRandomSlope <- function(object, ...) {
171	!	c("slope" = "lm_rs_ind_rnd_slope")
172		}

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

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(
68	8x	unique(unlist(object@groups)) %in% names(data_list$subject_to_index)
69		)
70		)
71
72	8x	validate_time_grid(object@times)
73
74	8x	group_grid <- expand.grid(
75	8x	group = names(object@groups),
76	8x	time = object@times,
77	8x	stringsAsFactors = FALSE
78		)
79
80	8x	generator <- as.QuantityGenerator(object, data)
81
82	8x	select_indexes <- mapply(
83	8x	function(group, time) {
84	490x	correct_subject <- generator@subjects %in% object@groups[[group]]
85	490x	correct_time <- generator@times == time
86	490x	seq_along(correct_time)[correct_subject & correct_time]
87		},
88	8x	group_grid$group,
89	8x	group_grid$time,
90	8x	SIMPLIFY = FALSE
91		)
92	8x	names(select_indexes) <- NULL
93
94	8x	QuantityCollapser(
95	8x	times = group_grid$time,
96	8x	groups = group_grid$group,
97	8x	indexes = select_indexes
98		)
99		}
100
101		#' @export
102		#'
103		#' @returns A named `list` representation of the object.
104		as.list.GridGrouped <- function(x, ...) {
105	1x	x@groups
106		}
107
108		#' @rdname coalesceGridTime
109		#' @export
110		coalesceGridTime.GridGrouped <- function(object, times, ...) {
111	10x	if (is.null(object@times)) {
112	2x	object <- GridGrouped(
113	2x	groups = object@groups,
114	2x	times = times
115		)
116		}
117	10x	object
118		}

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		#' @typed a: Prior
22		#' Prior distribution for the scale parameter `a`.
23		#' @typed b: Prior
24		#' Prior distribution for the shape parameter `b`.
25		#' @typed beta: Prior
26		#' Prior distribution for covariates coefficients `beta`.
27		#'
28		#' @export
29		#'
30		#' @returns A `SurvivalLogLogistic` object.
31		#'
32		#' @examples
33		#' SurvivalLogLogistic()
34		SurvivalLogLogistic <- function(
35		a = prior_lognormal(log(0.1), 5),
36		b = prior_gamma(2, 5),
37		beta = prior_normal(0, 2)
38		) {
39	9x	a <- set_limits(a, lower = getOption("jmpost.double_eps"))
40	9x	b <- set_limits(b, lower = getOption("jmpost.double_eps"))
41
42	9x	.SurvivalLogLogistic(
43	9x	SurvivalModel(
44	9x	name = "Log-Logistic",
45	9x	stan = StanModule("sm-loglogistic/model.stan"),
46	9x	parameters = ParameterList(
47	9x	Parameter(name = "sm_loglogis_a", prior = a, size = 1),
48	9x	Parameter(name = "sm_loglogis_b", prior = b, size = 1),
49	9x	Parameter(
50	9x	name = "beta_os_cov",
51	9x	prior = beta,
52	9x	size = "p_os_cov_design"
53		)
54		)
55		)
56		)
57		}

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		#'
20		#' @returns A `JointModelSamples` object.
21		#'
22		#' @examples
23		#' \dontrun{
24		#' # JointModelSamples objects are returned by sampleStanModel().
25		#' fit <- sampleStanModel(JointModel(LongitudinalGSF()), data = joint_data)
26		#' }
27		.JointModelSamples <- setClass(
28		"JointModelSamples",
29		slots = c(
30		model = "JointModel",
31		data = "DataJoint",
32		results = "CmdStanMCMC"
33		)
34		)
35
36
37		#' @rdname generateQuantities
38		#' @typed generator: QuantityGenerator
39		#' object that specifies which subjects and time points
40		#' to calculate the quantities at
41		#' @typed type: character
42		#' type of quantities to be generated, must be either "survival" or
43		#' "longitudinal".
44		#' @export
45		generateQuantities.JointModelSamples <- function(object, generator, type, ...) {
46	37x	data <- as_stan_list(object@data) \|>
47	37x	append(as_stan_list(object@model@parameters)) \|>
48	37x	append(as_stan_list(
49	37x	generator,
50	37x	data = object@data,
51	37x	model = object@model
52		))
53
54	36x	stanobj <- as.StanModule(object, generator = generator, type = type)
55	36x	model <- compileStanModel(stanobj)
56
57	36x	devnull <- utils::capture.output(
58	36x	results <- model$generate_quantities(
59	36x	data = data,
60	36x	fitted_params = object@results
61		)
62		)
63	36x	return(results)
64		}
65
66
67		#' `JointModelSamples` -> `StanModule`
68		#'
69		#' Converts a `JointModelSamples` object into a `StanModule` object ensuring
70		#' that the resulting `StanModule` object is able to generate post sampling
71		#' quantities.
72		#'
73		#' @inheritParams generateQuantities
74		#' @export
75		#'
76		#' @returns A `StanModule` object.
77		as.StanModule.JointModelSamples <- function(object, generator, type, ...) {
78	44x	assert_that(
79	44x	is(generator, "QuantityGenerator"),
80	44x	length(type) == 1,
81	44x	type %in% c("survival", "longitudinal")
82		)
83
84	44x	quant_stanobj <- read_stan("base/quantities.stan") \|>
85	44x	decorated_render(
86	44x	include_gq_longitudinal_idv = (type == "longitudinal") &
87	44x	is(generator, "QuantityGeneratorSubject"),
88	44x	include_gq_longitudinal_pop = (type == "longitudinal") &
89	44x	is(generator, "QuantityGeneratorPopulation"),
90	44x	include_gq_survival_idv = (type == "survival") &
91	44x	is(generator, "QuantityGeneratorSubject"),
92	44x	include_gq_survival_pred = (type == "survival") &
93	44x	is(generator, "QuantityGeneratorPrediction")
94		) \|>
95	44x	StanModule()
96
97	44x	stanobj <- Reduce(
98	44x	merge,
99	44x	list(
100	44x	as.StanModule(object@model),
101	44x	enableGQ(object@model),
102	44x	quant_stanobj
103		)
104		)
105	44x	stanobj
106		}
107
108
109		#' `JointModelSamples` -> Printable `Character`
110		#'
111		#' Converts [`JointModelSamples`] object into a printable string.
112		#'
113		#' @typed object: JointModelSamples
114		#' samples as drawn from a [`JointModel`].
115		#' @family JointModelSamples
116		#' @typed indent: numeric
117		#' how much white space to prefix the print string with.
118		#' @keywords internal
119		#' @export
120		#'
121		#' @returns A character vector suitable for printing.
122		as_print_string.JointModelSamples <- function(object, indent = 1, ...) {
123	18x	sizes <- vapply(
124	18x	cmdstanr::as.CmdStanMCMC(object)$metadata()[["stan_variable_sizes"]],
125	18x	\(x) {
126	280x	if (length(x) == 1 && x == 1) {
127	88x	return("")
128		}
129	192x	paste0("[", paste(x, collapse = ", "), "]")
130		},
131	18x	character(1)
132		)
133	18x	variable_string <- paste0(
134		" ",
135	18x	cmdstanr::as.CmdStanMCMC(object)$metadata()[["stan_variables"]],
136	18x	sizes
137		)
138	18x	template <- c(
139	18x	"JointModelSamples Object with:",
140		"",
141	18x	" # of samples per chain = %d",
142	18x	" # of chains = %d",
143		"",
144	18x	" Variables:",
145	18x	variable_string[order(variable_string)]
146		)
147	18x	pad <- rep(" ", indent) \|> paste(collapse = "")
148	18x	template_padded <- paste(pad, template)
149	18x	sprintf(
150	18x	paste(template_padded, collapse = "\n"),
151	18x	cmdstanr::as.CmdStanMCMC(object)$metadata()$iter_sampling,
152	18x	cmdstanr::as.CmdStanMCMC(object)$num_chains()
153		)
154		}
155
156		#' @rdname show-object
157		#' @export
158		setMethod(
159		f = "show",
160		signature = "JointModelSamples",
161		definition = function(object) {
162	18x	string <- as_print_string(object)
163	18x	cat("\n", string, "\n\n")
164		}
165		)
166
167
168		#' @rdname as.CmdStanMCMC
169		#'
170		#' @returns The underlying `cmdstanr::CmdStanMCMC` object.
171		as.CmdStanMCMC.JointModelSamples <- function(object, ...) {
172	84x	return(object@results)
173		}
174
175
176		#' Save a `JointModelSamples` object to a file.
177		#'
178		#' This function is just a wrapper around `saveRDS` that saves the object to a file
179		#' ensuring that all of the Stan samples are correctly stored. Note that as
180		#' `cmdstanr` objects store their samples as a csv file the samples may be lost
181		#' if you call `saveRDS` directly on the object.
182		#'
183		#' @typed object: JointModelSamples
184		#' the object to save.
185		#' @typed file: character
186		#' the file to save the object to.
187		#' @typed ...: ANY
188		#' additional arguments to [`saveRDS`].
189		#'
190		#' @family saveObject
191		#'
192		#' @export
193		#'
194		#' @returns Invisibly returns `NULL` after saving the object.
195		saveObject.JointModelSamples <- function(object, file, ...) {
196	1x	object@results$draws()
197	1x	try(object@results$sampler_diagnostics(), silent = TRUE)
198	1x	try(object@results$init(), silent = TRUE)
199	1x	try(object@results$profiles(), silent = TRUE)
200	1x	saveRDS(object, file, ...)
201		}
202
203		#' Extract shrinkage factors from a `JointModelSamples` object
204		#'
205		#' @typed object: JointModelSamples
206		#' the object to extract shrinkage factors from.
207		#' @param ... not used.
208		#' @return the shrinkage factors correctly named after the covariate names.
209		#'
210		#' @family shrinkage
211		#' @export
212		shrinkage.JointModelSamples <- function(object, ...) {
213	1x	draws <- cmdstanr::as.CmdStanMCMC(object)$draws()
214	1x	assert_class(object@data@survival, "DataSurvival")
215	1x	covariate_names <- covariates(object@data@survival)
216	1x	assert_subset(
217	1x	"prior_shrinkage_factors_beta_os_cov[1]",
218	1x	posterior::variables(draws)
219		)
220	1x	shrinkage_draws <- subset(
221	1x	draws,
222	1x	variable = "prior_shrinkage_factors_beta_os_cov"
223		)
224	1x	shrinkage_vars <- posterior::variables(shrinkage_draws)
225	1x	assert_true(identical(length(covariate_names), length(shrinkage_vars)))
226	1x	posterior::variables(shrinkage_draws) <- covariate_names
227	1x	shrinkage_draws
228		}

1		#' @include SurvivalModel.R
2		NULL
3
4
5		#' `SurvivalGamma`
6		#'
7		#' This class extends the general [`SurvivalModel`] class for using the
8		#' Gamma survival model.
9		#'
10		#' @exportClass SurvivalGamma
11		.SurvivalGamma <- setClass(
12		Class = "SurvivalGamma",
13		contains = "SurvivalModel"
14		)
15
16		# SurvivalGamma-constructors ----
17
18		#' @rdname SurvivalGamma-class
19		#'
20		#' @typed k: Prior
21		#' for the shape `k`.
22		#' @typed theta: Prior
23		#' for the scale `theta`.
24		#' @typed beta: Prior
25		#' for covariates coefficients `beta`.
26		#'
27		#' @export
28		#'
29		#' @returns A `SurvivalGamma` object.
30		#'
31		#' @examples
32		#' SurvivalGamma()
33		SurvivalGamma <- function(
34		k = prior_gamma(2, 0.5),
35		theta = prior_gamma(2, 0.5),
36		beta = prior_normal(0, 2)
37		) {
38	4x	k <- set_limits(k, lower = getOption("jmpost.double_eps"))
39	4x	theta <- set_limits(theta, lower = getOption("jmpost.double_eps"))
40
41	4x	.SurvivalGamma(
42	4x	SurvivalModel(
43	4x	name = "Gamma",
44	4x	stan = StanModule(x = "sm-gamma/model.stan"),
45	4x	parameters = ParameterList(
46	4x	Parameter(name = "sm_gamma_k", prior = k, size = 1),
47	4x	Parameter(name = "sm_gamma_theta", prior = theta, size = 1),
48	4x	Parameter(
49	4x	name = "beta_os_cov",
50	4x	prior = beta,
51	4x	size = "p_os_cov_design"
52		)
53		)
54		)
55		)
56		}

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

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

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		#'
25		#' @returns A `SurvivalModel` object.
26		#'
27		#' @examples
28		#' SurvivalModel()
29		SurvivalModel <- function(
30		stan = StanModule(),
31		parameters = ParameterList(),
32		name = "<Unnamed>",
33		...
34		) {
35	63x	base_stan <- read_stan("base/survival.stan")
36	63x	stan_full <- decorated_render(
37	63x	.x = base_stan,
38	63x	stan = add_missing_stan_blocks(as.list(stan))
39		)
40	63x	.SurvivalModel(
41	63x	StanModel(
42	63x	name = name,
43	63x	stan = StanModule(stan_full),
44	63x	parameters = parameters,
45		...
46		)
47		)
48		}
49
50		#' @export
51		#'
52		#' @returns A character vector suitable for printing.
53		as_print_string.SurvivalModel <- function(object, ...) {
54	15x	string <- sprintf(
55	15x	"\n%s Survival Model with parameters:\n%s\n\n",
56	15x	object@name,
57	15x	paste(" ", as_print_string(object@parameters)) \|>
58	15x	paste(collapse = "\n")
59		)
60	15x	return(string)
61		}

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

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

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		#' @typed lambda: Prior
22		#' for the exponential rate `lambda`.
23		#' @typed beta: Prior
24		#' for covariates coefficients `beta`.
25		#'
26		#' @export
27		#'
28		#' @returns A `SurvivalExponential` object.
29		#'
30		#' @examples
31		#' SurvivalExponential()
32		SurvivalExponential <- function(
33		lambda = prior_gamma(2, 5),
34		beta = prior_normal(0, 2)
35		) {
36	26x	lambda <- set_limits(lambda, lower = getOption("jmpost.double_eps"))
37	26x	.SurvivalExponential(
38	26x	SurvivalModel(
39	26x	name = "Exponential",
40	26x	stan = StanModule("sm-exponential/model.stan"),
41	26x	parameters = ParameterList(
42	26x	Parameter(name = "sm_exp_lambda", prior = lambda, size = 1),
43	26x	Parameter(
44	26x	name = "beta_os_cov",
45	26x	prior = beta,
46	26x	size = "p_os_cov_design"
47		)
48		)
49		)
50		)
51		}

1		#' Abstract Simulation Class for Longitudinal Data
2		#'
3		#' @typed times: numeric
4		#' 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		#'
13		#' @examples
14		#' SimLongitudinal(times = c(0, 30, 60))
15		.SimLongitudinal <- setClass(
16		"SimLongitudinal",
17		slots = list(
18		times = "numeric"
19		)
20		)
21
22		#' @rdname SimLongitudinal-class
23		#' @export
24		#'
25		#' @returns A `SimLongitudinal` object.
26		SimLongitudinal <- function(times = seq(0, 100, 50)) {
27	!	.SimLongitudinal(times = times)
28		}
29
30
31		#' @rdname show-object
32		#' @export
33		setMethod(
34		f = "show",
35		signature = "SimLongitudinal",
36		definition = function(object) {
37	4x	x <- sprintf("\nA %s Object\n\n", as_print_string(object))
38	4x	cat(x)
39	4x	return(object)
40		}
41		)
42
43		#' @rdname as_print_string
44		#' @exportS3Method NULL
45		as_print_string.SimLongitudinal <- function(object, ...) {
46	!	return("SimLongitudinal")
47		}