#' Density plot type
#' 
#' @md
#' @description Type function for density plots.
#' @inheritParams stats::density
#' @param bw the smoothing \code{\link[stats:bw.nrd]{bandwidth}} to be used,
#'   see \code{\link[stats]{density}} for details and options.
#' @param kernel a character string giving the smoothing kernel to be used. This
#'   must partially match one of `"gaussian"`, `"rectangular"`, `"triangular"`,
#'   `"epanechnikov"`, `"biweight"`, `"cosine"` or `"optcosine"`, with default
#'   `"gaussian"`, and may be abbreviated to a unique prefix (single letter).
#'
#'   `"cosine"` is smoother than `"optcosine"`, which is the usual 'cosine'
#'   kernel in the literature and almost MSE-efficient. However, `"cosine"` is
#'   the version used by S.
#' @inheritParams type_ribbon
#' @param joint.bw character string indicating whether (and how) the smoothing
#'   bandwidth should be computed from the joint data distribution when there
#'   are multiple subgroups. The options are `"mean"` (the default), `"full"`,
#'   and `"none"`. Also accepts a logical argument, where `TRUE` maps to
#'   `"mean"` and `FALSE` maps to `"none"`. See the "Bandwidth selection"
#'   section below for a discussion of practical considerations.
#' @inherit stats::density details
#' @param echo.bw where to report the smoothing bandwidth (and the number of
#'   observations behind it), since neither is visible from the curve itself.
#'   The default `FALSE` reports nothing. `TRUE` is shorthand for `"sub"`, and
#'   any combination of `"sub"` (subtitle), `"cap"` (caption), and `"cat"`
#'   (printed to the console, with `"print"` accepted as a synonym) may be
#'   given. A destination the user has already labelled is
#'   left alone. Bandwidths shared across subgroups are reported once and named
#'   as joint; individual bandwidths are reported per group, truncated after
#'   three so the label stays legible.
#'
#' @section Bandwidth selection: While the choice of smoothing bandwidth will
#'   always stand to affect a density visualization, it gains an added
#'   importance when multiple densities are drawn simultaneously (e.g., for
#'   subgroups with respect to `by` or `facet`). Allowing each subgroup to
#'   compute its own separate bandwidth independently offers greater flexibility
#'   in capturing the unique characteristics of each subgroup, particularly when
#'   distributions differ substantially in location and/or scale. However, this
#'   approach may overemphasize small random variations and make it harder to
#'   visually compare densities across subgroups. Hence, it is often useful to
#'   employ the same ("joint") bandwidth across all subgroups. The following
#'   strategies are available via the `joint.bw` argument:
#' 
#'   - The default `joint.bw = "mean"` first computes the individual bandwidths
#'   for each group but then computes their mean, weighted by the number of
#'   observations in each group. This will work well when all groups have
#'   similar amounts of scatter (similar variances), even when they have
#'   potentially rather different locations. The weighted averaging stabilizes
#'   potential fluctuations in the individual bandwidths, especially when some
#'   subgroups are rather small.
#' 
#'   - Alternatively, `joint.bw = "full"` can be used to compute the joint
#'   bandwidth from the full joint distribution (merging all groups). This will
#'   yield an even more robust bandwidth, especially when the groups overlap
#'   substantially (i.e., have similar locations and scales). However, it may
#'   lead to too large bandwidths and thus too much smoothing, especially when
#'   the locations of the groups differ substantially.
#' 
#'   - Finally, `joint.bw = "none"` disables the joint bandwidth so that each
#'   group just employs its individual bandwidth. This is often the best choice
#'   if the amounts of scatter differ substantially between the groups, thus
#'   necessitating different amounts of smoothing.
#' @section Titles: This tinyplot method for density plots differs from the base
#'   \code{\link[stats]{plot.density}} function in its treatment of titles. The
#'   x-axis title displays only the variable name, omitting details about the
#'   number of observations and smoothing bandwidth. Additionally, the main
#'   title is left blank by default for a cleaner appearance. Use `echo.bw` to
#'   put those details back in the subtitle, a caption, or the console.
#' @examples
#' # "density" type convenience string
#' tinyplot(~Sepal.Length, data = iris, type = "density")
#' 
#' # grouped density example
#' tinyplot(~Sepal.Length | Species, data = iris, type = "density")
#' 
#' # use `bg = "by"` (or, equivalent `fill = "by"`) to get filled densities
#' tinyplot(~Sepal.Length | Species, data = iris, type = "density", fill = "by")
#' 
#' # use `type_density()` to pass extra arguments for customization
#' tinyplot(
#'   ~Sepal.Length | Species, data = iris,
#'   type = type_density(bw = "SJ"),
#'   main = "Bandwidth computed using Sheather & Jones (1991)"
#' )
#' 
#' # The default for grouped density plots is to use the mean of the
#' # individual subgroup bandwidths (weighted by group size) as the
#' # joint bandwidth. Alternatively, the bandwidth from the "full"
#' # data or separate individual bandwidths ("none") can be used.
#' tinyplot(~Sepal.Length | Species, data = iris,
#'     ylim = c(0, 1.25), type = "density")        # mean (default)
#' tinyplot_add(joint.bw = "full", lty = 2)        # full data
#' tinyplot_add(joint.bw = "none", lty = 3)        # none (individual)
#' legend("topright", c("Mean", "Full", "None"), lty = 1:3, bty = "n", title = "Joint BW")
#'
#' # neither the bandwidth nor the number of observations is visible from the
#' # curve, so `echo.bw` can report them: in the subtitle, a caption, or the
#' # console
#' tinyplot(~Sepal.Length | Species, data = iris,
#'     type = type_density(echo.bw = "sub"))
#'
#' # individual bandwidths are listed per group
#' tinyplot(~Sepal.Length | Species, data = iris,
#'     type = type_density(joint.bw = "none", echo.bw = "sub"))
#'
#' @importFrom stats density weighted.mean
#' @importFrom stats bw.SJ bw.bcv bw.nrd bw.nrd0 bw.ucv 
#' @export
type_density = function(
        bw = "nrd0",
        joint.bw =  c("mean", "full", "none"),
        adjust = 1,
        kernel = c("gaussian", "epanechnikov", "rectangular", "triangular", "biweight", "cosine", "optcosine"),
        n = 512,
        # more args from density here?
        echo.bw = FALSE,
        alpha = NULL
    ) {
    kernel = match.arg(kernel, c("gaussian", "epanechnikov", "rectangular", "triangular", "biweight", "cosine", "optcosine"))
    if (is.logical(joint.bw)) {
        joint.bw = ifelse(joint.bw, "mean", "none")
    }
    joint.bw = match.arg(joint.bw, c("mean", "full", "none"))
    echo.bw = match_echo_bw(echo.bw)
    out = list(
        data = data_density(bw = bw, adjust = adjust, kernel = kernel, n = n,
                            joint.bw = joint.bw, echo.bw = echo.bw,
                            alpha = alpha),
        draw = NULL,
        name = "density"
    )
    class(out) = "tinyplot_type"
    return(out)
}

