#!/usr/bin/env Rscript
utils::globalVariables(".progress")
###############################################################################
# MaAsLin3 visualizations

# Copyright (c) 2025 Harvard School of Public Health

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
###############################################################################

# MaAsLin 3 ggplot theme based on Nature journal requirements
nature_theme <- function(x_axis_labels, y_label) {
    # set default text format based on categorical and length
    angle <- NULL
    hjust <- NULL
    size <- 8
    if (max(nchar(x_axis_labels), na.rm = TRUE) > 5) {
        angle <- 45
        hjust <- 1
        size <- 6
    }
    axis_title_size <- 10
    if (nchar(y_label) > 15) {
        axis_title_size <- 8
    }
    if (nchar(y_label) > 25) {
        axis_title_size <- 6
    }
    return (
        ggplot2::theme_bw() + ggplot2::theme(
            axis.text.x = ggplot2::element_text(
                size = size,
                vjust = 1,
                hjust = hjust,
                angle = angle
            ),
            axis.text.y = ggplot2::element_text(size = 8, hjust = 1),
            axis.title = ggplot2::element_text(size = axis_title_size),
            plot.title = ggplot2::element_text(size = 7, face = 'bold'),
            legend.title = ggplot2::element_text(size = 6, face = 'bold'),
            legend.text = ggplot2::element_text(size = 6),
            axis.line = ggplot2::element_line(colour = 'black', 
                                            linewidth = .25),
            axis.line.x = ggplot2::element_line(colour = 'black', 
                                                linewidth = .25),
            axis.line.y = ggplot2::element_line(colour = 'black', 
                                                linewidth = .25),
            panel.border = ggplot2::element_blank(),
            panel.grid.major = ggplot2::element_blank(),
            panel.grid.minor = ggplot2::element_blank()
        )
    )
}

# Pre-process MaAsLin 3 results for plotting
preprocess_merged_results <- function(merged_results) {
    # Disregard abundance-induced-prevalence errors in plotting
    merged_results$error[grepl("Prevalence association possibly induced", 
                                merged_results$error)] <- NA
    # Disregard small random effect group warning
    merged_results$error[grepl("<4 average observations per random effect", 
        merged_results$error)] <- NA
    
    merged_results <-
        merged_results[is.na(merged_results$error) &
                        !is.na(merged_results$qval_individual) &
                        !is.na(merged_results$coef),]
    
    if (nrow(merged_results) == 0) {
        logging::loginfo(paste(
            "No associations were without errors. 
                No summary plot generated."
        ))
        return(NULL)
    }
    merged_results$model <-
        data.table::fifelse(merged_results$model == 'linear', 'Abundance', 'Prevalence')
    merged_results$full_metadata_name <-
        data.table::fifelse(
            merged_results$metadata == merged_results$value,
            merged_results$metadata,
            paste0(merged_results$metadata, ' ', merged_results$value)
        )
    return(merged_results)
}

