server-scripts/qhtcp-workflow/apps/r/calculate_interaction_zscores.R

suppressMessages({
  library("ggplot2")
  library("plotly")
  library("htmlwidgets")
  library("dplyr")
  library("rlang")
  library("ggthemes")
  library("data.table")
  library("gridExtra")
  library("future")
  library("furrr")
  library("purrr")
})

# These parallelization libraries are very noisy
suppressPackageStartupMessages({
  library("future")
  library("furrr")
  library("purrr")
})

options(warn = 2)

# Constants for configuration
plot_width <- 14
plot_height <- 9
base_size <- 14

parse_arguments <- function() {
  args <- if (interactive()) {
    c(
      "/home/bryan/documents/develop/hartmanlab/qhtcp-workflow/out/20240116_jhartman2_DoxoHLD/20240116_jhartman2_DoxoHLD",
      "/home/bryan/documents/develop/hartmanlab/qhtcp-workflow/apps/r/SGD_features.tab",
      "/home/bryan/documents/develop/hartmanlab/qhtcp-workflow/out/20240116_jhartman2_DoxoHLD/easy/20240116_jhartman2_DoxoHLD/results_std.txt",
      "/home/bryan/documents/develop/hartmanlab/qhtcp-workflow/out/20240116_jhartman2_DoxoHLD/20240822_jhartman2_DoxoHLD/exp1",
      "Experiment 1: Doxo versus HLD",
      3,
      "/home/bryan/documents/develop/hartmanlab/qhtcp-workflow/out/20240116_jhartman2_DoxoHLD/20240822_jhartman2_DoxoHLD/exp2",
      "Experiment 2: HLD versus Doxo",
      3
    )
  } else {
    commandArgs(trailingOnly = TRUE)
  }

  out_dir <- normalizePath(args[1], mustWork = FALSE)
  sgd_gene_list <- normalizePath(args[2], mustWork = FALSE)
  easy_results_file <- normalizePath(args[3], mustWork = FALSE)

  # The remaining arguments should be in groups of 3
  exp_args <- args[-(1:3)]
  if (length(exp_args) %% 3 != 0) {
    stop("Experiment arguments should be in groups of 3: path, name, sd.")
  }

  # Extract the experiments into a list
  experiments <- list()
  for (i in seq(1, length(exp_args), by = 3)) {
    exp_name <- exp_args[i + 1]
    experiments[[exp_name]] <- list(
      path = normalizePath(exp_args[i], mustWork = FALSE),
      sd = as.numeric(exp_args[i + 2])
    )
  }

  # Extract the trailing number from each path
  trailing_numbers <- sapply(experiments, function(x) {
    path <- x$path
    nums <- gsub("[^0-9]", "", basename(path))
    as.integer(nums)
  })

  # Sort the experiments based on the trailing numbers
  sorted_experiments <- experiments[order(trailing_numbers)]

  list(
    out_dir = out_dir,
    sgd_gene_list = sgd_gene_list,
    easy_results_file = easy_results_file,
    experiments = sorted_experiments
  )
}

args <- parse_arguments()

# Should we keep output in exp dirs or combine in the study output dir?
# dir.create(file.path(args$out_dir, "zscores"), showWarnings = FALSE)
# dir.create(file.path(args$out_dir, "zscores", "qc"), showWarnings = FALSE)

theme_publication <- function(base_size = 14, base_family = "sans", legend_position = "bottom") {
  theme_foundation <- ggthemes::theme_foundation(base_size = base_size, base_family = base_family)
  
  theme_foundation %+replace%
    theme(
      plot.title = element_text(face = "bold", size = rel(1.6), hjust = 0.5),
      text = element_text(),
      panel.background = element_blank(),
      plot.background = element_blank(),
      panel.border = element_blank(),
      axis.title = element_text(face = "bold", size = rel(1.4)),
      axis.title.y = element_text(angle = 90, vjust = 2),
      # axis.title.x = element_text(vjust = -0.2), # TODO this causes errors
      axis.text = element_text(size = rel(1.2)),
      axis.line = element_line(colour = "black"),
      # axis.ticks = element_line(),
      panel.grid.major = element_line(colour = "#f0f0f0"),
      panel.grid.minor = element_blank(),
      legend.key = element_rect(colour = NA),
      legend.position = legend_position,
      legend.direction = ifelse(legend_position == "right", "vertical", "horizontal"),
      # legend.key.size = unit(0.5, "cm"),
      legend.spacing = unit(0, "cm"),
      legend.title = element_text(face = "italic", size = rel(1.3)),
      legend.text = element_text(size = rel(1.2)),
      plot.margin = unit(c(10, 5, 5, 5), "mm")
      # strip.background = element_rect(colour = "#f0f0f0", fill = "#f0f0f0"),
      # strip.text = element_text(face = "bold")
    )
}

scale_fill_publication <- function(...) {
  discrete_scale("fill", "Publication", manual_pal(values = c(
    "#386cb0", "#fdb462", "#7fc97f", "#ef3b2c", "#662506",
    "#a6cee3", "#fb9a99", "#984ea3", "#ffff33"
  )), ...)
}

scale_colour_publication <- function(...) {
  discrete_scale("colour", "Publication", manual_pal(values = c(
    "#386cb0", "#fdb462", "#7fc97f", "#ef3b2c", "#662506",
    "#a6cee3", "#fb9a99", "#984ea3", "#ffff33"
  )), ...)
}

# Load the initial dataframe from the easy_results_file
load_and_filter_data <- function(easy_results_file, sd = 3) {
  df <- read.delim(easy_results_file, skip = 2, as.is = TRUE, row.names = 1, strip.white = TRUE)

  df <- df %>%
    filter(!(.[[1]] %in% c("", "Scan"))) %>%
    filter(!is.na(ORF) & ORF != "" & !Gene %in% c("BLANK", "Blank", "blank") & Drug != "BMH21") %>%
    # Rename columns
    rename(L = l, num = Num., AUC = AUC96, scan = Scan, last_bg = LstBackgrd, first_bg = X1stBackgrd) %>%
    mutate(
      across(c(Col, Row, num, L, K, r, scan, AUC, last_bg, first_bg), as.numeric),
      delta_bg = last_bg - first_bg,
      delta_bg_tolerance = mean(delta_bg, na.rm = TRUE) + (sd * sd(delta_bg, na.rm = TRUE)),
      NG = if_else(L == 0 & !is.na(L), 1, 0),
      DB = if_else(delta_bg >= delta_bg_tolerance, 1, 0),
      SM = 0,
      OrfRep = if_else(ORF == "YDL227C", "YDL227C", OrfRep), # should these be hardcoded?
      conc_num = as.numeric(gsub("[^0-9\\.]", "", Conc)),
      conc_num_factor = as.numeric(as.factor(conc_num)) - 1, # for legacy purposes
      conc_num_factor_factor = as.factor(conc_num)
    )

  return(df)
}

# Update Gene names using the SGD gene list
update_gene_names <- function(df, sgd_gene_list) {
  # Load SGD gene list
  genes <- read.delim(file = sgd_gene_list,
    quote = "", header = FALSE,
    colClasses = c(rep("NULL", 3), rep("character", 2), rep("NULL", 11)))
  
  # Create a named vector for mapping ORF to GeneName
  gene_map <- setNames(genes$V5, genes$V4)
  # Vectorized match to find the GeneName from gene_map
  mapped_genes <- gene_map[df$ORF]
  # Replace NAs in mapped_genes with original Gene names (preserves existing Gene names if ORF is not found)
  updated_genes <- ifelse(is.na(mapped_genes) | df$OrfRep == "YDL227C", df$Gene, mapped_genes)
  # Ensure Gene is not left blank or incorrectly updated to "OCT1"
  df <- df %>%
    mutate(Gene = ifelse(updated_genes == "" | updated_genes == "OCT1", OrfRep, updated_genes))
  
  return(df)
}

