plot_utils.R

# Map census data 
#'
#' @param census_data, dataframe of census data for specified variable of interest 
#' @param conus_sf, sf of conus states outline
#' @param leg_title, character string for legend title
#' @param outfile_path, outfile path for pngs 
#' @param viz_config_df `data.frame` width, height, counties outline color, conus outline color, background color, font nam, and font size 
#' @param viz_config_pal `data.frame` assign colors for postively and negatively correlated dimensions for census maps
#' @param width, set figure width dimension
#' @param height, set figure height dimension
#' @param percent_leg if else statement where if TRUE, apply 0-100 legend, otherwise retain 0 - max of variable name
#' @param dollar_leg if else statement where if TRUE, apply $ to the estimate column to display cost in USD on map
#' @param font_size, set font size 
#' @param barheight, set colorbar bar height
#' @param barwidth, set colorbar bar width 
#' @param lim_vals, set legend limits
#' @param break_vals, set legend breaks 
#' @param low_ramp_col, assign color to low end of color ramp for legend 
plot_census_map <- function(census_data, conus_sf, leg_title, outfile_path, var, 
                            percent_leg, viz_config_df, viz_config_pal, width, height, 
                            font_size, barwidth, barheight, lim_vals, break_vals,
                            dollar_leg, low_ramp_col){
  
  font_legend <- viz_config_df$load_font
  font_add_google(font_legend)
  showtext_opts(dpi = 300, regular.wt = 200, bold.wt = 700)
  showtext_auto(enable = TRUE)
  
  census_map <- census_data |> 
    ggplot(aes(fill = .data[[var]])) + 
    geom_sf(color = viz_config_df$counties_outline_col,
            linewidth = 0.05) +
    geom_sf(data = conus_sf,
            fill = NA,
            color = viz_config_df$conus_outline_col,
            linewidth = 0.5,
            linetype = "solid") + 
    theme_void() +
    theme(text = element_text(family = viz_config_df$font_legend, size = font_size),
          legend.margin = margin(t = 5, b = 5), legend.position = 'bottom', legend.title.align = 0.5) +
    guides(fill = guide_colorbar(
           title.position = "top",
           title.theme = element_text(face = 'bold', family = viz_config_df$font_legend, size = font_size),
           direction = "horizontal",
           position = "bottom",
           barwidth = barwidth, 
           barheight = barheight 
           )) 
  
  if (percent_leg == FALSE) {
    census_map <- census_map +     
      scale_fill_gradientn(
        colors = colorRampPalette(c(low_ramp_col, viz_config_pal))(100), 
        name = leg_title,
        limits = lim_vals,
        labels = if (dollar_leg) scales::label_dollar() else scales::comma,
        na.value = "#F5F5F5"
      ) 
  
 } else {
   census_map <- census_map +     
     scale_fill_gradientn(
       colors = colorRampPalette(c("#F8F9FF", viz_config_pal))(100), 
       name = leg_title,
       limits = lim_vals,  # c(0, 100),
       breaks = break_vals, # c(0, 25, 50, 75, 100),
       labels = if (dollar_leg) scales::label_dollar() else function(x) paste0(x, "%"),
       na.value="#F5F5F5"
     )
 }

  
  background_color = "white"
  plot_margin =  0.025
  
  # cowplot to get map sizes larger
  canvas <- grid::rectGrob(
    x = 0, y = 0,
    width = width, height = height,
    gp = grid::gpar(fill = background_color, alpha = 1, col = background_color
    )
  )
      census_legend <-  get_plot_component(census_map, 'guide-box', return_all = TRUE)
    # compose final plot
    final_map <- ggdraw(ylim = c(0,1),
                        xlim = c(0,1)) +
      # White background
      draw_grob(canvas,
                x = 0, y = 1,
                height = height, width = width,
                hjust = 0, vjust = 1) +
      # Add main plot
      draw_plot(census_map + theme(legend.position="none"),
                x = -0.01,
                y = 0.08,
                height = 0.98,
                width = (1-plot_margin)*1.03) +
      # Add legend 
      draw_plot(census_legend[[3]],
                x = 0.48,
                y = 0.02,
                height = 0.09 ,
                width = 0.1 - plot_margin) 
  
ggsave(outfile_path, final_map, width = width, height = height, dpi = viz_config_df$dpi, bg = viz_config_df$bg_col, units = "in")
  
  return(outfile_path)
}


# Map raster data 
#'
#' @param raster_data .tif file of raster data for specified variable of interest 
#' # notes: update colorpalette according to Elmera's request

