#' @include utilities.R
NULL
#' Scatter Plot with Marginal Histograms
#' @description Create a scatter plot with marginal histograms, density plots or
#'  box plots.
#' @inheritParams ggscatter
#' @param main.plot.size the width of the main plot. Default is 2.
#' @param margin.plot.size the width of the marginal plot. Default is 1.
#' @param group a grouping variable. Change points color and shape by groups if
#'  the options \code{color} and \code{shape} are missing. Should be also
#'  specified when you want to create a marginal box plot that is grouped.
#' @param margin.plot the type of the marginal plot. Default is "hist".
#' @param margin.params parameters to be applied to the marginal plots.
#' @param margin.ggtheme the theme of the marginal plot. Default is
#'  \code{\link[ggplot2:ggtheme]{theme_void}()}.
#' @param margin.space logical value. If TRUE, adds space between the main plot
#'  and the marginal plot.
#' @param bins Number of histogram bins. Defaults to 30. Pick a better value that
#'  fit to your data.
#' @param linetype line type ("solid", "dashed", ...)
#' @param legend specify the legend position. Allowed values include: "top",
#'  "bottom", "left", "right".
#' @param ggtheme the theme to be used for the scatter plot. Default is
#'  \code{\link{theme_pubr}()}.
#' @param print logical value. If \code{TRUE} (default), print the plot.
#' @param ... other arguments passed to the function \code{\link{ggscatter}()}.
#' @return an object of class \code{ggscatterhist}, which is a list of ggplots,
#'  including the following elements: \itemize{\item sp: main scatter plot;
#'  \item xplot: marginal x-axis plot; \item yplot: marginal y-axis plot. }.
#'
#'  User can modify each of plot before printing.
#'
#' @examples
#' # Basic scatter plot with marginal density plot
#' ggscatterhist(iris,
#'   x = "Sepal.Length", y = "Sepal.Width",
#'   color = "#00AFBB",
#'   margin.params = list(fill = "lightgray")
#' )
#'
#'
#' # Grouped data
#' ggscatterhist(
#'   iris,
#'   x = "Sepal.Length", y = "Sepal.Width",
#'   color = "Species", size = 3, alpha = 0.6,
#'   palette = c("#00AFBB", "#E7B800", "#FC4E07"),
#'   margin.params = list(fill = "Species", color = "black", size = 0.2)
#' )
#'
#' # Use boxplot as marginal
#' ggscatterhist(
#'   iris,
#'   x = "Sepal.Length", y = "Sepal.Width",
#'   color = "Species", size = 3, alpha = 0.6,
#'   palette = c("#00AFBB", "#E7B800", "#FC4E07"),
#'   margin.plot = "boxplot",
#'   ggtheme = theme_bw()
#' )
#'
#' # Add vertical and horizontal line to a ggscatterhist
#' plots <- ggscatterhist(iris, x = "Sepal.Length", y = "Sepal.Width", print = FALSE)
#' plots$sp <- plots$sp +
#'   geom_hline(yintercept = 3, linetype = "dashed", color = "blue") +
#'   geom_vline(xintercept = 6, linetype = "dashed", color = "red")
#' plots
#'
#' @export
ggscatterhist <- function(
  data, x, y, group = NULL,
  color = "black", fill = NA, palette = NULL,
  shape = 19, size = 2,
  linetype = "solid", bins = 30,
  margin.plot = c("density", "histogram", "boxplot"),
  margin.params = list(),
  margin.ggtheme = theme_void(),
  margin.space = FALSE,
  main.plot.size = 2,
  margin.plot.size = 1,
  title = NULL, xlab = NULL, ylab = NULL,
  legend = "top", ggtheme = theme_pubr(),
  print = TRUE, ...
) {
  if (!has_cowplot_v0.9()) {
    warning("Install the latest developmental version of cowplot on github ",
      "to fully use all the features of ggscatterhist",
      call. = FALSE
    )
    # margin.space = TRUE
  }

  margin.plot <- match.arg(margin.plot)
  margin.params <- .check_margin_params(
    margin.params, data, color, fill, linetype
  )
  if (margin.plot == "histogram") {
    if (is.null(margin.params$position)) {
      margin.params$position <- "identity"
    }
    margin.params <- margin.params %>%
      .add_item(bins = bins)
  }

  if (!is.null(group)) {
    if (missing(color)) color <- group
    if (missing(shape)) shape <- group
  }

  . <- NULL

  sp <- ggscatter(
    data, x, y,
    color = color, fill = fill,
    palette = palette,
    shape = shape, size = size,
    xlab = xlab, ylab = ylab,
    ggtheme = ggtheme, title = title,
    legend = legend,
    ...
  )

  # Type of graphics to be added in the margins
  geomfunc <- switch(margin.plot,
    histogram = geom_histogram,
    density = geom_density,
    boxplot = geom_boxplot,
    geom_histogram
  )

  # Define the x and the y variables depending on
  # the geometry used in margins
  if (margin.plot %in% c("density", "histogram")) {
    xplot.x <- x
    xplot.y <- NULL
    yplot.x <- y
    yplot.y <- NULL
  } else if (margin.plot %in% c("boxplot")) {
    if (is.null(group)) {
      data <- data %>%
        mutate(.xgroupx. = factor(1))
      group <- ".xgroupx."
    }
    xplot.x <- group
    xplot.y <- x
    yplot.x <- group
    yplot.y <- y
  }

  # Create the different marginal plot
  xplot <- ggplot() + margin.params %>%
    .add_item(
      geomfunc = geomfunc, data = data,
      x = xplot.x, y = xplot.y, alpha = 0.7
    ) %>%
    do.call(geom_exec, .)
  xplot <- set_palette(xplot, palette)

  yplot <- ggplot() + margin.params %>%
    .add_item(
      geomfunc = geomfunc, data = data,
      x = yplot.x, y = yplot.y, alpha = 0.7
    ) %>%
    do.call(geom_exec, .)
  yplot <- set_palette(yplot, palette)

  # Align the marginal plot axes with the main scatter plot's data ranges.
  # The margins are built from the raw data, so any transformation that changes
  # the scatter limits (e.g. ellipse = TRUE, position jitter, explicit
  # xlim/ylim) was not reflected in the margins, making them misaligned with the
  # scatter (#220, #420). Anchor each margin's DATA axis to the scatter's panel
  # range: set that scale's expansion to zero (the scatter range already
  # includes expansion) and zoom with coord_*(). The other (count/density) axis
  # keeps its default expansion. Only linear axes are aligned (see
  # .get_scatter_ranges): forcing a transformed scatter range onto a raw-data
  # margin would push it off-screen.
  sp.ranges <- .get_scatter_ranges(sp)
  # Flip the marginal plots
  if (margin.plot %in% c("density", "histogram")) {
    if (sp.ranges$x_identity && !is.null(sp.ranges$x)) {
      xplot <- xplot +
        scale_x_continuous(expand = expansion(0)) +
        coord_cartesian(xlim = sp.ranges$x)
    }
    if (sp.ranges$y_identity && !is.null(sp.ranges$y)) {
      yplot <- yplot +
        scale_x_continuous(expand = expansion(0)) +
        coord_flip(xlim = sp.ranges$y)
    } else {
      yplot <- yplot + coord_flip()
    }
  } else if (margin.plot %in% c("boxplot")) {
    xplot <- xplot + coord_flip()
  }

  # Cleaning the plot
  .legend <- get_legend(sp)
  sp <- sp + theme(plot.margin = grid::unit(c(0, 0, 0.25, 0.25), "cm"))
  xplot <- xplot + margin.ggtheme + clean_theme() + rremove("legend") +
    theme(plot.margin = grid::unit(c(0, 0, 0, 0), "cm"))
  yplot <- yplot + margin.ggtheme + clean_theme() + rremove("legend") +
    theme(plot.margin = grid::unit(c(0, 0, 0, 0), "cm"))

  plots <- list(sp = sp, xplot = xplot, yplot = yplot)
  class(plots) <- c("ggscatterhist", "list")
  if (print) {
    res <- print(
      plots,
      margin.space = margin.space,
      main.plot.size = main.plot.size, margin.plot.size = margin.plot.size,
      title = title, legend = legend
    )
  }
  invisible(plots)
}

