plot_helper_functions.R

# This script contains helper functions for plotting data for cov-spectrum

library(lubridate)
library(ggplot2)
library(RColorBrewer)

#' Get an appropriate date scale, depending on the date range of the data
#' @param data The data to plot
#' @param date_colname The column name for the data data
#' @param max_breaks Maximum number of breaks for the scale. Note: will be ignored if number of years > max_breaks.
#' @return A date scale of the type scale_x_date()
get_date_scale <- function(data, date_colname = "date", max_breaks = 10) {
  # Declare break options for the scale in increasing order of coarseness
  break_options <- c("1 day", "1 week", "1 month", "2 months", "1 year") # last is maximum regardless of max_breaks
  break_units <- c("days", "weeks", "months", "months", "years")
  break_labels <- c("%d %b\n%Y", "%d %b\n%Y", "%b\n%Y", "%b\n%Y", "%Y")

  # Increase coarseness of scale if too many breaks
  n_breaks <- max_breaks + 1
  i <- 0
  while (n_breaks > max_breaks & i < length(break_options)) {
    i <- i + 1
    n_breaks <- as.numeric(
      lubridate::as.period(difftime(
        max(data[[date_colname]]), min(data[[date_colname]])
      )),
      break_units[i]
    )
    if (break_options[i] == "2 months") {
      n_breaks <- n_breaks / 2
    }
  }

  # Define scale
  scale <- scale_x_date(date_breaks = break_options[i], date_labels = break_labels[i])
  return(scale)
}

#' Get the appropriate geometry and aesthetic specifications for the plot type
#' @param plot_type One of "line" or "bar"
#' @return A list of geometry (`geom_` functions) and aesthetic specifications
get_bar_vs_line_specs <- function(plot_type) {
  if (plot_type == "line") {
    bar_line_specs <- list(
      "plot_type" = "line",
      "geom_bar_or_line" = geom_line,
      "geom_errbar_or_ribbon" = geom_ribbon,
      "alpha_estimate" = 1,
      "alpha_uncertainty" = 0.5,
      "color" = diverging_colors[1],
      "fill" = diverging_colors[1]
    )
  } else if (plot_type == "bar") {
    bar_line_specs <- list(
      "plot_type" = "bar",
      "geom_bar_or_line" = geom_col,
      "geom_errbar_or_ribbon" = geom_errorbar,
      "alpha_estimate" = 0.5,
      "alpha_uncertainty" = 1,
      "color" = NA,
      "fill" = diverging_colors[1]
    )
  }
  return(bar_line_specs)
}

#' Get a diverging color scale that will always be the same for the same set of values
#' @param locations List of unique vlues
#' @param aesthetics Argument to scale_color_manual(). Character string or vector of character strings listing the name(s) of the aesthetic(s) that this scale works with.
#' @param max_char_label Maximum number of characters for a legend label. Will truncate longer labels with "..."
#' @return A color scale of the type scale_color_manual()
get_color_scale <- function(values, aesthetics = "color", max_char_label = 15) {
  n_values <- length(values)
  if (n_values == 0) {
    stop("No values specified for color scale")
  } else if (n_values <= length(diverging_colors)) {
    colors <- diverging_colors[1:n_values]
  } else {
    colors <- get_unlimited_colors(n = n_values)
  }
  names(colors) <- sort(values)
  labels <- truncate_labels(labels = sort(values), max_char_label = max_char_label)
  color_scale <- scale_color_manual(values = colors, labels = labels, aesthetics = aesthetics)
  return(color_scale)
}

#' Truncate long labels with "..."
#' @param labels A vector of strings. Labels for a plot.
#' @param max_char_label The longest allowable label.
#' @return A vector of strings, some possibly truncated with "..."
truncate_labels <- function(labels, max_char_label) {
  for (i in seq_len(length(labels))) {
    label <- labels[i]
    if (nchar(label) > max_char_label) {
      label <- paste0(c(strsplit(label, "")[[1]][1:(max_char_label - 3)], "..."), collapse = "")
      labels[i] <- label
    }
  }
  return(labels)
}