plot_raster <- function(in_raster, conus_sf, conus_inner, outfile_path, 
                        viz_config_df, viz_config_pal, width, height,
                        font_size, barwidth, barheight,
                        low_ramp_col){
  # in_raster = p2_pop_density_processed, conus_sf = p2_conus_sf_proj, conus_inner = p2_conus_inner, viz_config_pal = p0_viz_config_pal$living_conditions
  # width = p0_viz_config_df$width_desktop, height = p0_viz_config_df$height_desktop, font_size  = p0_viz_config_df$font_size_desktop, barwidth = 20, barheight = 1,
  #outfile_path = "3_visualize/out/pop_density_rast_2020_en.png", low_ramp_col = "#eef0ff", 
  raster_data <- rast(in_raster)
  
  # plot
  font_legend <- viz_config_df$load_font
  font_add_google(font_legend)
  showtext_opts(dpi = 300, regular.wt = 200, bold.wt = 700)
  showtext_auto(enable = TRUE)
  
  # doesn't work with cowplot::get_plot_component()
  legend_label <- expression("Population per 1 km"^2) #bquote('Population per 1'~km^2)
  
  (pop_map <- ggplot() +
      #geom_spatraster(data = pop_usa_proj)+ 
      geom_sf(data = conus_sf, fill = low_ramp_col, color = low_ramp_col, linewidth = 0.5) +
      geom_spatraster(data = raster_data)+ #x = x, y = y, 
      geom_sf(data = conus_inner, fill = NA, color = viz_config_df$counties_outline_col, linewidth = 0.5) +
      scale_fill_gradientn(
        colors = colorRampPalette(c(low_ramp_col, viz_config_pal))(100), 
        name = legend_label,
        #limits = lim_vals,  # c(0, 100),
        #breaks = break_vals, # c(0, 25, 50, 75, 100),
        #labels = if (dollar_leg) scales::label_dollar() else function(x) paste0(x, "%"),
        labels = c(1, 10, 100, '1k', '10k', '100k'),
        na.value=NA
      )+
      #scale_fill_viridis_c(na.value = NA, option = 'mako', breaks = c(0, 1, 2, 3, 4, 5), direction = -1,
                           # leave a little room for the NAs and Inf (which are 0s)
                           # to be darker than the -1 values
       #                    labels = c(1, 10, 100, '1k', '10k', '100k'), begin = 0.05) +
      labs(fill = legend_label) + #"Population per 1 km^2"
      theme(plot.background = element_rect(fill = "white", color = "white"),
            panel.background = element_rect(fill = "white", color = "white"),
            axis.title = element_blank(), 
            axis.text = element_blank(), 
            axis.ticks = element_blank(), 
            panel.grid = element_blank(),
            legend.background = element_blank(),
            legend.direction = "horizontal",
            legend.position = "bottom",
            legend.title.align = 0.5) +
      guides(fill = guide_colorbar(
        title.position = "top",
        title.theme = element_text(face = 'bold', family = font_legend, size = font_size),
        direction = "horizontal",
        position = "bottom",
        barwidth = barwidth, 
        barheight = barheight 
      ))
    # guides(fill = guide_colorbar(
    #    direction = "horizontal",
    #    title = "Population (per km2)",
    #    barwidth = 15, 
    #    title.position = "top",
    #    title.theme = element_text(color = "white"),
    #    label.theme = element_text(color = "white")))
  )
  
  background_color = "white"
  plot_margin =  0.025
  
  # cowplot to get map sizes larger
  canvas <- grid::rectGrob(
    x = 0, y = 0,
    width = 6, height = 6,
    gp = grid::gpar(fill = background_color, alpha = 1, col = background_color
    )
  )
  pop_map_legend <-  get_plot_component(pop_map, 'guide-box', return_all = TRUE)
  # compose final plot
  final_map <- ggdraw(ylim = c(0,1),
                      xlim = c(0,1)) +
    # White background
    draw_grob(canvas,
              x = 0, y = 1,
              height = 6, width = 6,
              hjust = 0, vjust = 1) +
    # Add main plot
    draw_plot(pop_map + theme(legend.position="none"),
              x = -0.01,
              y = 0.08,
              height = 0.98,
              width = (1-plot_margin)*1.03) +
    # Add legend 
    draw_plot(pop_map_legend[[3]],
              x = 0.48,
              y = 0.02,
              height = 0.09 ,
              width = 0.1 - plot_margin) 
  
  ggsave(outfile_path, final_map, width = width, height = height, dpi = viz_config_df$dpi, bg = viz_config_df$bg_col, units = "in")
  
}