#' @method print ggscatterhist
#' @param x an object of class \code{ggscatterhist}.
#' @rdname ggscatterhist
#' @export
print.ggscatterhist <- function(x, margin.space = FALSE, main.plot.size = 2,
                                margin.plot.size = 1, title = NULL, legend = "top", ...) {
  sp <- x$sp
  xplot <- x$xplot
  yplot <- x$yplot
  .legend <- get_legend(sp)

  if (margin.space) {
    common.legend <- FALSE
    if (!is.null(.legend)) common.legend <- TRUE
    sp <- sp + theme(
      plot.title = element_blank(),
      plot.subtitle = element_blank()
    )

    fig <- ggarrange(
      xplot, NULL, sp, yplot,
      ncol = 2, nrow = 2,
      align = "hv", widths = c(main.plot.size, margin.plot.size),
      heights = c(margin.plot.size, main.plot.size),
      common.legend = common.legend, legend = legend
    )
    if (!is.null(title)) {
      fig <- annotate_figure(
        fig,
        top = text_grob(title, color = "black", size = 13, face = "bold")
      )
    }
  } else {
    if (!is.null(title)) sp <- sp + ggtitle(title)
    fig <- .insert_xaxis_grob(sp, xplot, grid::unit(margin.plot.size / 5, "null"), position = "top")
    fig <- .insert_yaxis_grob(fig, yplot, grid::unit(margin.plot.size / 5, "null"), position = "right")
    fig <- cowplot::ggdraw(fig)
  }
  print(fig)
  invisible(fig)
}


