RSFIA: FIA Tree Mortality Analysis

#' @title Plots mortality data as a boxplot of various types
#'
#' @description Can be used to plot multiple mortality types by section / province /
#' forest type, can drop IQR outliers, etc
#'
#' This is a plotting 'workhorse' wrapped by other functions in
#' to make final figures
#'
#' @param plot_type one of the mortality variables
#' @param by one of 'section', 'province', or 'forest type'
#' @param zero_as_NA logical flag, convert zeroes in mort data to NAs?
#' @param drop_NA logical flag, drop by groups that are all NA?
#' @param IQR_outliers logical flag, drop outliers as detected by IQR?
#' @param IQR_outliers which IQR group to plot? can be background or high
#' @param add_sample_size logical tag, add sample size to x-axis labels?
#' @param ftype Forest type to subset by
#' @param clear_plots logical to pass to PrepPlorEnvir
#' @param prnt Print summary statistics
#' @param flipped Flip the ggplot axis?
#'
#' @details
#'
#' Available mortality variables: TODO this
#'
#' @examples
#' PlotMortBoxplot(plot_type = 'non_fire', by = 'forest_type',
#' drop_NA = F, zero_as_NA = F)
#' @export
PlotMortBoxplot <- function(mort_df = NULL,
                            plot_type = 'all', by = 'section', zero_as_NA = F,
                            drop_NA = F, IQR_outliers = F, IQR_type = 'high',
                            ftype = NULL, clear_plots = T, prnt = F, flipped = F) {

  # Prepare plot environment
  if (is.null(mort_df)) {
    mort_df <- PrepPlotEnvir(ret = 'plots', clear = clear_plots)
  } else {
    PrepPlotEnvir(ret = F, clear = clear_plots)
  }
  ytag <- '10 Year Mortality Rate'
  if (length(ftype) > 0) {
    mort_df <- mort_df[which(mort_df$forest_type %in% ftype), ]
  }
  if (IQR_outliers) {
    IQR_groups <- list()
    if (by == 'mort_type') {
      stop('not implemented')
    }
    b <- which(c('forest_type', 'section', 'province') == by)
    var <- switch(b, 'forest_type', 'Cleland_section', 'Cleland_province')
    IQR_groups <- CalcCategoricalMortByIQR(
      x = mort_df$mort_rate, group = mort_df[[var]]
    )[, 4]
    if (IQR_type == 'background') {
      mort_df <- mort_df[!IQR_groups, ]
      ytag <- paste0(ytag, ', Background')
    }
    if (IQR_type == 'high') {
      mort_df <- mort_df[IQR_groups, ]
      ytag <- paste0(ytag, ', Outlier Events')
    }
  }
  # ggplot() + geom_bar(data = test2, aes(Section, fill = Mort_Category)) + coord_flip()
  mortality <- list()
  if (by == 'mort_type') {
    plot_type <- 'mort_type'
    message('Setting plot_type to mort_type')
    # Overall; Overall except harvest; overall except harvest + fire;
    # overall except harvest, fire, insects
    mort0 <- c('unique_plot_key', 'LAT', 'LON')
    mort1 <- c('mort_rate', 'non_harv_mort_rate', 'non_harv_fire_mort_rate',
                   'non_harv_fire_beet_mort_rate', 'Harvest_rate_wrt_ntree')
    mort_df <- mort_df[, c(mort0, mort1)]
    mort_df <- reshape2::melt(mort_df, id.vars = mort0, measure.vars = mort1)
    u <- 0
    for (i in mort1) {
      u <- u + 1
      if (i == 'mort_rate') {
        i <- '^mort_rate$'
      }
      rep_pat <- switch(u, 'Total', 'Non-Harvest', 'Non-Harvest/Fire', 'Remainder',
                        'Harvest')
      mort_df$variable <- gsub(pattern = i, replacement = rep_pat, x = mort_df$variable)
    }
    colnames(mort_df)[c(4, 5)] <- c('category', 'mortality')
    mortality <- mort_df$mortality
    df <- mort_df[, c(5, 4, 1, 2, 3)]
    acol <- 2
  }
  if (by != 'mort_type') {
    var <- which(c('all', 'log', 'fire', 'insect', 'diff',
                   'non_fire', 'non_harvest', 'non_harvest_fire') == plot_type)
    b <- switch(var, 'mort_rate', 'Harvest_rate_wrt_ntree', 'Harvest_rate_wrt_ntree',
                'Insect_rate_wrt_ntree', 'non_harv_fire_beet_mort_rate',
                'non_fire_mort_rate', 'non_harv_mort_rate', 'non_harv_fire_mort_rate')
    mortality <- mort_df[[b]]
  }

  if (length(mortality) == 0) stop('bad plot_type input')
  if (by != 'mort_type') {
    df <- data.frame(mortality, mort_df$Cleland_section, mort_df$Cleland_province,
                     mort_df$forest_type,
                     stringsAsFactors = F)
    df[, 2] <- KeyClelandCode(df[, 2], lvl = 'section')
    df[, 3] <- KeyClelandCode(df[, 3], lvl = 'province')
    if (by == 'section') {
      acol <- 2
    } else if (by == 'province') {
      acol <- 3
    } else if (by == 'forest_type') {
      acol <- 4
    } else {
      stop('bad by input')
    }
  }
  if (length(mortality) == 0) stop('bad plot_type input')

  # Zeros to NAs option:
  if (zero_as_NA) {
    df[, 1] <- ifelse(df[, 1] < 0.0001, NA, df[, 1])
  }
  # NA column exclude/fix option:
  NA_cols <- aggregate(df[, 1], by = list(df[, acol]),
                       FUN = function(x) all(is.na(x)))
  NA_cols <- NA_cols$Group.1[which(NA_cols[, 'x'])]
  if (drop_NA) {
    df <- df[which(!(df[, acol] %in% NA_cols)), ]
  } else {
    df[which(df[, acol] %in% NA_cols), 1] <- 0
  }
  # Add sample size:
  acol_new <- df[, acol]
  for (i in 1:length(unique(df[, acol]))) {
    i0 <- unique(df[, acol])[i]
    acol_tbl <- as.data.frame(table(df[, acol]))
    i_tag <- acol_tbl$Freq[match(i0, acol_tbl$Var1)]
    rep_tag <- paste0(i0, ' (', i_tag, ')')
    if (i0 == 'Northern Rockies') {
      i0 <- paste0(i0, '$')
    }
    if (i0 == 'Sierra Nevada') {
      i0 <- paste0(i0, '$')
    }
    acol_new <- gsub(pattern = i0, replacement = rep_tag, x = acol_new)
  }
  if (by != 'mort_type') {
    df[, acol] <- acol_new
  } else {
    message(paste0('Sample size; ', i_tag))
    samp_size <- i_tag
  }

  # Plot code:
  out_plot <- ggplot()
  if (by == 'section') {
    df_out <- aggregate(df[, 1], by = list(df[, acol]), FUN = mean)
    df_out[, 2] <- round(df_out[, 2] * 100, 2)
    colnames(df_out) <- c('Section', '10-Year Mortality Rate')
    if (prnt) {
      print(df_out)
    }

    out_plot <- out_plot + theme_bw() +
      geom_boxplot(aes(x = df[, acol], y = df[, 1], fill = df[, 3])) +
      #theme(axis.text.x = element_text(angle = -45, hjust = 0)) +
      theme(plot.margin = unit(c(1, 2.5, 1, 1), 'cm')) +
      xlab('') + ylab(ytag) +
      guides(fill = guide_legend(title = "Province")) +
      scale_fill_brewer(palette = "Set2") +
      theme(legend.text = element_text(size = 8))

  } else if (by == 'province') {
    out_plot <- out_plot + theme_bw() +
      geom_boxplot(aes(x = df[, acol], y = df[, 1])) +
      #theme(axis.text.x = element_text(angle = -45, hjust = 0)) +
      theme(plot.margin = unit(c(1, 2.5, 1, 1), 'cm')) +
      xlab('') + ylab(ytag) +
      theme(legend.text = element_text(size = 8))

  } else if (by == 'forest_type') {
    df_out <- aggregate(df[, 1], by = list(df[, acol]), FUN = mean)
    df_out[, 2] <- round(df_out[, 2] * 100, 2)
    colnames(df_out) <- c('Forest Type', '10-Year Mortality Rate')
    if (prnt) {
      print(df_out)
    }

    out_plot <- out_plot + theme_bw() +
      geom_boxplot(aes(x = df[, acol], y = df[, 1], fill = df[, 3])) +
      #theme(axis.text.x = element_text(angle = -45, hjust = 0)) +
      theme(plot.margin = unit(c(1, 2.5, 1, 1), 'cm')) +
      xlab('') + ylab(ytag) +
      guides(fill = guide_legend(title = 'Province')) +
      scale_fill_brewer(palette = 'Set2') +
      theme(legend.text = element_text(size = 8))

  } else if (by == 'mort_type') {
    df_out <- aggregate(df[, 1], by = list(df[, acol]), FUN = mean)
    df_out[, 2] <- round(df_out[, 2] * 100, 2)
    colnames(df_out) <- c('Mortality Source', '10-Year Mortality Rate')
    if (prnt) {
      print(df_out)
    }

    xtag <- paste0(ftype, ', n = ', samp_size)
    out_plot <- out_plot + theme_bw() +
      geom_boxplot(aes(x = df[, acol], y = df[, 1])) +
      #theme(axis.text.x = element_text(angle = -45, hjust = 0)) +
      theme(plot.margin = unit(c(0.5, 1, 0.5, 1), 'cm')) +
      xlab(xtag) + ylab(ytag) +
      theme(legend.text = element_text(size = 8))

  } else {
    stop('by must be either province, section, or forest type')
  }

  if (flipped) {
    out_plot <- out_plot + coord_flip()
  } else {
    out_plot <- out_plot + theme(axis.text.x = element_text(angle = -45, hjust = 0))
  }

  return(out_plot)
}
#' @describeIn PlotMortBoxplot Wrapper for multi-panel mort boxplot
#' @export
PlotOutlierMortBoxplot <- function(lab_pos = -0.5) {
  plot0 <- PlotMortBoxplot(plot_type = 'all', IQR_outliers = T, IQR_type = 'high',
                           drop_NA = F, zero_as_NA = F, flipped = F)
  txt0 <- textGrob('(b)', gp = gpar(fontsize = 18))
  txt0_ymin <- 1 * lab_pos
  plot0 <- plot0 + coord_flip(clip = 'off', ylim = c(0, 1)) +
    annotation_custom(txt0, xmin = 38.5, ymax = txt0_ymin)

  plot1 <- PlotMortBoxplot(plot_type = 'all', IQR_outliers = T, IQR_type = 'background',
                           drop_NA = F, zero_as_NA = F, flipped = F)
  txt1 <- textGrob('(a)', gp = gpar(fontsize = 18))
  txt1_ymin <- 0.65 * lab_pos
  plot1 <- plot1 + coord_flip(clip = 'off', ylim = c(0, 0.65)) +
    annotation_custom(txt1, xmin = 38.5, ymax = txt1_ymin)

  Multiplot(plot1, plot0)
}
#' @describeIn PlotMortBoxplot Wrapper for multi-panel ftype boxplot
#' @export
PlotOutlierForestBoxplot <- function(lab_pos = -0.5) {
  mort_df <- BinUncommonForestType()[[1]]
  mort_df <- SubDominantForestType(mort_df = mort_df, sub = F, n_plot_cutoff = 400)
  message('Aggregating forest types by most common province...')
  mort_df <- AggForTypeByProv(mort_df = mort_df)

  # bottom panel:
  plot0 <- PlotMortBoxplot(mort_df = mort_df,
                           plot_type = 'all', by = 'forest_type',
                           IQR_outliers = T, IQR_type = 'high',
                           drop_NA = F, zero_as_NA = F, flipped = F)
  txt0 <- textGrob('(b)', gp = gpar(fontsize = 18))
  txt0_ymin <- 1 * lab_pos
  plot0 <- plot0 + coord_flip(clip = 'off', ylim = c(0, 1)) +
    annotation_custom(txt0, xmin = 38.5, ymax = txt0_ymin)

  # top panel:
  plot1 <- PlotMortBoxplot(mort_df = mort_df,
                           plot_type = 'all', by = 'forest_type',
                           IQR_outliers = T, IQR_type = 'background',
                           drop_NA = F, zero_as_NA = F, flipped = F)
  txt1 <- textGrob('(a)', gp = gpar(fontsize = 18))
  txt1_ymin <- 0.65 * lab_pos
  plot1 <- plot1 + coord_flip(clip = 'off', ylim = c(0, 0.65)) +
    annotation_custom(txt1, xmin = 38.5, ymax = txt1_ymin)

  Multiplot(plot1, plot0)
  invisible()
}
#' @describeIn PlotMortBoxplot Wrapper for dominant-forest boxplot
#' @family plot_wrappers
#' @export
PlotDomForBoxplot <- function(lab_pos = -0.5) {
  mort_df <- SubDominantForestType(sub = T, n_plot_cutoff = 400)

  plot0 <- PlotMortBoxplot(mort_df = mort_df,
                           plot_type = 'all', by = 'forest_type',
                           IQR_outliers = T, IQR_type = 'high',
                           drop_NA = F, zero_as_NA = F, flipped = F)
  txt0 <- textGrob('(b)', gp = gpar(fontsize = 18))
  txt0_ymin <- 1 * lab_pos
  plot0 <- plot0 + coord_cartesian(clip = 'off', ylim = c(0, 1)) +
    annotation_custom(txt0, xmin = 38.5, ymax = txt0_ymin) +
    theme(axis.text.x = element_text(angle = -45, hjust = 0))

  plot1 <- PlotMortBoxplot(mort_df = mort_df,
                           plot_type = 'all', by = 'forest_type',
                           IQR_outliers = T, IQR_type = 'background',
                           drop_NA = F, zero_as_NA = F, flipped = F)
  txt1 <- textGrob('(a)', gp = gpar(fontsize = 18))
  txt1_ymin <- 0.65 * lab_pos
  plot1 <- plot1 + coord_cartesian(clip = 'off', ylim = c(0, 0.65)) +
    annotation_custom(txt1, xmin = 38.5, ymax = txt1_ymin) +
    theme(axis.text.x = element_text(angle = -45, hjust = 0))

  Multiplot(plot1, plot0)
  invisible()
}