calculate_summary_stats <- function(df, variables, group_vars) {
  
  summary_stats <- df %>%
    group_by(across(all_of(group_vars))) %>%
    summarise(
      N = n(),
      across(all_of(variables),
        list(
          mean = ~ mean(.x, na.rm = TRUE),
          median = ~ median(.x, na.rm = TRUE),
          max = ~ ifelse(all(is.na(.x)), NA, max(.x, na.rm = TRUE)),
          min = ~ ifelse(all(is.na(.x)), NA, min(.x, na.rm = TRUE)),
          sd = ~ sd(.x, na.rm = TRUE),
          se = ~ sd(.x, na.rm = TRUE) / sqrt(n() - 1)
        ),
        .names = "{.fn}_{.col}"
      ),
      .groups = "drop"
    )
  
  # Create a cleaned version of df that doesn't overlap with summary_stats
  cleaned_df <- df %>%
    select(-any_of(setdiff(intersect(names(df), names(summary_stats)), group_vars)))
  
  df_joined <- left_join(cleaned_df, summary_stats, by = group_vars)

  return(list(summary_stats = summary_stats, df_with_stats = df_joined))
}

calculate_interaction_scores <- function(df, max_conc, bg_stats,
  group_vars = c("OrfRep", "Gene", "num")) {

  # Calculate total concentration variables
  total_conc_num <- length(unique(df$conc_num))

  calculations <- df %>%
    group_by(across(all_of(group_vars))) %>%
    mutate(
      NG = sum(NG, na.rm = TRUE),
      DB = sum(DB, na.rm = TRUE),
      SM = sum(SM, na.rm = TRUE),
      num_non_removed_concs = total_conc_num - sum(DB, na.rm = TRUE) - 1,

      # Calculate raw data
      Raw_Shift_L = first(mean_L) - bg_stats$mean_L,
      Raw_Shift_K = first(mean_K) - bg_stats$mean_K,
      Raw_Shift_r = first(mean_r) - bg_stats$mean_r,
      Raw_Shift_AUC = first(mean_AUC) - bg_stats$mean_AUC,
      Z_Shift_L = first(Raw_Shift_L) / bg_stats$sd_L,
      Z_Shift_K = first(Raw_Shift_K) / bg_stats$sd_K,
      Z_Shift_r = first(Raw_Shift_r) / bg_stats$sd_r,
      Z_Shift_AUC = first(Raw_Shift_AUC) / bg_stats$sd_AUC,
      Exp_L = WT_L + Raw_Shift_L,
      Exp_K = WT_K + Raw_Shift_K,
      Exp_r = WT_r + Raw_Shift_r,
      Exp_AUC = WT_AUC + Raw_Shift_AUC,
      Delta_L = mean_L - Exp_L,
      Delta_K = mean_K - Exp_K,
      Delta_r = mean_r - Exp_r,
      Delta_AUC = mean_AUC - Exp_AUC,
      Delta_L = if_else(NG == 1, mean_L - WT_L, Delta_L),
      Delta_K = if_else(NG == 1, mean_K - WT_K, Delta_K),
      Delta_r = if_else(NG == 1, mean_r - WT_r, Delta_r),
      Delta_AUC = if_else(NG == 1, mean_AUC - WT_AUC, Delta_AUC),
      Delta_L = if_else(SM == 1, mean_L - WT_L, Delta_L),

      # Calculate Z-scores
      Zscore_L = Delta_L / WT_sd_L,
      Zscore_K = Delta_K / WT_sd_K,
      Zscore_r = Delta_r / WT_sd_r,
      Zscore_AUC = Delta_AUC / WT_sd_AUC,

      # Fit linear models and store in list-columns
      gene_lm_L = list(lm(Delta_L ~ conc_num_factor, data = pick(everything()))),
      gene_lm_K = list(lm(Delta_K ~ conc_num_factor, data = pick(everything()))),
      gene_lm_r = list(lm(Delta_r ~ conc_num_factor, data = pick(everything()))),
      gene_lm_AUC = list(lm(Delta_AUC ~ conc_num_factor, data = pick(everything()))),

      # Extract coefficients using purrr::map_dbl
      lm_intercept_L = map_dbl(gene_lm_L, ~ coef(.x)[1]),
      lm_slope_L = map_dbl(gene_lm_L, ~ coef(.x)[2]),
      lm_intercept_K = map_dbl(gene_lm_K, ~ coef(.x)[1]),
      lm_slope_K = map_dbl(gene_lm_K, ~ coef(.x)[2]),
      lm_intercept_r = map_dbl(gene_lm_r, ~ coef(.x)[1]),
      lm_slope_r = map_dbl(gene_lm_r, ~ coef(.x)[2]),
      lm_intercept_AUC = map_dbl(gene_lm_AUC, ~ coef(.x)[1]),
      lm_slope_AUC = map_dbl(gene_lm_AUC, ~ coef(.x)[2]),

      # Calculate lm_Score_* based on coefficients
      lm_Score_L = max_conc * lm_slope_L + lm_intercept_L,
      lm_Score_K = max_conc * lm_slope_K + lm_intercept_K,
      lm_Score_r = max_conc * lm_slope_r + lm_intercept_r,
      lm_Score_AUC = max_conc * lm_slope_AUC + lm_intercept_AUC,

      # Calculate R-squared values
      R_Squared_L = map_dbl(gene_lm_L, ~ summary(.x)$r.squared),
      R_Squared_K = map_dbl(gene_lm_K, ~ summary(.x)$r.squared),
      R_Squared_r = map_dbl(gene_lm_r, ~ summary(.x)$r.squared),
      R_Squared_AUC = map_dbl(gene_lm_AUC, ~ summary(.x)$r.squared)
    ) %>%
    ungroup()

  # Calculate overall mean and SD for lm_Score_* variables
  lm_means_sds <- calculations %>%
    summarise(
      lm_mean_L = mean(lm_Score_L, na.rm = TRUE),
      lm_sd_L = sd(lm_Score_L, na.rm = TRUE),
      lm_mean_K = mean(lm_Score_K, na.rm = TRUE),
      lm_sd_K = sd(lm_Score_K, na.rm = TRUE),
      lm_mean_r = mean(lm_Score_r, na.rm = TRUE),
      lm_sd_r = sd(lm_Score_r, na.rm = TRUE),
      lm_mean_AUC = mean(lm_Score_AUC, na.rm = TRUE),
      lm_sd_AUC = sd(lm_Score_AUC, na.rm = TRUE)
    )

  calculations <- calculations %>%
    mutate(
      Z_lm_L = (lm_Score_L - lm_means_sds$lm_mean_L) / lm_means_sds$lm_sd_L,
      Z_lm_K = (lm_Score_K - lm_means_sds$lm_mean_K) / lm_means_sds$lm_sd_K,
      Z_lm_r = (lm_Score_r - lm_means_sds$lm_mean_r) / lm_means_sds$lm_sd_r,
      Z_lm_AUC = (lm_Score_AUC - lm_means_sds$lm_mean_AUC) / lm_means_sds$lm_sd_AUC
    )

  # Summarize some of the stats
  interactions <- calculations %>%
    group_by(across(all_of(group_vars))) %>%
    mutate(
      # Calculate raw shifts
      Raw_Shift_L = first(Raw_Shift_L),
      Raw_Shift_K = first(Raw_Shift_K),
      Raw_Shift_r = first(Raw_Shift_r),
      Raw_Shift_AUC = first(Raw_Shift_AUC),

      # Calculate Z-shifts
      Z_Shift_L = first(Z_Shift_L),
      Z_Shift_K = first(Z_Shift_K),
      Z_Shift_r = first(Z_Shift_r),
      Z_Shift_AUC = first(Z_Shift_AUC),

      # Sum Z-scores
      Sum_Z_Score_L = sum(Zscore_L),
      Sum_Z_Score_K = sum(Zscore_K),
      Sum_Z_Score_r = sum(Zscore_r),
      Sum_Z_Score_AUC = sum(Zscore_AUC),

      # Calculate Average Z-scores
      Avg_Zscore_L = Sum_Z_Score_L / num_non_removed_concs,
      Avg_Zscore_K = Sum_Z_Score_K / num_non_removed_concs,
      Avg_Zscore_r = Sum_Z_Score_r / (total_conc_num - 1),
      Avg_Zscore_AUC = Sum_Z_Score_AUC / (total_conc_num - 1)
    ) %>%
    arrange(desc(Z_lm_L), desc(NG)) %>%
    ungroup()

  # Declare column order for output
  calculations <- calculations %>%
    select(
      "OrfRep", "Gene", "num", "N",
      "conc_num", "conc_num_factor", "conc_num_factor_factor",
      "mean_L", "mean_K", "mean_r", "mean_AUC",
      "median_L", "median_K", "median_r", "median_AUC",
      "sd_L", "sd_K", "sd_r", "sd_AUC",
      "se_L", "se_K", "se_r", "se_AUC",
      "Raw_Shift_L", "Raw_Shift_K", "Raw_Shift_r", "Raw_Shift_AUC",
      "Z_Shift_L", "Z_Shift_K", "Z_Shift_r", "Z_Shift_AUC",
      "WT_L", "WT_K", "WT_r", "WT_AUC",
      "WT_sd_L", "WT_sd_K", "WT_sd_r", "WT_sd_AUC",
      "Exp_L", "Exp_K", "Exp_r", "Exp_AUC",
      "Delta_L", "Delta_K", "Delta_r", "Delta_AUC",
      "Zscore_L", "Zscore_K", "Zscore_r", "Zscore_AUC",
      "NG", "SM", "DB"
    )

  interactions <- interactions %>%
    select(
      "OrfRep", "Gene", "num", "NG", "DB", "SM",
      "conc_num", "conc_num_factor", "conc_num_factor_factor",
      "Raw_Shift_L", "Raw_Shift_K", "Raw_Shift_r", "Raw_Shift_AUC",
      "Z_Shift_L", "Z_Shift_K", "Z_Shift_r", "Z_Shift_AUC",
      "Sum_Z_Score_L", "Sum_Z_Score_K", "Sum_Z_Score_r", "Sum_Z_Score_AUC",
      "Avg_Zscore_L", "Avg_Zscore_K", "Avg_Zscore_r", "Avg_Zscore_AUC",
      "lm_Score_L", "lm_Score_K", "lm_Score_r", "lm_Score_AUC",
      "R_Squared_L", "R_Squared_K", "R_Squared_r", "R_Squared_AUC",
      "Z_lm_L", "Z_lm_K", "Z_lm_r", "Z_lm_AUC"
    )

  # Clean the original dataframe by removing overlapping columns
  cleaned_df <- df %>%
    select(-any_of(
      setdiff(intersect(names(df), names(calculations)),
      c("OrfRep", "Gene", "num", "conc_num", "conc_num_factor", "conc_num_factor_factor"))))

  # Join the original dataframe with calculations
  df_with_calculations <- left_join(cleaned_df, calculations,
    by = c("OrfRep", "Gene", "num", "conc_num", "conc_num_factor", "conc_num_factor_factor"))

  # Remove overlapping columns between df_with_calculations and interactions
  df_with_calculations_clean <- df_with_calculations %>%
    select(-any_of(
      setdiff(intersect(names(df_with_calculations), names(interactions)),
      c("OrfRep", "Gene", "num", "conc_num", "conc_num_factor"))))

  # Join with interactions to create the full dataset
  full_data <- left_join(df_with_calculations_clean, interactions,
    by = c("OrfRep", "Gene", "num", "conc_num", "conc_num_factor", "conc_num_factor_factor"))

  return(list(
    calculations = calculations,
    interactions = interactions,
    full_data = full_data
  ))
}

