# Range keys --------------------------------------------------------------

#' Range keys
#'
#' @description
#' These functions are helper functions for working with ranged data as keys
#' in guides. They all share the goal creating of a guide key, but have
#' different methods:
#'
#' * `key_range_auto()` is a function factory whose functions make an attempt
#'   to infer ranges from the scale's labels.
#' * `key_range_manual()` uses user-provided vectors to set ranges.
#' * `key_range_map()` makes mappings from a `<data.frame>` to set ranges.
#' * `key_range_rle()` uses run-length encoding to determine the start and
#'   end of runs (blocks of repeated data values).
#'
#' @param sep A `<character[1]>` giving a [regular expression][base::regex] to
#'   use for splitting labels provided by the scale using
#'   [`strsplit()`][base::strsplit]. Defaults to splitting on any
#'   non-alphanumeric character.
#' @param reverse A `<logical[1]>` which if `FALSE` (default) treats the first
#'   labels as the inner labels and the last labels as the outer labels.
#'   If `TRUE`, thee first labels are treated as the outer labels and the last
#'   labels are treated as the inner labels.
#' @param start,end A vector that can be interpreted by the scale, giving the
#'   start and end positions of each range respectively.
#' @param name A `<character>` or list of expressions
#' @param level An `<integer>` giving the depth of each range to avoid overlaps
#'   between different ranges. When `level` is smaller than 1, no brackets are
#'   drawn.
#' @param data A `<data.frame>` or similar object coerced by
#'   [`fortify()`][ggplot2::fortify] to a `<data.frame>`, in which the `mapping`
#'   argument is evaluated.
#' @param ...
#' The `...` parameter has two purposes.
#' 1. In `key_range_map()` it is [`<data-masking>`][rlang::topic-data-mask].
#'   A set of mappings similar to those provided to [`aes()`][ggplot2::aes],
#'   which will be evaluated in the `data` argument. These *must* contain
#'   `start` and `end` mappings.
#' 2. In other keys, `...` can be used to transfer graphical properties to the
#'   individual ranges of a guide. For example, using `colour = "blue"` will
#'   draw parts of the guides associated with ranges in blue. There is a shallow
#'   hierarchy in that `text_colour`, `line_colour`, `rect_colour` and
#'   `point_colour` are the specific properties for elements, but all inherit
#'   from the main `colour` setting. Likewise, `size`, `linewidth`, `linetype`
#'   and `fill` have specific variants for elements.
#' @param x A `<vector[n]>` for which to determine run-starts and run-ends.
#' @param call A [call][rlang::topic-error-call] to display in messages.
#'
#' @details
#' The `level` variable is optional and when missing, the guides use an
#' algorithm similar to `IRanges::disjointBins()` to avoid overlaps.
#'
#' The `key_range_auto()` does *not* work with expression labels.
#'
#' @name key_range
#' @family keys
#' @return
#' For `key_range_auto()` a function. For `key_range_manual()` and
#' `key_range_map()` a `<data.frame>` with the `<key_range>` class.
#'
#' @examples
#' # Example scale
#' template <- scale_x_discrete(limits = c("A 1", "B 1", "C&1", "D&2", "E&2"))
#'
#' # By default, splits on all non-alphanumeric characters
#' auto <- key_range_auto()
#' auto(template)
#'
#' # Only split on a specific character
#' auto <- key_range_auto(sep = "&")
#' auto(template)
#'
#' # Treating the letters as outer labels and numbers as inner labels
#' auto <- key_range_auto(reverse = TRUE)
#' auto(template)
#'
#' # Providing custom values
#' key_range_manual(
#'   start = c(1, 5,  10),
#'   end   = c(4, 15, 11),
#'   level = c(0, 2, 1),
#'   name  = c("A", "B", "C")
#' )
#'
#' # Values from a <data.frame>
#' key_range_map(presidential, start = start, end = end, name = name)
#'
#' # Values from run length encoding
#' key_range_rle(c("AB", "AB", "C", "DEF", "DEF", "DEF"))
NULL

#' @rdname key_range
#' @export
key_range_auto <- function(sep = "[^[:alnum:]]+", reverse = FALSE, ...) {
  check_string(sep)
  check_bool(reverse)
  force(sep)
  force(reverse)
  dots <- extra_args(...)
  call <- current_call()
  fun <- function(scale, aesthetic = NULL) {
    range_from_label(
      scale = scale, aesthetic = aesthetic,
      sep = sep, reverse = reverse, extra_args = dots,
      call = call
    )
  }
  class(fun) <- union("key_range_auto_function", class(fun))
  fun
}

#' @rdname key_range
#' @export
key_range_manual <- function(start, end, name = NULL, level = NULL, ..., call = NULL) {
  df <- data_frame0(
    start = start, end = end, .label = name, .level = level,
    !!!extra_args(...),
    .error_call = call %||% current_call()
  )
  check_columns(df, c("start", "end"))
  class(df) <- c("key_range", "key_guide", class(df))
  df
}

