#' A class for an object to contain the proportions of arriving patients to be

#' allocated to each arm of the trial.

#' 

#' @slot recruit_arm_prevalence Numeric vector of the expected proportion 

#' of patients eligible for each recruitment arm.

#' @slot recruit_arm_prevalence_start Numeric vector of the initial proportions

#' of patients eligible for each recruitment arm.

#' @slot recruit_arm_names Character vector of the names of the recruitment 

#' arms.

#' @slot shared_control TRUE if using a shared control arm for all 

#' experimental arms.

#' @slot ctrl_ratio Proportion of patients assigned to control

#' @slot treatment_arm_ids Named list of lists of recruitment arms by 

#' treatment arm.

#' @slot treatment_arm_ids_start Named list of lists of the initial 

#' configuration of recruitment arms by treatment arm.

#' @slot recruit_arm_id Automatically generated integer vector of the ID 

#' numbers of the recruitment arms.

#' @slot treatment_counts Automatically generated named integer vector of 

#' number of recruitment arms recruiting to each treatment arm.

#' @slot treatment_arm_struct Automatically generated logical matrix of 

#' treatment arms by recruitment arms.

#' @slot treatment_arm_struct_start Automatically generated logical matrix of 

#' the initial configuration of treatment arms by recruitment arms.

#' @slot experimental_arm_prevalence Automatically generated matrix of 

#' prevalences of treatment arms by recruitment arms

#' 

#' @param props_df Dataframe of expected biomarker prevalences for the 

#' regions, with one column `category` containing names for the 

#' biomarkers, and one column per region.

#' @param arms_ls List of lists of recruitment arms which recruit to

#' each treatment arm.

#' @param centres_df Dataframe containing columns `site`, the index

#' number of each site; `start_month`, the month in which that site

#' is expected to start recruiting; `mean_rate`, the expected number 

#' of patients from that site per month; `region`, the index of the

#' region the site is in (should be in the same order as the columns

#' in `props_df`); and an optional column `site_cap`, if there is a 

#' recruitment cap on any of the sites.

#' @param precision For the Dirichlet model of biomarker prevalences, 

#' variability decreases as precision increases. Defaults to 10.

#' @param shared_control TRUE if all experimental arms share one 

#' control arm; FALSE if they each have separate control arms.

#' @param ctrl_ratio Proportion of patients assigned to control

#' @param fixed_region_prevalences TRUE if biomarker prevalences 

#' should be considered to be identical for all sites within a 

#' region; FALSE if they should be drawn from a Dirichlet distribution

#' with a mean of the specified prevalence.

#' 

#' @name trial_structure

#' 

#' @import S7

#' 