generate_and_save_plots <- function(out_dir, filename, plot_configs) {
  message("Generating ", filename, ".pdf and ", filename, ".html")

  static_plots <- list()
  plotly_plots <- list()

  for (i in seq_along(plot_configs$plots)) {
    config <- plot_configs$plots[[i]]
    df <- config$df

    message("Processing plot ", i, ": ", config$title)
    message("Plot type: ", config$plot_type)

    if (is.null(df)) {
      message("Dataframe for plot ", i, " is NULL.")
      next
    }

    # Define aes_mapping, ensuring y_var is only used when it's not NULL
    aes_mapping <- switch(config$plot_type,
      "bar" = if (!is.null(config$color_var)) {
        aes(x = .data[[config$x_var]], fill = .data[[config$color_var]], color = .data[[config$color_var]])
      } else {
        aes(x = .data[[config$x_var]])
      },
      "density" = if (!is.null(config$color_var)) {
        aes(x = .data[[config$x_var]], color = .data[[config$color_var]])
      } else {
        aes(x = .data[[config$x_var]])
      },
      # For other plot types, only include y_var if it's not NULL
      if (!is.null(config$y_var) && !is.null(config$color_var)) {
        aes(x = .data[[config$x_var]], y = .data[[config$y_var]], color = .data[[config$color_var]])
      } else if (!is.null(config$y_var)) {
        aes(x = .data[[config$x_var]], y = .data[[config$y_var]])
      } else {
        aes(x = .data[[config$x_var]])  # no y_var needed for density and bar plots
      }
    )

    # Apply theme_publication with legend_position
    plot <- ggplot(df, aes_mapping) + theme_publication(legend_position = config$legend_position)

    # Apply appropriate plot function
    plot <- switch(config$plot_type,
      "scatter" = generate_scatter_plot(plot, config),
      "box" = generate_box_plot(plot, config),
      "density" = plot + geom_density(),
      "bar" = plot + geom_bar(stat = "count"), # count occurrences
      plot  # default (unused)
    )

    # Add titles and labels
    if (!is.null(config$title)) plot <- plot + ggtitle(config$title)
    if (!is.null(config$x_label)) plot <- plot + xlab(config$x_label)
    if (!is.null(config$y_label)) plot <- plot + ylab(config$y_label)
    if (!is.null(config$coord_cartesian)) plot <- plot + coord_cartesian(ylim = config$coord_cartesian)

    # Convert ggplot to plotly, skipping subplot
    if (!is.null(config$tooltip_vars)) {
      plotly_plot <- suppressWarnings(plotly::ggplotly(plot, tooltip = config$tooltip_vars))
    } else {
      plotly_plot <- suppressWarnings(plotly::ggplotly(plot))
    }

    # Adjust legend position in plotly
    if (!is.null(config$legend_position) && config$legend_position == "bottom") {
      plotly_plot <- plotly_plot %>% layout(legend = list(orientation = "h"))
      plotly_plot <- plotly_plot + theme(legend.direction = NULL)
    }

    # Add static and interactive plots to lists
    static_plots[[i]] <- plot
    plotly_plots[[i]] <- plotly_plot
  }

  # Save static PDF plots
  pdf(file.path(out_dir, paste0(filename, ".pdf")), width = 16, height = 9)

  if (is.null(plot_configs$grid_layout)) {
    # Print each plot on a new page if grid_layout is not set
    for (plot in static_plots) {
      print(plot)
    }
  } else {
    # Use grid.arrange if grid_layout is set
    grid_nrow <- ifelse(is.null(plot_configs$grid_layout$nrow), length(plot_configs$plots), plot_configs$grid_layout$nrow)
    grid_ncol <- ifelse(is.null(plot_configs$grid_layout$ncol), 1, plot_configs$grid_layout$ncol)
    grid.arrange(grobs = static_plots, ncol = grid_ncol, nrow = grid_nrow)
  }

  dev.off()

  # Save individual interactive HTML plots without subplot
  for (i in seq_along(plotly_plots)) {
    html_file <- file.path(out_dir, paste0(filename, "_plot_", i, ".html"))
    message("Saving HTML plot ", i, ": ", html_file)
    htmlwidgets::saveWidget(plotly_plots[[i]], file = html_file, selfcontained = TRUE)
  }
}