has_cowplot_v0.9 <- function() {
  vv <- as.character(utils::packageVersion("cowplot"))
  cc <- utils::compareVersion(vv, "0.8.0.8") > 0
  cc
}

# Extract the final x and y coordinate ranges (including expansion) of a built
# scatter plot, plus whether each axis uses an identity (linear) transformation.
# Used to align the marginal plots in ggscatterhist() (#220, #420). Alignment is
# applied only to linear axes: a transformed scatter axis (log, sqrt, ...) lives
# in different units than the raw-data margins, so forcing its range onto the
# margin would push the margin off-screen. Falls back gracefully across ggplot2
# panel_params / scale layouts.
.get_scatter_ranges <- function(sp) {
  b <- ggplot2::ggplot_build(sp)
  pp <- b$layout$panel_params[[1]]
  axis_range <- function(axis) {
    r <- pp[[paste0(axis, ".range")]]
    if (is.null(r) && !is.null(pp[[axis]])) r <- pp[[axis]]$continuous_range
    r
  }
  axis_is_identity <- function(axis) {
    sc <- b$layout[[paste0("panel_scales_", axis)]][[1]]
    if (is.null(sc)) return(TRUE)
    tr <- tryCatch(sc$get_transformation(), error = function(e) NULL)
    if (is.null(tr)) tr <- tryCatch(sc$trans, error = function(e) NULL)
    # Treat "can't determine" as linear (the common case); a real transform
    # reports its name (e.g. "log-10", "sqrt").
    is.null(tr) || is.null(tr$name) || identical(tr$name, "identity")
  }
  list(
    x = axis_range("x"), y = axis_range("y"),
    x_identity = axis_is_identity("x"), y_identity = axis_is_identity("y")
  )
}

.check_margin_params <- function(params, data, color = "black", fill = NA, linetype = "solid") {
  if (is.null(params$color)) {
    if (color %in% colnames(data)) {
      col.val <- dplyr::pull(data, color)
      if (!is.factor(col.val)) {
        params$color <- "black"
      } else {
        params$color <- color
      }
    }
  }

  if (is.null(params$fill)) {
    if (fill %in% colnames(data)) {
      fill.val <- .select_vec(data, fill)
      if (!is.factor(fill.val)) {
        params$fill <- "black"
      } else {
        params$fill <- fill
      }
    }
  }
  if (is.null(params$linetype)) params$linetype <- linetype
  params
}

# Helper functions to insert marginal plots
# ::::::::::::::::::::::::::::::::::::::::::::::::::::
# Use cowplot::insert_xaxis_grob and cowplot::insert_yaxis_grob,
# when 0.9 stable version released

.insert_xaxis_grob <- function(
  plot, grob, height = grid::unit(0.2, "null"),
  position = c("top", "bottom")
) {
  if (inherits(grob, "ggplot")) {
    grob <- .get_panel(grob)
  }

  gt <- .plot_to_gtable(plot)
  pp <- gt$layout[gt$layout$name == "panel", ]

  if (position[1] == "top") {
    g <- gtable::gtable_add_rows(gt, height, pp$t - 1)
    g <- gtable::gtable_add_grob(
      g, grob, pp$t, pp$l, pp$t,
      pp$r,
      clip = "inherit", name = "xaxis-grob-t"
    )
  } else {
    g <- gtable::gtable_add_rows(gt, height, pp$b)
    g <- gtable::gtable_add_grob(
      g, grob, pp$b + 1, pp$l,
      pp$b + 1, pp$r,
      clip = "inherit", name = "xaxis-grob-b"
    )
  }
}

.insert_yaxis_grob <- function(
  plot, grob, width = grid::unit(0.2, "null"),
  position = c("right", "left")
) {
  if (inherits(grob, "ggplot")) {
    grob <- .get_panel(grob)
  }

  gt <- .plot_to_gtable(plot)
  pp <- gt$layout[gt$layout$name == "panel", ]

  if (position[1] == "right") {
    g <- gtable::gtable_add_cols(gt, width, pp$r)
    g <- gtable::gtable_add_grob(
      g, grob, pp$t, pp$r + 1,
      pp$b, pp$r + 1,
      clip = "inherit", name = "yaxis-grob-r"
    )
  } else {
    g <- gtable::gtable_add_cols(gt, width, pp$l - 1)
    g <- gtable::gtable_add_grob(
      g, grob, pp$t, pp$l, pp$b,
      pp$l,
      clip = "inherit", name = "yaxis-grob-l"
    )
  }
}


.plot_to_gtable <- function(plot) {
  if (inherits(plot, "gtable")) {
    return(plot)
  } else {
    ggplot2::ggplotGrob(plot)
  }
}

.get_panel <- function(plot) {
  gt <- .plot_to_gtable(plot)
  panelIndex <- which(gt$layout$name == "panel")
  panel <- gt$grobs[[panelIndex]]
}