trial_structure <- S7::new_class("trial_structure",

  package = "biomkrAccrual",

  properties = list(

    # These need explicitly setting

    recruit_arm_prevalence = S7::class_double,

    recruit_arm_prevalence_start = S7::class_double,

    recruit_arm_names = S7::class_character,

    shared_control = S7::class_logical,

    ctrl_ratio = S7::class_vector,

    treatment_arm_ids = S7::class_list,

    treatment_arm_ids_start = S7::class_list,

    # These are generated from existing properties at the time they execute

    recruit_arm_ids = S7::new_property(

      getter = function(self) seq_len(nrow(self@recruit_arm_prevalence))

    ),

    treatment_counts = S7::new_property(

      getter = function(self) length(self@treatment_arm_ids)

    ),

    # Will inherit class matrix despite S7 silliness

    treatment_arm_struct = S7::new_property(

      getter = function(self) {

        get_matrix_struct(self@treatment_arm_ids, self@recruit_arm_prevalence)

      }

    ),

    treatment_arm_struct_start = S7::new_property(

      getter = function(self) {

        get_matrix_struct(

          self@treatment_arm_ids_start, self@recruit_arm_prevalence_start

        )

      }

    ),

    # array[recruit arms, treat arms, prevalence set]

    experimental_arm_prevalence = S7::new_property(

      getter = 

        function(self) {

          get_array_prevalence(

            self@treatment_arm_struct, 

            self@recruit_arm_prevalence,

            self@shared_control,

            self@ctrl_ratio

          )

        }

    )

  ),

  # Make new instance by calling trial_structure(props_df, arms_ls)

  constructor = function(

    props_df = S7::class_missing, 

    arms_ls = S7::class_missing,

    centres_df = S7::class_missing,

    precision = S7::class_missing,

    shared_control = S7::class_missing,

    ctrl_ratio = S7::class_missing,

    fixed_region_prevalences = S7::class_missing

  ) {

    # Create the object and populate it

    S7::new_object(

      # Parent class

      S7::S7_object(),

      recruit_arm_names = props_df$category, 

      recruit_arm_prevalence = 

        get_recruit_arm_prevalence(

          props_df, centres_df, precision, fixed_region_prevalences

        ),

      recruit_arm_prevalence_start = 

        get_recruit_arm_prevalence(

          props_df, centres_df, precision, fixed_region_prevalences

        ),

      shared_control = shared_control,

      ctrl_ratio = ctrl_ratio,

      treatment_arm_ids = arms_ls,

      treatment_arm_ids_start = arms_ls

    )

  },

  validator = function(self) {

    if (!is.numeric(self@recruit_arm_prevalence[, -1])) {

      "Recruitment arm prevalences must be numbers"

    } else if (

      length(self@recruit_arm_names) != nrow(self@recruit_arm_prevalence)

    ) {

      "Different number of recruitment arm names supplied than prevalences"

    } else if (

      !(all(sapply(

        self@treatment_arm_ids, 

        function(v) is.integer(unlist(v)) || is.na(v)

      )))

    ) {

      "Elements from the treatment arm list should be integer vectors or NA"

    } 

  }

)

#' Dirichet regression model for biomarker prevalence

#' 

#' For each site, draws from the Dirichlet distribution using the

#' expected prevalences for the region. 

#' expected prevalences by region

#' 

#' @param props_df Dataframe with one row per biomarker.  Has one column 

#' `category`, containing names of the biomarkers, plus one column for 

#' each region, containing the expected biomarker prevalences for the regions.

#' @param centres_df Dataframe with one row per site, including one column

#' `region`, containing the index number for the region for each site.

#' Indices are assumed to be in the same order as the columns in `props_df`.

#' @param precision Variability decreases as precision increases.

#' @param fixed_region_prevalences TRUE if biomarker prevalences 

#' should be considered to be identical for all sites within a 

#' region; FALSE if they should be drawn from a Dirichlet distribution

#' with a mean of the specified prevalence.

#'   

#' @return Matrix of prevalences with one column per site and one 

#' row per biomarker; each column sums to 1.

#' 

#' @importFrom checkmate assert_names assert_data_frame assert_atomic_vector 

#' assert_integerish assert_numeric

#' 