# Make the coefficient plot part of the summary figure
make_coef_plot <- function(merged_results_sig, 
                        coef_plot_vars,
                        max_significance,
                        median_comparison_prevalence,
                        median_comparison_abundance,
                        median_df,
                        plot_threshold = 10) {
    coef_plot_data <-
        merged_results_sig[merged_results_sig$full_metadata_name %in% 
                            coef_plot_vars,]
    
    # Limit plotted coefficients to median +/- 10 times distance to quartiles
    get_lo = function(x, pt = plot_threshold) {
        med_x = collapse::fmedian(x)
        
        q_x = collapse::fquantile(x, .25)
        
        med_x - pt * (med_x - q_x)
    }
    
    get_hi = function(x, pt = plot_threshold) {
        med_x = collapse::fmedian(x)
        
        q_x = collapse::fquantile(x, .75)
        
        med_x + pt * (q_x - med_x)
    }
    
    quantile_df <- coef_plot_data |>
        collapse::fgroup_by("full_metadata_name") |>
        collapse::collapv(by = "full_metadata_name",
                          FUN = list(get_lo, get_hi),
                          cols = "coef") |> 
        collapse::fungroup() |> 
        collapse::fselect("full_metadata_name", 
                          "lower_q" = "get_lo.coef", 
                          "upper_q" = "get_hi.coef") |> 
        data.frame(check.names = FALSE)
        
    rownames(quantile_df) <- quantile_df$full_metadata_name
    
    # Make sure insignificant coefficients don't distort the plot
    coef_plot_data <-
        coef_plot_data[coef_plot_data$qval_individual < 
                        max_significance |
                        (coef_plot_data$coef > quantile_df[
                            coef_plot_data$full_metadata_name, 
                            'lower_q'] &
                                coef_plot_data$coef < quantile_df[
                                    coef_plot_data$full_metadata_name, 
                                    'upper_q']),]
    
    # Choose breaks for plot
    custom_break_fun <- function(n) {
        return(function(x) {
            extended_breaks <- scales::breaks_extended(n)(x)
            if (max(x) > 0) {
                extended_breaks <- extended_breaks[
                    extended_breaks <= max(x) * 0.9]
            } else {
                extended_breaks <- extended_breaks[
                    extended_breaks <= max(x) * 1.1]
            }
            if (min(x) > 0) {
                extended_breaks <- extended_breaks[
                    extended_breaks >= min(x) * 1.1]
            } else {
                extended_breaks <- extended_breaks[
                    extended_breaks >= min(x) * 0.9]
            }
            extended_breaks
        })
    }
    
    # Start plot
    p1 <-
        ggplot2::ggplot(coef_plot_data,
                        ggplot2::aes(x = .data$coef, y = .data$feature))
    
    # Show nulls the coefficients were compared against
    if (median_comparison_prevalence |
        median_comparison_abundance) {
        p1 <- p1 +
            ggplot2::guides(linetype = ggplot2::guide_legend(
                title = 'Null hypothesis', order = 1),
            ) +
            ggplot2::geom_vline(
                data = median_df[median_df$full_metadata_name %in% 
                                    coef_plot_vars,],
                ggplot2::aes(
                    xintercept = .data$median_val,
                    linetype = .data$model
                ),
                color = "darkgray"
            ) +
            ggplot2::scale_linetype_manual(values = c("Prevalence" = 
                                                        "dashed", 
                                                    "Abundance" = 
                                                        "solid"))
    } else {
        p1 <- p1 +
            ggplot2::geom_vline(
                ggplot2::aes(xintercept = 0),
                color = "darkgray",
                linetype = 'dashed'
            )
    }
    
    # Q-value color scale
    scale_fill_gradient_limits <-
        c(min(max_significance, 10 ^ floor(log10(
            min(coef_plot_data$qval_individual)
        ))), 1)
    if (min(coef_plot_data$qval_individual) < max_significance) {
        scale_fill_gradient_breaks <-
            c(10 ^ floor(log10(
                min(coef_plot_data$qval_individual)
            )), max_significance, 1)
    } else {
        scale_fill_gradient_breaks <- c(max_significance, 1)
    }
    if (min(coef_plot_data$qval_individual) < max_significance) {
        scale_fill_gradient_labels <-
            c(paste0("1e", floor(log10(
                min(coef_plot_data$qval_individual)
            ))),
            paste0(max_significance),
            "1")
    } else {
        scale_fill_gradient_labels <- c(paste0(max_significance),
                                        "1")
    }
    
    # Create the whole plot
    p1 <- p1 +
        ggplot2::geom_errorbar(
            ggplot2::aes(
                xmin = .data$coef - .data$stderr,
                xmax = .data$coef + .data$stderr
            ),
            width = 0.2
        ) +
        ggplot2::geom_point(
            data = coef_plot_data[coef_plot_data$model == 
                                    'Prevalence',],
            ggplot2::aes(
                shape = .data$model,
                fill = .data$qval_individual
            ),
            size = 4.5,
            color = "black"
        ) +
        ggplot2::scale_fill_gradient(
            low = "#008B8B",
            high = "white",
            limits = scale_fill_gradient_limits,
            breaks = scale_fill_gradient_breaks,
            labels = scale_fill_gradient_labels,
            transform = scales::pseudo_log_trans(sigma = 0.001),
            name = bquote("Prevalence" ~ P["FDR"])
        ) +
        ggnewscale::new_scale_fill() +
        ggplot2::geom_point(
            data = coef_plot_data[coef_plot_data$model == 'Abundance',],
            ggplot2::aes(
                shape = .data$model,
                fill = .data$qval_individual
            ),
            size = 4.5,
            color = "black"
        ) +
        ggplot2::scale_fill_gradient(
            low = "#8B008B",
            high = "white",
            limits = scale_fill_gradient_limits,
            breaks = scale_fill_gradient_breaks,
            labels = scale_fill_gradient_labels,
            transform = scales::pseudo_log_trans(sigma = 0.001),
            name = bquote("Abundance" ~ P["FDR"])
        ) +
        ggplot2::scale_x_continuous(
            breaks = custom_break_fun(n = 6),
            limits = c(
                min(coef_plot_data$coef) - 
                    quantile(coef_plot_data$stderr, 0.8),
                max(coef_plot_data$coef) + 
                    quantile(coef_plot_data$stderr, 0.8)
            )
        ) +
        ggplot2::scale_shape_manual(name = "Association", values =
                                        c(21, 24)) +
        ggplot2::guides(shape = ggplot2::guide_legend(order = 2), ) +
        ggplot2::labs(x = expression(paste(beta, " coefficient")),  
                    y = "Feature") +
        ggplot2::theme_bw() +
        ggplot2::theme(
            axis.title = ggplot2::element_text(size = 16),
            axis.text.y = ggplot2::element_text(size = 14),
            axis.text.x = ggplot2::element_text(size = 14),
            legend.title = ggplot2::element_text(size = 16),
            legend.text = ggplot2::element_text(size = 14, 
                                                face = "plain"),
            legend.position = "right",
            legend.background = ggplot2::element_rect(
                fill = "transparent"),
            panel.spacing = ggplot2::unit(0, "lines"),
            panel.grid.minor = ggplot2::element_blank(),
            strip.text = ggplot2::element_text(size = 14),
            strip.background = ggplot2::element_rect(
                fill = "transparent")
        ) +
        ggplot2::facet_wrap(
            ~ factor(full_metadata_name, 
                    levels = unique(coef_plot_vars)),
            scales = 'free_x',
            ncol = length(coef_plot_vars)
        )
    return(p1)
}

# Make the heatmap plot part of the summary figure
make_heatmap_plot <- function(merged_results_sig,
                            heatmap_vars,
                            max_significance,
                            median_comparison_prevalence,
                            median_comparison_abundance) {
    # Create column for significance star annotation
    merged_results_sig$sig_star <-
        cut(
            merged_results_sig$qval_individual,
            breaks = c(-Inf, max_significance / 10, max_significance, Inf),
            label = c("**", "*", "")
        )
    
    # Bin coefficients into categories
    coefficient_thresh <-
        round(max(abs(
            quantile(merged_results_sig$coef, c(0.1, 0.9))
        )) / 10, 1) * 5
    if (coefficient_thresh == 0) {
        coefficient_thresh <- 0.5
    }
    coef_breaks <-
        c(
            -coefficient_thresh,-coefficient_thresh / 2,
            0,
            coefficient_thresh / 2,
            coefficient_thresh,
            Inf
        )
    threshold_set <-
        c(
            paste0("(-Inf,", -1 * coefficient_thresh, "]"),
            paste0(
                "(",-1 * coefficient_thresh,
                ",",-1 / 2 * coefficient_thresh,
                "]"
            ),
            paste0("(", -1 / 2 * coefficient_thresh, ",0]"),
            paste0("(0,", 1 / 2 * coefficient_thresh, "]"),
            paste0(
                "(",
                1 / 2 * coefficient_thresh,
                ",",
                1 * coefficient_thresh,
                "]"
            ),
            paste0("(", 1 * coefficient_thresh, ",Inf)")
        )
    
    threshold_indices <-
        vapply(merged_results_sig$coef, function(value) {
            which(value < coef_breaks)[1]
        }, FUN.VALUE = 0)
    
    merged_results_sig <- merged_results_sig |>  
        collapse::fmutate(coef_cat = threshold_set[threshold_indices])
    
    merged_results_sig$coef_cat <-
        factor(merged_results_sig$coef_cat, levels = threshold_set)
    
    # Create color scheme
    scale_fill_values <-
        rev((RColorBrewer::brewer.pal(n = 6, name = "RdBu")))
    names(scale_fill_values) <- threshold_set
    
    # Create grid for plotting
    heatmap_data <-
        merged_results_sig[merged_results_sig$full_metadata_name %in% 
                            heatmap_vars,]
    
    grid <- expand.grid(
        feature = unique(heatmap_data$feature),
        full_metadata_name = unique(heatmap_data$full_metadata_name),
        model = unique(heatmap_data$model)
    )
    if (length(unique(grid$model)) == 2) {
        grid$model <- factor(grid$model, levels = c('Prevalence', 'Abundance'))
    }
    heatmap_data <-
        merge(
            grid,
            heatmap_data,
            by = c("feature", "full_metadata_name", "model"),
            all.x = TRUE
        )
    heatmap_data$coef[is.na(heatmap_data$coef)] <- NA
    
    # Create whole plot
    p2 <-
        ggplot2::ggplot(heatmap_data,
                        ggplot2::aes(
                            x = factor(.data$full_metadata_name, 
                                    unique(heatmap_vars)),
                            y = .data$feature
                        )) +
        ggplot2::geom_tile(
            data = heatmap_data,
            ggplot2::aes(fill = .data$coef_cat),
            colour = "white",
            linewidth = 0.2
        ) +
        ggplot2::scale_fill_manual(
            name = "Beta coefficient",
            na.value = "#EEEEEE",
            values = scale_fill_values
        ) +
        ggplot2::geom_text(
            ggplot2::aes(
                label = .data$sig_star,
                color = .data$sig_star
            ),
            size = 6,
            vjust = 0.75,
            hjust = 0.5,
            key_glyph = ggplot2::draw_key_blank
        ) +
        ggplot2::scale_color_manual(
            name = bquote("Covariates" ~ P["FDR"]),
            breaks = c("*", "**", ""),
            values = c("black", "black", "black"),
            labels = c(paste0("* < ", round(
                max_significance, 3
            )), paste0(
                "** < ", round(max_significance / 10, 5)
            ), ""),
            drop = FALSE
        ) +
        ggplot2::labs(x = '',
                    y = "Feature",
                    caption = "") +
        ggplot2::theme_bw() +
        ggplot2::theme(
            axis.title = ggplot2::element_text(size = 16),
            axis.text.x = ggplot2::element_text(
                size = 14,
                angle = 90,
                vjust = 0.5,
                hjust = 1
            ),
            legend.title = ggplot2::element_text(size = 16),
            legend.text = ggplot2::element_text(size = 14, 
                                                face = "plain"),
            legend.position = "right",
            legend.background = ggplot2::element_rect(
                fill = "transparent"),
            panel.spacing = ggplot2::unit(0, "lines"),
            panel.grid.minor = ggplot2::element_blank(),
            strip.text = ggplot2::element_text(size = 14),
            strip.background = ggplot2::element_rect(
                fill = "transparent")
        ) +
        ggplot2::guides(
            fill = ggplot2::guide_legend(order = 1),
            color = ggplot2::guide_legend(order = 2),
        ) +
        ggplot2::facet_grid(~ model, labeller = ggplot2::labeller(
            model = c("abundance" = "Abundance", 
                    "prevalence" = "Prevalence")
        ))
    return(p2)
}

