
#' @title RColorBrewer Color Ramp for EpiModel Plots
#'
#' @description Returns a vector of colors consistent with a high-brightness set
#'              of colors from an `RColorBrewer` palette.
#'
#' @param plt `RColorBrewer` palette from [`RColorBrewer::brewer.pal`].
#' @param n Number of colors to return.
#' @param delete.lights If TRUE, delete the lightest colors from the color
#'        palette; this helps with plotting in many high-contrast palettes.
#'
#' @details
#' `RColorBrewer` provides easy access to helpful color palettes, but the
#' built-in palettes are limited to the set of colors in the existing palette.
#' This function expands the palette size to any number of colors by filling
#' in the gaps. Also, colors within the "div" and "seq" set of palettes whose
#' colors are very light (close to white) are deleted by default for better
#' visualization of plots.
#'
#' @return
#' A vector of length equal to `n` with a range of color values consistent
#' with an RColorBrewer color palette.
#'
#' @seealso [RColorBrewer::RColorBrewer]
#' @keywords colorUtils internal
#' @export
#'
#' @examples
#' # Shows a 100-color ramp for 4 RColorBrewer palettes
#' par(mfrow = c(2, 2), mar=c(1, 1, 2, 1))
#' pals <- c("Spectral", "Greys", "Blues", "Set1")
#' for (i in seq_along(pals)) {
#'  plot(1:100, 1:100, type = "n", axes = FALSE, main = pals[i])
#'  abline(v = 1:100, lwd = 6, col = brewer_ramp(100, pals[i]))
#' }
#'
brewer_ramp <- function(n, plt, delete.lights = TRUE) {

  if (n < 1) {
    stop("n must be a positive integer")
  }

  bpi <- brewer.pal.info
  if (!(plt %in% row.names(bpi))) {
    stop("plt must match an RColorBrewer palette name. See
         RColorBrewer::brewer.pal.info",
         .call = FALSE)
  }

  pltmax <- bpi[row.names(bpi) == plt, ]$maxcolors
  pltcat <- bpi[row.names(bpi) == plt, ]$category

  if (pltcat == "div") {
    if (delete.lights == TRUE) {
      colors <- brewer.pal(pltmax, plt)[-c(4:7)]
    } else {
      colors <- brewer.pal(pltmax, plt)
    }
  }
  if (pltcat == "qual") {
    colors <- brewer.pal(pltmax, plt)
  }
  if (pltcat == "seq") {
    if (delete.lights == TRUE) {
      colors <- rev(brewer.pal(pltmax, plt)[-c(1:3)])
    } else {
      colors <- rev(brewer.pal(pltmax, plt))
    }
  }
  if (plt == "Set1") {
    colors <- brewer.pal(9, "Set1")[-6]
  }

  pal <- colorRampPalette(colors)

  return(pal(n))
}

#' @title Delete Elements from Attribute List
#'
#' @description Deletes elements from the main attribute list.
#'
#' @param dat Either the `netsim_dat` class object passed through
#'            `netsim` simulations, or the main data object passed through
#'            `icm` simulations.
#' @param ids ID numbers to delete from the list.
#'
#' @inherit recovery.net return
#'
#' @export
#' @keywords internal
delete_attr <- function(dat, ids) {
  remove_node_attr(dat, ids)
}


#' @title Stable Sampling Function
#'
#' @description Provides a sampling function useful for dynamic simulations, in
#'              which the length of the input vector may be multiple lengths and
#'              the size of the sample may be 0.
#'
#' @param x Either a vector of one or more elements from which to choose, or a
#'        positive integer.
#' @param size Non-negative integer giving the number of items to choose.
#' @param replace Should sampling be with replacement?
#' @param prob Vector of probability weights for obtaining the elements of the
#'        vector being sampled.
#'
#' @return A vector containing the sampled value(s).
#'
#' @export
#' @keywords internal
ssample <- function(x, size, replace = FALSE, prob = NULL) {

  if (length(x) > 1) {
    return(sample(x, size, replace, prob))
  }

  if (length(x) == 1 && size > 0) {
    return(x)
  }

  if (length(x) == 1 && size == 0) {
    return(NULL)
  }

}