get_recruit_arm_prevalence <- function(

  props_df, centres_df, precision, fixed_region_prevalences

) {

  # Check format of fixed_region_prevalences

  checkmate::assert_logical(

    fixed_region_prevalences, 

    len = 1, 

    any.missing = FALSE, 

    null.ok = FALSE

  )

  # If fixed_region_prevalences is FALSE, check contents of 

  # precision; otherwise should be NULL

  if (fixed_region_prevalences) {

    assert_null(precision)

  } else {

    # Check format and content of precision

    checkmate::assert_numeric(

      precision,

      lower = 10^-7,

      finite = TRUE,

      len = 1,

      any.missing = FALSE,

      null.ok = FALSE

    )

  }

  # Check format and content of centres_df

  checkmate::assert_data_frame(

    centres_df,

    types = "numeric",

    any.missing = FALSE,

    min.cols = 5,

    max.cols = 6,

    min.rows = 1,

    col.names = "named",

    null.ok = FALSE

  )

  checkmate::assert_names(

    names(centres_df),

    subset.of = c(

      "site", "start_month", "mean_rate", "region", "site_cap", "start_week"

    ),

    must.include = c(

      "site", "start_month", "mean_rate", "region", "start_week"

    )

  )

  sites_in_region <- centres_df$region

  checkmate::assert_atomic_vector(

    sites_in_region,

    any.missing = FALSE,

    min.len = 1,

    max.len = 10^4

  )

  checkmate::assert_integerish(

    sites_in_region,

    lower = 1,

    upper = 10^4,

    any.missing = FALSE,

    null.ok = FALSE

  )

  # Check format and content of props_df

  checkmate::assert_data_frame(

    props_df,

    types = c("numeric", "character"),

    any.missing = FALSE,

    # number of regions + "category"

    min.cols = max(sites_in_region) + 1,

    min.rows = 2,

    null.ok = FALSE

  )

  checkmate::assert_names(names(props_df), must.include = "category")

  # Any column that isn't "category" is assumed to be a region -

  # allows for named regions

  region_prevalence <- 

    props_df[, grep("^category$", names(props_df), invert = TRUE)]

  checkmate::assert_numeric(

    as.matrix(region_prevalence),

    lower = 0,

    upper = 1,

    finite = TRUE

  )

  if (fixed_region_prevalences) {

    # Use region prevalences unchanged

    recruit_arm_prevalence <- as.matrix(

      region_prevalence[, sites_in_region]

    )

    # Scale columns to sum to 1

    recruit_arm_prevalence <- sweep(

      recruit_arm_prevalence,

      2,

      colSums(recruit_arm_prevalence),

      FUN = "/"

    )

  } else {

    # Draw from Dirichlet distribution

    recruit_arm_prevalence <- do_dirichlet_draws(

      region_prevalence, sites_in_region, precision

    )

  }

  return(recruit_arm_prevalence)

}

#' Draws from dirichlet regression model for biomarker prevalences

#' to create the prevalence matrix.

#' 

#' @param region_prevalence Dataframe with one column for each 

#' region and one row for each biomarker, containing prevalences as 

#' probabilities.

#' @param sites_in_region Vector with a region index number for each

#' site, the index determined by the order of the columns in

#' `region_prevalence`.

#' @param precision Variability decreases as precision increases.

#'  

do_dirichlet_draws <- function(region_prevalence, sites_in_region, precision) {

  # Predeclare prevalence matrix

  recruit_arm_prevalence_mx <- matrix(

    data = NA, 

    ncol = length(sites_in_region),

    nrow = nrow(region_prevalence)

  )

  # Draw prevalences for sites in each region in turn

  for (region in unique(sites_in_region)) {

    # Sites in region

    site_indices <- which(sites_in_region == region)

    bio_prevalence <- rdirichlet_alt(

      n = length(site_indices),

      mu = region_prevalence[, region],

      phi = precision

    )

    # rdirichlet_alt produces the transpose of what we want,

    # for consistency with other implementations of rdirichlet

    recruit_arm_prevalence_mx[, site_indices] <- t(bio_prevalence)

  }

  return(recruit_arm_prevalence_mx)

}

#' Converts trial structure and prevalence information into matrix form

#' 

#' @param arms_ls List of lists of recruitment arms which recruit to

#' each treatment arm.

#' @param recruit_arm_prevalence Matrix of prevalences with one row per 

#' biomarker and one column per site; each column sums to 1.

#' 

#' @return Logical matrix with one row per biomarker and one column per

#' treatment arm; TRUE where the treatment recruits from that biomarker.

#' 

get_matrix_struct <- function(arms_ls, recruit_arm_prevalence) {

  # Predeclare matrix as no_biomarkers * no_treatments

  no_treats <- length(arms_ls)

  no_biomarkers <- nrow(recruit_arm_prevalence)

  # This one is logical, to avoid rounding errors

  arm_structure_mx <- 

    matrix(FALSE, max(unlist(arms_ls), no_biomarkers, na.rm = TRUE), no_treats)

  # Loop, changing to TRUE for arms including that treatment

  for (icol in seq_len(no_treats)) {

    if (any(!is.na(arms_ls[[icol]]))) {

      arm_structure_mx[unlist(arms_ls[[icol]]), icol] <- TRUE

    }

  }

  return(arm_structure_mx)

}