# MaAsLin3 summary_plot function for overall view of associations
maaslin3_summary_plot <-
    function(merged_results,
            summary_plot_file,
            figures_folder,
            first_n = 30,
            max_significance = 0.1,
            coef_plot_vars = NULL,
            heatmap_vars = NULL,
            median_comparison_abundance = FALSE,
            median_comparison_prevalence = FALSE,
            balanced = FALSE,
            save_plots_rds = FALSE) {
        ret_plots <- list()
        if (first_n > 200) {
            logging::logerror(
                paste(
                    "At most 200 features can be plotted in the heatmap. 
                    Please choose a smaller first_n."
                )
            )
            return()
        }
        
        merged_results <- preprocess_merged_results(merged_results)
        
        if (is.null(merged_results) || nrow(merged_results) == 0) {
            return()
        }
        
        median_df <- merged_results |> 
            collapse::fselect(  c("full_metadata_name", "model", "null_hypothesis")) |> 
            collapse::fgroup_by(c("full_metadata_name", "model")) |> 
            collapse::fmedian() |> 
            collapse::frename("median_val" = "null_hypothesis")

        # Check variables can be plotted
        if (!is.null(coef_plot_vars) | !is.null(heatmap_vars)) {
            if (any(
                !c(coef_plot_vars, heatmap_vars) %in% 
                unique(merged_results$full_metadata_name)
            )) {
                logging::loginfo(
                    paste0(
                        "The following specified variables were not 
                        found in the associations: ",
                        paste0(setdiff(
                            c(coef_plot_vars, heatmap_vars),
                            unique(merged_results$full_metadata_name)
                        ), collapse = ', '),
                        collapse = ''
                    )
                )
                logging::loginfo(paste0(
                    "Available associations: ",
                    paste0(
                        unique(merged_results$full_metadata_name),
                        collapse = ', '
                    ),
                    collapse = ''
                ))
                return(NULL)
            }
            merged_results <-
                merged_results[merged_results$full_metadata_name %in%
                                c(coef_plot_vars, heatmap_vars),]
        }
        
        # Subset associations for plotting
        merged_results_joint_only <- merged_results |> 
            collapse::fselect(c('feature', 'qval_joint', 'full_metadata_name')) |> 
            collapse::funique() |> 
            collapse::roworderv("qval_joint")
        
        if (length(unique(merged_results_joint_only$feature)) < first_n) {
            first_n <- length(unique(merged_results_joint_only$feature))
            signif_taxa <-
                unique(merged_results_joint_only$feature)[seq(first_n)]
        } else {
            # If balanced is turned on but there are not
            # coefs choosen, error out
            if (balanced){
                if (is.null(coef_plot_vars)){
                    logging::logerror(
                        paste(
                            "Balanced plotting requires 
                            you set the variables you 
                            want to plot using
                            the parameter coef_plot_vars"
                        )
                    )
                    return()
                } else {
                    # grab the first N feature where 
                    # N=N/(length of coef_plot_var) to
                    # plot the coef plot
                    first_n_per <- ceiling(first_n/length(coef_plot_vars))
                    
                    signif_taxa <- merged_results_joint_only |>
                        collapse::roworderv(cols = c("full_metadata_name", 
                                                     "qval_joint")) |> 
                        collapse::fslicev("full_metadata_name",
                                          n = first_n_per) |> 
                        collapse::get_elem("feature") |> 
                        collapse::funique()
                }
            } else {
                signif_taxa <-
                unique(merged_results_joint_only$feature)[seq(first_n)]
            }
        }

        
        merged_results_sig <- merged_results |> 
            collapse::fsubset(merged_results$feature %in% signif_taxa)
        # TODO: consider collapse::`%iin%`
        
        # Order features
        ord_feature <-
            with(merged_results_sig, reorder(feature, coef))
        ord_feature <- levels(ord_feature)
        
        merged_results_sig$feature <-
            factor(merged_results_sig$feature, levels = ord_feature)
        
        # Choose variables for plotting if not set
        if (is.null(coef_plot_vars)) {
            
            mean_log_qval <- merged_results_sig |> 
                collapse::collapv(by   = "full_metadata_name",
                                  cols = "qval_joint",
                                  FUN  = \(x) collapse::fmean(log(x), na.rm = TRUE)) |> 
                collapse::fselect("full_metadata_name", 
                                  "mean_value" = "qval_joint")
                
            coef_plot_vars <-
                mean_log_qval$full_metadata_name[
                    order(mean_log_qval$mean_value)]
            
            coef_plot_vars <- setdiff(coef_plot_vars, heatmap_vars)
            
            if (length(coef_plot_vars) > 0) {
                coef_plot_vars <-
                    coef_plot_vars[seq(min(2, length(coef_plot_vars)))]
            }
        }
        
        # Choose variables for plotting if not set
        if (is.null(heatmap_vars)) {
            
            mean_log_qval <- merged_results_sig |> 
                collapse::collapv(by   = "full_metadata_name",
                                  cols = "qval_joint",
                                  FUN  = \(x) collapse::fmean(log(x), na.rm = TRUE)) |> 
                collapse::fselect("full_metadata_name", 
                                  "mean_value" = "qval_joint")
            
            heatmap_vars <-
                mean_log_qval$full_metadata_name[
                    order(mean_log_qval$mean_value)]
            
            heatmap_vars <- setdiff(heatmap_vars, coef_plot_vars)
        }
        
        # Create coefficient plot
        if (length(coef_plot_vars) > 0 &
            sum(merged_results_sig$full_metadata_name %in% 
                coef_plot_vars) >= 1) {
            if (balanced) {
                plot_thres <- 5
            } else {
                plot_thres <- 10
            }
            p1 <- make_coef_plot(merged_results_sig, 
                                coef_plot_vars,
                                max_significance,
                                median_comparison_prevalence,
                                median_comparison_abundance,
                                median_df,
                                plot_threshold = plot_thres)
            
        } else {
            p1 <- NULL
        }
        
        # Create heatmap plot
        if (length(heatmap_vars) > 0 &
            sum(merged_results_sig$full_metadata_name %in% heatmap_vars) >= 1) {
            
            p2 <- make_heatmap_plot(merged_results_sig, 
                                    heatmap_vars,
                                    max_significance,
                                    median_comparison_prevalence,
                                    median_comparison_abundance)
            
            if (!is.null(p1)) {
                p2 <- p2 + ggplot2::theme(
                    axis.text.y = ggplot2::element_blank(),
                    axis.title.y = ggplot2::element_blank(),
                    axis.ticks.y = ggplot2::element_blank(),
                )
            }
            
        } else {
            p2 <- NULL
        }
        
        # Combine plots
        if (!is.null(p1) & !is.null(p2)) {
            final_plot <- patchwork::wrap_plots(
                p1,
                p2,
                ncol = 3,
                widths = c(max(
                    0, length(coef_plot_vars) * (max(15, max(
                        nchar(as.character(coef_plot_vars))
                    ))) / 15 - 2
                ) + 2,
                max(0, length(
                    heatmap_vars
                ) / 4 - 2) + 2,
                0.5),
                guides = 'collect'
            )
        } else if (is.null(p1) & !is.null(p2)) {
            final_plot <- p2
        } else if (!is.null(p1) & is.null(p2)) {
            final_plot <- p1
        } else {
            final_plot <- NULL
        }
        ret_plots[["coefficient"]] <- p1
        ret_plots[["heat"]] <- p2
        ret_plots[["final"]] <- final_plot
        
        # Save plot
        if (!is.null(final_plot)) {
            height_out <-
                9.5 + max(first_n / 5 - 5, 0) + max(nchar(c(
                    as.character(coef_plot_vars),
                    as.character(heatmap_vars)
                ))) / 10
            width_out <-
                5 + max(nchar(merged_results$feature)) / 12 +
                (length(coef_plot_vars) * (max(20, max(
                    nchar(as.character(coef_plot_vars))
                ))) / 20) * 2.5 +
                length(heatmap_vars) * 0.25
            
            # Silence ggplot warnings
            tryCatch({
                withCallingHandlers({
                    ggplot2::ggsave(
                        summary_plot_file,
                        plot = final_plot,
                        height = height_out,
                        width = width_out
                    )
                },
                warning = function(w) {
                    invokeRestart("muffleWarning")
                })
            })
            
            png_file <- file.path(figures_folder, "summary_plot.png")
            
            tryCatch({
                withCallingHandlers({
                    ggplot2::ggsave(png_file,
                                    plot = final_plot,
                                    height = height_out,
                                    width = width_out)
                },
                warning = function(w) {
                    invokeRestart("muffleWarning")
                })
            })
            if (save_plots_rds) {
                saveRDS(final_plot,
                        file = paste(figures_folder,
                            "/" , "summary_plot_gg.RDS", sep = ""))
            }
        }
        return(ret_plots)
    }