#' Get a geometry to show the data
#' @param bar_line_specs Geometry and aesthetic specifications for the plot type, generated by get_bar_vs_line_specs()
#' @return A ggplot geometry to show the data (geom_line or geom_bar)
get_bar_line_geom <- function(bar_line_specs) {
  suppressWarnings({
    bar_line_geom <- bar_line_specs$geom_bar_or_line(
      alpha = bar_line_specs$alpha_estimate,
      size = line_bar_size,  # defined in R/plot_shared_elements.R
      color = bar_line_specs$color,
      fill = bar_line_specs$fill
    )
  })  # suppress warning that fill is not applicable for line geom
  return(bar_line_geom)
}

#' Get a geometry to show uncertainty
#' @param data The data to plot (used to see if uncertainty columns present)
#' @param bar_line_specs Geometry and aesthetic specifications for the plot type, generated by get_bar_vs_line_specs()
#' @param fill_var A string. Variable name to set colors by if color aesthetic being used.
#' @param fill_color A string. Color name if color aesthetic not being used.
#' @param low_pattern A string. Pattern to uniquely identify column for lower uncertainty bound by.
#' @param high_pattern A string. Pattern to uniquely identify column for high uncertainty bound by.
#' @return A ggplot geometry to show uncertainty (may be a blank geometry for no uncertainty)
get_uncertainty_geom <- function(data, bar_line_specs, fill_var = NULL, fill_color = "black", low_pattern = "CILow", high_pattern = "CIHigh") {
  low_var <- colnames(data)[grepl(pattern = low_pattern, x = colnames(data))]
  high_var <- colnames(data)[grepl(pattern = high_pattern, x = colnames(data))]
  if (length(low_var) == 1 & length(high_var) == 1) {
    if (!is.null(fill_var)) {
      uncertainty_geom <- bar_line_specs$geom_errbar_or_ribbon(
        aes(ymin = .data[[low_var]], ymax = .data[[high_var]], fill = .data[[fill_var]]),
        alpha = bar_line_specs$alpha_uncertainty,
        linetype = case_when(bar_line_specs$plot_type == "line" ~ 0, T ~ 1)
      )
    } else {
      suppressWarnings({
        uncertainty_geom <- bar_line_specs$geom_errbar_or_ribbon(
          aes(ymin = .data[[low_var]], ymax = .data[[high_var]]),
          alpha = bar_line_specs$alpha_uncertainty,
          fill = fill_color,
          color = fill_color,
          linetype = case_when(bar_line_specs$plot_type == "line" ~ 0, T ~ 1)
        )
      })  # suppress warning that one of fill or color is not applicable, depending on whether the geom is ribbon or errbar
    }
  } else {
    uncertainty_geom <- geom_blank()
  }
  return(uncertainty_geom)
}

#' Get plot title and subtitle depending on plot type and metadata
#' @param config Configuration information.
#' @param metadata The request metadata.
#' @Return A named list. The title and subtitle for the plot.
get_titles <- function(config, metadata) {
  # If plot is a collection, give collection info
  if ("collection" %in% names(metadata)) {
    collection_title <- truncate_labels(labels = metadata$collection$title, max_char_label = 70)
    subtitle <- paste0("Collection #", metadata$collection$id, " '", collection_title, "'\nmaintained by ", metadata$collection$maintainer)
  } else {
    subtitle <- NULL
  }

  if ("variant" %in% names(metadata)) {
    prefix <- truncate_labels(labels = metadata$variant, max_char_label = 35)
  } else {
    prefix <- ""
  }

  if ("location" %in% names(metadata)) {
    suffix <- paste0(" in ", metadata$location)
  } else {
    suffix <- ""
  }

  if (config$plotName == "estimated-cases") {
    title <- paste0("Estimated cases of ", prefix, suffix)
  } else if (startsWith(x = config$plotName, prefix = "sequences-over-time")) {
    if (prefix == "") {
      title <- paste0("Sequences over time", suffix)
    } else {
      title <- paste0(prefix, " sequences over time", suffix)
    }
  } else {
    title <- NULL
  }

  return(list("title" = title, "subtitle" = subtitle))
}