#' Make array with the prevalences by treatment arm, recruitment arm and 

#' prevalence set

#' @param arm_structure_mx Logical matrix of recruitment arms by treatment arm

#' @param recruit_arm_prevalence Data frame with sets of prevalences, in 

#' columns, for each arm (rows) 

#' @param shared_control TRUE if all treatment arms share the

#' same control arm; FALSE if each treatment arm has its own 

#' control. Defaults to TRUE.

#' @param ctrl_ratio Ratio of patients assigned to treatment versus control

#' 

#' @return arm_prevalence_ar Array of prevalences, recruitment arms * 

#' (treatment arms + control arms) * prevalence sets

#' 

get_array_prevalence <- function(

  arm_structure_mx, recruit_arm_prevalence, shared_control, ctrl_ratio

) {

  no_treatments <- ncol(arm_structure_mx)

  no_recruit_arms <- nrow(arm_structure_mx)

  no_regions <- ncol(recruit_arm_prevalence)

  scale_factor <- 

    ifelse(

      sapply(

        rowSums(arm_structure_mx), 

        function(x) isTRUE(all.equal(x, 0, tolerance = 1e-6))

      ),

      1, 

      rowSums(arm_structure_mx)

    )

  # List of matrices by centre 

  ## Total proportion recruited for centre on each open arm 

  ## - more efficient to fix later

  prev_ls <- lapply(

    seq(no_regions),

    function(i) recruit_arm_prevalence[, i] * arm_structure_mx / scale_factor

  )

  # Apply control configuration

  if (shared_control) {

    prev_ls <- lapply(

      prev_ls,

      function(mx) cbind(mx * ctrl_ratio[1], rowSums(mx) * ctrl_ratio[2])

    )

  } else {

    prev_ls <- lapply(

      prev_ls,

      function(mx) cbind(mx * ctrl_ratio[1], mx * ctrl_ratio[2])

    )

  }

  # Make into array

  arm_prevalence_ar <- array(

    data = do.call(cbind, prev_ls),

    dim = c(dim(prev_ls[[1]]), length(prev_ls))

  )

  return(arm_prevalence_ar)

}

#' Close off treatment arms.

#' In the list of treatment arm IDs, replaces the vector of recruitment

#' arm IDs with NA.

#' If any recruitment arms are closed, sets their prevalence to zero but 

#' does not recalculate prevalence vector, on the assumption that recruitment

#' from aites will fall because no patients with those characteristics will 

#' be recruited from that point.

#' Class object will automatically generate new trial structure and 

#' prevalence matrices.

#' 

#' @param structure_obj An object of class `trial_structure`.

#' @param arms Vector or scalar of integer arm ID numbers to remove

#' 

remove_treat_arms <- S7::new_generic("remove_treat_arms", "structure_obj")

S7::method(remove_treat_arms, trial_structure) <- function(

  structure_obj,

  arms

) {

  # Mark treatment arms as removed using NA; 

  # automatic getter for treatment_arm_struct does the rest

  structure_obj@treatment_arm_ids[arms] <- NA_integer_

  # Set prevalence to 0 for any recruitment arms which now have no

  # experimental arms to recruit to

  no_exp_arms <- rowSums(structure_obj@treatment_arm_struct) < 1

  structure_obj@recruit_arm_prevalence[no_exp_arms, ] <- 0

  return(structure_obj)

}

#' Check whether an object is of class "trial_structure".

#' 

#' @param x Object to test

#' 

#' @return TRUE or FALSE

#' 

#' @importFrom S7 new_generic method class_any

#' @export

#' 

is.trial_structure <- S7::new_generic("is.trial_structure", "x")

S7::method(is.trial_structure, S7::class_any) <- function(x) {

  inherits(x, "biomkrAccrual::trial_structure")

}