# Create a scatterplot for abundance vs. continuous associations
make_scatterplot <- function(joined_features_metadata_abun,
                            metadata_name,
                            feature_name,
                            normalization,
                            transformation,
                            coef_val,
                            qval,
                            N_nonzero,
                            N_total,
                            results_value) {
    match.arg(normalization, c('Total sum scaling', 'Center log ratio', 'None'))
    match.arg(transformation, c('Log base 2', 'Pseudo-log base 2', 'None'))

    temp_plot <-
        ggplot2::ggplot(data = joined_features_metadata_abun,
                        ggplot2::aes(
                            as.numeric(.data$metadata),
                            as.numeric(.data$feature_abun)
                        )) +
        ggplot2::geom_point(
            fill = 'darkolivegreen4',
            color = 'black',
            alpha = .5,
            shape = 21,
            size = 1,
            stroke = 0.15
        ) +
        ggplot2::scale_x_continuous(limits = c(
            min(joined_features_metadata_abun['metadata']),
            max(joined_features_metadata_abun['metadata'])
        )) +
        ggplot2::scale_y_continuous(
            limits = c(
                min(joined_features_metadata_abun[
                    'feature_abun']),
                max(joined_features_metadata_abun[
                    'feature_abun'])
            ),
            expand = ggplot2::expansion(mult = c(0, 0.2))
        ) +
        ggplot2::stat_smooth(
            method = "glm",
            formula = 'y ~ x',
            linewidth = 0.5,
            color = 'blue',
            na.rm = TRUE
        ) +
        ggplot2::guides(alpha = 'none') +
        ggplot2::xlab(metadata_name) +
        ggplot2::ylab(
            paste0(
                feature_name,
                '\n(Normalization: ',
                normalization,
                ', Transformation: ',
                transformation,
                ')'
            )
        ) +
        nature_theme(
            joined_features_metadata_abun['metadata'],
            paste0(
                feature_name,
                '\n(Normalization: ',
                normalization,
                ', Transformation: ',
                transformation,
                ')'
            )
        ) +
        ggplot2::annotate(
            geom = "text",
            x = Inf,
            y = Inf,
            hjust = 1,
            vjust = 1,
            label = sprintf(
            "FDR: %s\nCoefficient (in full model): %sN: %s\nN (not zero): %s",
                formatC(qval, format = "e", digits = 1),
                if (is.na(coef_val)) 'NA' else formatC(
                    coef_val,
                    format = "e",
                    digits = 1
                ),
                formatC(
                    N_total,
                    format = 'f',
                    digits = 0
                ),
                formatC(
                    N_nonzero,
                    format = 'f',
                    digits = 0
                )
            ) ,
            color = "black",
            size = 2,
            fontface = "italic"
        )
    return(temp_plot)
}