#' @title Add New Epidemiology Variables
#'
#' @description Inspired by `dplyr::mutate`, `mutate_epi` adds new
#'              variables to the epidemiological and related variables within
#'              simulated model objects of any class in `EpiModel`.
#'
#' @param x An `EpiModel` object of class `dcm`, `icm`, or
#'        `netsim`.
#' @param ... Name-value pairs of expressions (see examples below).
#'
#' @return The updated `EpiModel` object of class `dcm`, `icm`,
#'         or `netsim`.
#'
#' @export
#'
#' @examples
#' # DCM example
#' param <- param.dcm(inf.prob = 0.2, act.rate = 0.25)
#' init <- init.dcm(s.num = 500, i.num = 1)
#' control <- control.dcm(type = "SI", nsteps = 500)
#' mod1 <- dcm(param, init, control)
#' mod1 <- mutate_epi(mod1, prev = i.num/num)
#' plot(mod1, y = "prev")
#'
#' \donttest{
#' # Network model example
#' nw <- network_initialize(n = 100)
#' nw <- set_vertex_attribute(nw, "group", rep(1:2, each = 50))
#' formation <- ~edges
#' target.stats <- 50
#' coef.diss <- dissolution_coefs(dissolution = ~offset(edges), duration = 20)
#' est1 <- netest(nw, formation, target.stats, coef.diss, verbose = FALSE)
#'
#' param <- param.net(inf.prob = 0.3, inf.prob.g2 = 0.15)
#' init <- init.net(i.num = 1, i.num.g2 = 0)
#' control <- control.net(type = "SI", nsteps = 10, nsims = 3,
#'                        verbose = FALSE)
#' mod1 <- netsim(est1, param, init, control)
#' mod1
#'
#' # Add the prevalences to the dataset
#' mod1 <- mutate_epi(mod1, i.prev = i.num / num,
#'                          i.prev.g2 = i.num.g2 / num.g2)
#' plot(mod1, y = c("i.prev", "i.prev.g2"), qnts = 0.5, legend = TRUE)
#'
#' # Add incidence rate per 100 person years (assume time step = 1 week)
#' mod1 <- mutate_epi(mod1, ir100 = 5200*(si.flow + si.flow.g2) /
#'                                       (s.num + s.num.g2))
#' as.data.frame(mod1)
#' as.data.frame(mod1, out = "mean")
#' }
#'
mutate_epi <- function(x, ...) {

  dt <- lazy_dots(...)
  ndat <- lazy_eval(dt, x$epi)

  not.df <- which(sapply(ndat, class) != "data.frame")
  if (length(not.df) > 0) {
    existing_names <- names(x$epi[[1]])
    ncols <- length(existing_names)
    for (jj in not.df) {
      val <- rep(ndat[jj][[1]], length.out = x$control$nsteps)
      ndat[jj][[1]] <- data.frame(replicate(ncols, val))
      names(ndat[[jj]]) <- existing_names
    }
  }

  x$epi <- c(x$epi, ndat)
  return(x)

}

#' @title Apportion Using the Largest Remainder Method
#'
#' @description Apportions a vector of values given a specified frequency
#'              distribution of those values such that the length of the output
#'              is robust to rounding and other instabilities.
#'
#' @param vector.length Length for the output vector.
#' @param values Values for the output vector.
#' @param proportions Proportion distribution with one number for each value.
#'        This must sum to 1.
#' @param shuffled If `TRUE`, randomly shuffle the order of the vector.
#'
#' @return A vector of length `vector.length` containing the apportioned
#'         values from `values`.
#'
#' @export
#'
#' @examples
#' \dontrun{
#' ## Example 1: Without rounding
#' apportioned_vec_1 <- apportion_lr(4, c(1, 2, 3, 4, 5),
#'                                      c(0.25, 0, 0.25, 0.25, 0.25))
#'
#' ## Example 2: With rounding
#' apportioned_vec_2 <- apportion_lr(5, c(1, 2, 3, 4, 5),
#'                                      c(0.21, 0, 0.29, 0.25, 0.25))
#' }
#'
apportion_lr <- function(vector.length, values,
                         proportions, shuffled = FALSE) {

  if (vector.length != round(vector.length)) {
    stop("argument vector.length must be a positive integer")
  }
  if (vector.length <= 0) {
    stop("argument vector.length must be a positive integer")
  }
  if (is.vector(values) == FALSE) {
    stop("argument values must be a vector")
  }
  if (!(length(proportions) == length(values) && round(sum(proportions),
                                                       10) == 1) &&
        (!(length(proportions) == length(values) - 1 && round(sum(proportions),
                                                              10) <= 1 &&
             round(sum(proportions), 10) >= 0))) {
    stop("error in proportions length or proportions sum")
  }

  if (length(proportions) == length(values) - 1) {
    proportions <- c(proportions, 1 - round(sum(proportions), 10))
  }
  result <- rep(NA, vector.length)
  exp.nums <- proportions * vector.length
  counts <- floor(exp.nums)
  remainders <- exp.nums - counts
  leftovers <- vector.length - sum(counts)
  if (leftovers > 0) {
    additions <- order(remainders, decreasing = TRUE)[1:leftovers]
    counts[additions]   <- counts[additions] + 1
  }
  result <- rep(values, counts)
  if (shuffled == TRUE) {
    result <- sample(result, length(result))
  }

  return(result)
}