#' @rdname key_range
#' @export
key_range_map <- function(data, ..., call = NULL) {
  mapping <- aes(!!!enquos(...))

  df <- eval_aes(
    data, mapping,
    required = c("start", "end"),
    optional = c("name", "level", .element_params),
    call = call %||% current_call(),
    arg_mapping = "mapping", arg_data = "data"
  )

  df <- rename(
    df, c("name", "level", .element_params),
    c(".label", ".level", paste0(".", .element_params))
  )
  class(df) <- c("key_range", "key_guide", class(df))
  df
}

#' @rdname key_range
#' @export
key_range_rle <- function(x, ..., call = NULL) {
  rle <- new_rle(x)
  key_range_manual(
    start = as_mapped_discrete(rle$start),
    end   = as_mapped_discrete(rle$end),
    name  = rle$key,
    level = 1L,
    ...,
    call = call %||% current_call()
  )
}

# Extractor ---------------------------------------------------------------

range_extract_key <- function(
  scale, aesthetic, key,
  drop_zero = TRUE, pad_discrete = 0.0,  oob = "squish",
  ...
) {
  if (is.function(key)) {
    key <- key(scale, aesthetic)
  }
  if (is.null(key$.level)) {
    key <- disjoin_ranges(key)
  }

  # Mark discrete variables separately for start and end
  pad_start <- -1.0 * is_discrete(key$start) * pad_discrete
  pad_end   <- +1.0 * is_discrete(key$end)   * pad_discrete

  map <- aesthetic %in% c("x", "y")
  key$start <- scale_transform(key$start, scale, map = map, "start")
  key$end   <- scale_transform(key$end,   scale, map = map, "end")

  # Backtransform AsIs variables
  limits <- scale$continuous_range %||% scale$get_limits()
  key$start <- descale(key$start, to = limits)
  key$end   <- descale(key$end,   to = limits)

  # Sort starts and ends
  key[c("start", "end")] <- list(
    start = pmin(key$start, key$end),
    end   = pmax(key$start, key$end)
  )

  # Mark ranges where no brackets should be drawn
  key$.draw <- TRUE
  if (!isFALSE(drop_zero)) {
    key$.draw <- abs(key$end - key$start) > sqrt(.Machine$double.eps)
  }
  key$.draw <- key$.draw & key$.level > 0L

  # Apply padding for discrete variables
  extend <- pad_discrete
  if (scale$is_discrete() && !is.null(extend)) {
    key$start <- key$start + pad_start
    key$end   <- key$end   + pad_end
  }

  # Apply out-of-bounds rules
  range_oob(key, oob, limits)
}


# Helpers -----------------------------------------------------------------

## Out-of-bounds utilities ------------------------------------------------

range_oob <- function(ranges, method, limits) {
  limits <- sort(limits)
  ranges <- switch(
    method,
    squish = range_squish(ranges, limits),
    censor = range_censor(ranges, limits),
    ranges
  )
  vec_slice(ranges, !is.na(ranges$.draw))
}

range_squish <- function(ranges, limits) {
  start  <- ranges$start
  end    <- ranges$end
  oob_start <- is_oob(start, limits)
  oob_end   <- is_oob(end,   limits)
  keep <- !oob_start | !oob_end | (start < limits[1L] & end > limits[2L])
  ranges$.draw[!keep] <- NA
  ranges$start <- pmin(pmax(ranges$start, limits[1L]), limits[2L])
  ranges$end   <- pmin(pmax(ranges$end,   limits[1L]), limits[2L])
  ranges
}

range_censor <- function(ranges, limits) {
  oob_start <- is_oob(ranges$start, limits)
  oob_end   <- is_oob(ranges$end, limits)
  keep      <- !oob_start & !oob_end
  ranges$.draw[!keep] <- NA
  ranges
}

## Other helpers ----------------------------------------------------------