# Create boxplot for abundance vs. categorical associations
make_boxplot_lm <- function(joined_features_metadata_abun,
                            metadata_name,
                            feature_name,
                            normalization,
                            transformation,
                            coef_val,
                            qval,
                            N_nonzero,
                            N_total,
                            results_value) {
    match.arg(normalization, c('Total sum scaling', 'Center log ratio', 'None'))
    match.arg(transformation, c('Log base 2', 'Pseudo-log base 2', 'None'))
    
    temp_plot <-
        ggplot2::ggplot(data = joined_features_metadata_abun,
                        ggplot2::aes(.data$metadata, 
                                    .data$feature_abun)) +
        ggplot2::geom_boxplot(
            ggplot2::aes(fill = .data$metadata),
            outlier.alpha = 0.0,
            na.rm = TRUE,
            alpha = .5,
            show.legend = FALSE
        ) +
        ggplot2::geom_point(
            ggplot2::aes(fill = .data$metadata),
            alpha = 0.75 ,
            size = 1,
            shape = 21,
            stroke = 0.15,
            color = 'black',
            position = ggplot2::position_jitterdodge()
        ) +
        ggplot2::scale_fill_brewer(palette = "Spectral") +
        ggplot2::scale_y_continuous(
            expand = ggplot2::expansion(mult = c(0, 0.2)))
    
    temp_plot <- temp_plot +
        nature_theme(
            as.character(
                joined_features_metadata_abun$metadata),
            paste0(
                feature_name,
                '\n(Normalization: ',
                normalization,
                ', Transformation: ',
                transformation,
                ')'
            )
        ) +
        ggplot2::theme(
            panel.grid.major = ggplot2::element_blank(),
            panel.grid.minor = ggplot2::element_blank(),
            panel.background = ggplot2::element_blank(),
            axis.line = ggplot2::element_line(colour = "black")
        ) +
        ggplot2::xlab(metadata_name) +
        ggplot2::ylab(
            paste0(
                feature_name,
                '\n(Normalization: ',
                normalization,
                ', Transformation: ',
                transformation,
                ')'
            )
        ) +
        ggplot2::theme(legend.position = "none") +
        ggplot2::annotate(
            geom = "text",
            x = Inf,
            y = Inf,
            hjust = 1,
            vjust = 1,
            label = sprintf(
                "Value: %s\nFDR: %s\nCoefficient (in full model): %s",
                paste0(results_value, collapse = ', '),
                paste0(
                    formatC(qval, format = "e", digits = 1),
                    collapse = ', '
                ),
                paste0(
                    vapply(coef_val, function(cv) {
                        if (is.na(cv)) 'NA' else formatC(cv, format = "e", digits = 1)
                    }, character(1)),
                    collapse = ', '
                )
            ) ,
            color = "black",
            size = 2,
            fontface = "italic"
        )
    return(temp_plot)
}

# Create plots for abundance associations
make_lm_plot <- function(this_signif_association,
                        joined_features_metadata,
                        metadata,
                        metadata_name,
                        feature_name,
                        normalization,
                        transformation,
                        feature_specific_covariate_name,
                        feature_specific_covariate) {
    match.arg(normalization, c('Total sum scaling', 'Center log ratio', 'None'))
    match.arg(transformation, c('Log base 2', 'Pseudo-log base 2', 'None'))

    coef_val <-
        this_signif_association[
            this_signif_association$model == 'linear',]$coef
    qval <-
        this_signif_association[
            this_signif_association$model == 'linear',]$qval_individual
    N_nonzero <-
        this_signif_association[
            this_signif_association$model == 'linear',]$N_not_zero
    N_total <-
        this_signif_association[
            this_signif_association$model == 'linear',]$N
    results_value <-
        this_signif_association[
            this_signif_association$model == 'linear',]$value
    
    joined_features_metadata_abun <-
        joined_features_metadata[
            !is.na(joined_features_metadata$feature_abun),]
    if (is.numeric(joined_features_metadata_abun$metadata) &
        length(unique(
            joined_features_metadata_abun$metadata)) > 1) {
        logging::loginfo(
            "Creating scatter plot for continuous 
                        data (linear), %s vs %s",
            metadata_name,
            feature_name
        )
        temp_plot <- make_scatterplot(joined_features_metadata_abun,
                                    metadata_name,
                                    feature_name,
                                    normalization,
                                    transformation,
                                    coef_val,
                                    qval,
                                    N_nonzero,
                                    N_total,
                                    results_value)
    } else {
        x_axis_label_names <- 
            unique(joined_features_metadata_abun$metadata)
        
        sorted_fixed_order <-
            order(match(results_value, levels(x_axis_label_names)))
        coef_val <- coef_val[sorted_fixed_order]
        qval <- qval[sorted_fixed_order]
        N_nonzero <- N_nonzero[sorted_fixed_order]
        N_total <- N_total[sorted_fixed_order]
        results_value <-
            results_value[sorted_fixed_order]
        
        if (!is.null(feature_specific_covariate_name)) {
            if (metadata_name == feature_specific_covariate_name) {
                renamed_levels <-
                    as.character(levels(
                        feature_specific_covariate[, feature_name]))
            } else {
                renamed_levels <- as.character(
                    levels(metadata[, metadata_name]))
            }
        } else {
            renamed_levels <- as.character(
                levels(metadata[, metadata_name]))
        }
        
        if (length(renamed_levels) == 0) {
            renamed_levels <- x_axis_label_names
        }
        for (name in x_axis_label_names) {
            total <-
                length(which(
                    joined_features_metadata_abun$metadata == name
                ))
            new_n <-
                paste(name, " (n = ", total, ")", sep = "")
            levels(joined_features_metadata_abun[, 'metadata'])[
                levels(
                    joined_features_metadata_abun[, 'metadata']) == 
                    name] <-
                new_n
            renamed_levels <-
                replace(renamed_levels,
                        renamed_levels == name,
                        new_n)
        }
        
        logging::loginfo(
            "Creating box plot for categorical data (linear), 
                        %s vs %s",
            metadata_name,
            feature_name
        )
        
        temp_plot <- make_boxplot_lm(joined_features_metadata_abun,
                                    metadata_name,
                                    feature_name,
                                    normalization,
                                    transformation,
                                    coef_val,
                                    qval,
                                    N_nonzero,
                                    N_total,
                                    results_value)
    }
    return(temp_plot)
}