generate_scatter_plot <- function(plot, config) {

  # Define the points
  shape <- if (!is.null(config$shape)) config$shape else 3
  size <- if (!is.null(config$size)) config$size else 1.5
  position <-
    if (!is.null(config$position) && config$position == "jitter") {
      position_jitter(width = 0.1, height = 0)
    } else {
      "identity"
    }

  plot <- plot + geom_point(
    shape = shape,
    size = size,
    position = position
  )

  if (!is.null(config$cyan_points) && config$cyan_points) {
    plot <- plot + geom_point(
      aes(x = .data[[config$x_var]], y = .data[[config$y_var]]),
      color = "cyan",
      shape = 3,
      size = 0.5
    )
  }
  
  # Add Smooth Line if specified
  if (!is.null(config$smooth) && config$smooth) {
    smooth_color <- if (!is.null(config$smooth_color)) config$smooth_color else "blue"
    
    if (!is.null(config$lm_line)) {
      plot <- plot +
        geom_abline(
          intercept = config$lm_line$intercept,
          slope = config$lm_line$slope,
          color = smooth_color
        )
    } else {
      plot <- plot +
        geom_smooth(
          method = "lm",
          se = FALSE,
          color = smooth_color
        )
    }
  }
  
  # Add SD Bands if specified
  if (!is.null(config$sd_band)) {
    plot <- plot +
      annotate(
        "rect",
        xmin = -Inf, xmax = Inf,
        ymin = config$sd_band, ymax = Inf,
        fill = ifelse(!is.null(config$fill_positive), config$fill_positive, "#542788"),
        alpha = ifelse(!is.null(config$alpha_positive), config$alpha_positive, 0.3)
      ) +
      annotate(
        "rect",
        xmin = -Inf, xmax = Inf,
        ymin = -config$sd_band, ymax = -Inf,
        fill = ifelse(!is.null(config$fill_negative), config$fill_negative, "orange"),
        alpha = ifelse(!is.null(config$alpha_negative), config$alpha_negative, 0.3)
      ) +
      geom_hline(
        yintercept = c(-config$sd_band, config$sd_band),
        color = ifelse(!is.null(config$hl_color), config$hl_color, "gray")
      )
  }

  # Add Rectangles if specified
  if (!is.null(config$rectangles)) {
    for (rect in config$rectangles) {
      plot <- plot + annotate(
        "rect",
        xmin = rect$xmin,
        xmax = rect$xmax,
        ymin = rect$ymin,
        ymax = rect$ymax,
        fill = ifelse(is.null(rect$fill), NA, rect$fill),
        color = ifelse(is.null(rect$color), "black", rect$color),
        alpha = ifelse(is.null(rect$alpha), 0.1, rect$alpha)
      )
    }
  }
  
  # Add error bars if specified
  if (!is.null(config$error_bar) && config$error_bar && !is.null(config$y_var)) {
    if (!is.null(config$error_bar_params)) {
      plot <- plot +
        geom_errorbar(
          aes(
            ymin = config$error_bar_params$ymin,
            ymax = config$error_bar_params$ymax
          ),
          alpha = 0.3,
          linewidth = 0.5
        )
    } else {
      y_mean_col <- paste0("mean_", config$y_var)
      y_sd_col <- paste0("sd_", config$y_var)
      
      plot <- plot +
        geom_errorbar(
          aes(
            ymin = !!sym(y_mean_col) - !!sym(y_sd_col),
            ymax = !!sym(y_mean_col) + !!sym(y_sd_col)
          ),
          alpha = 0.3,
          linewidth = 0.5
        )
    }
  }
  
  # Customize X-axis if specified
  if (!is.null(config$x_breaks) && !is.null(config$x_labels) && !is.null(config$x_label)) {
    plot <- plot +
      scale_x_discrete(
        name = config$x_label,
        breaks = config$x_breaks,
        labels = config$x_labels
      )
  }
  
  # Set Y-axis limits if specified
  if (!is.null(config$ylim_vals)) {
    plot <- plot + scale_y_continuous(limits = config$ylim_vals)
  }
  
  # Add annotations if specified
  if (!is.null(config$annotations)) {
    for (annotation in config$annotations) {
      plot <- plot +
        annotate(
          "text",
          x = annotation$x,
          y = annotation$y,
          label = annotation$label,
          hjust = ifelse(is.null(annotation$hjust), 0.5, annotation$hjust),
          vjust = ifelse(is.null(annotation$vjust), 0.5, annotation$vjust),
          size = ifelse(is.null(annotation$size), 6, annotation$size),
          color = ifelse(is.null(annotation$color), "black", annotation$color)
        )
    }
  }
  
  return(plot)
}

generate_box_plot <- function(plot, config) {
  # Convert x_var to a factor within aes mapping
  plot <- plot + geom_boxplot(aes(x = factor(.data[[config$x_var]])))

  # Apply scale_x_discrete for breaks, labels, and axis label if provided
  if (!is.null(config$x_breaks) && !is.null(config$x_labels) && !is.null(config$x_label)) {
    plot <- plot + scale_x_discrete(
      name = config$x_label,
      breaks = config$x_breaks,
      labels = config$x_labels
    )
  }

  return(plot)
}

generate_plate_analysis_plot_configs <- function(variables, df_before = NULL, df_after = NULL,
  plot_type = "scatter", stages = c("before", "after")) {
  plots <- list()
  
  for (var in variables) {
    for (stage in stages) {
      df_plot <- if (stage == "before") df_before else df_after

      # Check for non-finite values in the y-variable
      df_plot_filtered <- df_plot %>% filter(is.finite(!!sym(var)))

      # Adjust settings based on plot_type
      config <- list(
        df = df_plot_filtered,
        x_var = "scan",
        y_var = var,
        plot_type = plot_type,
        title = paste("Plate analysis by Drug Conc for", var, stage, "quality control"),
        color_var = "conc_num_factor_factor",
        position = if (plot_type == "scatter") "jitter" else NULL,
        size = 0.2,
        error_bar = (plot_type == "scatter")
      )

      # Add config to plots list
      plots <- append(plots, list(config))
    }
  }

  return(list(grid_layout = list(ncol = 1, nrow = length(plots)), plots = plots))
}

generate_interaction_plot_configs <- function(df, limits_map = NULL, plot_type = "reference") {
  if (is.null(limits_map)) {
    limits_map <- list(
      L = c(0, 130),
      K = c(-20, 160),
      r = c(0, 1),
      AUC = c(0, 12500),
      Delta_L = c(-60, 60),
      Delta_K = c(-60, 60),
      Delta_r = c(-0.6, 0.6),
      Delta_AUC = c(-6000, 6000)
    )
  }

  group_vars <- if (plot_type == "reference") c("OrfRep", "Gene", "num") else c("OrfRep", "Gene")
  df_filtered <- df %>%
    mutate(OrfRepCombined = if (plot_type == "reference") paste(OrfRep, Gene, num, sep = "_") else paste(OrfRep, Gene, sep = "_"))

  # Separate the plots into two groups: overall variables and delta comparisons
  overall_plots <- list()
  delta_plots <- list()

  for (var in c("L", "K", "r", "AUC")) {
    y_limits <- limits_map[[var]]

    config <- list(
      df = df_filtered,
      plot_type = "scatter",
      x_var = "conc_num_factor_factor",
      y_var = var,
      x_label = unique(df_filtered$Drug)[1],
      title = sprintf("Scatter RF for %s with SD", var),
      coord_cartesian = y_limits,
      error_bar = TRUE,
      x_breaks = unique(df_filtered$conc_num_factor_factor),
      x_labels = as.character(unique(df_filtered$conc_num)),
      position = "jitter",
      smooth = TRUE
    )
    overall_plots <- append(overall_plots, list(config))
  }

  unique_groups <- df_filtered %>% select(all_of(group_vars)) %>% distinct()

  for (i in seq_len(nrow(unique_groups))) {
    group <- unique_groups[i, ]
    group_data <- df_filtered %>% semi_join(group, by = group_vars)

    OrfRep <- as.character(group$OrfRep)
    Gene <- if ("Gene" %in% names(group)) as.character(group$Gene) else ""
    num <- if ("num" %in% names(group)) as.character(group$num) else ""

    for (var in c("Delta_L", "Delta_K", "Delta_r", "Delta_AUC")) {
      y_limits <- limits_map[[var]]
      y_span <- y_limits[2] - y_limits[1]

      # Error bars
      WT_sd_var <- paste0("WT_sd_", sub("Delta_", "", var))
      WT_sd_value <- group_data[[WT_sd_var]][1]
      error_bar_ymin <- 0 - (2 * WT_sd_value)
      error_bar_ymax <- 0 + (2 * WT_sd_value)

      # Annotations
      Z_Shift_value <- round(group_data[[paste0("Z_Shift_", sub("Delta_", "", var))]][1], 2)
      Z_lm_value <- round(group_data[[paste0("Z_lm_", sub("Delta_", "", var))]][1], 2)
      NG_value <- group_data$NG[1]
      DB_value <- group_data$DB[1]
      SM_value <- group_data$SM[1]

      annotations <- list(
        list(x = 1, y = y_limits[2] - 0.2 * y_span, label = paste("ZShift =", Z_Shift_value)),
        list(x = 1, y = y_limits[2] - 0.3 * y_span, label = paste("lm ZScore =", Z_lm_value)),
        list(x = 1, y = y_limits[1] + 0.2 * y_span, label = paste("NG =", NG_value)),
        list(x = 1, y = y_limits[1] + 0.1 * y_span, label = paste("DB =", DB_value)),
        list(x = 1, y = y_limits[1], label = paste("SM =", SM_value))
      )

      config <- list(
        df = group_data,
        plot_type = "scatter",
        x_var = "conc_num_factor_factor",
        y_var = var,
        x_label = unique(group_data$Drug)[1],
        title = paste(OrfRep, Gene, sep = "      "),
        coord_cartesian = y_limits,
        annotations = annotations,
        error_bar = TRUE,
        error_bar_params = list(
          ymin = error_bar_ymin,
          ymax = error_bar_ymax
        ),
        smooth = TRUE,
        x_breaks = unique(group_data$conc_num_factor_factor),
        x_labels = as.character(unique(group_data$conc_num)),
        ylim_vals = y_limits
      )
      delta_plots <- append(delta_plots, list(config))
    }
  }

  return(list(
    overall_plots = list(grid_layout = list(ncol = 2, nrow = 2), plots = overall_plots),
    delta_plots = list(grid_layout = list(ncol = 4, nrow = 3), plots = delta_plots)
  ))
}

