R/facet_color.R

In deeptime: Plotting Tools for Anyone Working in Deep Time

#' Lay out panels in a grid with colored strips
#'
#' `facet_grid_color` behaves similarly to [ggplot2::facet_grid()] in that it
#' forms a matrix of panels defined by row and column faceting variables. The
#' main difference is that it also allows the user to specify the background and
#' label colors of the individual facet strips using the `colors` and
#' `lab_colors` arguments. If you have only one variable with many levels, try
#' [facet_wrap_color()].
#'
#' `facet_grid_geo(...)` is an alias of `facet_grid_color()` with the default of
#' `colors` set to `stages`.
#'
#' @param colors Specifies which colors to use to replace the strip backgrounds.
#'   Either A) a function that returns a color for a given strip label, B) the
#'   character name of a function that does the same, C) a named character
#'   vector with names matching strip labels and values indicating the desired
#'   colors, or D) a data.frame representing a lookup table with columns named
#'   "name" (matching strip labels) and "color" (indicating desired colors). If
#'   the function returns `NA`, the default background color will be used.
#' @param lab_colors Specifies which colors to use for the strip labels. Either
#'   A) a function that returns a color for a given strip label, B) the
#'   character name of a function that does the same, C) a named character
#'   vector with names matching strip labels and values indicating the desired
#'   colors, D) a data.frame representing a lookup table with columns named
#'   "name" (matching strip labels) and "lab_color" (indicating desired colors),
#'   or E) "auto" (the default), which set the labels to black or white,
#'   whichever has better contrast with the background color, based on
#'   [recommendations by the International Telecommunication Union](https://www.itu.int/rec/R-REC-BT.601-7-201103-I/en).
#'   If the function returns `NA`, the default label color will be used.
#' @inheritParams ggplot2::facet_grid
#' @importFrom ggplot2 ggproto FacetGrid ggproto_parent
#' @importFrom rlang arg_match0 is_function
#' @family faceting functions
#' @export
#'
#' @examples
#' library(ggplot2)
#' df <- data.frame(x = 1:10, y = 1:10, period = c("Permian", "Triassic"))
#' ggplot(df) +
#'   geom_point(aes(x, y)) +
#'   facet_grid_color(cols = vars(period), colors = periods)
facet_grid_color <- function(colors, rows = NULL, cols = NULL,
                             scales = "fixed", space = "fixed", shrink = TRUE,
                             labeller = "label_value", lab_colors = "auto",
                             as.table = TRUE, switch = NULL,
                             drop = TRUE, margins = FALSE,
                             axes = "margins", axis.labels = "all") {
  colors <- convert_colors(colors)
  lab_colors <- convert_lab_colors(lab_colors)

  # function and arguments copied from ggplot 3.5.0

  # Should become a warning in a future release
  if (is.logical(cols)) {
    margins <- cols
    cols <- NULL
  }

  scales <- arg_match0(scales %||% "fixed", c("fixed", "free_x",
                                              "free_y", "free"))
  free <- list(
    x = any(scales %in% c("free_x", "free")),
    y = any(scales %in% c("free_y", "free"))
  )

  space <- arg_match0(space %||% "fixed", c("fixed", "free_x",
                                            "free_y", "free"))
  space_free <- list(
    x = any(space %in% c("free_x", "free")),
    y = any(space %in% c("free_y", "free"))
  )

  draw_axes <- arg_match0(axes, c("margins", "all_x", "all_y", "all"))
  draw_axes <- list(
    x = any(draw_axes %in% c("all_x", "all")),
    y = any(draw_axes %in% c("all_y", "all"))
  )

  # Omitting labels is special-cased internally, so even when no internal axes
  # are to be drawn, register as labelled.
  axis_labels <- arg_match0(axis.labels, c("margins", "all_x",
                                           "all_y", "all"))
  axis_labels <- list(
    x = !draw_axes$x || any(axis_labels %in% c("all_x", "all")),
    y = !draw_axes$y || any(axis_labels %in% c("all_y", "all"))
  )

  if (!is.null(switch)) {
    arg_match0(switch, c("both", "x", "y"))
  }

  facets_list <- grid_as_facets_list(rows, cols)

  # Check for deprecated labellers
  labeller <- match.fun(labeller)

  params <- list(rows = facets_list$rows, cols = facets_list$cols,
                 margins = margins, free = free, space_free = space_free,
                 labeller = labeller, colors = colors, lab_colors = lab_colors,
                 as.table = as.table, switch = switch, drop = drop,
                 draw_axes = draw_axes, axis_labels = axis_labels)

  ggproto(NULL, FacetGridColor,
          shrink = shrink,
          params = params
  )
}

