#' 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")

}