## Normalize the echo.bw argument to a character vector of destinations.
## FALSE/NULL disables it; TRUE is shorthand for the subtitle. "print" is
## accepted as a synonym for "cat".
match_echo_bw = function(echo.bw) {
    if (is.null(echo.bw) || isFALSE(echo.bw)) return(character(0))
    if (isTRUE(echo.bw)) return("sub")
    if (!is.character(echo.bw)) {
        stop("`echo.bw` must be TRUE, FALSE, or a character vector.",
             call. = FALSE)
    }
    echo.bw[echo.bw == "print"] = "cat"
    valid = c("sub", "cap", "cat")
    bad = setdiff(echo.bw, valid)
    if (length(bad)) {
        stop(
            "`echo.bw` must be one or more of ",
            paste(sprintf('"%s"', valid), collapse = ", "),
            ' (or "print" for "cat"); got ',
            paste(sprintf('"%s"', bad), collapse = ", "), ".",
            call. = FALSE
        )
    }
    unique(echo.bw)
}

## Collapse a vector for display, following the convention used elsewhere in
## tinyplot: a lone value prints bare, several print as a bracketed list, and
## more than three are truncated so the label stays readable.
format_echo_vec = function(x, numeric = TRUE) {
    x = if (numeric) sprintf("%.4g", x) else as.character(x)
    if (length(x) == 1L) return(x)
    if (length(x) > 3L) x = c(x[1:3], "...")
    paste0("[", paste(x, collapse = ", "), "]")
}