#' @export
#' @rdname facet_grid_color
facet_grid_geo <- function(colors = stages, rows = NULL, cols = NULL,
                           scales = "fixed", space = "fixed", shrink = TRUE,
                           labeller = "label_value", lab_colors = "auto",
                           as.table = TRUE, switch = NULL,
                           drop = TRUE, margins = FALSE,
                           axes = "margins", axis.labels = "all") {
  facet_grid_color(colors = colors, rows = rows, cols = cols,
                   scales = scales, space = space, shrink = shrink,
                   labeller = labeller, lab_colors = lab_colors,
                   as.table = as.table, switch = switch,
                   drop = drop, margins = margins,
                   axes = axes, axis.labels = axis.labels)
}

grid_as_facets_list <- function(...) {
  asNamespace("ggplot2")$grid_as_facets_list(...)
}

#' @importFrom rlang is_function
convert_colors <- function(colors) {
  # convert colors to a function
  if (!is_function(colors)) {
    if (is.character(colors) && !is.null(names(colors))) {
      name <- names(colors)
      color <- unname(colors)
      colors <- function(x) {
        if (x %in% name) color[which(x == name)[1]] else NA
      }
    } else if (is.character(colors) && length(colors) == 1) {
      colors <- match.fun(colors)
    } else if (is.data.frame(colors)) {
      if (all(c("name", "color") %in% names(colors))) {
        name <- colors$name
        color <- colors$color
        colors <- function(x) {
          if (x %in% name) color[which(x == name)[1]] else NA
        }
      } else {
        cli::cli_abort("If using a data.frame for `colors`, the data.frame must
                       have columns named 'name' and 'color'.")
      }
    } else {
      cli::cli_abort("Invalid type for `colors`; only functions, function names,
                     named character vectors, and data.frames are allowed.")
    }
  }
  return(colors)
}

#' @importFrom rlang is_function
convert_lab_colors <- function(lab_colors) {
  # convert colors to a function
  if (!is_function(lab_colors)) {
    if (is.character(lab_colors) && lab_colors == "auto") {
        return(lab_colors)
    } else if (is.character(lab_colors) && !is.null(names(lab_colors))) {
      name <- names(lab_colors)
      color <- unname(lab_colors)
      lab_colors <- function(x) {
        if (x %in% name) color[which(x == name)[1]] else NA
      }
    } else if (is.character(lab_colors) && length(lab_colors) == 1) {
      lab_colors <- match.fun(lab_colors)
    } else if (is.data.frame(lab_colors)) {
      if (all(c("name", "lab_color") %in% names(lab_colors))) {
        name <- lab_colors$name
        color <- lab_colors$lab_color
        lab_colors <- function(x) {
          if (x %in% name) color[which(x == name)[1]] else NA
        }
      } else {
        cli::cli_abort("If using a data.frame for `lab_colors`, the data.frame
                       must have columns named 'name' and 'lab_color'.")
      }
    } else {
      cli::cli_abort("Invalid type for `lab_colors`; only functions, function
                     names, named character vectors, data.frames, and \"auto\"
                     are allowed.")
    }
  }
  return(lab_colors)
}

#' @rdname facet_grid_color
#' @format NULL
#' @usage NULL
#' @export
#' @importFrom ggplot2 ggproto FacetGrid ggproto_parent
FacetGridColor <- ggproto("FacetGridColor", FacetGrid,
  draw_panels = function(panels, layout, x_scales, y_scales, ranges, coord,
                         data, theme, params, self) {
    panel_table <-
      ggproto_parent(FacetGrid, self)$draw_panels(panels, layout,
                                                  x_scales, y_scales, ranges,
                                                  coord, data, theme, params)
    strips <- grep("strip", panel_table$layout$name)
    for (i in strips) {
      label <-
        panel_table$grobs[[i]]$grobs[[1]]$children[[2]]$children[[1]]$label
      fill <- tryCatch(params$colors(label), error = function(e) NA)
      if (!is.na(fill)) {
        panel_table$grobs[[i]]$grobs[[1]]$children[[1]]$gp$fill <- fill
        if (!is_function(params$lab_colors) && params$lab_colors == "auto") {
          panel_table$grobs[[i]]$grobs[[1]]$children[[2]]$children[[1]]$
            gp$col <- white_or_black(fill)
        }
      }
      if (is_function(params$lab_colors)) {
        color <- tryCatch(params$lab_colors(label), error = function(e) NA)
        if (!is.na(color)) {
          panel_table$grobs[[i]]$grobs[[1]]$children[[2]]$children[[1]]$
            gp$col <- color
        }
      }
    }
    panel_table
  }
)

#' Wrap a 1d ribbon of panels into 2d with colored strips
#'
#' `facet_wrap_color` behaves similarly to [ggplot2::facet_wrap()] in that it
#' wraps a 1d sequence of panels into 2d. The main difference is that it also
#' allows the user to specify the background and label colors of the individual
#' facet strips using the `colors` and `lab_colors` arguments. This is generally
#' a better use of screen space than [facet_grid_color()] because most displays
#' are roughly rectangular.
#'
#' `facet_wrap_geo(...)` is an alias of `facet_wrap_color()` with the default of
#' `colors` set to `stages`.
#'
#' @param colors Specifies which colors to use to replace the strip backgrounds.
#'   Either A) a function that returns a color for a given strip label, B) the
#'   character name of a function that does the same, C) a named character
#'   vector with names matching strip labels and values indicating the desired
#'   colors, or D) a data.frame representing a lookup table with columns named
#'   "name" (matching strip labels) and "color" (indicating desired colors). If
#'   the function returns `NA`, the default background color will be used.
#' @param lab_colors Specifies which colors to use for the strip labels. Either
#'   A) a function that returns a color for a given strip label, B) the
#'   character name of a function that does the same, C) a named character
#'   vector with names matching strip labels and values indicating the desired
#'   colors, D) a data.frame representing a lookup table with columns named
#'   "name" (matching strip labels) and "lab_color" (indicating desired colors),
#'   or E) "auto" (the default), which set the labels to black or white,
#'   whichever has better contrast with the background color, based on
#'   [recommendations by the International Telecommunication Union](https://www.itu.int/rec/R-REC-BT.601-7-201103-I/en).
#'   If the function returns `NA`, the default label color will be used.
#' @inheritParams ggplot2::facet_wrap
#' @importFrom ggplot2 ggproto FacetWrap ggproto_parent
#' @importFrom rlang arg_match0
#' @importFrom utils packageVersion
#' @family faceting functions
#' @export
#'
#' @examples
#' library(ggplot2)
#' df <- data.frame(x = 1:10, y = 1:10, period = c("Permian", "Triassic"))
#' ggplot(df) +
#'   geom_point(aes(x, y)) +
#'   facet_wrap_color(vars(period), colors = periods)
facet_wrap_color <- function(facets, colors, nrow = NULL, ncol = NULL,
                             scales = "fixed", shrink = TRUE,
                             labeller = "label_value", lab_colors = "auto",
                             as.table = TRUE, drop = TRUE,
                             dir = "h", strip.position = "top",
                             axes = "margins", axis.labels = "all") {
  colors <- convert_colors(colors)
  lab_colors <- convert_lab_colors(lab_colors)

  # function and arguments copied from ggplot 3.5.0
  scales <- arg_match0(scales %||% "fixed", c("fixed", "free_x",
                                              "free_y", "free"))
  scales <- arg_match0(scales %||% "fixed", c("fixed", "free_x", "free_y", "free"))
  if (packageVersion("ggplot2") > "3.5.2") {
    dir <- arg_match0(dir, c("h", "v", "lt", "tl", "lb", "bl", "rt", "tr", "rb", "br"))
    if (nchar(dir) == 1) {
      dir <- base::switch(
        dir,
        h = if (as.table) "lt" else "lb",
        v = if (as.table) "tl" else "tr"
      )
    }
  } else {
    dir <- arg_match0(dir, c("h", "v"))
  }

  free <- list(
    x = any(scales %in% c("free_x", "free")),
    y = any(scales %in% c("free_y", "free"))
  )

  # If scales are free, always draw the axes
  draw_axes <- arg_match0(axes, c("margins", "all_x", "all_y", "all"))
  draw_axes <- list(
    x = free$x || any(draw_axes %in% c("all_x", "all")),
    y = free$y || any(draw_axes %in% c("all_y", "all"))
  )

  # Omitting labels is special-cased internally, so only omit labels if
  # scales are not free and the axis is to be drawn
  axis_labels <- arg_match0(axis.labels, c("margins", "all_x", "all_y", "all"))
  axis_labels <- list(
    x = free$x || !draw_axes$x || any(axis_labels %in% c("all_x", "all")),
    y = free$y || !draw_axes$y || any(axis_labels %in% c("all_y", "all"))
  )

  # Check for deprecated labellers
  labeller <- match.fun(labeller)

  # Flatten all facets dimensions into a single one
  facets <- ggplot2::facet_wrap(facets = facets)$params$facets

  strip.position <- arg_match0(strip.position, c("top", "bottom",
                                                 "left", "right"))

  check_number_whole(ncol, allow_null = TRUE, min = 1)
  check_number_whole(nrow, allow_null = TRUE, min = 1)

  if (identical(dir, "v")) {
    # swap
    tmp <- ncol
    ncol <- nrow
    nrow <- tmp
  }

  params <- list(
    facets = facets,
    free = free,
    as.table = as.table,
    strip.position = strip.position,
    drop = drop,
    ncol = ncol,
    nrow = nrow,
    labeller = labeller,
    colors = colors,
    lab_colors = lab_colors,
    dir = dir,
    draw_axes = draw_axes,
    axis_labels = axis_labels
  )

  ggproto(NULL, FacetWrapColor,
          shrink = shrink,
          params = params
  )
}

#' @export
#' @rdname facet_wrap_color
facet_wrap_geo <- function(facets, colors = stages, nrow = NULL, ncol = NULL,
                           scales = "fixed", shrink = TRUE,
                           labeller = "label_value", lab_colors = "auto",
                           as.table = TRUE, drop = TRUE,
                           dir = "h", strip.position = "top",
                           axes = "margins", axis.labels = "all") {
  facet_wrap_color(facets = facets, colors = colors, nrow = nrow,
                   ncol = ncol, scales = scales, shrink = shrink,
                   labeller = labeller, lab_colors = lab_colors,
                   as.table = as.table, drop = drop,
                   dir = dir, strip.position = strip.position,
                   axes = axes, axis.labels = axis.labels)
}

check_number_whole <- function(...) {
  asNamespace("rlang")$check_number_whole(...)
}

#' @rdname facet_wrap_color
#' @format NULL
#' @usage NULL
#' @export
#' @importFrom ggplot2 ggproto FacetWrap ggproto_parent
#' @importFrom rlang is_function
FacetWrapColor <- ggproto("FacetWrapColor", FacetWrap,
  draw_panels = function(panels, layout, x_scales, y_scales, ranges, coord,
                         data, theme, params, self) {
    panel_table <-
      ggproto_parent(FacetWrap, self)$draw_panels(panels, layout,
                                                  x_scales, y_scales, ranges,
                                                  coord, data, theme, params)
    strips <- grep("strip", panel_table$layout$name)
    for (i in strips) {
      label <-
        panel_table$grobs[[i]]$grobs[[1]]$children[[2]]$children[[1]]$label
      fill <- tryCatch(params$colors(label), error = function(e) NA)
      if (!is.na(fill)) {
        panel_table$grobs[[i]]$grobs[[1]]$children[[1]]$gp$fill <- fill
        if (!is_function(params$lab_colors) && params$lab_colors == "auto") {
          panel_table$grobs[[i]]$grobs[[1]]$children[[2]]$children[[1]]$
            gp$col <- white_or_black(fill)
        }
      }
      if (is_function(params$lab_colors)) {
        color <- tryCatch(params$lab_colors(label), error = function(e) NA)
        if (!is.na(color)) {
          panel_table$grobs[[i]]$grobs[[1]]$children[[2]]$children[[1]]$
            gp$col <- color
        }
      }
    }
    panel_table
  }
)

new_grid_facets <- function(...) {
  asNamespace("ggh4x")$new_grid_facets(...)
}

assert_strip <- function(...) {
  asNamespace("ggh4x")$assert_strip(...)
}

#' Layout panels in a grid with nested colored strips
#'
#' `facet_nested_color` behaves similarly to [ggh4x::facet_nested()] in that it
#' forms a matrix of panels defined by row and column faceting variables and
#' nests grouped facets. The main difference is that it also allows the user to
#' specify the background and label colors of the individual facet strips using
#' the `colors` and `lab_colors` arguments.
#'
#' `facet_nested_geo(...)` is an alias of `facet_nested_color()` with the
#' default of `colors` set to `rbind(periods, stages)`.
#'
#' @inheritParams ggh4x::facet_nested
#' @inheritParams facet_grid_color
#' @inherit ggh4x::facet_nested details
#' @importFrom ggplot2 element_line element_blank
#' @importFrom ggh4x strip_nested
#' @importFrom grid unit
#' @family faceting functions
#' @export
#' @examples
#' library(ggplot2)
#' df <- data.frame(x = 1:10, y = 1:10,
#'                  period = factor(c("Permian", "Triassic", "Jurassic",
#'                                    "Cretaceous", "Paleogene"),
#'                                  levels = c("Permian", "Triassic",
#'                                             "Jurassic", "Cretaceous",
#'                                             "Paleogene")),
#'                  era = factor(c("Paleozoic", "Mesozoic", "Mesozoic",
#'                                 "Mesozoic", "Cenozoic"),
#'                                 levels = c("Paleozoic", "Mesozoic",
#'                                            "Cenozoic")))
#' ggplot(df, aes(x, y)) +
#'   geom_point() +
#'   facet_nested_color(~ era + period, colors = rbind(periods, eras))
facet_nested_color <- function(colors, rows = NULL, cols = NULL,
                               scales = "fixed", space = "fixed",
                               axes = "margins", remove_labels = "none",
                               independent = "none", shrink = TRUE,
                               labeller = "label_value", lab_colors = "auto",
                               as.table = TRUE, switch = NULL, drop = TRUE,
                               margins = FALSE,
                               nest_line = element_line(inherit.blank = TRUE),
                               solo_line = FALSE, resect = unit(0, "mm"),
                               render_empty = TRUE,
                               strip = strip_nested(), bleed = NULL) {
  colors <- convert_colors(colors)
  lab_colors <- convert_lab_colors(lab_colors)

  # copied from ggh4x v0.2.8
  strip <- assert_strip(strip)
  if (!is.null(bleed)) {
    lifecycle::deprecate_warn(
      when = "0.2.0",
      what = "facet_nested(bleed)",
      details = paste0("The `bleed` argument should be set in the ",
                       "`strip_nested()` function instead.")
    )
    strip$params$bleed <- isTRUE(bleed)
  }
  # Convert logical to elements for backward compatibility
  if (isTRUE(nest_line)) {
    nest_line <- element_line()
  }
  if (isFALSE(nest_line)) {
    nest_line <- element_blank()
  }
  if (!inherits(nest_line, c("element_line", "element_blank"))) {
    cli::cli_abort(
      "The {.arg nest_line} argument must be {.cls element_blank} or inherit \\
      from {.cls element_line}."
    )
  }

  params <- list(nest_line = nest_line, solo_line = isTRUE(solo_line),
                 resect = resect, colors = colors, lab_colors = lab_colors)

  new_grid_facets(
    rows, cols,
    scales, space, axes, remove_labels, independent,
    shrink, labeller, as.table, switch,
    drop, margins, render_empty, strip,
    params = params,
    super = FacetNestedColor
  )
}

#' @export
#' @rdname facet_nested_color
facet_nested_geo <- function(colors = rbind(periods, stages),
                             rows = NULL, cols = NULL,
                             scales = "fixed", space = "fixed",
                             axes = "margins", remove_labels = "none",
                             independent = "none", shrink = TRUE,
                             labeller = "label_value", lab_colors = "auto",
                             as.table = TRUE, switch = NULL, drop = TRUE,
                             margins = FALSE,
                             nest_line = element_line(inherit.blank = TRUE),
                             solo_line = FALSE, resect = unit(0, "mm"),
                             render_empty = TRUE,
                             strip = strip_nested(), bleed = NULL) {
  facet_nested_color(colors = colors, rows = rows, cols = cols,
                     scales = scales, space = space, axes = axes,
                     remove_labels = remove_labels,
                     independent = independent, shrink = shrink,
                     labeller = labeller, lab_colors = lab_colors,
                     as.table = as.table, switch = switch, drop = drop,
                     margins = margins, nest_line = nest_line,
                     solo_line = solo_line, resect = resect,
                     render_empty = render_empty, strip = strip,
                     bleed = bleed)
}

#' @rdname facet_nested_color
#' @format NULL
#' @usage NULL
#' @export
#' @importFrom ggplot2 ggproto ggproto_parent
#' @importFrom ggh4x FacetNested
#' @importFrom rlang is_function
FacetNestedColor <- ggproto("FacetNestedColor", FacetNested,
  draw_panels = function(panels, layout, x_scales, y_scales, ranges, coord,
                         data, theme, params, self) {
    panel_table <-
      ggproto_parent(FacetNested, self)$draw_panels(panels, layout,
                                                    x_scales, y_scales, ranges,
                                                    coord, data, theme, params)
    strips <- grep("strip", panel_table$layout$name)
    for (i in strips) {
      label <-
        panel_table$grobs[[i]]$grobs[[1]]$children[[2]]$children[[1]]$label
      fill <- tryCatch(params$colors(label), error = function(e) NA)
      if (!is.na(fill)) {
        panel_table$grobs[[i]]$grobs[[1]]$children[[1]]$gp$fill <- fill
        if (!is_function(params$lab_colors) && params$lab_colors == "auto") {
          panel_table$grobs[[i]]$grobs[[1]]$children[[2]]$children[[1]]$
            gp$col <- white_or_black(fill)
        }
      }
      if (is_function(params$lab_colors)) {
        color <- tryCatch(params$lab_colors(label), error = function(e) NA)
        if (!is.na(color)) {
          panel_table$grobs[[i]]$grobs[[1]]$children[[2]]$children[[1]]$
            gp$col <- color
        }
      }
    }
    panel_table
  }
)

new_wrap_facets <- function(...) {
  asNamespace("ggh4x")$new_wrap_facets(...)
}

#' Ribbon of panels with nested colored strips
#'
#' `facet_nested_wrap_color` behaves similarly to [ggh4x::facet_nested_wrap()]
#' in that it wraps a sequence of panels onto a two-dimensional layout, and
#' nests grouped facets where possible.. The main difference is that it also
#' allows the user to specify the background and label colors of the individual
#' facet strips using the `colors` and `lab_colors` arguments.
#'
#' `facet_nested_wrap_geo(...)` is an alias of `facet_nested_wrap_color()` with the
#' default of `colors` set to `rbind(periods, stages)`.
#'
#' @inheritParams ggh4x::facet_nested_wrap
#' @inheritParams facet_wrap_color
#' @inherit ggh4x::facet_nested_wrap details
#' @importFrom ggplot2 element_line element_blank
#' @importFrom ggh4x strip_nested
#' @importFrom grid unit
#' @family faceting functions
#' @export
#' @examples
#' library(ggplot2)
#' df <- data.frame(x = 1:10, y = 1:10,
#'                  period = factor(c("Permian", "Triassic", "Jurassic",
#'                                    "Cretaceous", "Paleogene"),
#'                                  levels = c("Permian", "Triassic",
#'                                             "Jurassic", "Cretaceous",
#'                                             "Paleogene")),
#'                  era = factor(c("Paleozoic", "Mesozoic", "Mesozoic",
#'                                 "Mesozoic", "Cenozoic"),
#'                                 levels = c("Paleozoic", "Mesozoic",
#'                                            "Cenozoic")))
#' ggplot(df, aes(x, y)) +
#'   geom_point() +
#'   facet_nested_wrap_color(~ era + period, colors = rbind(periods, eras))
facet_nested_wrap_color <- function(
    facets, colors, nrow = NULL, ncol = NULL,
    scales = "fixed", axes = "margins",
    remove_labels = "none",
    shrink = TRUE, labeller = "label_value", lab_colors = "auto",
    as.table = TRUE, drop = TRUE,
    dir = "h", strip.position = "top",
    nest_line = element_line(inherit.blank = TRUE),
    solo_line = FALSE,
    resect = unit(0, "mm"),
    trim_blank = TRUE,
    strip = strip_nested(),
    bleed = NULL
) {
  colors <- convert_colors(colors)
  lab_colors <- convert_lab_colors(lab_colors)

  # copied from ggh4x v0.2.8
  strip <- assert_strip(strip)
  if (!is.null(bleed)) {
    lifecycle::deprecate_warn(
      when = "0.2.0",
      what = "facet_nested_wrap(bleed)",
      details = paste0("The `bleed` argument should be set in the ",
                       "`strip_nested()` function instead.")
    )
    strip$params$bleed <- isTRUE(bleed)
  }
  # Convert logical to elements for backward compatibility
  if (isTRUE(nest_line)) {
    nest_line <- element_line()
  }
  if (isFALSE(nest_line)) {
    nest_line <- element_blank()
  }
  if (!inherits(nest_line, c("element_line", "element_blank"))) {
    cli::cli_abort(
      "The {.arg nest_line} argument must be {.cls element_blank} or inherit \\
      from {.cls element_line}."
    )
  }
  params <- list(nest_line = nest_line, solo_line = isTRUE(solo_line),
                 resect = resect, colors = colors, lab_colors = lab_colors)
  new_wrap_facets(
    facets, nrow, ncol,
    scales, axes, remove_labels,
    shrink, labeller,
    as.table, drop, dir,
    strip.position, strip,
    trim_blank, params,
    super = FacetNestedWrapColor
  )
}

#' @export
#' @rdname facet_nested_wrap_color
facet_nested_wrap_geo <- function(
    facets, colors = rbind(periods, stages), nrow = NULL, ncol = NULL,
    scales = "fixed", axes = "margins",
    remove_labels = "none",
    shrink = TRUE, labeller = "label_value", lab_colors = "auto",
    as.table = TRUE, drop = TRUE,
    dir = "h", strip.position = "top",
    nest_line = element_line(inherit.blank = TRUE),
    solo_line = FALSE,
    resect = unit(0, "mm"),
    trim_blank = TRUE,
    strip = strip_nested(),
    bleed = NULL
) {
  facet_nested_wrap_color(facets = facets, colors = colors, nrow = nrow,
                          ncol = ncol, scales = scales, axes = axes,
                          remove_labels = remove_labels, shrink = shrink,
                          labeller = labeller, lab_colors = lab_colors,
                          as.table = as.table, drop = drop,
                          dir = dir, strip.position = strip.position,
                          nest_line = nest_line, solo_line = solo_line,
                          resect = resect, trim_blank = trim_blank,
                          strip = strip, bleed = bleed)
}

#' @rdname facet_nested_wrap_color
#' @format NULL
#' @usage NULL
#' @export
#' @importFrom ggplot2 ggproto ggproto_parent
#' @importFrom ggh4x FacetNestedWrap
#' @importFrom rlang is_function
FacetNestedWrapColor <- ggproto("FacetNestedWrapColor", FacetNestedWrap,
  draw_panels = function(panels, layout, x_scales, y_scales, ranges, coord,
                         data, theme, params, self) {
    panel_table <-
      ggproto_parent(FacetNestedWrap, self)$draw_panels(panels, layout,
                                                        x_scales, y_scales,
                                                        ranges, coord, data,
                                                        theme, params)
    strips <- grep("strip", panel_table$layout$name)
    for (i in strips) {
      label <-
        panel_table$grobs[[i]]$grobs[[1]]$children[[2]]$children[[1]]$label
      fill <- tryCatch(params$colors(label), error = function(e) NA)
      if (!is.na(fill)) {
        panel_table$grobs[[i]]$grobs[[1]]$children[[1]]$gp$fill <- fill
        if (!is_function(params$lab_colors) && params$lab_colors == "auto") {
          panel_table$grobs[[i]]$grobs[[1]]$children[[2]]$children[[1]]$
            gp$col <- white_or_black(fill)
        }
      }
      if (is_function(params$lab_colors)) {
        color <- tryCatch(params$lab_colors(label), error = function(e) NA)
        if (!is.na(color)) {
          panel_table$grobs[[i]]$grobs[[1]]$children[[2]]$children[[1]]$
            gp$col <- color
        }
      }
    }
    panel_table
  }
)