# Code for adjusting the boundaries of the kernal estimates used in the z-curve graphs. Copied from github.
# Copied from: https://github.com/echasnovski/ggplot2/blob/geom_density-bounds/R/stat-density.r

#' @param bw The smoothing bandwidth to be used.
#'   If numeric, the standard deviation of the smoothing kernel.
#'   If character, a rule to choose the bandwidth, as listed in
#'   [stats::bw.nrd()].
#' @param adjust A multiplicate bandwidth adjustment. This makes it possible
#'    to adjust the bandwidth while still using the a bandwidth estimator.
#'    For example, `adjust = 1/2` means use half of the default bandwidth.
#' @param kernel Kernel. See list of available kernels in [density()].
#' @param n number of equally spaced points at which the density is to be
#'   estimated, should be a power of two, see [density()] for
#'   details
#' @param trim If `FALSE`, the default, each density is computed on the
#'   full range of the data. If `TRUE`, each density is computed over the
#'   range of that group: this typically means the estimated x values will
#'   not line-up, and hence you won't be able to stack density values.
#'   This parameter only matters if you are displaying multiple densities in
#'   one plot or if you are manually adjusting the scale limits.
#' @param bounds Known lower and upper bounds for estimated data. Default
#'   `c(-Inf, Inf)` means that there are no (finite) bounds.
#' @section Computed variables:
#' \describe{
#'   \item{density}{density estimate}
#'   \item{count}{density * number of points - useful for stacked density
#'      plots}
#'   \item{scaled}{density estimate, scaled to maximum of 1}
#'   \item{ndensity}{alias for `scaled`, to mirror the syntax of
#'    [`stat_bin()`]}
#' }
#' @export
#' @rdname geom_density
stat_density <- function(mapping = NULL, data = NULL,
                         geom = "area", position = "stack",
                         ...,
                         bw = "nrd0",
                         adjust = 1,
                         kernel = "gaussian",
                         n = 512,
                         trim = FALSE,
                         na.rm = FALSE,
                         bounds = c(-Inf, Inf),
                         orientation = NA,
                         show.legend = NA,
                         inherit.aes = TRUE) {
  
  layer(
    data = data,
    mapping = mapping,
    stat = StatDensity,
    geom = geom,
    position = position,
    show.legend = show.legend,
    inherit.aes = inherit.aes,
    params = list(
      bw = bw,
      adjust = adjust,
      kernel = kernel,
      n = n,
      trim = trim,
      na.rm = na.rm,
      bounds = bounds,
      orientation = orientation,
      ...
    )
  )
}

#' @rdname ggplot2-ggproto
#' @format NULL
#' @usage NULL
#' @export
StatDensity <- ggproto("StatDensity", Stat,
                       required_aes = "x|y",
                       
                       default_aes = aes(x = after_stat(density), y = after_stat(density), fill = NA, weight = NULL),
                       
                       setup_params = function(data, params) {
                         params$flipped_aes <- has_flipped_aes(data, params, main_is_orthogonal = FALSE, main_is_continuous = TRUE)
                         
                         has_x <- !(is.null(data$x) && is.null(params$x))
                         has_y <- !(is.null(data$y) && is.null(params$y))
                         if (!has_x && !has_y) {
                           abort("stat_density() requires an x or y aesthetic.")
                         }
                         
                         params
                       },
                       
                       extra_params = c("na.rm", "orientation"),
                       
                       compute_group = function(data, scales, bw = "nrd0", adjust = 1, kernel = "gaussian",
                                                n = 512, trim = FALSE, na.rm = FALSE, bounds = c(-Inf, Inf),
                                                flipped_aes = FALSE) {
                         data <- flip_data(data, flipped_aes)
                         if (trim) {
                           range <- range(data$x, na.rm = TRUE)
                         } else {
                           range <- scales[[flipped_names(flipped_aes)$x]]$dimension()
                         }
                         
                         density <- compute_density(data$x, data$weight, from = range[1],
                                                    to = range[2], bw = bw, adjust = adjust, kernel = kernel, n = n,
                                                    bounds = bounds)
                         density$flipped_aes <- flipped_aes
                         flip_data(density, flipped_aes)
                       }
                       
)

compute_density <- function(x, w, from, to, bw = "nrd0", adjust = 1,
                            kernel = "gaussian", n = 512,
                            bounds = c(-Inf, Inf)) {
  nx <- length(x)
  if (is.null(w)) {
    w <- rep(1 / nx, nx)
  } else {
    w <- w / sum(w)
  }
  
  # if less than 2 points return data frame of NAs and a warning
  if (nx < 2) {
    warn("Groups with fewer than two data points have been dropped.")
    return(new_data_frame(list(
      x = NA_real_,
      density = NA_real_,
      scaled = NA_real_,
      ndensity = NA_real_,
      count = NA_real_,
      n = NA_integer_
    ), n = 1))
  }
  
  if (all(is.infinite(bounds))) {
    dens <- stats::density(x, weights = w, bw = bw, adjust = adjust,
                           kernel = kernel, n = n, from = from, to = to)
  } else {
    dens <- stats::density(x, weights = w, bw = bw, adjust = adjust,
                           kernel = kernel, n = n)
    
    dens <- reflect_density(dens = dens, bounds = bounds, from = from, to = to)
  }
  
  new_data_frame(list(
    x = dens$x,
    density = dens$y,
    scaled =  dens$y / max(dens$y, na.rm = TRUE),
    ndensity = dens$y / max(dens$y, na.rm = TRUE),
    count =   dens$y * nx,
    n = nx
  ), n = length(dens$x))
}

reflect_density <- function(dens, bounds, from, to) {
  if (all(is.infinite(bounds))) {
    return(dens)
  }
  
  f_dens <- stats::approxfun(
    x = dens$x, y = dens$y, method = "linear", yleft = 0, yright = 0
  )
  
  out_x <- intersection_grid(dens$x, bounds, from, to)
  out_y <- f_dens(out_x) + f_dens(bounds[1] + (bounds[1] - out_x)) +
    f_dens(bounds[2] + (bounds[2] - out_x))
  
  list(x = out_x, y = out_y)
}

intersection_grid <- function(grid, bounds, from, to) {
  left <- max(from, bounds[1])
  right <- min(to, bounds[2])
  
  seq(from = left, to = right, length.out = length(grid))
}