# Create boxplot for prevalence vs. continuous associations
make_boxplot_logistic <- function(joined_features_metadata_prev,
                                metadata_name,
                                feature_name,
                                normalization,
                                transformation,
                                coef_val,
                                qval,
                                N_nonzero,
                                N_total,
                                results_value) {
    match.arg(normalization, c('Total sum scaling', 'Center log ratio', 'None'))
    match.arg(transformation, c('Log base 2', 'Pseudo-log base 2', 'None'))
    
    temp_plot <-
        ggplot2::ggplot(data = joined_features_metadata_prev,
                        ggplot2::aes(.data$feature_abun, 
                                    .data$metadata)) +
        ggplot2::geom_boxplot(
            ggplot2::aes(fill = .data$feature_abun),
            outlier.alpha = 0.0,
            na.rm = TRUE,
            alpha = .5,
            show.legend = FALSE
        ) +
        ggplot2::geom_point(
            ggplot2::aes(fill = .data$feature_abun),
            alpha = 0.75 ,
            size = 1,
            shape = 21,
            stroke = 0.15,
            color = 'black',
            position = ggplot2::position_jitterdodge()
        ) +
        ggplot2::scale_fill_brewer(palette = "Spectral") +
        ggplot2::scale_x_discrete(
            expand = ggplot2::expansion(mult = c(0, 0.7)))
    
    temp_plot <- temp_plot +
        nature_theme(metadata_name,
                    joined_features_metadata_prev[
                        'feature_abun']) +
        ggplot2::theme(
            panel.grid.major = ggplot2::element_blank(),
            panel.grid.minor = ggplot2::element_blank(),
            panel.background = ggplot2::element_blank(),
            axis.line = ggplot2::element_line(colour = "black")
        ) +
        ggplot2::xlab(feature_name) +
        ggplot2::ylab(metadata_name) +
        ggplot2::theme(legend.position = "none") +
        ggplot2::annotate(
            geom = "text",
            x = Inf,
            y = Inf,
            hjust = 1,
            vjust = 1,
            label = sprintf(
            "FDR: %s\nCoefficient (in full model): %s\nN: %s\nN (not zero): %s",
                formatC(qval, format = "e", digits = 1),
                if (is.na(coef_val)) 'NA' else formatC(
                    coef_val,
                    format = "e",
                    digits = 1
                ),
                formatC(
                    N_total,
                    format = 'f',
                    digits = 0
                ),
                formatC(
                    N_nonzero,
                    format = 'f',
                    digits = 0
                )
            ) ,
            color = "black",
            size = 2,
            fontface = "italic"
        ) +
        ggplot2::coord_flip()
    return(temp_plot)
}

# Create grid for prevalence vs. categorical associations
make_tile_plot <- function(joined_features_metadata_prev,
                        metadata_name,
                        feature_name,
                        normalization,
                        transformation,
                        coef_val,
                        qval,
                        N_nonzero,
                        N_total,
                        results_value) {
    match.arg(normalization, c('Total sum scaling', 'Center log ratio', 'None'))
    match.arg(transformation, c('Log base 2', 'Pseudo-log base 2', 'None'))
    
    count_df <- joined_features_metadata_prev |> 
        collapse::fcountv(c("feature_abun",
                            "metadata")) |> 
        collapse::roworderv(c("feature_abun",
                              "metadata")) |> 
        collapse::frename("count" = "N")
    
    x_vals <-
        unique(joined_features_metadata_prev$feature_abun)
    y_vals <-
        unique(joined_features_metadata_prev$metadata)
    complete_grid <-
        expand.grid(feature_abun = x_vals,
                    metadata = y_vals)
    
    table_df <- collapse::join(complete_grid,
                               count_df,
                               verbose = FALSE,
                               on = c("feature_abun", "metadata")) 
    
    table_df$count = collapse::replace_na(table_df$count, value = 0)
    
    temp_plot <-
        ggplot2::ggplot(table_df,
                        ggplot2::aes(
                            x = .data$metadata,
                            y = .data$feature_abun
                        )) +
        ggplot2::geom_tile(
            ggplot2::aes(fill = .data$count),
            color = "white",
            lwd = 1.5,
            linetype = 1
        ) +
        ggplot2::geom_text(
            ggplot2::aes(label = .data$count),
            color = "black",
            size = 32 / nrow(table_df)
        ) +
        ggplot2::coord_fixed(ratio = 0.5) +
        ggplot2::scale_fill_gradient2(
            low = "#075AFF",
            mid = "#FFFFCC",
            high = "#FF0000"
        ) +
        ggplot2::scale_y_discrete(
            expand = ggplot2::expansion(mult = c(0, 1 + nrow(
                table_df
            ) / 6))) +
        ggplot2::theme(
            panel.background = ggplot2::element_blank(),
            legend.position = "none"
        )
    
    temp_plot <- temp_plot +
        nature_theme(as.character(table_df$metadata),
                    joined_features_metadata_prev[
                        'feature_abun']) +
        ggplot2::theme(
            panel.grid.major = ggplot2::element_blank(),
            panel.grid.minor = ggplot2::element_blank(),
            panel.background = ggplot2::element_blank(),
            axis.line = ggplot2::element_line(colour = "black")
        ) +
        ggplot2::xlab(metadata_name) +
        ggplot2::ylab(feature_name) +
        ggplot2::theme(legend.position = "none") +
        ggplot2::annotate(
            geom = "text",
            x = Inf,
            y = Inf,
            hjust = 1,
            vjust = 1,
            label = sprintf(
                "Value: %s\nFDR: %s\nCoefficient (in full model): %s",
                paste0(results_value, collapse = ', '),
                paste0(
                    formatC(qval, format = "e", digits = 1),
                    collapse = ', '
                ),
                paste0(
                    vapply(coef_val, function(cv) {
                        if (is.na(cv)) 'NA' else formatC(cv, format = "e", digits = 1)
                    }, character(1)),
                    collapse = ', '
                )
            ) ,
            color = "black",
            size = 2,
            fontface = "italic"
        )
    return(temp_plot)
}