generate_rank_plot_configs <- function(df, variables, is_lm = FALSE, adjust = FALSE, overlap_color = FALSE) {
  sd_bands <- c(1, 2, 3)
  avg_zscore_cols <- paste0("Avg_Zscore_", variables)
  z_lm_cols <- paste0("Z_lm_", variables)
  rank_avg_zscore_cols <- paste0("Rank_", variables)
  rank_z_lm_cols <- paste0("Rank_lm_", variables)

  configs <- list()

  # Adjust values if necessary
  if (adjust) {
    df <- df %>%
      mutate(
        across(all_of(avg_zscore_cols), ~ifelse(is.na(.), 0.001, .)),
        across(all_of(z_lm_cols), ~ifelse(is.na(.), 0.001, .))
      )
  }

  # Calculate rank columns for Avg_Zscore and Z_lm columns
  df_ranked <- df %>%
    mutate(across(all_of(avg_zscore_cols), ~rank(., na.last = "keep"), .names = paste0("Rank_", avg_zscore_cols))) %>%
    mutate(across(all_of(z_lm_cols), ~rank(., na.last = "keep"), .names = paste0("Rank_lm_", z_lm_cols)))

  # Generate plots for SD-based L and K variables
  for (variable in c("L", "K")) {
    rank_var <- if (is_lm) paste0("Rank_lm_", variable) else paste0("Rank_", variable)
    zscore_var <- if (is_lm) paste0("Z_lm_", variable) else paste0("Avg_Zscore_", variable)
    y_label <- if (is_lm) paste("Int Z score", variable) else paste("Avg Z score", variable)

    for (sd_band in sd_bands) {
      num_enhancers <- sum(df_ranked[[zscore_var]] >= sd_band, na.rm = TRUE)
      num_suppressors <- sum(df_ranked[[zscore_var]] <= -sd_band, na.rm = TRUE)

      # Plot with annotations
      configs[[length(configs) + 1]] <- list(
        df = df_ranked,
        x_var = rank_var,
        y_var = zscore_var,
        plot_type = "scatter",
        title = paste(y_label, "vs. Rank for", variable, "above", sd_band, "SD"),
        sd_band = sd_band,
        fill_positive = "#542788",
        fill_negative = "orange",
        alpha_positive = 0.3,
        alpha_negative = 0.3,
        annotations = list(
          list(
            x = median(df_ranked[[rank_var]], na.rm = TRUE),
            y = max(df_ranked[[zscore_var]], na.rm = TRUE) * 0.9,
            label = paste("Deletion Enhancers =", num_enhancers)
          ),
          list(
            x = median(df_ranked[[rank_var]], na.rm = TRUE),
            y = min(df_ranked[[zscore_var]], na.rm = TRUE) * 0.9,
            label = paste("Deletion Suppressors =", num_suppressors)
          )
        ),
        shape = 3,
        size = 0.1,
        y_label = y_label,
        x_label = "Rank",
        legend_position = "none"
      )

      # Plot without annotations
      configs[[length(configs) + 1]] <- list(
        df = df_ranked,
        x_var = rank_var,
        y_var = zscore_var,
        plot_type = "scatter",
        title = paste(y_label, "vs. Rank for", variable, "above", sd_band, "SD No Annotations"),
        sd_band = sd_band,
        fill_positive = "#542788",
        fill_negative = "orange",
        alpha_positive = 0.3,
        alpha_negative = 0.3,
        annotations = NULL,
        shape = 3,
        size = 0.1,
        y_label = y_label,
        x_label = "Rank",
        legend_position = "none"
      )
    }
  }

  # Generate Avg ZScore and Rank Avg ZScore plots for each variable
  for (variable in variables) {
    for (plot_type in c("Avg Zscore vs lm", "Rank Avg Zscore vs lm")) {
      title <- paste(plot_type, variable)

      # Define specific variables based on plot type
      x_var <- if (plot_type == "Avg Zscore vs lm") paste0("Avg_Zscore_", variable) else paste0("Rank_", variable)
      y_var <- if (plot_type == "Avg Zscore vs lm") paste0("Z_lm_", variable) else paste0("Rank_lm_", variable)

      # Fit the linear model
      lm_model <- lm(as.formula(paste(y_var, "~", x_var)), data = df_ranked)
      intercept <- coef(lm_model)[1]
      slope <- coef(lm_model)[2]
      r_squared <- summary(lm_model)$r.squared

      annotations <- list(
        list(
          x = mean(range(df_ranked[[x_var]], na.rm = TRUE)),
          y = mean(range(df_ranked[[y_var]], na.rm = TRUE)),
          label = paste("R-squared =", round(r_squared, 2)),
          hjust = 0.5,
          vjust = 1,
          size = 5
        )
      )

      rectangles <- if (plot_type == "Avg Zscore vs lm") {
        list(list(xmin = -2, xmax = 2, ymin = -2, ymax = 2, fill = NA, color = "grey20", alpha = 0.1))
      } else {
        NULL
      }

      configs[[length(configs) + 1]] <- list(
        df = df_ranked,
        x_var = x_var,
        y_var = y_var,
        plot_type = "scatter",
        title = title,
        annotations = annotations,
        shape = 3,
        size = 0.25,
        smooth = TRUE,
        smooth_color = "black",
        lm_line = list(intercept = intercept, slope = slope),
        legend_position = "right",
        color_var = if (overlap_color) "Overlap" else NULL,
        x_label = x_var,
        y_label = y_var,
        rectangles = rectangles
      )
    }
  }

  return(list(grid_layout = list(ncol = 3, nrow = 2), plots = configs))
}