range_from_label <- function(
  scale, aesthetic = NULL, sep =  "[^[:alnum:]]+", reverse = FALSE,
  extra_args = list(), call = caller_env()
) {
  # Extract a standard key from the scale
  aesthetic <- aesthetic %||% scale$aesthetics[1L]
  key <- Guide$extract_key(scale, aesthetic)

  # Reject expressions, as we cannot split these
  if (!is.character(key$.label)) {
    type <- obj_type_friendly(key$.label)
    cli::cli_abort(
      c("Cannot split the guide's {.field label}.",
        i = "It must be a {.cls character} vector, not {type}."),
      call = call
    )
  }

  # Split labels
  labels <- strsplit(key$.label, sep)

  # Pad label vectors if necessary
  lengths <- lengths(labels)
  max_len <- max(lengths)
  if (any(lengths != max_len)) {
    cli::cli_warn(c(
      paste0("Not all {.field labels} in {.fn key_range_auto} can be split ",
             "into equal lengths."),
      i = "Is \"{sep}\" the correct {.arg sep} argument?"
    ))
    labels <- lapply(labels, pad, length = max_len)
  }
  labels <- do.call(rbind, labels)

  if (isTRUE(reverse)) {
    labels <- labels[, rev(seq_len(ncol(labels)))]
  }

  key$.label <- labels[, 1L, drop = TRUE]
  labels <- labels[, -1L, drop = FALSE]

  # Set first series of unbracketed labels
  value  <- key[[1L]]
  key <- data_frame0(
    start = value, end = value, .label = key$.label, .level = 0L
  )
  if (is_empty(labels)) {
    return(data_frame0(key, !!!extra_args))
  }
  ranges <- apply(labels, 2L, function(labs) {
    rle   <- vec_unrep(labs)
    start <- cumsum(rle$times) - rle$times + 1L
    data_frame0(
      start = value[start],
      end   = value[cumsum(rle$times)],
      .label = labs[start]
    )
  })
  nrows  <- list_sizes(ranges)
  ranges <- vec_c(!!!ranges)
  ranges$.level <- rep.int(seq_along(nrows), nrows)
  range <- vec_slice(ranges, !is.na(ranges$.label))
  df <- vec_rbind(key, range)
  df <- data_frame0(df, !!!extra_args)
  class(df) <- c("key_range", "key_guide", class(df))
  df
}

justify_ranges <- function(key, levels, element, level_elements) {

  if (is_blank(element)) {
    return(key)
  }

  ends <- intersect(c("thetaend", "xend", "yend"), names(key))
  if (length(ends) < 1L) {
    return(key)
  }
  starts <- gsub("end$", "", ends[1L])

  just_name <- switch(ends[1L], yend = "vjust", "hjust")
  just <- element[[just_name]] %||% 0.5

  if (!is.null(level_elements)) {
    just <- map_dbl(level_elements, function(x) x[[just_name]] %||% just)
    just <- just[match(key$.level, levels)]
  }

  key[[starts]] <- switch(
    ends[1L],
    thetaend = justify_range(key$theta, key$thetaend, just, theta = TRUE),
    xend     = justify_range(key$x, key$xend, just),
    yend     = justify_range(key$y, key$yend, just)
  )

  key
}

justify_range <- function(start, end, just, theta = FALSE) {
  if (theta) {
    add <- end < start
    end[add] <- end[add] + 2.0 * pi
  }
  (end - start) * just + start
}

disjoin_ranges <- function(ranges) {

  n_ranges <- nrow(ranges)
  if (n_ranges < 2L) {
    ranges$.level <- rep(1L, nrow(ranges))
    return(ranges)
  }

  # Sort and extract
  ranges <- ranges[order(ranges$start, ranges$end), , drop = FALSE]
  ranges <- vec_slice(ranges, order(ranges$start, ranges$end))
  starts <- ranges$start
  ends   <- ranges$end

  # Initialise first range
  end_tracker <- ends[1L]
  bin <- rep(NA_integer_, nrow(ranges))
  bin[1L] <- 1L

  # Find bins
  for (range_id in seq_len(n_ranges)[-1L]) {
    candidate <- which(end_tracker < starts[range_id])
    if (length(candidate) > 0L) {
      # If there is room in this bin, update this bin
      ans <- candidate[1L]
      end_tracker[ans] <- ends[range_id]
    } else {
      # Register new bin
      end_tracker <- c(end_tracker, ends[range_id])
      ans <- length(end_tracker)
    }
    bin[range_id] <- ans
  }

  ranges$.level <- bin
  ranges
}

extract_range_params <- function(scale, params, ...) {
  params$position <- params$position %|W|% NULL
  params$limits   <- scale$get_limits()

  new_names <- c("start", "end")
  aesthetic <- params$aesthetic
  if (aesthetic %in% c("x", "y")) {
    new_names  <- paste0(aesthetic, c("", "end"))
  } else if (is_theta(params$position)) {
    new_names <- c("x", "xend")
  }
  params$key <- rename(params$key, c("start", "end"), new_names)
  params
}

setup_range_params <- function(params) {
  if (params$aesthetic %in% c("x", "y")) {
    # parameters are already transformed
    return(params)
  }

  limits   <- params$limits %||% c(0.0, 1.0)
  other    <- switch(params$position, bottom = , left = 1L, 0L)

  if (!is_empty(params$key)) {
    key <- params$key

    start <- guide_rescale(key$start, limits)
    end   <- guide_rescale(key$end, limits)

    key <- switch(
      params$direction,
      horizontal = replace_null(key, x = start, xend = end, y = other),
      vertical   = replace_null(key, y = start, yend = end, x = other),
      key
    )

    params$key <- key
  }

  if (!is_empty(params$decor)) {
    decor <- params$decor
    value <- guide_rescale(decor[[params$aesthetic]], limits)
    decor <- switch(
      params$direction,
      horizontal = replace_null(decor, x = value, y = other),
      vertical   = replace_null(decor, y = value, x = other),
      decor
    )

    params$decor <- decor
  }

  params
}
