#' Segment keys
#'
#' @description
#' These functions are helper functions for working with segment data as keys
#' in guides. They all share the goal of creating a guide key, but have
#' different methods:
#'
#' * `key_segment_manual()` directly uses user-provided vectors to set segments.
#' * `key_segment_map()` makes mappings from a `<data.frame>` to set segments.
#' * `key_dendro()` is a specialty case for coercing dendrogram data to
#'   segments. Be aware that setting the key alone cannot affect the scale
#'   limits, and will give misleading results when used incorrectly!
#'
#' @param value,value_end A vector that is interpreted to be along the scale
#'   that the guide codifies.
#' @param oppo,oppo_end A vector that is interpreted to be orthogonal to the
#'   `value` and `value_end` variables.
#' @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 dendro A data structure that can be coerced to a dendrogram through
#'   the [`as.dendrogram()`][stats::as.dendrogram()] function. When `NULL`
#'   (default) an attempt is made to search for such data in the scale.
#' @param type A string, either `"rectangle"` or `"triangle"`, indicating the
#'   shape of edges between nodes of the dendrogram.
#' @param ...
#' The `...` parameter has two purposes.
#' 1. In `key_segments_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
#'   `value` and `oppo` mappings.
#' 2. In other keys, `...` can be used to transfer graphical properties to the
#'   individual breaks of a guide. For example, using `colour = "blue"` will
#'   draw parts of the guides associated with breaks in blue. There is a shallow
#'   hierarchy in that `line_colour` is the specific property for segment
#'   elements, but others inherit from the main `colour` setting. Likewise,
#'   `linewidth` and `linetype` have specific variants for line elements.
#' @param .call A [call][rlang::topic-error-call] to display in messages.
#'
#' @export
#' @name key_segments
#' @family keys
#' @return
#' For `key_segments_manual()` and `key_segments_map()`, a `<data.frame>` with
#' the `<key_range>` class.
#'
#' @examples
#' # Giving vectors directly
#' key_segment_manual(
#'   value = 0:1, value_end = 2:3,
#'   oppo  = 1:0, oppo_end  = 3:2
#' )
#'
#' # Taking columns of a data frame
#' data <- data.frame(x = 0:1, y = 1:0, xend = 2:3, yend = 3:2)
#' key_segment_map(data, value = x, oppo = y, value_end = xend, oppo_end = yend)
#'
#' # Using dendrogram data
#' clust <- hclust(dist(USArrests), "ave")
#' key_dendro(clust)(scale_x_discrete())
key_segment_manual <- function(value, oppo, value_end = value,
                               oppo_end = oppo, ...) {
  df <- data_frame0(
    value = value, oppo = oppo,
    value_end = value_end, oppo_end = oppo_end,
    !!!extra_args(..., .valid_args = .line_params),
    .error_call = current_env()
  )
  check_columns(df, c("value", "oppo"))
  class(df) <- c("key_segment", "key_guide", class(df))
  df
}

#' @rdname key_segments
#' @export
key_segment_map <- function(data, ..., .call = caller_env()) {

  mapping <- aes(!!!enquos(...))

  df <- eval_aes(
    data, mapping,
    required = c("value", "oppo"),
    optional = c("value_end", "oppo_end", .line_params),
    call = .call, arg_mapping = "mapping", arg_data = "data"
  )

  df <- rename(df, .line_params, paste0(".", .line_params))
  class(df) <- c("key_segment", "key_guide", class(df))
  df

}

#' @rdname key_segments
#' @export
key_dendro <- function(dendro = NULL, type = "rectangle", ..., .call = NULL) {
  call <- .call %||% current_call()
  dots <- extra_args(..., .valid_args = .line_params)
  force(dendro)
  function(scale, aesthetic = NULL, ...) {
    key <- extract_dendro(scale$scale$clust %||% dendro, type = type)
    data_frame0(key, !!!dots, .error_call = call)
  }
}

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

segment_extract_key <- function(scale, aesthetic, key, ...) {
  key <- standard_extract_key(scale, aesthetic, key, ...)

  # Backtransform AsIs variables
  range <- scale$continuous_range %||% scale$get_limits()
  key$value     <- descale(key$value, range)
  key$value_end <- descale(key$value_end, range)

  remove_vars <- character()

  # Interleave values
  value_vars <- c("value", "value_end")
  if (all(value_vars %in% names(key))) {
    value <- vec_interleave(key$value, key$value_end)
    remove_vars <- c(remove_vars, value_vars)
  }

  # Interleave opposites
  oppo_vars <- c("oppo", "oppo_end")
  if (all(oppo_vars %in% names(key))) {
    oppo <- vec_interleave(key$oppo, key$oppo_end)
    remove_vars <- c(remove_vars, oppo_vars)
  }

  # Reconstruct key
  key[remove_vars] <- NULL
  new_key <- data_frame0(value = value, oppo = oppo)
  i <- rep(vec_seq_along(key), each = 2L)
  new_key[names(key)] <- vec_slice(key, i)
  new_key$group <- new_key$group %||% i
  new_key$oppo  <- rescale(new_key$oppo, from = range(new_key$oppo, 0.0))

  # Normalise key column names
  if (aesthetic %in% c("x", "y")) {
    new_names <- union(aesthetic, c("x", "y"))
    new_key <- rename(new_key, c("value", "oppo"), new_names)
  } else {
    new_key <- rename(new_key, "value", aesthetic)
    new_key$.value <- new_key[[aesthetic]]
  }
  new_key
}