generate_correlation_plot_configs <- function(df, highlight_cyan = FALSE) {
  relationships <- list(
    list(x = "Z_lm_L", y = "Z_lm_K", label = "Interaction L vs. Interaction K"),
    list(x = "Z_lm_L", y = "Z_lm_r", label = "Interaction L vs. Interaction r"),
    list(x = "Z_lm_L", y = "Z_lm_AUC", label = "Interaction L vs. Interaction AUC"),
    list(x = "Z_lm_K", y = "Z_lm_r", label = "Interaction K vs. Interaction r"),
    list(x = "Z_lm_K", y = "Z_lm_AUC", label = "Interaction K vs. Interaction AUC"),
    list(x = "Z_lm_r", y = "Z_lm_AUC", label = "Interaction r vs. Interaction AUC")
  )

  plots <- list()

  for (rel in relationships) {
    lm_model <- lm(as.formula(paste(rel$y, "~", rel$x)), data = df)
    r_squared <- summary(lm_model)$r.squared

    config <- list(
      df = df,
      x_var = rel$x,
      y_var = rel$y,
      plot_type = "scatter",
      title = rel$label,
      annotations = list(
        list(x = mean(df[[rel$x]], na.rm = TRUE),
             y = mean(df[[rel$y]], na.rm = TRUE),
             label = paste("R-squared =", round(r_squared, 3)))
      ),
      smooth = TRUE,
      smooth_color = "tomato3",
      lm_line = list(intercept = coef(lm_model)[1], slope = coef(lm_model)[2]),
      shape = 3,
      size = 0.5,
      color_var = "Overlap",
      cyan_points = highlight_cyan
    )

    plots <- append(plots, list(config))
  }

  return(list(grid_layout = list(ncol = 3, nrow = 2), plots = plots))
}