# Create plots for prevalence associations
make_logistic_plot <- function(this_signif_association,
                            joined_features_metadata,
                            metadata,
                            metadata_name,
                            feature_name,
                            normalization,
                            transformation,
                            feature_specific_covariate_name,
                            feature_specific_covariate) {
    match.arg(normalization, c('Total sum scaling', 'Center log ratio', 'None'))
    match.arg(transformation, c('Log base 2', 'Pseudo-log base 2', 'None'))
    
    coef_val <-
        this_signif_association[
            this_signif_association$model == 'logistic',]$coef
    qval <-
        this_signif_association[
            this_signif_association$model == 'logistic',
        ]$qval_individual
    N_nonzero <-
        this_signif_association[
            this_signif_association$model == 'logistic',]$N_not_zero
    N_total <-
        this_signif_association[
            this_signif_association$model == 'logistic',]$N
    results_value <-
        this_signif_association[
            this_signif_association$model == 'logistic',]$value
    
    joined_features_metadata_prev <-
        joined_features_metadata
    joined_features_metadata_prev$feature_abun <-
        data.table::fifelse(
            is.na(joined_features_metadata_prev$feature_abun),
            'Absent',
            'Present'
        )
    
    joined_features_metadata_prev$feature_abun <-
        factor(
            joined_features_metadata_prev$feature_abun,
            levels = c('Present', 'Absent')
        )
    
    if (is.numeric(joined_features_metadata_prev$metadata) &
        length(unique(
            joined_features_metadata_prev$metadata)) > 1) {
        logging::loginfo(
            "Creating boxplot for continuous data (logistic), %s vs %s",
            metadata_name,
            feature_name
        )
        temp_plot <- make_boxplot_logistic(joined_features_metadata_prev,
                                        metadata_name,
                                        feature_name,
                                        normalization,
                                        transformation,
                                        coef_val,
                                        qval,
                                        N_nonzero,
                                        N_total,
                                        results_value)
        
    } else {
        joined_features_metadata_prev$feature_abun <-
            factor(
                joined_features_metadata_prev$feature_abun,
                levels = c('Absent', 'Present')
            )
        
        x_axis_label_names <-
            unique(joined_features_metadata_prev$metadata)
        
        sorted_fixed_order <-
            order(match(results_value, levels(x_axis_label_names)))
        coef_val <- coef_val[sorted_fixed_order]
        qval <- qval[sorted_fixed_order]
        N_nonzero <- N_nonzero[sorted_fixed_order]
        N_total <- N_total[sorted_fixed_order]
        results_value <-
            results_value[sorted_fixed_order]
        
        if (!is.null(feature_specific_covariate_name)) {
            if (metadata_name == feature_specific_covariate_name) {
                renamed_levels <-
                    as.character(levels(
                        feature_specific_covariate[, feature_name]))
            } else {
                renamed_levels <- 
                    as.character(levels(metadata[, metadata_name]))
            }
        } else {
            renamed_levels <- 
                as.character(levels(metadata[, metadata_name]))
        }
        
        if (length(renamed_levels) == 0) {
            renamed_levels <- x_axis_label_names
        }
        for (name in x_axis_label_names) {
            mean_abun <- mean(
                joined_features_metadata_prev$feature_abun[
                    joined_features_metadata_prev$metadata == 
                        name] == 'Present',
                na.rm = TRUE
            )
            new_n <-
                paste(name,
                    " (p = ",
                    round(mean_abun, 2) * 100,
                    "%)",
                    sep = "")
            levels(joined_features_metadata_prev[, 'metadata'])[
                levels(joined_features_metadata_prev[, 'metadata']) 
                == name] <-
                new_n
            renamed_levels <-
                replace(renamed_levels,
                        renamed_levels == name,
                        new_n)
        }
        
        logging::loginfo(
            "Creating tile plot for categorical data (logistic), %s vs %s",
            metadata_name,
            feature_name
        )
        temp_plot <- make_tile_plot(joined_features_metadata_prev,
                                    metadata_name,
                                    feature_name,
                                    normalization,
                                    transformation,
                                    coef_val,
                                    qval,
                                    N_nonzero,
                                    N_total,
                                    results_value)
    }
    return(temp_plot)
}