# Dendrogram utilities ----------------------------------------------------

# Simplified version of `stats:::plotNode`.
# It only looks for the segments and ignores labels and most other attributes.
extract_dendro <- function(tree, type = "rectangle") {

  # Check arguments
  whole_tree <- tree <- as_dendrogram(tree)

  # Initialise stuff
  depth <- 0L
  llimit <- list()
  x1 <- i <- 1L
  x2 <- number_of_members(tree)
  KK <- kk <- integer()

  mtx <- setup_segments(stats::nobs(tree), type)
  col_order <- attr(mtx, "order", exact = TRUE)
  row_index <- attr(mtx, "offset", exact = TRUE)
  row_step  <- length(row_index)

  repeat {
    depth <- depth + 1L
    inner <- !stats::is.leaf(tree) && x1 != x2

    node <- node_limit(x1, x2, tree)
    llimit[[depth]] <- node$limit

    ymax <- attr(tree, "height")
    xmax <- node$x

    if (inner) {
      for (k in seq_along(tree)) {
        child <- tree[[k]]

        ymin <- attr(child, "height") %||% 0.0
        xmin <- node$limit[k] + (attr(child, "midpoint") %||% 0.0)

        mtx[i + row_index, ] <- c(xmax, ymax, xmin, ymin)[col_order]
        i <- i + row_step
      }
      if (length(tree) > 0L) {
        KK[depth] <- length(tree)
        kk[depth] <- 1L
        x1 <- node$limit[1L]
        x2 <- node$limit[2L]
        tree <- tree[[1L]]
      }
    } else {
      repeat {
        # We climb back up the tree until we have unvisited children
        depth <- depth - 1L
        if (!depth || kk[depth] < KK[depth]) {
          break
        }
      }
      if (!depth) {
        break
      }
      length(kk) <- depth
      kk[depth] <- k <- kk[depth] + 1L
      x1 <- llimit[[depth]][k]
      x2 <- llimit[[depth]][k + 1L]
      tree <- whole_tree[[kk]]
    }
  }
  as.data.frame(mtx)
}

setup_segments <- function(n_obs, type = "triangle") {
  type <- arg_match0(type, c("rectangle", "triangle"))
  if (type == "triangle") {
    n_segments <- 2L * n_obs - 2L
    i_increment <- 1L
    i_offset <- 0L
    coord_order <- c(1L, 2L, 3L, 4L)
  } else {
    n_segments <- 4L * n_obs - 4L
    i_increment <- 2L
    i_offset <- 0L:1L
    # xmax -> xmin -> ymax -> ymax -> xmin -> xmin -> ymax -> ymin
    coord_order <- c(1L, 3L, 2L, 2L, 3L, 3L, 2L, 4L)
  }

  mtx <- matrix(NA_real_, n_segments, ncol = 4L)

  attributes(mtx) <- list(
    offset = i_offset,
    increment = i_increment,
    order = coord_order,
    dim = dim(mtx)
  )
  colnames(mtx) <- c("value", "oppo", "value_end", "oppo_end")
  mtx
}

# Copy of `stats:::.memberDend()`
number_of_members <- function(tree) {
  attr(tree, "x.member") %||% attr(tree, "members") %||% 1L
}

# Simplified version of `stats:::plotNodeLimit`,
# It has `center = FALSE` build-in.
node_limit <- function(x1, x2, subtree) {
  inner <- !stats::is.leaf(subtree) && x1 != x2
  if (inner) {
    K <- length(subtree)
    limit <- integer(K)
    xx1 <- x1
    for (k in 1L:K) {
      xx1 <- xx1 + number_of_members(subtree[[k]])
      limit[k] <- xx1
    }
  } else {
    limit <- x2
  }
  limit <- c(x1, limit)
  mid <- attr(subtree, "midpoint")
  center <- inner && !is.numeric(mid)
  x <- if (center) mean(c(x1, x2)) else x1 + (mid %||% 0L)
  list(x = x, limit = limit)
}

as_dendrogram <- function(tree) {
  # Check arguments
  try_fetch(
    stats::as.dendrogram(tree),
    error = function(cnd) {
      cli::cli_abort(
        "Could not find or coerce {.arg dendro} argument.",
        parent = cnd
      )
    }
  )
}