data_density = function(bw = "nrd0", adjust = 1, kernel = "gaussian", n = 512,
                        joint.bw = "none", echo.bw = character(0),
                        alpha = NULL) {
    fun = function(settings, ...) {
        env2env(settings, environment(), c("by", "bg", "facet", "ylab", "col", "ribbon.alpha", "datapoints"))
        ribbon.alpha = if (is.null(alpha)) .tpar[["ribbon.alpha"]] else (alpha)
        has_weights = !is.null(datapoints[["weights"]])
        if (has_weights) settings$weights_used = TRUE
        
        if (is.null(ylab)) ylab = "Density"
        
        datapoints = split(datapoints, list(datapoints$by, datapoints$facet))
        datapoints = Filter(function(k) nrow(k) > 0, datapoints)
        
        if (joint.bw == "none" || is.numeric(bw)) {
            dens_bw = bw
        } else {
            if (joint.bw == "mean") {
                # Use weighted mean of subgroup bandwidths
                bws = sapply(datapoints, function(dat) bw_fun(kernel = bw, dat$x))
                ws = sapply(datapoints, nrow)
                dens_bw = weighted.mean(bws, ws)
            } else if (joint.bw == "full") {
                dens_bw = bw_fun(kernel = bw, unlist(sapply(datapoints, `[[`, "x")))
            }
        }
        
        dens = lapply(datapoints, function(dat) {
            wts = if (has_weights) dat[["weights"]] / sum(dat[["weights"]]) else NULL
            density(dat$x, bw = dens_bw, kernel = kernel, n = n, weights = wts)
        })

        if (length(echo.bw)) {
            bws = vapply(dens, `[[`, numeric(1), "bw")
            ns = vapply(dens, `[[`, numeric(1), "n")
            # A shared bandwidth is reported once and named as joint; individual
            # bandwidths are reported per group, in the order the groups plot.
            shared = length(bws) > 1L && isTRUE(all.equal(min(bws), max(bws)))
            bw_txt = format_echo_vec(if (shared) bws[1] else bws)
            bw_lab = if (shared) "Joint Bandwidth" else "Bandwidth"
            echo_txt = paste0(
                "N = ", format_echo_vec(ns, numeric = FALSE),
                "   ", bw_lab, " = ", bw_txt
            )
            if ("cat" %in% echo.bw) cat(echo_txt, "\n", sep = "")
            # Never overwrite a label the user has set themselves.
            for (dest in intersect(c("sub", "cap"), echo.bw)) {
                if (is.null(settings[[dest]])) settings[[dest]] = echo_txt
            }
        }

        datapoints = Map(function(dat, d) {
            data.frame(
                by = dat$by[1], # already split
                facet = dat$facet[1], # already split
                y = d$y,
                x = d$x
            )
        }, datapoints, dens)
        datapoints = do.call(rbind, datapoints)
        datapoints$ymax = datapoints$y
        datapoints$ymin = rep.int(0, nrow(datapoints))
        
        # flags for legend and fill
        dtype = if (!is.null(bg)) "ribbon" else "l"

        type = dtype
        by = if (length(unique(datapoints$by)) == 1) by else datapoints$by
        facet = if (length(unique(datapoints$facet)) == 1) facet else datapoints$facet
        
        # legend customizations (only for filled density plots)
        if (!is.null(bg)) {
            settings$legend_args[["pch"]] = settings$legend_args[["pch"]] %||% 22
            settings$legend_args[["pt.cex"]] = settings$legend_args[["pt.cex"]] %||% 3.5
            settings$legend_args[["pt.lwd"]] = settings$legend_args[["pt.lwd"]] %||% par("lwd")
            settings$legend_args[["lty"]] = settings$legend_args[["lty"]] %||% 0
            settings$legend_args[["y.intersp"]] = settings$legend_args[["y.intersp"]] %||% 1.25
            settings$legend_args[["seg.len"]] = settings$legend_args[["seg.len"]] %||% 1.25
        }
        
        env2env(environment(), settings, c(
            "ylab",
            "type",
            "ribbon.alpha",
            "datapoints",
            "by",
            "facet"
        ))
    }
    return(fun)
}