# Create individual plots for significant associations
maaslin3_association_plots <-
    function(merged_results,
             metadata,
             features,
             max_significance = 0.1,
             figures_folder,
             max_pngs = 10,
             normalization,
             transform,
             feature_specific_covariate = NULL,
             feature_specific_covariate_name = NULL,
             feature_specific_covariate_record = NULL,
             save_plots_rds = FALSE) {
        
        
        match.arg(normalization, c("TSS", "CLR", "NONE"))
        
        match.arg(transform, c("LOG", "PLOG", "NONE"))
        
        # Disregard abundance-induced-prevalence errors in plotting
        merged_results$error[grepl("Prevalence association possibly induced", 
                                   merged_results$error)] <- NA
        
        # Disregard small random effect group warning
        merged_results$error[grepl("<4 average observations per random effect", 
                                   merged_results$error)] <- NA
        
        new_name_normalization <-
            c('Total sum scaling', 'Center log ratio', 'None')
        
        names(new_name_normalization) <- c("TSS", "CLR", "NONE")
        
        normalization <- new_name_normalization[normalization]
        
        new_name_transformation <-
            c('Log base 2', 'Pseudo-log base 2', 'None')
        
        names(new_name_transformation) <- c("LOG", 'PLOG', "NONE")
        
        transformation <- new_name_transformation[transform]
        
        merged_results <-
            merged_results[is.na(merged_results$error) &
                               !is.na(merged_results$qval_individual) &
                               merged_results$qval_individual < 
                               max_significance,]
        
        if (nrow(merged_results) == 0) {
            logging::loginfo(paste("All associations had errors 
                                or were insignificant."))
            return(NULL)
        }
        
        merged_results <-
            merged_results[order(merged_results$qval_individual),]
        
        logging::loginfo(
            paste(
                "Plotting associations from most",
                "to least significant,",
                "grouped by metadata"
            )
        )
        
        saved_plots <- list()
        
        features_by_metadata <-
            unique(merged_results[, c('feature', 'metadata', 'model')])
        
        names(features_by_metadata)[2] = "meta_var"
        
        assoc_num = min(nrow(features_by_metadata), max_pngs)
        
        to_map = features_by_metadata[seq_len(assoc_num),]
        
        to_map$fv = lapply(to_map$feature, \(x) {
            features[[x]]
        })
        
        stat_df = collapse::join(to_map,
                                 merged_results, 
                                 on = c("feature", 
                                        "meta_var" = "metadata",
                                        "model"),
                                 how = "left",
                                 verbose = FALSE) 
        
        stat_df = stat_df |> 
            collapse::fmutate(split_var = paste(stat_df$feature, 
                                                stat_df$meta_var,
                                                stat_df$model,
                                                sep = "_")) 
        
        stat_list = split(stat_df , 
                          stat_df$split_var) 
        # split() sorts the resultant list by f :(
        
        assoc_df = data.frame(split_var = names(stat_list),
                              assoc_stats = I(stat_list))
        
        to_map = to_map |> 
            collapse::fmutate(split_var = paste(to_map$feature, 
                                                to_map$meta_var,
                                                to_map$model,
                                                sep = "_")) |> 
            collapse::join(assoc_df, on = "split_var",
                           how = "left",
                           verbose = FALSE)
        
        collapse::get_vars(to_map, "split_var") <- NULL
        
        # Some associations have multiple rows of statistics to attach. A join +
        # tidyr::nest() would be simpler, but we don't have tidyr.
        
        # # V Attach association stats here so we don't have to pass
        # # around/re-join merged_results at each iteration.
        
        small_meta = metadata |> 
            collapse::fselect(collapse::funique(features_by_metadata$meta_var))
        
        association_plots_folder <-
            file.path(figures_folder, 'association_plots')
        
        if (!file.exists(association_plots_folder)) {
            dir.create(association_plots_folder)
        }
        
        arg_list = list(meta = small_meta,
                        feat_rn = rownames(features),
                        fscn = feature_specific_covariate_name,
                        fsc = feature_specific_covariate,
                        normalization = normalization,
                        transformation = transformation,
                        save_plots_rds = save_plots_rds,
                        ap_dir = association_plots_folder)
        
        if (mirai::daemons_set()) {
            mirai::everywhere({},
                              make_lm_plot = make_lm_plot,
                              make_logistic_plot = make_logistic_plot)
            
            plot_list <- mirai::mirai_map(to_map, 
                                   plot_one_assoc,
                                   .args = arg_list)[.progress]
        } else {
            plot_list <- mapply(plot_one_assoc,
                                to_map$feature,  #life without purrr::pmap :(
                                to_map$meta_var, 
                                to_map$model, 
                                to_map$fv, 
                                to_map$assoc_stats,
                                MoreArgs = arg_list,
                                SIMPLIFY = FALSE,
                                USE.NAMES = FALSE)
        }
        
        to_map$pl <- plot_list
        
        if (save_plots_rds) {
            # save lists for each metadata variable
             if (save_plots_rds) {
                
                split_plots = split(plot_list,
                                    to_map$meta_var)  
                
                ap_dir = association_plots_folder
                save_plot_list_fun = \(x, y) {
                    saveRDS(x, 
                            file = file.path(ap_dir, 
                                             paste0(make.names(y), 
                                                    "_gg_associations.RDS")))
                }
                
                if (mirai::daemons_set()) {
                    # I don't think parallelization helps much here, might be
                    # dependent on disk. TODO: remove?
                    plot_df = data.frame(x = I(split_plots),
                                         y = names(split_plots))
                    
                    mirai::mirai_map(.x = plot_df,
                              .f = save_plot_list_fun,
                              ap_dir = association_plots_folder)[.progress]
                } else {
                    mapply(split_plots,
                           names(split_plots),
                           FUN = save_plot_list_fun)
                }
             }
        } else {
            return(invisible())
        }
        
        names(plot_list) = paste(sep = "_",
                                 to_map$meta_var, 
                                 to_map$feature,
                                 to_map$model)
            
        return(plot_list)
    }

plot_one_assoc = function(feature, meta_var, model, fv, assoc_stats,
                          meta,
                          feat_rn,
                          fscn, # feature_specific_covariate_name
                          fsc,
                          normalization,
                          transformation,
                          save_plots_rds,
                          ap_dir) {
    
    feature_name <- feature
    
    feature_abun <- data.frame(
        sample = feat_rn,
        feature_abun = fv, # features[, feature_name], 
        check.names = FALSE)
    
    metadata_name <- meta_var # features_by_metadata[row_num, 'metadata']
    
    if (!is.null(fscn)) {
        if (metadata_name == fscn) {
            metadata_sub <-
                data.frame(
                    sample = rownames(fsc),
                    metadata = fsc[, feature_name],
                    check.names = FALSE
                )
        } else {
            metadata_sub <- data.frame(sample = rownames(meta),
                                       metadata = meta[, metadata_name], 
                                       check.names = FALSE)
        }
    } else {
        metadata_sub <- data.frame(sample = rownames(meta),
                                   metadata = meta[, metadata_name], 
                                   check.names = FALSE)
    }
    
    joined_features_metadata <-
        collapse::join(feature_abun, 
                       metadata_sub,
                       how = "inner",
                       verbose = FALSE,
                       on = c('sample'))
    
    model_name <- model # features_by_metadata[row_num, 'model']
    
    this_signif_association = assoc_stats
    
    if ('linear' == model_name) {
        temp_plot <- make_lm_plot(this_signif_association,
                                            joined_features_metadata,
                                            meta,
                                            metadata_name,
                                            feature_name,
                                            normalization,
                                            transformation,
                                            fscn,
                                            fsc)
    }
    
    if ('logistic' == model_name) {
        temp_plot <- make_logistic_plot(this_signif_association,
                                                  joined_features_metadata,
                                                  meta,
                                                  metadata_name,
                                                  feature_name,
                                                  normalization,
                                                  transformation,
                                                  fscn,
                                                  fsc)
    }
    
    out_dir = file.path(ap_dir,
                        make.names(meta_var), 
                        model)
    
    if (!dir.exists(out_dir)) dir.create(out_dir, recursive = TRUE)
    
    png_file = file.path(out_dir, 
                         paste0(make.names(meta_var),
                                "_",
                                make.names(feature_name),
                                "_",
                                model, ".png"))
    
    height <- max(960, 18 * max(nchar(unlist(strsplit(
        temp_plot$labels$y, '\n')))))
    
    tryCatch({
        withCallingHandlers({
            ggplot2::ggsave(filename = png_file, 
                            plot = temp_plot, 
                            dpi = 600, 
                            width = 960 / 300, 
                            height = height / 300)
        }, warning = function(w) { 
            invokeRestart("muffleWarning") })
    })
    
    if (save_plots_rds) {
        return(temp_plot)
    } else {
        invisible()
    }
}