main <- function() {
  lapply(names(args$experiments), function(exp_name) {
    exp <- args$experiments[[exp_name]]
    exp_path <- exp$path
    exp_sd <- exp$sd
    out_dir <- file.path(exp_path, "zscores")
    out_dir_qc <- file.path(exp_path, "zscores", "qc")
    dir.create(out_dir, recursive = TRUE, showWarnings = FALSE)
    dir.create(out_dir_qc, recursive = TRUE, showWarnings = FALSE)

    summary_vars <- c("L", "K", "r", "AUC", "delta_bg") # fields to filter and calculate summary stats across
    interaction_vars <- c("L", "K", "r", "AUC") # fields to calculate interaction z-scores
    
    message("Loading and filtering data for experiment: ", exp_name)
    df <- load_and_filter_data(args$easy_results_file, sd = exp_sd) %>%
      update_gene_names(args$sgd_gene_list) %>%
      as_tibble()

    # Filter rows above delta background tolerance
    df_above_tolerance <- df %>% filter(DB == 1)
    df_na <- df %>% mutate(across(all_of(summary_vars), ~ ifelse(DB == 1, NA, .))) # formerly X
    df_no_zeros <- df_na %>% filter(L > 0) # formerly X_noZero
    
    # Save some constants
    max_conc <- max(df$conc_num_factor)
    l_half_median <- (median(df_above_tolerance$L, na.rm = TRUE)) / 2
    k_half_median <- (median(df_above_tolerance$K, na.rm = TRUE)) / 2

    message("Calculating summary statistics before quality control")
    df_stats <- calculate_summary_stats(
      df = df,
      variables = summary_vars,
      group_vars = c("conc_num", "conc_num_factor", "conc_num_factor_factor"))$df_with_stats

    message("Calculating summary statistics after quality control")
    ss <- calculate_summary_stats(
      df = df_na,
      variables = summary_vars,
      group_vars = c("conc_num", "conc_num_factor", "conc_num_factor_factor"))
    df_na_ss <- ss$summary_stats
    df_na_stats <- ss$df_with_stats
    write.csv(df_na_ss, file = file.path(out_dir, "summary_stats_all_strains.csv"), row.names = FALSE)
    # For plotting (ggplot warns on NAs)
    df_na_stats_filtered <- df_na_stats %>% filter(if_all(all_of(summary_vars), is.finite))

    df_na_stats <- df_na_stats %>%
      mutate(
        WT_L = mean_L,
        WT_K = mean_K,
        WT_r = mean_r,
        WT_AUC = mean_AUC,
        WT_sd_L = sd_L,
        WT_sd_K = sd_K,
        WT_sd_r = sd_r,
        WT_sd_AUC = sd_AUC
      )

    # Pull the background means and standard deviations from zero concentration for interactions
    bg_stats <- df_na_stats %>%
      filter(conc_num == 0) %>%
      summarise(
        mean_L = first(mean_L),
        mean_K = first(mean_K),
        mean_r = first(mean_r),
        mean_AUC = first(mean_AUC),
        sd_L = first(sd_L),
        sd_K = first(sd_K),
        sd_r = first(sd_r),
        sd_AUC = first(sd_AUC)
      )

    message("Calculating summary statistics after quality control excluding zero values")
    df_no_zeros_stats <- calculate_summary_stats(
      df = df_no_zeros,
      variables = summary_vars,
      group_vars = c("conc_num", "conc_num_factor", "conc_num_factor_factor")
    )$df_with_stats

    message("Filtering by 2SD of K")
    df_na_within_2sd_k <- df_na_stats %>%
      filter(K >= (mean_K - 2 * sd_K) & K <= (mean_K + 2 * sd_K))
    df_na_outside_2sd_k <- df_na_stats %>%
      filter(K < (mean_K - 2 * sd_K) | K > (mean_K + 2 * sd_K))

    message("Calculating summary statistics for L within 2SD of K")
    # TODO We're omitting the original z_max calculation, not sure if needed?
    ss <- calculate_summary_stats(df_na_within_2sd_k, "L", group_vars = c("conc_num", "conc_num_factor", "conc_num_factor_factor"))$summary_stats
    write.csv(ss,
      file = file.path(out_dir_qc, "max_observed_L_vals_for_spots_within_2sd_K.csv"),
      row.names = FALSE)
    
    message("Calculating summary statistics for L outside 2SD of K")
    ss <- calculate_summary_stats(df_na_outside_2sd_k, "L", group_vars = c("conc_num", "conc_num_factor", "conc_num_factor_factor"))
    df_na_l_outside_2sd_k_stats <- ss$df_with_stats
    write.csv(ss$summary_stats,
      file = file.path(out_dir, "max_observed_L_vals_for_spots_outside_2sd_K.csv"),
      row.names = FALSE)
    
    # Each plots list corresponds to a file
    l_vs_k_plot_configs <- list(
      plots = list(
        list(
          df = df,
          x_var = "L",
          y_var = "K",
          plot_type = "scatter",
          tooltip_vars = c("OrfRep", "Gene", "delta_bg"),
          title = "Raw L vs K before quality control",
          color_var = "conc_num_factor_factor",
          error_bar = FALSE,
          legend_position = "right"
        )
      )
    )

    frequency_delta_bg_plot_configs <- list(
      plots = list(
        list(
          df = df_stats,
          x_var = "delta_bg",
          y_var = NULL,
          plot_type = "density",
          title = "Density plot for Delta Background by [Drug] (All Data)",
          color_var = "conc_num_factor_factor",
          x_label = "Delta Background",
          y_label = "Density",
          error_bar = FALSE,
          legend_position = "right"
        ),
        list(
          df = df_stats,
          x_var = "delta_bg",
          y_var = NULL,
          plot_type = "bar",
          title = "Bar plot for Delta Background by [Drug] (All Data)",
          color_var = "conc_num_factor_factor",
          x_label = "Delta Background",
          y_label = "Count",
          error_bar = FALSE,
          legend_position = "right"
        )
      )
    )

    above_threshold_plot_configs <- list(
      plots = list(
        list(
          df = df_above_tolerance,
          x_var = "L",
          y_var = "K",
          plot_type = "scatter",
          tooltip_vars = c("OrfRep", "Gene", "delta_bg"),
          title = paste("Raw L vs K for strains above Delta Background threshold of",
            round(df_above_tolerance$delta_bg_tolerance[[1]], 3), "or above"),
          color_var = "conc_num_factor_factor",
          position = "jitter",
          annotations = list(
            list(
              x = l_half_median,
              y = k_half_median,
              label = paste("# strains above Delta Background tolerance =", nrow(df_above_tolerance))
            )
          ),
          error_bar = FALSE,
          legend_position = "right"
        )
      )
    )

    plate_analysis_plot_configs <- generate_plate_analysis_plot_configs(
      variables = summary_vars,
      df_before = df_stats,
      df_after = df_na_stats_filtered
    )

    plate_analysis_boxplot_configs <- generate_plate_analysis_plot_configs(
      variables = summary_vars,
      df_before = df_stats,
      df_after = df_na_stats_filtered,
      plot_type = "box"
    )

    plate_analysis_no_zeros_plot_configs <- generate_plate_analysis_plot_configs(
      variables = summary_vars,
      stages = c("after"),  # Only after QC
      df_after = df_no_zeros_stats
    )

    plate_analysis_no_zeros_boxplot_configs <- generate_plate_analysis_plot_configs(
      variables = summary_vars,
      stages = c("after"),  # Only after QC
      df_after = df_no_zeros_stats,
      plot_type = "box"
    )

    l_outside_2sd_k_plot_configs <- list(
      plots = list(
        list(
          df = df_na_l_outside_2sd_k_stats,
          x_var = "L",
          y_var = "K",
          plot_type = "scatter",
          title = "Raw L vs K for strains falling outside 2SD of the K mean at each Conc",
          color_var = "conc_num_factor_factor",
          position = "jitter",  # Apply jitter for better visibility
          tooltip_vars = c("OrfRep", "Gene", "delta_bg"),
          annotations = list(
            list(
              x = mean(df_na_l_outside_2sd_k_stats$L, na.rm = TRUE),
              y = mean(df_na_l_outside_2sd_k_stats$K, na.rm = TRUE),
              label = paste("Total strains:", nrow(df_na_l_outside_2sd_k_stats))
            )
          ),
          error_bar = FALSE,  # No error bars for this plot
          legend_position = "right"
        )
      )
    )

    delta_bg_outside_2sd_k_plot_configs <- list(
      plots = list(
        list(
          df = df_na_l_outside_2sd_k_stats,
          x_var = "delta_bg",
          y_var = "K",
          plot_type = "scatter",
          title = "Delta Background vs K for strains falling outside 2SD of the K mean at each Conc",
          color_var = "conc_num_factor_factor",
          position = "jitter",  # Apply jitter for better visibility
          tooltip_vars = c("OrfRep", "Gene", "delta_bg"),
          annotations = list(
            list(
              x = mean(df_na_l_outside_2sd_k_stats$delta_bg, na.rm = TRUE),
              y = mean(df_na_l_outside_2sd_k_stats$K, na.rm = TRUE),
              label = paste("Total strains:", nrow(df_na_l_outside_2sd_k_stats))
            )
          ),
          error_bar = FALSE,  # No error bars for this plot
          legend_position = "right"
        )
      )
    )

    message("Generating quality control plots in parallel")
    # # future::plan(future::multicore, workers = parallel::detectCores())
    future::plan(future::multisession, workers = 3) # generate 3 plots in parallel

    plot_configs <- list(
      list(out_dir = out_dir_qc, filename = "L_vs_K_before_quality_control",
        plot_configs = l_vs_k_plot_configs),
      list(out_dir = out_dir_qc, filename = "frequency_delta_background",
        plot_configs = frequency_delta_bg_plot_configs),
      list(out_dir = out_dir_qc, filename = "L_vs_K_above_threshold",
        plot_configs = above_threshold_plot_configs),
      list(out_dir = out_dir_qc, filename = "plate_analysis",
        plot_configs = plate_analysis_plot_configs),
      list(out_dir = out_dir_qc, filename = "plate_analysis_boxplots",
        plot_configs = plate_analysis_boxplot_configs),
      list(out_dir = out_dir_qc, filename = "plate_analysis_no_zeros",
        plot_configs = plate_analysis_no_zeros_plot_configs),
      list(out_dir = out_dir_qc, filename = "plate_analysis_no_zeros_boxplots",
        plot_configs = plate_analysis_no_zeros_boxplot_configs),
      list(out_dir = out_dir_qc, filename = "L_vs_K_for_strains_2SD_outside_mean_K",
        plot_configs = l_outside_2sd_k_plot_configs),
      list(out_dir = out_dir_qc, filename = "delta_background_vs_K_for_strains_2sd_outside_mean_K",
        plot_configs = delta_bg_outside_2sd_k_plot_configs)
    )

    # Generating quality control plots in parallel
    furrr::future_map(plot_configs, function(config) {
      generate_and_save_plots(config$out_dir, config$filename, config$plot_configs)
    }, .options = furrr_options(seed = TRUE))

    # Process background strains
    bg_strains <- c("YDL227C")
    lapply(bg_strains, function(strain) {
      
      message("Processing background strain: ", strain)
      
      # Handle missing data by setting zero values to NA
      # and then removing any rows with NA in L col
      df_bg <- df_na %>%
        filter(OrfRep == strain) %>%
        mutate(
          L = if_else(L == 0, NA, L),
          K = if_else(K == 0, NA, K),
          r = if_else(r == 0, NA, r),
          AUC = if_else(AUC == 0, NA, AUC)
        ) %>%
        filter(!is.na(L))
      
      # Recalculate summary statistics for the background strain
      message("Calculating summary statistics for background strain")
      ss_bg <- calculate_summary_stats(df_bg, summary_vars, group_vars = c("OrfRep", "conc_num", "conc_num_factor", "conc_num_factor_factor"))
      summary_stats_bg <- ss_bg$summary_stats
      write.csv(summary_stats_bg,
        file = file.path(out_dir, paste0("summary_stats_background_strain_", strain, ".csv")),
        row.names = FALSE)
      
      # Set the missing values to the highest theoretical value at each drug conc for L
      # Leave other values as 0 for the max/min
      df_reference <- df_na_stats %>% # formerly X2_RF
        filter(OrfRep == strain) %>%
        filter(!is.na(L)) %>%
        group_by(conc_num, conc_num_factor, conc_num_factor_factor) %>%
        mutate(
          max_l_theoretical = max(max_L, na.rm = TRUE),
          L = ifelse(L == 0 & !is.na(L) & conc_num > 0, max_l_theoretical, L),
          SM = ifelse(L >= max_l_theoretical & !is.na(L) & conc_num > 0, 1, 0),
          L = ifelse(L >= max_l_theoretical & !is.na(L) & conc_num > 0, max_l_theoretical, L)) %>%
        ungroup()

      # Ditto for deletion strains
      df_deletion <- df_na_stats %>% # formerly X2
        filter(OrfRep != strain) %>%
        filter(!is.na(L)) %>%
        mutate(SM = 0) %>%
        group_by(conc_num, conc_num_factor, conc_num_factor_factor) %>%
        mutate(
          max_l_theoretical = max(max_L, na.rm = TRUE),
          L = ifelse(L == 0 & !is.na(L) & conc_num > 0, max_l_theoretical, L),
          SM = ifelse(L >= max_l_theoretical & !is.na(L) & conc_num > 0, 1, SM),
          L = ifelse(L >= max_l_theoretical & !is.na(L) & conc_num > 0, max_l_theoretical, L)) %>%
        ungroup()

      message("Calculating reference strain interaction scores")
      df_reference_stats <- calculate_summary_stats(
        df = df_reference,
        variables = interaction_vars,
        group_vars = c("OrfRep", "Gene", "num", "conc_num", "conc_num_factor", "conc_num_factor_factor")
        )$df_with_stats
      reference_results <- calculate_interaction_scores(df_reference_stats, max_conc, bg_stats, group_vars = c("OrfRep", "Gene", "num"))
      zscore_calculations_reference <- reference_results$calculations
      zscore_interactions_reference <- reference_results$interactions
      zscore_interactions_reference_joined <- reference_results$full_data

      message("Calculating deletion strain(s) interactions scores")
      df_deletion_stats <- calculate_summary_stats(
        df = df_deletion,
        variables = interaction_vars,
        group_vars = c("OrfRep", "Gene", "conc_num", "conc_num_factor", "conc_num_factor_factor")
        )$df_with_stats
      deletion_results <- calculate_interaction_scores(df_deletion_stats, max_conc, bg_stats, group_vars = c("OrfRep", "Gene"))
      zscore_calculations <- deletion_results$calculations
      zscore_interactions <- deletion_results$interactions
      zscore_interactions_joined <- deletion_results$full_data

      # Writing Z-Scores to file
      write.csv(zscore_calculations_reference, file = file.path(out_dir, "zscore_calculations_reference.csv"), row.names = FALSE)
      write.csv(zscore_calculations, file = file.path(out_dir, "zscore_calculations.csv"), row.names = FALSE)
      write.csv(zscore_interactions_reference, file = file.path(out_dir, "zscore_interactions_reference.csv"), row.names = FALSE)
      write.csv(zscore_interactions, file = file.path(out_dir, "zscore_interactions.csv"), row.names = FALSE)

      # Create interaction plots
      message("Generating reference interaction plots")
      reference_plot_configs <- generate_interaction_plot_configs(zscore_interactions_reference_joined, plot_type = "reference")
      generate_and_save_plots(out_dir, "interaction_plots_reference", reference_plot_configs)

      message("Generating deletion interaction plots")
      deletion_plot_configs <- generate_interaction_plot_configs(zscore_interactions_joined, plot_type = "deletion")
      generate_and_save_plots(out_dir, "interaction_plots", deletion_plot_configs)

      # Define conditions for enhancers and suppressors
      # TODO Add to study config?
      threshold <- 2
      enhancer_condition_L <- zscore_interactions$Avg_Zscore_L >= threshold
      suppressor_condition_L <- zscore_interactions$Avg_Zscore_L <= -threshold
      enhancer_condition_K <- zscore_interactions$Avg_Zscore_K >= threshold
      suppressor_condition_K <- zscore_interactions$Avg_Zscore_K <= -threshold
      
      # Subset data
      enhancers_L <- zscore_interactions[enhancer_condition_L, ]
      suppressors_L <- zscore_interactions[suppressor_condition_L, ]
      enhancers_K <- zscore_interactions[enhancer_condition_K, ]
      suppressors_K <- zscore_interactions[suppressor_condition_K, ]
      
      # Save enhancers and suppressors
      message("Writing enhancer/suppressor csv files")
      write.csv(enhancers_L, file = file.path(out_dir, "zscore_interactions_deletion_enhancers_L.csv"), row.names = FALSE)
      write.csv(suppressors_L, file = file.path(out_dir, "zscore_interactions_deletion_suppressors_L.csv"), row.names = FALSE)
      write.csv(enhancers_K, file = file.path(out_dir, "zscore_interactions_deletion_enhancers_K.csv"), row.names = FALSE)
      write.csv(suppressors_K, file = file.path(out_dir, "zscore_interactions_deletion_suppressors_K.csv"), row.names = FALSE)
      
      # Combine conditions for enhancers and suppressors
      enhancers_and_suppressors_L <- zscore_interactions[enhancer_condition_L | suppressor_condition_L, ]
      enhancers_and_suppressors_K <- zscore_interactions[enhancer_condition_K | suppressor_condition_K, ]
      
      # Save combined enhancers and suppressors
      write.csv(enhancers_and_suppressors_L,
        file = file.path(out_dir, "zscore_interactions_deletion_enhancers_and_suppressors_L.csv"), row.names = FALSE)
      write.csv(enhancers_and_suppressors_K,
        file = file.path(out_dir, "zscore_interaction_deletion_enhancers_and_suppressors_K.csv"), row.names = FALSE)
      
      # Handle linear model based enhancers and suppressors
      lm_threshold <- 2 # TODO add to study config?
      enhancers_lm_L <- zscore_interactions[zscore_interactions$Z_lm_L >= lm_threshold, ]
      suppressors_lm_L <- zscore_interactions[zscore_interactions$Z_lm_L <= -lm_threshold, ]
      enhancers_lm_K <- zscore_interactions[zscore_interactions$Z_lm_K >= lm_threshold, ]
      suppressors_lm_K <- zscore_interactions[zscore_interactions$Z_lm_K <= -lm_threshold, ]
      
      # Save linear model based enhancers and suppressors
      message("Writing linear model enhancer/suppressor csv files")
      write.csv(enhancers_lm_L,
        file = file.path(out_dir, "zscore_interactions_deletion_enhancers_lm_L.csv"), row.names = FALSE)
      write.csv(suppressors_lm_L,
        file = file.path(out_dir, "zscore_interactions_deletion_suppressors_lm_L.csv"), row.names = FALSE)
      write.csv(enhancers_lm_K,
        file = file.path(out_dir, "zscore_interactions_deletion_enhancers_lm_K.csv"), row.names = FALSE)
      write.csv(suppressors_lm_K,
        file = file.path(out_dir, "zscore_interactions_deletion_suppressors_lm_K.csv"), row.names = FALSE)

      message("Generating rank plots")
      rank_plot_configs <- generate_rank_plot_configs(
        df = zscore_interactions_joined,
        variables = interaction_vars,
        is_lm = FALSE,
        adjust = TRUE
      )
      generate_and_save_plots(out_dir = out_dir, filename = "rank_plots",
        plot_configs = rank_plot_configs)

      message("Generating ranked linear model plots")
      rank_lm_plot_configs <- generate_rank_plot_configs(
        df = zscore_interactions_joined,
        variables = interaction_vars,
        is_lm = TRUE,
        adjust = TRUE
      )
      generate_and_save_plots(out_dir = out_dir, filename = "rank_plots_lm",
        plot_configs = rank_lm_plot_configs)

      message("Filtering and reranking plots")
      interaction_threshold <- 2 # TODO add to study config?
      # Formerly X_NArm
      zscore_interactions_filtered <- zscore_interactions_joined %>%
        filter(!is.na(Z_lm_L) & !is.na(Avg_Zscore_L)) %>%
        mutate(
          Overlap = case_when(
            Z_lm_L >= interaction_threshold & Avg_Zscore_L >= interaction_threshold ~ "Deletion Enhancer Both",
            Z_lm_L <= -interaction_threshold & Avg_Zscore_L <= -interaction_threshold ~ "Deletion Suppressor Both",
            Z_lm_L >= interaction_threshold & Avg_Zscore_L <= interaction_threshold ~ "Deletion Enhancer lm only",
            Z_lm_L <= interaction_threshold & Avg_Zscore_L >= interaction_threshold ~ "Deletion Enhancer Avg Zscore only",
            Z_lm_L <= -interaction_threshold & Avg_Zscore_L >= -interaction_threshold ~ "Deletion Suppressor lm only",
            Z_lm_L >= -interaction_threshold & Avg_Zscore_L <= -interaction_threshold ~ "Deletion Suppressor Avg Zscore only",
            Z_lm_L >= interaction_threshold & Avg_Zscore_L <= -interaction_threshold ~ "Deletion Enhancer lm, Deletion Suppressor Avg Z score",
            Z_lm_L <= -interaction_threshold & Avg_Zscore_L >= interaction_threshold ~ "Deletion Suppressor lm, Deletion Enhancer Avg Z score",
            TRUE ~ "No Effect"
          ),
          lm_R_squared_L = summary(lm(Z_lm_L ~ Avg_Zscore_L))$r.squared,
          lm_R_squared_K = summary(lm(Z_lm_K ~ Avg_Zscore_K))$r.squared,
          lm_R_squared_r = summary(lm(Z_lm_r ~ Avg_Zscore_r))$r.squared,
          lm_R_squared_AUC = summary(lm(Z_lm_AUC ~ Avg_Zscore_AUC))$r.squared
        )

      message("Generating filtered ranked plots")
      rank_plot_filtered_configs <- generate_rank_plot_configs(
        df = zscore_interactions_filtered,
        variables = interaction_vars,
        is_lm = FALSE,
        adjust = FALSE,
        overlap_color = TRUE
      )
      generate_and_save_plots(
        out_dir = out_dir,
        filename = "RankPlots_na_rm",
        plot_configs = rank_plot_filtered_configs)

      message("Generating filtered ranked linear model plots")
      rank_plot_lm_filtered_configs <- generate_rank_plot_configs(
        df = zscore_interactions_filtered,
        variables = interaction_vars,
        is_lm = TRUE,
        adjust = FALSE,
        overlap_color = TRUE
      )
      generate_and_save_plots(
        out_dir = out_dir,
        filename = "rank_plots_lm_na_rm",
        plot_configs = rank_plot_lm_filtered_configs)

      message("Generating correlation curve parameter pair plots")
      correlation_plot_configs <- generate_correlation_plot_configs(zscore_interactions_filtered)
      generate_and_save_plots(
        out_dir = out_dir,
        filename = "correlation_cpps",
        plot_configs = correlation_plot_configs,
      )
    })
  })
}
main()