# Plotting results



#==============================================================================>
# Helper functions

#---------------------------------------
labels_10x <- function(x) {
  text <- str_replace_all(scientific_format()(x), ".*e", "10^") %>%
    str_replace_all("[+]", "")
  ifelse(x == 0, "0", parse(text = text))
}



#---------------------------------------
breaks_10x <- function(x) {
  breaks <- 10^(0:10)
  index <- which( (breaks > x[1]) & (breaks < x[2]) )
  #index <- c( max(index[1]-1, 1), index )
  #index <- c(index, min(index[length(index)]+1, length(breaks)))
  index <- sort(unique(index))
  n <- length(index)
  if ( n > 5 ) index <- index[seq(1, n, by = 2)]
  breaks[index]
}



#==============================================================================>
# Main plotting functions



#------------------------------------------------------------------------------
#' @export
plot_boxes <- function(d, x, y, fill = NULL, log10 = FALSE, 
                       points = TRUE, alpha = 1, facet = NULL, ncol = 3) {

  # Modifying d columns to keep code the same
  d$x <- d[[x]]
  d$y <- d[[y]]

  if ( log10 ) {
    d <- d %>%
      mutate(y = ifelse(y <= 1e-1, NA, y)) %>%
      filter(is.finite(y))   
  }

  if (! is.null(fill) ) {
    d$fill <- d[[fill]] 
  }

  if (! is.null(facet) ) {
    d$group <- d[[facet]] 
  }

  # Basic plot
  p <- ggplot(d) +
    geom_boxplot(aes(x = x, y = y, fill = fill), outlier.size = 0 ) +
    scale_fill_brewer(fill, palette = "Dark2") +
    xlab(x) + 
    ylab(y)

  # Adding points if required
  if ( points ) {
    p <- p + 
      geom_point(aes(x = x, y = y, group = fill), alpha = alpha,
                 position = position_dodge(width = 0.75))
  }

  # Faceting if required
  if (! is.null(facet) ) {
    p <- p + facet_wrap(~ group, ncol = ncol, scale = "free_y")
  }

  # Log scale if required
  if ( log10 ) {
    p <- p + scale_y_log10( breaks = breaks_10x, labels = labels_10x)
  }

  p + theme_bw(18)
}



#------------------------------------------------------------------------------
#' @export
add_comparisons <- function(p, d, x, x.split, y, reference = "reference", 
                            y.offset = 0.5, tick.height = 0.1) {

  # Modifying d columns to keep code the same
  d$x <- d[[x]]
  d$y <- d[[y]]
  d$x.split <- d[[x.split]] 
  d$reference <- d[[reference]]

  # Expanding x values into xstart and xstop
  level.names <- levels(d$x.split)
  n.levels <- length(level.names) 
  level.rank <- 1:n.levels
  names(level.rank) <- level.names

  f_offset <- function(x) (x - 1)/(n.levels - 1)*0.5 - 0.25 

  d <- d %>%
    mutate(x = as.numeric(factor(d$x)),
           start = as.numeric(d$reference),
           start = f_offset(start),
           end = as.numeric(d$x.split),
           end = f_offset(end),
           xstart = x + start,
           xend = x + end,
           xmid = (xstart + xend)/2,
           level = rank(level.rank[reference] + level.rank[x.split]),
           y = y + (level)^y.offset*y,
           ymin = (1 - tick.height)*y) %>%
    filter(label != "")

  # Plotting
  p <- p +
    geom_segment(data = d, aes(x = xstart, xend = xend, y = y, yend = y),
                 size = 1, lineend = "round") +
    geom_segment(data = d, aes(x = xstart, xend = xstart, y = y, yend = ymin),
                 size = 1, lineend = "round") +
    geom_segment(data = d, aes(x = xend, xend = xend, y = y, yend = ymin),
                 size = 1, lineend = "round") +
    geom_text(data = d, aes(x = xmid, y = y, label = label), 
              vjust = 0.3, size = 6)

  p
}



#------------------------------------------------------------------------------
#' @export
plot_bars <- function(d, x, fill = NULL, facet = NULL, ncol = 3, position = "stack") {

  # Modifying d columns to keep code the same
  d$x <- d[[x]]

  if (! is.null(fill) ) {
    d$fill <- d[[fill]] 
  }

  if (! is.null(facet) ) {
    d$group <- d[[facet]] 
  }

  # Basic plot
  p <- ggplot(d, aes(x = x, fill = fill)) +
    stat_count(position = position) +
    scale_fill_brewer(fill, palette = "Dark2") +
    xlab(x)

  # Faceting if required
  if (! is.null(facet) ) {
    p <- p + facet_wrap(~ group, ncol = ncol, scale = "free_y")
  }

  p + theme_bw(18)
}



#------------------------------------------------------------------------------
#' @export
plot_points <- function(d, x, y, colour = NULL, size = NULL, log10 = FALSE,  
                        facet = NULL, ncol = 3) {

  # Modifying d columns to keep code the same
  d$x <- d[[x]]
  d$y <- d[[y]]

  if (! is.null(colour) ) {
    d$colour <- d[[colour]] 
  }

  if (! is.null(facet) ) {
    d$group <- d[[facet]] 
  }

  # Basic plot
  p <- ggplot(d, aes(x = x, y = y, colour = colour)) +
    geom_point(size = size) +
    scale_colour_brewer(colour, palette = "Dark2") +
    xlab(x)

  # Faceting if required
  if (! is.null(facet) ) {
    p <- p + facet_wrap(~ group, ncol = ncol, scale = "free_y")
  }

  # Log scale if required
  if ( log10 ) {
    p <- p + scale_y_log10( breaks = breaks_10x, labels = labels_10x)
  }

  p + theme_bw(18)
}