--- a
+++ b/R/recruit.R
@@ -0,0 +1,554 @@
+#' A class to contain the accrued patient allocations,
+#' by site, treatment/control arm and week
+#' 
+#' @slot accrual 3-D array with axes site, experimental arm
+#' and week
+#' @slot target_arm_size Number of subjects required for each treatment arm.
+#' @slot target_control Number of subjects required for control arm(s).
+#' @slot target_interim Number of subjects required for treatment arm at 
+#' interim analysis.
+#' @slot accrual_period Number of weeks in recruitment period.
+#' @slot interim_period Number of weeks to recruit for interim analysis.
+#' @slot phase_changes Vector of week numbers when arms closed
+#' @slot site_closures Vector of weeks sites closed; NA indicates open
+#' @slot week Current recruitment week
+#' @slot active_arms Vector of indices of open arms
+#' @slot active_sites Vector of indices of open sites
+#' @slot shared_control TRUE if a shared control arm is being used, else FALSE
+#' @slot site_in_region Vector of indices for which set of expected 
+#' prevalences each site should use 
+#' @slot site_cap Vector of maximum number of patients for each site
+#' @slot site_mean_rate Vector of expected recruitment rates for each site
+#' @slot site_rate Vector of recruitment rates for each site, drawn from a 
+#' gamma distribution
+#' @slot start_week Vector of weeks that each site opens recruitment
+#' @slot site_index Vector of index numbers for each site
+#' @slot treatment_arm_ids Named list of lists of recruitment arms by 
+#' treatment arm.
+#' 
+#' @param treatment_arm_ids Named list of lists of recruitment arms by 
+#' treatment arm.
+#' @param shared_control TRUE if all experimental arms share one control arm;
+#' FALSE if each has their own
+#' @param target_arm_size Number of subjects required for each treatment arm.
+#' @param target_control Number of subjects required for control arm(s).
+#' @param target_interim Number of subjects required for treatment arm at 
+#' interim analysis.
+#' @param accrual_period Number of weeks in recruitment period.
+#' @param interim_period Number of weeks to recruit for interim analysis.
+#' @param centres_df Dataframe with columns "site", "start_month", "mean_rate", 
+#' "region" and "site_cap"
+#' 
+#' @usage accrual(treatment_arm_ids, shared_control, target_arm_size, 
+#' target_control, target_interim, accrual_period, interim_period, centres_df)
+#' @name accrual
+#'
+#' @export
+#' 
+accrual <- S7::new_class("accrual",
+  package = "biomkrAccrual",
+  properties = list(
+    accrual = S7::class_integer,
+    target_arm_size = S7::class_integer,
+    target_control = S7::class_integer,
+    target_interim = S7::class_integer,
+    accrual_period = S7::class_integer,
+    interim_period = S7::class_integer,
+    control_ratio = S7::class_double,
+    phase_changes = S7::class_integer,
+    site_closures = S7::class_integer,
+    week = S7::class_integer,
+    active_arms = S7::class_integer,
+    active_sites = S7::class_integer,
+    shared_control = S7::class_logical,
+    site_in_region = S7::class_integer,
+    site_cap = S7::class_integer,
+    site_mean_rate = S7::class_double,
+    site_rate = S7::class_double,
+    site_start_week = S7::class_integer,
+    site_index = S7::class_integer,
+    treatment_arm_ids = S7::class_list
+  ),
+  constructor = function(
+    treatment_arm_ids = S7::class_missing,
+    shared_control = S7::class_missing,
+    target_arm_size = S7::class_missing,
+    target_control = S7::class_missing,
+    target_interim = S7::class_missing,
+    accrual_period = S7::class_missing,
+    interim_period = S7::class_missing,
+    control_ratio = S7::class_missing,
+    centres_df = S7::class_missing
+  ) {
+    # Create the object and populate it
+    S7::new_object(
+      # Parent class
+      S7::S7_object(),
+      # weeks * treatments * no_centres
+      accrual = array(
+        # Defaults to 0 for no patients recruited
+        as.integer(0), 
+        dim = c(
+          # Max. weeks
+          accrual_period,
+          # No. experimental arms including control
+          length(treatment_arm_ids) + 
+            ifelse(shared_control, 1, length(treatment_arm_ids)),
+          # No. sites 
+          length(unique(centres_df$site))
+        ),
+        dimnames = list(
+          Weeks = NULL,
+          Arms = c(
+            names(treatment_arm_ids), 
+            if (shared_control) {
+              "Control"
+            } else {
+              paste("C", names(treatment_arm_ids), sep = "-")
+            }
+          ),
+          Centres = c(paste("Centre", unique(centres_df$site)))
+        )
+      ),
+      target_arm_size = as.integer(target_arm_size),
+      target_control = as.integer(target_control),
+      target_interim = as.integer(target_interim),
+      accrual_period = accrual_period,
+      interim_period = interim_period,
+      control_ratio = control_ratio,
+      phase_changes = rep(NA_integer_, length(treatment_arm_ids)),
+      site_closures = rep(NA_integer_, length(unique(centres_df$site))),
+      week = as.integer(1),
+      active_arms = seq_len(length(treatment_arm_ids)),
+      active_sites = seq_len(length(unique(centres_df$site))),
+      shared_control = shared_control,
+      site_in_region = as.integer(centres_df$region),
+      site_cap = as.integer(centres_df$site_cap),
+      site_mean_rate = as.numeric(centres_df$mean_rate),
+      site_rate = NA_real_,
+      site_start_week = as.integer(centres_df$start_week),
+      site_index = as.integer(centres_df$site),
+      treatment_arm_ids = treatment_arm_ids
+    )
+  }
+)
+
+
+#' Sum accrual array by site, to enable checking against site caps
+#' @param obj Object of class "accrual"
+#' @return vector of total accrual by recruitment site
+#' 
+#' @export 
+#' 
+site_sums <- S7::new_generic("site_sums", "obj")
+S7::method(site_sums, accrual) <- function(obj) {
+  # Sum across dimensions 1:dims
+  colSums(obj@accrual, dims = 2)
+}
+
+
+# Sum accrual array by experimental arm (including control).
+#' 
+#' @param x Accrual array with dimensions Weeks, Arms and Centres.
+#' @export
+treat_sums <- function(x, ...) {
+  UseMethod("treat_sums", x)
+}
+
+#' Sum accrual array by experimental arm (including control).
+#' 
+#' @param x Accrual array with dimensions Weeks, Arms and Centres.
+#' @param no_treat_arms Number of treatment arms (as opposed to control 
+#' arms).
+#' @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 control_total Logical; if TRUE return single total for all 
+#' control arms (not used if `shared_control` is TRUE); defaults to FALSE.
+#' 
+#' @return vector of total accrual by experimental arm.
+#' 
+#' @export
+#' 
+treat_sums.array <- function(
+  x, 
+  control_total = FALSE,
+  no_treat_arms,
+  shared_control = TRUE,
+  na.rm = TRUE
+) {
+  # Permute the array so that the first dimension is the
+  # dimension you want to get sums for (experimental arms)
+  arm_sums <-
+    as.integer(colSums(rowSums(x, dims = 2, na.rm = na.rm)))
+
+  # If want total for control arms rather than separate values
+  if (!shared_control && control_total) {
+    arm_sums <- c(
+      # Treatment arms
+      arm_sums[seq_along(no_treat_arms)],
+      # Control
+      sum(arm_sums[seq(no_treat_arms + 1, length.out = no_treat_arms)])
+    )
+  }
+
+  return(arm_sums)
+}
+
+
+#' Sum accrual array element of accrual object, by experimental 
+#' arm (including control).
+#' 
+#' @name treat_sums
+#' 
+#' @param x Object of class `accrual`. 
+#' @param control_total Logical; if TRUE return single total for all 
+#' control arms
+#' 
+#' @return vector of total accrual by experimental arm
+#' 
+#' @export
+#' 
+`treat_sums.biomkrAccrual::accrual` <- function(
+  x,
+  control_total = FALSE,
+  na.rm = TRUE
+) {
+
+  # Call treat_sums.array() on accrual array element
+  treat_sums(
+    x@accrual,
+    control_total, 
+    length(x@phase_changes), 
+    x@shared_control,
+    na.rm 
+  )
+}
+
+
+#' Select arms to cap for single site, or sites to cap for single arm
+#' @param population Vector of allocations of patients this week
+#' @param captotal Value of cap for site or arm (as appropriate)
+#' @return Vector of arms to cap (can include multiples of the same arm)
+#' 
+do_choose_cap <- function(population, captotal) {
+  if (length(population) == captotal) { 
+    # Use whole population if correct length
+    capped <- population
+  } else {  
+    # Sample      
+    capped <- sample(population, size = captotal)
+  }
+
+  return(capped)
+}
+
+
+#' Get no. weeks from months
+#' Currently 4 week month, fix later using lubridate
+#' @param months Duration in months
+#' @return Duration in weeks
+#' 
+get_weeks <- function(months) {
+  as.integer(round(months * 4, 0))
+}
+
+
+#' Method to increment site rates by gamma-distributed rates of sites
+#' opening an accrual pathway in the current week.
+#' @param obj An object of type "accrual"
+#' @return Modified object with new site rates
+#' 
+#' @importFrom stats rgamma
+#' 
+set_site_rates <- S7::new_generic("site_start_rates", "obj")
+S7::method(set_site_rates, accrual) <- function(obj, fixed_site_rates) {
+
+  # If this is the first time calling this, initialise with rate 0
+  if (any(is.na(obj@site_rate))) {
+    obj@site_rate <- rep(0, dim(obj@accrual)[3])
+  }
+
+  indices <- which(obj@site_start_week == obj@week)
+
+  if (length(indices) > 0) {
+
+    # mean_rates are in recruitment per month, convert to weeks
+    if (fixed_site_rates) {
+      rates <- obj@site_mean_rate(indices) / 4
+    } else {
+      var_lambda <- 0.25
+      rates <- 0.25 * stats::rgamma(
+        n = length(indices),
+        shape = obj@site_mean_rate[indices]^2 / var_lambda,
+        # Per week not per month
+        rate = obj@site_mean_rate[indices] / var_lambda
+      )
+    }
+
+    # When multiple recruitment sources at a given site, want them
+    # to stack
+    obj@site_rate[obj@site_index[indices]] <- 
+      obj@site_rate[obj@site_index[indices]] + rates
+
+  }
+
+  return(obj)
+}
+
+
+#' Implement site cap on a week's accrual. 
+#' @param obj Accrual object
+#' @param site_cap Maximum number of patients per site
+#' 
+#' @return Modified accrual object with capped week's accrual and 
+#' with any capped sites removed from active_sites
+#' 
+apply_site_cap <- S7::new_generic("apply_site_cap", "obj")
+S7::method(apply_site_cap, accrual) <- function(obj) {
+  # If any sites exceed their cap, remove accrual from randomly
+  # selected arms until sites are at cap 
+  # Represents sites closing during the week
+  site_captotal <- site_sums(obj) - obj@site_cap
+
+  if (any(site_captotal > 0)) {
+    # Loop over sites which are above the cap
+    for (site in which(site_captotal > 0)) {
+      # Vector of instances of populated arms in week's accrual, 
+      # including control, e.g. c(1, 1, 2, 4)
+      population <- unlist(sapply(
+        which(obj@accrual[obj@week, , site] > 0), 
+        function(j) rep(j, obj@accrual[obj@week, j, site])
+      ))
+
+      # Randomly select population instances to remove, leaving
+      # enough to max out the cap
+      if (length(population) > 0) {
+        capped <- do_choose_cap(population, site_captotal[site])
+
+        # Remove those instances
+        for (i in capped) {
+          obj@accrual[obj@week, i, site] <- 
+            as.integer(obj@accrual[obj@week, i, site] - 1)
+        }
+      }
+    }
+  }
+
+  # Don't remove sites from active_sites here, because they may
+  # be reinstated during arm capping
+  
+  return(obj)
+}
+
+
+#' Implement arm cap on week's accrual to experimental arms
+#' @param accrual_obj Object of class `accrual`
+#' @param struct_obj Object of class `trial_structure`
+#' 
+#' @return Modified accrual object with capped week's accrual
+#' 
+#' @importFrom rlang abort
+#' 
+apply_arm_cap <- 
+  S7::new_generic("apply_arm_cap", c("accrual_obj", "struct_obj"))
+S7::method(apply_arm_cap, list(accrual, trial_structure)) <- 
+  function(accrual_obj, struct_obj) {
+    
+    # Get totals for experimental arms (dropping control)
+    arm_sums <- 
+      treat_sums(accrual_obj)[seq_len(length(accrual_obj@phase_changes))]
+
+    # Compare with cap
+    arm_captotal <- arm_sums - accrual_obj@target_arm_size
+
+    # Inactive arms can be at cap but not exceed it
+    if (any(arm_captotal[-accrual_obj@active_arms] > 0)) {
+      rlang::abort(paste(
+        "Inactive arm exceeded cap:", 
+        which(arm_captotal[-accrual_obj@active_arms] > 0)
+      ))
+    }
+
+    ### Change references to active arms
+
+    # Active arm indices which exceed cap
+    active_tocap <- which(arm_captotal > 0)
+
+    if (length(active_tocap) > 0) {
+      for (arm in active_tocap) {
+        # Which sites recruited to that arm this week?
+        population <- as.integer(unlist(sapply(
+          which(accrual_obj@accrual[accrual_obj@week, arm, ] > 0), 
+          function(j) rep(j, accrual_obj@accrual[accrual_obj@week, arm, j])
+        )))
+
+        # Possible accruals to remove
+        if (length(population) > 0) {
+          capped <- do_choose_cap(population, arm_captotal[arm])
+
+          # Remove capped instances
+          for (i in capped) {
+            accrual_obj@accrual[accrual_obj@week, arm, i] <- 
+              as.integer(accrual_obj@accrual[accrual_obj@week, arm, i] - 1)
+          }
+        }
+      }
+    }
+
+    # Record closing week for capped arms
+    capped_arms <- arm_captotal[accrual_obj@active_arms] >= 0
+
+    accrual_obj@phase_changes[accrual_obj@active_arms[capped_arms]] <- 
+      accrual_obj@week
+    
+
+    # Update active_arms
+    accrual_obj@active_arms <- which(arm_captotal < 0)
+    # Also on the trial structure object
+    struct_obj <- remove_treat_arms(struct_obj, which(arm_captotal >= 0))
+
+    # Recheck site caps
+    site_captotal <- site_sums(accrual_obj) - accrual_obj@site_cap
+    
+    # Record closing week for capped sites
+    capped_sites <- site_captotal[accrual_obj@active_sites] >= 0
+    accrual_obj@site_closures[accrual_obj@active_sites[capped_sites]] <-
+      accrual_obj@week
+
+    # Update active sites
+    accrual_obj@active_sites <- which(site_captotal < 0)
+
+    return(list(accrual_obj, struct_obj)) 
+  }
+
+
+
+
+#' Randomise a week's expected accrual amongst the sites, according to 
+#' prevalence.
+#' @param accrual_obj An object of class `accrual`.
+#' @param struct_obj An object of class `trial_structure`.
+#' @param fixed_site_rates TRUE if centre recruitment rates should 
+#' be treated as exact; FALSE if they should be drawn from a gamma
+#' distribution with a mean of the specified rate.
+#' 
+#' @return Matrix of week's accrual by site and recruitment arm.
+#' 
+week_accrue <- S7::new_generic("week_accrue", c("accrual_obj", "struct_obj"))
+S7::method(week_accrue, list(accrual, trial_structure)) <- 
+  function(accrual_obj, struct_obj, fixed_site_rates) {
+
+    # Update the site rates
+    accrual_obj <- set_site_rates(accrual_obj, fixed_site_rates)
+
+    # Initialising (sites * experimental arms)
+    week_mx <- matrix(
+      as.integer(0), 
+      nrow = dim(accrual_obj@accrual)[2], 
+      ncol = dim(accrual_obj@accrual)[3]
+    )
+
+    week_acc <- as.integer(rep(0, dim(accrual_obj@accrual)[3]))
+
+    # Accrual per site is poisson distributed
+    week_acc[accrual_obj@active_sites] <- rpois(
+      n = length(accrual_obj@active_sites), 
+      lambda = accrual_obj@site_rate[accrual_obj@active_sites]
+    )
+
+    # Loop over recruiting sites
+    for (isite in which(week_acc > 0)) {
+
+      # Total probability for each experimental arm for the 
+      # relevant site prevalence set
+      probs <- colSums(struct_obj@experimental_arm_prevalence[
+        , , accrual_obj@site_in_region[isite]
+      ])
+
+
+      # Sample experimental arms according to probabilities
+      # Adding a dummy arm to take unassigned allocations due 
+      # to arm closure, representing reduced site recruitment
+      assigns <- sample(
+        seq_len(length(probs) + 1),
+        prob = c(probs, 1 - sum(probs)),
+        size = week_acc[isite],
+        replace = TRUE
+      )
+      
+      # Total assignments to each open arm plus dummy arm
+      assign_table <- table(assigns)
+      indices <- as.numeric(names(assign_table))
+    
+      # Drop the dummy arm if anything was assigned to it
+      if (max(indices) > length(probs)) {
+        assign_table <- assign_table[-length(assign_table)]
+        indices <- indices[-length(indices)]
+      }
+
+      # Increment week's assignment matrix
+      week_mx[indices, isite] <- 
+        week_mx[indices, isite] + as.vector(assign_table)
+    }
+
+    return(list(accrual_obj, week_mx))
+  }
+
+
+#' Generate accrual for a week and assign it to the 
+#' accrual object. Increments the property "week",
+#' which holds the next week number to accrue.
+#' @param accrual_obj An object of class "accrual"
+#' @param struct_obj An object of class "trial_structure"
+#' @param fixed_site_rates TRUE if expected site rate to be used; FALSE 
+#' draws the site rate from a gamma distribution
+#' 
+#' @return An object of class "accrual"
+#'
+accrue_week <- S7::new_generic("accrue_week", c("accrual_obj", "struct_obj"))
+S7::method(accrue_week, list(accrual, trial_structure)) <- 
+  function(accrual_obj, struct_obj, fixed_site_rates) {
+
+    # Should not get here if there aren't any but
+    if (length(accrual_obj@active_sites) > 0) {
+
+      week_acc_ls <- week_accrue(accrual_obj, struct_obj, fixed_site_rates)
+      accrual_obj <- week_acc_ls[[1]]
+      week_acc <- week_acc_ls[[2]]
+
+      # Assign the week's accrual to the object
+      accrual_obj@accrual[accrual_obj@week, , ] <-
+        week_acc
+
+      # Apply site cap
+      accrual_obj <- apply_site_cap(accrual_obj)
+
+      # Apply arm cap, and then adjust site cap appropriately
+      obj_list <- apply_arm_cap(accrual_obj, struct_obj)
+      accrual_obj <- obj_list[[1]]
+      struct_obj <- obj_list[[2]]
+
+    }
+
+    return(list(accrual_obj, struct_obj))
+  }
+
+
+
+#' Check whether an object is of class "accrual".
+#' 
+#' @param x Object to test
+#' 
+#' @return TRUE or FALSE
+#' 
+#' @importFrom S7 new_generic method class_any
+#' 
+#' @export
+#' 
+is.accrual <- S7::new_generic("is.accrual", "x")
+S7::method(is.accrual, S7::class_any) <- function(x) {
+  inherits(x, "biomkrAccrual::accrual")
+}
+
+